correct
28_DirichletMultinomialLogLikelihood.py
started
2026-05-11T15:54:55Z
timingkernelbench.timing
kernel mean / median / std
2.52 / — / 0.00 μs
ref mean / median / std
0.05 / — / 4.58e-04 μs
numerical precisioneval_kernel_against_ref
gpu memoryextended_metrics.measure_memory
kernel launches / fusiontorch.profiler
fusion ratio (ref / kernel)
2.25x
SOL (speed-of-light)nsight_hardware_counters
roofline / occupancynsight_hardware_counters
warp stalls (top 5)nsight
reference PyTorch source (the prompt the agent sees)
import torch
import torch.nn as nn
class Model(nn.Module):
"""
Log-likelihood under the Dirichlet-Multinomial (Polya) distribution.
Given count vectors n and concentration parameters alpha, computes
log p(n | alpha) using the log-gamma formulation. Core computation
in Bayesian topic models and mixed-membership models.
"""
def __init__(self, vocab_size):
super().__init__()
self.vocab_size = vocab_size
self.log_alpha = nn.Parameter(torch.zeros(vocab_size))
def forward(self, counts: torch.Tensor) -> torch.Tensor:
"""
Args:
counts: (B, V) – integer count vectors (as float for differentiability)
Returns:
log_lik: (B,) – log p(counts | alpha) for each sample
"""
alpha = torch.exp(self.log_alpha).unsqueeze(0) # (1, V)
N = counts.sum(-1) # (B,)
alpha_sum = alpha.sum(-1) # (1,)
# log p(n | alpha) = log Gamma(alpha_sum) - log Gamma(N + alpha_sum)
# + sum_v [ log Gamma(n_v + alpha_v) - log Gamma(alpha_v) ]
log_lik = (
torch.lgamma(alpha_sum)
- torch.lgamma(N + alpha_sum)
+ (torch.lgamma(counts + alpha) - torch.lgamma(alpha)).sum(-1)
)
return log_lik
vocab_size = 256
batch_size = 512
def get_inputs():
bs = popcorn_pri.jitter_int(batch_size)
counts = torch.randint(0, 10, (bs, vocab_size)).float()
return [counts]
def get_init_inputs():
return [vocab_size]
turn 03.8s turn · 3.8s LLM · 0.0s tools · tokens in=1,610 · out=278 · total=1,888
reasoning
Let me analyze this task carefully.
The model computes the log-likelihood under the Dirichlet-Multinomial distribution. The key computation involves:
1. `alpha = torch.exp(self.log_alpha)` - exponential of learnable parameters
2. `N = counts.sum(-1)` - sum of counts per batch
3. `alpha_sum = alpha.sum(-1)` - sum of alpha values
4. Log-gamma computations with these values
The forward pass:
```python
log_lik = (
torch.lgamma(alpha_sum)
- torch.lgamma(N + alpha_sum)
+ (torch.lgamma(counts + alpha) - torch.lgamma(alpha)).sum(-1)
)
```
Let me think about how to optimize this with a custom CUDA kernel:
1. The main computation is `lgamma` function calls, which are expensive
2. We need to compute lgamma for:
- `alpha_sum` (scalar, same for all batches)
- `N + alpha_sum` (per batch)
- `counts + alpha` (B × V elements)
- `alpha` (V elements, same for all batches)
Let me first get the GPU specs to understand what we're working with.
assistantI'll start by analyzing the task and getting GPU specifications to understand the hardware constraints.
call → get_gpu_specs
callget_gpu_specs
{}
result ← get_gpu_specs (ok)
okget_gpu_specs
GPU specs for NVIDIA H100 80GB HBM3:
total memory (runtime): 79.1 GB
spec entry: H100
GPU Architecture: Hopper
GPU Memory: 80GB
Memory Bandwidth: 3.35 TB/s
FP64 TFLOPS: 34
FP64 Tensor Core TFLOPS: 67
FP32 TFLOPS: 67
TF32 Tensor Core TFLOPS: 989 with sparsity
BFLOAT16 Tensore Core TFLOPS: 1979 with sparsity
FP16 Tensor Core TFLOPS: 1979 with sparsity
FP8 Tensor Core TFLOPS: 3958 with sparsity
INT8 Tensor Core TOPS: 3958 with sparsity
Register File Size: 64K 32-bit registers per SM
Maximum number of registers per thread: 255
Maximum number of thread blocks per SM: 32
Shared memory capacity per SM: 228 KB
Maximum shared memory per thread block: 227 KB
turn 1 · 1 tool call(s) · 89.2s turn · 14.1s LLM · tokens in=1,875 · out=1,341 · total=3,216 · 6,602 chars (click to expand)
turn 189.2s turn · 14.1s LLM · 75.2s tools · tokens in=1,875 · out=1,341 · total=3,216
reasoning
Now I understand the hardware. This is H100 with excellent FP32 performance (67 TFLOPS). The kernel needs to compute lgamma functions, which is compute-intensive.
Let me design an optimized CUDA kernel for this computation:
1. The formula is:
log_lik[i] = lgamma(alpha_sum) - lgamma(N[i] + alpha_sum) + sum_v(lgamma(counts[i,v] + alpha[v]) - lgamma(alpha[v]))
2. Key observations:
- alpha_sum is the same for all batch elements
- lgamma(alpha) is the same for all batch elements (can be precomputed)
- lgamma(alpha_sum) is a single value
- lgamma(N[i] + alpha_sum) varies per batch element
- lgamma(counts[i,v] + alpha[v]) varies per (batch, vocab) pair
3. Parallelization strategy:
- Assign each thread to a batch element
- Within each thread, compute the sum over vocab dimension
- Precompute lgamma(alpha[v]) and lgamma(alpha_sum) once
Let me write the kernel:
assistantNow let me analyze the computation and write an optimized CUDA kernel. The key formula involves lgamma operations which are compute-intensive. I'll create a kernel where each thread handles one batch element.
call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n// Fast lgamma approximation using Stirling's series for large values\n// For small values, fall back to device lgamma\n__device__ __forceinline__ float fast_lgamma(float x) {\n return lgammaf(x);\n}\n\n// Kernel: each thread handles one batch element\n// Computes log-likelihood under Dirichlet-Multinomial distribution\n__global__ void dirichlet_multinomial_loglik_kernel(\n const float* __restrict__ counts,\n const float* __restrict__ alpha,\n const float* __restrict__ lgamma_alpha, // precomputed lgamma(alpha)\n float lgamma_alpha_sum, // precomputed lgamma(alpha_sum)\n float alpha_sum,\n float* __restrict__ log_lik,\n int batch_size,\n int vocab_size\n) {\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n if (i >= batch_size) return;\n \n // Compute N = sum of counts for this batch element\n float N = 0.0f;\n float sum_lgamma_n_alpha = 0.0f;\n \n // Process vocab dimension - use vectorized loads if possible\n const float* counts_row = counts + i * vocab_size;\n \n // Unrolled loop for better ILP\n #pragma unroll 4\n for (int v = 0; v < vocab_size; v++) {\n float n_v = counts_row[v];\n N += n_v;\n // lgamma(n_v + alpha_v) - lgamma(alpha_v)\n sum_lgamma_n_alpha += fast_lgamma(n_v + alpha[v]) - lgamma_alpha[v];\n }\n \n // Compute final log-likelihood\n // log p(n | alpha) = lgamma(alpha_sum) - lgamma(N + alpha_sum) + sum\n log_lik[i] = lgamma_alpha_sum - fast_lgamma(N + alpha_sum) + sum_lgamma_n_alpha;\n}\n\n// Host function to launch kernel\ntorch::Tensor dirichlet_multinomial_forward(\n torch::Tensor counts,\n torch::Tensor log_alpha\n) {\n // Ensure contiguous input\n counts = counts.contiguous();\n log_alpha = log_alpha.contiguous();\n \n const int batch_size = counts.size(0);\n const int vocab_size = counts.size(1);\n \n // Compute alpha = exp(log_alpha)\n auto alpha = torch::empty({vocab_size}, counts.options());\n auto lgamma_alpha = torch::empty({vocab_size}, counts.options());\n \n // Compute alpha and lgamma(alpha) on device\n auto alpha_tensor = log_alpha.exp();\n auto lgamma_alpha_tensor = torch::lgamma(alpha_tensor);\n \n // Compute alpha_sum and lgamma(alpha_sum)\n float alpha_sum = alpha_tensor.sum().item<float>();\n float lgamma_alpha_sum = std::lgamma(alpha_sum);\n \n // Output tensor\n auto log_lik = torch::empty({batch_size}, counts.options());\n \n // Launch kernel\n const int threads = 256;\n const int blocks = (batch_size + threads - 1) / threads;\n \n dirichlet_multinomial_loglik_kernel<<<blocks, threads>>>(\n counts.data_ptr<float>(),\n alpha_tensor.data_ptr<float>(),\n lgamma_alpha_tensor.data_ptr<float>(),\n lgamma_alpha_sum,\n alpha_sum,\n log_lik.data_ptr<float>(),\n batch_size,\n vocab_size\n );\n \n return log_lik;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor dirichlet_multinomial_forward(torch::Tensor counts, torch::Tensor log_alpha);\n\"\"\"\n\n# Load the CUDA extension\ndirichlet_ext = load_inline(\n name='dirichlet_multinomial',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['dirichlet_multinomial_forward'],\n extra_cuda_cflags=['-O3', '--use_fast_math'],\n verbose=False\n)\n\n\nclass ModelNew(nn.Module):\n \"\"\"\n Log-likelihood under the Dirichlet-Multinomial (Polya) distribution.\n Given count vectors n and concentration parameters alpha, computes\n log p(n | alpha) using the log-gamma formulation. Core computation\n in Bayesian topic models and mixed-membership models.\n \"\"\"\n\n def __init__(self, vocab_size):\n super().__init__()\n self.vocab_size = vocab_size\n self.log_alpha = nn.Parameter(torch.zeros(vocab_size))\n\n def forward(self, counts: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Args:\n counts: (B, V) \u2013 integer count vectors (as float for differentiability)\n Returns:\n log_lik: (B,) \u2013 log p(counts | alpha) for each sample\n \"\"\"\n return dirichlet_ext.dirichlet_multinomial_forward(counts, self.log_alpha)"
}
result ← compile_kernel (ok)
okcompile_kernel
compile_kernel PASSED: kernel compiled without errors.
turn 277.9s turn · 9.4s LLM · 68.5s tools · tokens in=2,997 · out=1,083 · total=4,080
reasoning
The kernel compiled successfully. Now let me check correctness.
call → run_correctness
callrun_correctness
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n// Fast lgamma approximation using Stirling's series for large values\n// For small values, fall back to device lgamma\n__device__ __forceinline__ float fast_lgamma(float x) {\n return lgammaf(x);\n}\n\n// Kernel: each thread handles one batch element\n// Computes log-likelihood under Dirichlet-Multinomial distribution\n__global__ void dirichlet_multinomial_loglik_kernel(\n const float* __restrict__ counts,\n const float* __restrict__ alpha,\n const float* __restrict__ lgamma_alpha, // precomputed lgamma(alpha)\n float lgamma_alpha_sum, // precomputed lgamma(alpha_sum)\n float alpha_sum,\n float* __restrict__ log_lik,\n int batch_size,\n int vocab_size\n) {\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n if (i >= batch_size) return;\n \n // Compute N = sum of counts for this batch element\n float N = 0.0f;\n float sum_lgamma_n_alpha = 0.0f;\n \n // Process vocab dimension - use vectorized loads if possible\n const float* counts_row = counts + i * vocab_size;\n \n // Unrolled loop for better ILP\n #pragma unroll 4\n for (int v = 0; v < vocab_size; v++) {\n float n_v = counts_row[v];\n N += n_v;\n // lgamma(n_v + alpha_v) - lgamma(alpha_v)\n sum_lgamma_n_alpha += fast_lgamma(n_v + alpha[v]) - lgamma_alpha[v];\n }\n \n // Compute final log-likelihood\n // log p(n | alpha) = lgamma(alpha_sum) - lgamma(N + alpha_sum) + sum\n log_lik[i] = lgamma_alpha_sum - fast_lgamma(N + alpha_sum) + sum_lgamma_n_alpha;\n}\n\n// Host function to launch kernel\ntorch::Tensor dirichlet_multinomial_forward(\n torch::Tensor counts,\n torch::Tensor log_alpha\n) {\n // Ensure contiguous input\n counts = counts.contiguous();\n log_alpha = log_alpha.contiguous();\n \n const int batch_size = counts.size(0);\n const int vocab_size = counts.size(1);\n \n // Compute alpha = exp(log_alpha)\n auto alpha = torch::empty({vocab_size}, counts.options());\n auto lgamma_alpha = torch::empty({vocab_size}, counts.options());\n \n // Compute alpha and lgamma(alpha) on device\n auto alpha_tensor = log_alpha.exp();\n auto lgamma_alpha_tensor = torch::lgamma(alpha_tensor);\n \n // Compute alpha_sum and lgamma(alpha_sum)\n float alpha_sum = alpha_tensor.sum().item<float>();\n float lgamma_alpha_sum = std::lgamma(alpha_sum);\n \n // Output tensor\n auto log_lik = torch::empty({batch_size}, counts.options());\n \n // Launch kernel\n const int threads = 256;\n const int blocks = (batch_size + threads - 1) / threads;\n \n dirichlet_multinomial_loglik_kernel<<<blocks, threads>>>(\n counts.data_ptr<float>(),\n alpha_tensor.data_ptr<float>(),\n lgamma_alpha_tensor.data_ptr<float>(),\n lgamma_alpha_sum,\n alpha_sum,\n log_lik.data_ptr<float>(),\n batch_size,\n vocab_size\n );\n \n return log_lik;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor dirichlet_multinomial_forward(torch::Tensor counts, torch::Tensor log_alpha);\n\"\"\"\n\n# Load the CUDA extension\ndirichlet_ext = load_inline(\n name='dirichlet_multinomial',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['dirichlet_multinomial_forward'],\n extra_cuda_cflags=['-O3', '--use_fast_math'],\n verbose=False\n)\n\n\nclass ModelNew(nn.Module):\n \"\"\"\n Log-likelihood under the Dirichlet-Multinomial (Polya) distribution.\n Given count vectors n and concentration parameters alpha, computes\n log p(n | alpha) using the log-gamma formulation. Core computation\n in Bayesian topic models and mixed-membership models.\n \"\"\"\n\n def __init__(self, vocab_size):\n super().__init__()\n self.vocab_size = vocab_size\n self.log_alpha = nn.Parameter(torch.zeros(vocab_size))\n\n def forward(self, counts: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Args:\n counts: (B, V) \u2013 integer count vectors (as float for differentiability)\n Returns:\n log_lik: (B,) \u2013 log p(counts | alpha) for each sample\n \"\"\"\n return dirichlet_ext.dirichlet_multinomial_forward(counts, self.log_alpha)"
}
result ← run_correctness (ok)
okrun_correctness
run_correctness PASSED: (5 / 5) trials all matched the reference.
Numerical precision: max_abs_err=1.46e-03 mean_abs_err=5.55e-04 max_rel_err=2.42e-07
turn 3 · 1 tool call(s) · 91.5s turn · 12.9s LLM · tokens in=4,126 · out=1,373 · total=5,499 · 6,537 chars (click to expand)
turn 391.5s turn · 12.9s LLM · 78.6s tools · tokens in=4,126 · out=1,373 · total=5,499
reasoning
The kernel passes correctness checks. The max absolute error is 1.46e-03 which is well within the 1e-4 tolerance for fp32 (wait, let me check - the tolerance is atol=rtol=1e-4 for fp32, so 1.46e-03 might be too high? Let me check the actual output - it says "PASSED" with all trials matched, so it must be within tolerance. Actually, looking at the relative error of 2.42e-07, that's much smaller than the tolerance. The tolerance is rtol=1e-4, so the check is abs(a-b) <= atol + rtol * abs(b) = 1e-4 + 1e-4 * abs(b). For typical values around 1000, that would allow absolute errors up to ~0.1. So 1.46e-03 is fine.)
Now let me optimize the kernel further. Key optimizations:
1. Move alpha_sum and lgamma_alpha_sum computation to GPU
2. Use shared memory for alpha values if vocab_size is small enough
3. Better memory access patterns
Let me create a more optimized version:
assistantCorrectness passed. Now let me optimize further by moving more computation to GPU and improving memory access patterns:
call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n// Kernel: each thread handles one batch element\n// Computes log-likelihood under Dirichlet-Multinomial distribution\n__global__ void dirichlet_multinomial_loglik_kernel(\n const float* __restrict__ counts,\n const float* __restrict__ alpha,\n const float* __restrict__ lgamma_alpha,\n float lgamma_alpha_sum,\n float alpha_sum,\n float* __restrict__ log_lik,\n int batch_size,\n int vocab_size\n) {\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n if (i >= batch_size) return;\n \n // Compute N = sum of counts for this batch element\n float N = 0.0f;\n float sum_lgamma_n_alpha = 0.0f;\n \n // Process vocab dimension\n const float* counts_row = counts + i * vocab_size;\n \n // Use vectorized loads where possible (float4 for coalesced access)\n int v = 0;\n \n // Process in chunks of 4 for vectorized loads\n for (; v + 3 < vocab_size; v += 4) {\n float4 c = reinterpret_cast<const float4*>(counts_row + v)[0];\n float4 a = reinterpret_cast<const float4*>(alpha + v)[0];\n float4 lg = reinterpret_cast<const float4*>(lgamma_alpha + v)[0];\n \n N += c.x + c.y + c.z + c.w;\n sum_lgamma_n_alpha += lgammaf(c.x + a.x) - lg.x;\n sum_lgamma_n_alpha += lgammaf(c.y + a.y) - lg.y;\n sum_lgamma_n_alpha += lgammaf(c.z + a.z) - lg.z;\n sum_lgamma_n_alpha += lgammaf(c.w + a.w) - lg.w;\n }\n \n // Handle remaining elements\n for (; v < vocab_size; v++) {\n float n_v = counts_row[v];\n N += n_v;\n sum_lgamma_n_alpha += lgammaf(n_v + alpha[v]) - lgamma_alpha[v];\n }\n \n // Compute final log-likelihood\n log_lik[i] = lgamma_alpha_sum - lgammaf(N + alpha_sum) + sum_lgamma_n_alpha;\n}\n\n// Host function to launch kernel\ntorch::Tensor dirichlet_multinomial_forward(\n torch::Tensor counts,\n torch::Tensor log_alpha\n) {\n // Ensure contiguous input\n counts = counts.contiguous();\n log_alpha = log_alpha.contiguous();\n \n const int batch_size = counts.size(0);\n const int vocab_size = counts.size(1);\n \n // Compute alpha = exp(log_alpha)\n auto alpha_tensor = log_alpha.exp();\n auto lgamma_alpha_tensor = torch::lgamma(alpha_tensor);\n \n // Compute alpha_sum and lgamma(alpha_sum)\n float alpha_sum = alpha_tensor.sum().item<float>();\n float lgamma_alpha_sum = std::lgamma(alpha_sum);\n \n // Output tensor\n auto log_lik = torch::empty({batch_size}, counts.options());\n \n // Launch kernel\n const int threads = 256;\n const int blocks = (batch_size + threads - 1) / threads;\n \n dirichlet_multinomial_loglik_kernel<<<blocks, threads>>>(\n counts.data_ptr<float>(),\n alpha_tensor.data_ptr<float>(),\n lgamma_alpha_tensor.data_ptr<float>(),\n lgamma_alpha_sum,\n alpha_sum,\n log_lik.data_ptr<float>(),\n batch_size,\n vocab_size\n );\n \n return log_lik;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor dirichlet_multinomial_forward(torch::Tensor counts, torch::Tensor log_alpha);\n\"\"\"\n\n# Load the CUDA extension\ndirichlet_ext = load_inline(\n name='dirichlet_multinomial',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['dirichlet_multinomial_forward'],\n extra_cuda_cflags=['-O3', '--use_fast_math'],\n verbose=False\n)\n\n\nclass ModelNew(nn.Module):\n \"\"\"\n Log-likelihood under the Dirichlet-Multinomial (Polya) distribution.\n Given count vectors n and concentration parameters alpha, computes\n log p(n | alpha) using the log-gamma formulation. Core computation\n in Bayesian topic models and mixed-membership models.\n \"\"\"\n\n def __init__(self, vocab_size):\n super().__init__()\n self.vocab_size = vocab_size\n self.log_alpha = nn.Parameter(torch.zeros(vocab_size))\n\n def forward(self, counts: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Args:\n counts: (B, V) \u2013 integer count vectors (as float for differentiability)\n Returns:\n log_lik: (B,) \u2013 log p(counts | alpha) for each sample\n \"\"\"\n return dirichlet_ext.dirichlet_multinomial_forward(counts, self.log_alpha)"
}
result ← compile_kernel (ok)
okcompile_kernel
compile_kernel PASSED: kernel compiled without errors.
turn 410.3s turn · 10.2s LLM · 0.1s tools · tokens in=5,264 · out=1,118 · total=6,382
reasoning
Good, now let me test correctness of the vectorized version.
call → run_correctness
callrun_correctness
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n// Kernel: each thread handles one batch element\n// Computes log-likelihood under Dirichlet-Multinomial distribution\n__global__ void dirichlet_multinomial_loglik_kernel(\n const float* __restrict__ counts,\n const float* __restrict__ alpha,\n const float* __restrict__ lgamma_alpha,\n float lgamma_alpha_sum,\n float alpha_sum,\n float* __restrict__ log_lik,\n int batch_size,\n int vocab_size\n) {\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n if (i >= batch_size) return;\n \n // Compute N = sum of counts for this batch element\n float N = 0.0f;\n float sum_lgamma_n_alpha = 0.0f;\n \n // Process vocab dimension\n const float* counts_row = counts + i * vocab_size;\n \n // Use vectorized loads where possible (float4 for coalesced access)\n int v = 0;\n \n // Process in chunks of 4 for vectorized loads\n for (; v + 3 < vocab_size; v += 4) {\n float4 c = reinterpret_cast<const float4*>(counts_row + v)[0];\n float4 a = reinterpret_cast<const float4*>(alpha + v)[0];\n float4 lg = reinterpret_cast<const float4*>(lgamma_alpha + v)[0];\n \n N += c.x + c.y + c.z + c.w;\n sum_lgamma_n_alpha += lgammaf(c.x + a.x) - lg.x;\n sum_lgamma_n_alpha += lgammaf(c.y + a.y) - lg.y;\n sum_lgamma_n_alpha += lgammaf(c.z + a.z) - lg.z;\n sum_lgamma_n_alpha += lgammaf(c.w + a.w) - lg.w;\n }\n \n // Handle remaining elements\n for (; v < vocab_size; v++) {\n float n_v = counts_row[v];\n N += n_v;\n sum_lgamma_n_alpha += lgammaf(n_v + alpha[v]) - lgamma_alpha[v];\n }\n \n // Compute final log-likelihood\n log_lik[i] = lgamma_alpha_sum - lgammaf(N + alpha_sum) + sum_lgamma_n_alpha;\n}\n\n// Host function to launch kernel\ntorch::Tensor dirichlet_multinomial_forward(\n torch::Tensor counts,\n torch::Tensor log_alpha\n) {\n // Ensure contiguous input\n counts = counts.contiguous();\n log_alpha = log_alpha.contiguous();\n \n const int batch_size = counts.size(0);\n const int vocab_size = counts.size(1);\n \n // Compute alpha = exp(log_alpha)\n auto alpha_tensor = log_alpha.exp();\n auto lgamma_alpha_tensor = torch::lgamma(alpha_tensor);\n \n // Compute alpha_sum and lgamma(alpha_sum)\n float alpha_sum = alpha_tensor.sum().item<float>();\n float lgamma_alpha_sum = std::lgamma(alpha_sum);\n \n // Output tensor\n auto log_lik = torch::empty({batch_size}, counts.options());\n \n // Launch kernel\n const int threads = 256;\n const int blocks = (batch_size + threads - 1) / threads;\n \n dirichlet_multinomial_loglik_kernel<<<blocks, threads>>>(\n counts.data_ptr<float>(),\n alpha_tensor.data_ptr<float>(),\n lgamma_alpha_tensor.data_ptr<float>(),\n lgamma_alpha_sum,\n alpha_sum,\n log_lik.data_ptr<float>(),\n batch_size,\n vocab_size\n );\n \n return log_lik;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor dirichlet_multinomial_forward(torch::Tensor counts, torch::Tensor log_alpha);\n\"\"\"\n\n# Load the CUDA extension\ndirichlet_ext = load_inline(\n name='dirichlet_multinomial',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['dirichlet_multinomial_forward'],\n extra_cuda_cflags=['-O3', '--use_fast_math'],\n verbose=False\n)\n\n\nclass ModelNew(nn.Module):\n \"\"\"\n Log-likelihood under the Dirichlet-Multinomial (Polya) distribution.\n Given count vectors n and concentration parameters alpha, computes\n log p(n | alpha) using the log-gamma formulation. Core computation\n in Bayesian topic models and mixed-membership models.\n \"\"\"\n\n def __init__(self, vocab_size):\n super().__init__()\n self.vocab_size = vocab_size\n self.log_alpha = nn.Parameter(torch.zeros(vocab_size))\n\n def forward(self, counts: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Args:\n counts: (B, V) \u2013 integer count vectors (as float for differentiability)\n Returns:\n log_lik: (B,) \u2013 log p(counts | alpha) for each sample\n \"\"\"\n return dirichlet_ext.dirichlet_multinomial_forward(counts, self.log_alpha)"
}
result ← run_correctness (ok)
okrun_correctness
run_correctness PASSED: (5 / 5) trials all matched the reference.
Numerical precision: max_abs_err=1.46e-03 mean_abs_err=5.55e-04 max_rel_err=2.42e-07
turn 5 · 1 tool call(s) · 89.9s turn · 11.1s LLM · tokens in=6,426 · out=1,180 · total=7,606 · 6,009 chars (click to expand)
turn 589.9s turn · 11.1s LLM · 78.9s tools · tokens in=6,426 · out=1,180 · total=7,606
reasoning
Good, vectorized loads work. Now let me try using shared memory to cache alpha and lgamma_alpha to reduce redundant global memory accesses across threads. Since each thread accesses the same alpha values, caching in shared memory should help.
With vocab_size=256 and batch_size=512, I can:
- Use shared memory for alpha and lgamma_alpha (256 * 4 * 2 = 2KB)
- This allows faster access to alpha values
Let me also try using more threads per thread and a 2D grid structure for better memory access patterns.
assistantGood. Let me try using shared memory to cache the alpha values since they're accessed by all threads:
call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n// Kernel: each thread handles one batch element\n// Computes log-likelihood under Dirichlet-Multinomial distribution\n__global__ void dirichlet_multinomial_loglik_kernel(\n const float* __restrict__ counts,\n const float* __restrict__ alpha,\n const float* __restrict__ lgamma_alpha,\n float lgamma_alpha_sum,\n float alpha_sum,\n float* __restrict__ log_lik,\n int batch_size,\n int vocab_size\n) {\n // Shared memory for alpha and lgamma_alpha (accessed by all threads in block)\n extern __shared__ float shared_data[];\n float* s_alpha = shared_data;\n float* s_lgamma_alpha = shared_data + vocab_size;\n \n // Cooperatively load alpha and lgamma_alpha into shared memory\n for (int v = threadIdx.x; v < vocab_size; v += blockDim.x) {\n s_alpha[v] = alpha[v];\n s_lgamma_alpha[v] = lgamma_alpha[v];\n }\n __syncthreads();\n \n int i = blockIdx.x * blockDim.x + threadIdx.x;\n if (i >= batch_size) return;\n \n // Compute N = sum of counts for this batch element\n float N = 0.0f;\n float sum_lgamma_n_alpha = 0.0f;\n \n // Process vocab dimension\n const float* counts_row = counts + i * vocab_size;\n \n // Process vocab elements\n for (int v = 0; v < vocab_size; v++) {\n float n_v = counts_row[v];\n N += n_v;\n sum_lgamma_n_alpha += lgammaf(n_v + s_alpha[v]) - s_lgamma_alpha[v];\n }\n \n // Compute final log-likelihood\n log_lik[i] = lgamma_alpha_sum - lgammaf(N + alpha_sum) + sum_lgamma_n_alpha;\n}\n\n// Host function to launch kernel\ntorch::Tensor dirichlet_multinomial_forward(\n torch::Tensor counts,\n torch::Tensor log_alpha\n) {\n // Ensure contiguous input\n counts = counts.contiguous();\n log_alpha = log_alpha.contiguous();\n \n const int batch_size = counts.size(0);\n const int vocab_size = counts.size(1);\n \n // Compute alpha = exp(log_alpha)\n auto alpha_tensor = log_alpha.exp();\n auto lgamma_alpha_tensor = torch::lgamma(alpha_tensor);\n \n // Compute alpha_sum and lgamma(alpha_sum)\n float alpha_sum = alpha_tensor.sum().item<float>();\n float lgamma_alpha_sum = std::lgamma(alpha_sum);\n \n // Output tensor\n auto log_lik = torch::empty({batch_size}, counts.options());\n \n // Launch kernel with shared memory\n const int threads = 256;\n const int blocks = (batch_size + threads - 1) / threads;\n const size_t shared_mem_size = vocab_size * 2 * sizeof(float);\n \n dirichlet_multinomial_loglik_kernel<<<blocks, threads, shared_mem_size>>>(\n counts.data_ptr<float>(),\n alpha_tensor.data_ptr<float>(),\n lgamma_alpha_tensor.data_ptr<float>(),\n lgamma_alpha_sum,\n alpha_sum,\n log_lik.data_ptr<float>(),\n batch_size,\n vocab_size\n );\n \n return log_lik;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor dirichlet_multinomial_forward(torch::Tensor counts, torch::Tensor log_alpha);\n\"\"\"\n\n# Load the CUDA extension\ndirichlet_ext = load_inline(\n name='dirichlet_multinomial',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['dirichlet_multinomial_forward'],\n extra_cuda_cflags=['-O3', '--use_fast_math'],\n verbose=False\n)\n\n\nclass ModelNew(nn.Module):\n \"\"\"\n Log-likelihood under the Dirichlet-Multinomial (Polya) distribution.\n Given count vectors n and concentration parameters alpha, computes\n log p(n | alpha) using the log-gamma formulation. Core computation\n in Bayesian topic models and mixed-membership models.\n \"\"\"\n\n def __init__(self, vocab_size):\n super().__init__()\n self.vocab_size = vocab_size\n self.log_alpha = nn.Parameter(torch.zeros(vocab_size))\n\n def forward(self, counts: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Args:\n counts: (B, V) \u2013 integer count vectors (as float for differentiability)\n Returns:\n log_lik: (B,) \u2013 log p(counts | alpha) for each sample\n \"\"\"\n return dirichlet_ext.dirichlet_multinomial_forward(counts, self.log_alpha)"
}
result ← compile_kernel (ok)
okcompile_kernel
compile_kernel PASSED: kernel compiled without errors.
turn 6269.2s turn · 8.8s LLM · 260.3s tools · tokens in=7,505 · out=1,057 · total=8,562
reasoning
Now let me test correctness with the shared memory version.
call → run_correctness
callrun_correctness
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n// Kernel: each thread handles one batch element\n// Computes log-likelihood under Dirichlet-Multinomial distribution\n__global__ void dirichlet_multinomial_loglik_kernel(\n const float* __restrict__ counts,\n const float* __restrict__ alpha,\n const float* __restrict__ lgamma_alpha,\n float lgamma_alpha_sum,\n float alpha_sum,\n float* __restrict__ log_lik,\n int batch_size,\n int vocab_size\n) {\n // Shared memory for alpha and lgamma_alpha (accessed by all threads in block)\n extern __shared__ float shared_data[];\n float* s_alpha = shared_data;\n float* s_lgamma_alpha = shared_data + vocab_size;\n \n // Cooperatively load alpha and lgamma_alpha into shared memory\n for (int v = threadIdx.x; v < vocab_size; v += blockDim.x) {\n s_alpha[v] = alpha[v];\n s_lgamma_alpha[v] = lgamma_alpha[v];\n }\n __syncthreads();\n \n int i = blockIdx.x * blockDim.x + threadIdx.x;\n if (i >= batch_size) return;\n \n // Compute N = sum of counts for this batch element\n float N = 0.0f;\n float sum_lgamma_n_alpha = 0.0f;\n \n // Process vocab dimension\n const float* counts_row = counts + i * vocab_size;\n \n // Process vocab elements\n for (int v = 0; v < vocab_size; v++) {\n float n_v = counts_row[v];\n N += n_v;\n sum_lgamma_n_alpha += lgammaf(n_v + s_alpha[v]) - s_lgamma_alpha[v];\n }\n \n // Compute final log-likelihood\n log_lik[i] = lgamma_alpha_sum - lgammaf(N + alpha_sum) + sum_lgamma_n_alpha;\n}\n\n// Host function to launch kernel\ntorch::Tensor dirichlet_multinomial_forward(\n torch::Tensor counts,\n torch::Tensor log_alpha\n) {\n // Ensure contiguous input\n counts = counts.contiguous();\n log_alpha = log_alpha.contiguous();\n \n const int batch_size = counts.size(0);\n const int vocab_size = counts.size(1);\n \n // Compute alpha = exp(log_alpha)\n auto alpha_tensor = log_alpha.exp();\n auto lgamma_alpha_tensor = torch::lgamma(alpha_tensor);\n \n // Compute alpha_sum and lgamma(alpha_sum)\n float alpha_sum = alpha_tensor.sum().item<float>();\n float lgamma_alpha_sum = std::lgamma(alpha_sum);\n \n // Output tensor\n auto log_lik = torch::empty({batch_size}, counts.options());\n \n // Launch kernel with shared memory\n const int threads = 256;\n const int blocks = (batch_size + threads - 1) / threads;\n const size_t shared_mem_size = vocab_size * 2 * sizeof(float);\n \n dirichlet_multinomial_loglik_kernel<<<blocks, threads, shared_mem_size>>>(\n counts.data_ptr<float>(),\n alpha_tensor.data_ptr<float>(),\n lgamma_alpha_tensor.data_ptr<float>(),\n lgamma_alpha_sum,\n alpha_sum,\n log_lik.data_ptr<float>(),\n batch_size,\n vocab_size\n );\n \n return log_lik;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor dirichlet_multinomial_forward(torch::Tensor counts, torch::Tensor log_alpha);\n\"\"\"\n\n# Load the CUDA extension\ndirichlet_ext = load_inline(\n name='dirichlet_multinomial',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['dirichlet_multinomial_forward'],\n extra_cuda_cflags=['-O3', '--use_fast_math'],\n verbose=False\n)\n\n\nclass ModelNew(nn.Module):\n \"\"\"\n Log-likelihood under the Dirichlet-Multinomial (Polya) distribution.\n Given count vectors n and concentration parameters alpha, computes\n log p(n | alpha) using the log-gamma formulation. Core computation\n in Bayesian topic models and mixed-membership models.\n \"\"\"\n\n def __init__(self, vocab_size):\n super().__init__()\n self.vocab_size = vocab_size\n self.log_alpha = nn.Parameter(torch.zeros(vocab_size))\n\n def forward(self, counts: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Args:\n counts: (B, V) \u2013 integer count vectors (as float for differentiability)\n Returns:\n log_lik: (B,) \u2013 log p(counts | alpha) for each sample\n \"\"\"\n return dirichlet_ext.dirichlet_multinomial_forward(counts, self.log_alpha)"
}
result ← run_correctness (ok)
okrun_correctness
run_correctness PASSED: (5 / 5) trials all matched the reference.
Numerical precision: max_abs_err=1.46e-03 mean_abs_err=5.55e-04 max_rel_err=2.42e-07
turn 786.3s turn · 8.2s LLM · 78.2s tools · tokens in=8,608 · out=873 · total=9,481
reasoning
Correctness passed. Let me try combining shared memory with vectorized loads for the counts (the counts are not cached in shared memory, so vectorized loads would still help there). I should also check if the shared memory version helps or not.
Let me try another approach - use __ldg to cache the read-only access to alpha and lgamma_alpha in the read-only cache. This is simpler than shared memory and might be effective.
assistantLet me try using read-only cache for alpha arrays via `__ldg` intrinsic, which can be more effective than shared memory for broadcast reads:
call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n// Use texture cache for read-only data\n__device__ __forceinline__ float load_float(const float* ptr) {\n return __ldg(ptr);\n}\n\n// Kernel: each thread handles one batch element\n__global__ void dirichlet_multinomial_loglik_kernel(\n const float* __restrict__ counts,\n const float* __restrict__ alpha,\n const float* __restrict__ lgamma_alpha,\n float lgamma_alpha_sum,\n float alpha_sum,\n float* __restrict__ log_lik,\n int batch_size,\n int vocab_size\n) {\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n if (i >= batch_size) return;\n \n float N = 0.0f;\n float sum_lgamma_n_alpha = 0.0f;\n \n const float* counts_row = counts + i * vocab_size;\n \n // Use read-only cache for alpha and lgamma_alpha (broadcast pattern)\n for (int v = 0; v < vocab_size; v++) {\n float n_v = counts_row[v];\n float a_v = load_float(alpha + v);\n N += n_v;\n sum_lgamma_n_alpha += lgammaf(n_v + a_v) - load_float(lgamma_alpha + v);\n }\n \n log_lik[i] = lgamma_alpha_sum - lgammaf(N + alpha_sum) + sum_lgamma_n_alpha;\n}\n\ntorch::Tensor dirichlet_multinomial_forward(\n torch::Tensor counts,\n torch::Tensor log_alpha\n) {\n counts = counts.contiguous();\n log_alpha = log_alpha.contiguous();\n \n const int batch_size = counts.size(0);\n const int vocab_size = counts.size(1);\n \n auto alpha_tensor = log_alpha.exp();\n auto lgamma_alpha_tensor = torch::lgamma(alpha_tensor);\n \n float alpha_sum = alpha_tensor.sum().item<float>();\n float lgamma_alpha_sum = std::lgamma(alpha_sum);\n \n auto log_lik = torch::empty({batch_size}, counts.options());\n \n const int threads = 256;\n const int blocks = (batch_size + threads - 1) / threads;\n \n dirichlet_multinomial_loglik_kernel<<<blocks, threads>>>(\n counts.data_ptr<float>(),\n alpha_tensor.data_ptr<float>(),\n lgamma_alpha_tensor.data_ptr<float>(),\n lgamma_alpha_sum,\n alpha_sum,\n log_lik.data_ptr<float>(),\n batch_size,\n vocab_size\n );\n \n return log_lik;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor dirichlet_multinomial_forward(torch::Tensor counts, torch::Tensor log_alpha);\n\"\"\"\n\ndirichlet_ext = load_inline(\n name='dirichlet_multinomial',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['dirichlet_multinomial_forward'],\n extra_cuda_cflags=['-O3', '--use_fast_math'],\n verbose=False\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self, vocab_size):\n super().__init__()\n self.vocab_size = vocab_size\n self.log_alpha = nn.Parameter(torch.zeros(vocab_size))\n\n def forward(self, counts: torch.Tensor) -> torch.Tensor:\n return dirichlet_ext.dirichlet_multinomial_forward(counts, self.log_alpha)"
}
result ← compile_kernel (ok)
okcompile_kernel
compile_kernel PASSED: kernel compiled without errors.
turn 87.6s turn · 6.9s LLM · 0.7s tools · tokens in=9,406 · out=771 · total=10,177
reasoning
Let me check correctness and then I'll run a static check before submission.
call → run_correctness
callrun_correctness
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n// Use texture cache for read-only data\n__device__ __forceinline__ float load_float(const float* ptr) {\n return __ldg(ptr);\n}\n\n// Kernel: each thread handles one batch element\n__global__ void dirichlet_multinomial_loglik_kernel(\n const float* __restrict__ counts,\n const float* __restrict__ alpha,\n const float* __restrict__ lgamma_alpha,\n float lgamma_alpha_sum,\n float alpha_sum,\n float* __restrict__ log_lik,\n int batch_size,\n int vocab_size\n) {\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n if (i >= batch_size) return;\n \n float N = 0.0f;\n float sum_lgamma_n_alpha = 0.0f;\n \n const float* counts_row = counts + i * vocab_size;\n \n // Use read-only cache for alpha and lgamma_alpha (broadcast pattern)\n for (int v = 0; v < vocab_size; v++) {\n float n_v = counts_row[v];\n float a_v = load_float(alpha + v);\n N += n_v;\n sum_lgamma_n_alpha += lgammaf(n_v + a_v) - load_float(lgamma_alpha + v);\n }\n \n log_lik[i] = lgamma_alpha_sum - lgammaf(N + alpha_sum) + sum_lgamma_n_alpha;\n}\n\ntorch::Tensor dirichlet_multinomial_forward(\n torch::Tensor counts,\n torch::Tensor log_alpha\n) {\n counts = counts.contiguous();\n log_alpha = log_alpha.contiguous();\n \n const int batch_size = counts.size(0);\n const int vocab_size = counts.size(1);\n \n auto alpha_tensor = log_alpha.exp();\n auto lgamma_alpha_tensor = torch::lgamma(alpha_tensor);\n \n float alpha_sum = alpha_tensor.sum().item<float>();\n float lgamma_alpha_sum = std::lgamma(alpha_sum);\n \n auto log_lik = torch::empty({batch_size}, counts.options());\n \n const int threads = 256;\n const int blocks = (batch_size + threads - 1) / threads;\n \n dirichlet_multinomial_loglik_kernel<<<blocks, threads>>>(\n counts.data_ptr<float>(),\n alpha_tensor.data_ptr<float>(),\n lgamma_alpha_tensor.data_ptr<float>(),\n lgamma_alpha_sum,\n alpha_sum,\n log_lik.data_ptr<float>(),\n batch_size,\n vocab_size\n );\n \n return log_lik;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor dirichlet_multinomial_forward(torch::Tensor counts, torch::Tensor log_alpha);\n\"\"\"\n\ndirichlet_ext = load_inline(\n name='dirichlet_multinomial',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['dirichlet_multinomial_forward'],\n extra_cuda_cflags=['-O3', '--use_fast_math'],\n verbose=False\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self, vocab_size):\n super().__init__()\n self.vocab_size = vocab_size\n self.log_alpha = nn.Parameter(torch.zeros(vocab_size))\n\n def forward(self, counts: torch.Tensor) -> torch.Tensor:\n return dirichlet_ext.dirichlet_multinomial_forward(counts, self.log_alpha)"
}
result ← run_correctness (ok)
okrun_correctness
run_correctness PASSED: (5 / 5) trials all matched the reference.
Numerical precision: max_abs_err=1.46e-03 mean_abs_err=5.55e-04 max_rel_err=2.42e-07
turn 97.6s turn · 7.6s LLM · 0.0s tools · tokens in=10,251 · out=865 · total=11,116
reasoning
Let me submit the kernel now. It's passing correctness and I've tried several optimizations. The kernel uses:
1. One thread per batch element
2. Read-only cache (__ldg) for alpha arrays which are broadcast to all threads
3. Coalesced memory access for counts
4. Fast math compiler flag
The computation is lgamma-bound, so the main bottleneck is the lgamma function calls. The kernel should be faster than PyTorch's native implementation due to kernel fusion (avoiding multiple kernel launches for intermediate operations).
call → static_check
callstatic_check
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n// Use texture cache for read-only data\n__device__ __forceinline__ float load_float(const float* ptr) {\n return __ldg(ptr);\n}\n\n// Kernel: each thread handles one batch element\n__global__ void dirichlet_multinomial_loglik_kernel(\n const float* __restrict__ counts,\n const float* __restrict__ alpha,\n const float* __restrict__ lgamma_alpha,\n float lgamma_alpha_sum,\n float alpha_sum,\n float* __restrict__ log_lik,\n int batch_size,\n int vocab_size\n) {\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n if (i >= batch_size) return;\n \n float N = 0.0f;\n float sum_lgamma_n_alpha = 0.0f;\n \n const float* counts_row = counts + i * vocab_size;\n \n // Use read-only cache for alpha and cgamma_alpha (broadcast pattern)\n for (int v = 0; v < vocab_size; v++) {\n float n_v = counts_row[v];\n float a_v = load_float(alpha + v);\n N += n_v;\n sum_lgamma_n_alpha += lgammaf(n_v + a_v) - load_float(lgamma_alpha + v);\n }\n \n log_lik[i] = lgamma_alpha_sum - lgammaf(N + alpha_sum) + sum_lgamma_n_alpha;\n}\n\ntorch::Tensor dirichlet_multinomial_forward(\n torch::Tensor counts,\n torch::Tensor log_alpha\n) {\n counts = counts.contiguous();\n log_alpha = log_alpha.contiguous();\n \n const int batch_size = counts.size(0);\n const int vocab_size = counts.size(1);\n \n auto alpha_tensor = log_alpha.exp();\n auto lgamma_alpha_tensor = torch::lgamma(alpha_tensor);\n \n float alpha_sum = alpha_tensor.sum().item<float>();\n float lgamma_alpha_sum = std::lgamma(alpha_sum);\n \n auto log_lik = torch::empty({batch_size}, counts.options());\n \n const int threads = 256;\n const int blocks = (batch_size + threads - 1) / threads;\n \n dirichlet_multinomial_loglik_kernel<<<blocks, threads>>>(\n counts.data_ptr<float>(),\n alpha_tensor.data_ptr<float>(),\n lgamma_alpha_tensor.data_ptr<float>(),\n lgamma_alpha_sum,\n alpha_sum,\n log_lik.data_ptr<float>(),\n batch_size,\n vocab_size\n );\n \n return log_lik;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor dirichlet_multinomial_forward(torch::Tensor counts, torch::Tensor log_alpha);\n\"\"\"\n\ndirichlet_ext = load_inline(\n name='dirichlet_multinomial',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['dirichlet_multinomial_forward'],\n extra_cuda_cflags=['-O3', '--use_fast_math'],\n verbose=False\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self, vocab_size):\n super().__init__()\n self.vocab_size = vocab_size\n self.log_alpha = nn.Parameter(torch.zeros(vocab_size))\n\n def forward(self, counts: torch.Tensor) -> torch.Tensor:\n return dirichlet_ext.dirichlet_multinomial_forward(counts, self.log_alpha)"
}
result ← static_check (ok)
okstatic_check
static_check PASSED: no violations or warnings detected.
turn 10153.9s turn · 6.7s LLM · 147.2s tools · tokens in=11,052 · out=770 · total=11,822 · FINAL
reasoning
The kernel passes static check and correctness tests. Let me submit it now.
call → submit_kernel
callsubmit_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n// Use texture cache for read-only data\n__device__ __forceinline__ float load_float(const float* ptr) {\n return __ldg(ptr);\n}\n\n// Kernel: each thread handles one batch element\n__global__ void dirichlet_multinomial_loglik_kernel(\n const float* __restrict__ counts,\n const float* __restrict__ alpha,\n const float* __restrict__ lgamma_alpha,\n float lgamma_alpha_sum,\n float alpha_sum,\n float* __restrict__ log_lik,\n int batch_size,\n int vocab_size\n) {\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n if (i >= batch_size) return;\n \n float N = 0.0f;\n float sum_lgamma_n_alpha = 0.0f;\n \n const float* counts_row = counts + i * vocab_size;\n \n // Use read-only cache for alpha and lgamma_alpha (broadcast pattern)\n for (int v = 0; v < vocab_size; v++) {\n float n_v = counts_row[v];\n float a_v = load_float(alpha + v);\n N += n_v;\n sum_lgamma_n_alpha += lgammaf(n_v + a_v) - load_float(lgamma_alpha + v);\n }\n \n log_lik[i] = lgamma_alpha_sum - lgammaf(N + alpha_sum) + sum_lgamma_n_alpha;\n}\n\ntorch::Tensor dirichlet_multinomial_forward(\n torch::Tensor counts,\n torch::Tensor log_alpha\n) {\n counts = counts.contiguous();\n log_alpha = log_alpha.contiguous();\n \n const int batch_size = counts.size(0);\n const int vocab_size = counts.size(1);\n \n auto alpha_tensor = log_alpha.exp();\n auto lgamma_alpha_tensor = torch::lgamma(alpha_tensor);\n \n float alpha_sum = alpha_tensor.sum().item<float>();\n float lgamma_alpha_sum = std::lgamma(alpha_sum);\n \n auto log_lik = torch::empty({batch_size}, counts.options());\n \n const int threads = 256;\n const int blocks = (batch_size + threads - 1) / threads;\n \n dirichlet_multinomial_loglik_kernel<<<blocks, threads>>>(\n counts.data_ptr<float>(),\n alpha_tensor.data_ptr<float>(),\n lgamma_alpha_tensor.data_ptr<float>(),\n lgamma_alpha_sum,\n alpha_sum,\n log_lik.data_ptr<float>(),\n batch_size,\n vocab_size\n );\n \n return log_lik;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor dirichlet_multinomial_forward(torch::Tensor counts, torch::Tensor log_alpha);\n\"\"\"\n\ndirichlet_ext = load_inline(\n name='dirichlet_multinomial',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['dirichlet_multinomial_forward'],\n extra_cuda_cflags=['-O3', '--use_fast_math'],\n verbose=False\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self, vocab_size):\n super().__init__()\n self.vocab_size = vocab_size\n self.log_alpha = nn.Parameter(torch.zeros(vocab_size))\n\n def forward(self, counts: torch.Tensor) -> torch.Tensor:\n return dirichlet_ext.dirichlet_multinomial_forward(counts, self.log_alpha)"
}
result ← submit_kernel (ok)
oksubmit_kernel
submit_kernel PASSED: (5 / 5) correctness trials all passed.
Kernel runtime: 2.52 μs
Runtime stats: mean=2.52μs median=?μs std=0.00μs
Numerical precision: max_abs_err=1.46e-03 mean_abs_err=5.55e-04 max_rel_err=2.42e-07
Memory: 32.51 MB peak (ref: 33.51 MB, ratio: 0.9701x)
Kernel launches: 8 (ref: 18, fusion ratio: 2.25)
Energy: 242.2 mJ/run (ref: 0.0 mJ, ratio: -1x)
SOL score: 0.1807
final submitted kernel (103 lines, 2,960 chars) — full source, untruncated
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>
// Use texture cache for read-only data
__device__ __forceinline__ float load_float(const float* ptr) {
return __ldg(ptr);
}
// Kernel: each thread handles one batch element
__global__ void dirichlet_multinomial_loglik_kernel(
const float* __restrict__ counts,
const float* __restrict__ alpha,
const float* __restrict__ lgamma_alpha,
float lgamma_alpha_sum,
float alpha_sum,
float* __restrict__ log_lik,
int batch_size,
int vocab_size
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= batch_size) return;
float N = 0.0f;
float sum_lgamma_n_alpha = 0.0f;
const float* counts_row = counts + i * vocab_size;
// Use read-only cache for alpha and lgamma_alpha (broadcast pattern)
for (int v = 0; v < vocab_size; v++) {
float n_v = counts_row[v];
float a_v = load_float(alpha + v);
N += n_v;
sum_lgamma_n_alpha += lgammaf(n_v + a_v) - load_float(lgamma_alpha + v);
}
log_lik[i] = lgamma_alpha_sum - lgammaf(N + alpha_sum) + sum_lgamma_n_alpha;
}
torch::Tensor dirichlet_multinomial_forward(
torch::Tensor counts,
torch::Tensor log_alpha
) {
counts = counts.contiguous();
log_alpha = log_alpha.contiguous();
const int batch_size = counts.size(0);
const int vocab_size = counts.size(1);
auto alpha_tensor = log_alpha.exp();
auto lgamma_alpha_tensor = torch::lgamma(alpha_tensor);
float alpha_sum = alpha_tensor.sum().item<float>();
float lgamma_alpha_sum = std::lgamma(alpha_sum);
auto log_lik = torch::empty({batch_size}, counts.options());
const int threads = 256;
const int blocks = (batch_size + threads - 1) / threads;
dirichlet_multinomial_loglik_kernel<<<blocks, threads>>>(
counts.data_ptr<float>(),
alpha_tensor.data_ptr<float>(),
lgamma_alpha_tensor.data_ptr<float>(),
lgamma_alpha_sum,
alpha_sum,
log_lik.data_ptr<float>(),
batch_size,
vocab_size
);
return log_lik;
}
"""
cpp_source = """
torch::Tensor dirichlet_multinomial_forward(torch::Tensor counts, torch::Tensor log_alpha);
"""
dirichlet_ext = load_inline(
name='dirichlet_multinomial',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['dirichlet_multinomial_forward'],
extra_cuda_cflags=['-O3', '--use_fast_math'],
verbose=False
)
class ModelNew(nn.Module):
def __init__(self, vocab_size):
super().__init__()
self.vocab_size = vocab_size
self.log_alpha = nn.Parameter(torch.zeros(vocab_size))
def forward(self, counts: torch.Tensor) -> torch.Tensor:
return dirichlet_ext.dirichlet_multinomial_forward(counts, self.log_alpha)