reasoning
We need to write a ModelNew that implements L2 normalization along dim=1. We have large batch_size=32768 and dim=65535.
Input x is [32768, 65535]. The reference computes:
norm = torch.norm(x, p=2, dim=1, keepdim=True) => compute sqrt(sum_i x[i]^2) over dim=1. Output shape [32768, 1].
out = x / norm
This is memory bandwidth bound. We need to write a fused CUDA kernel that computes the L2 norm per row, then divides each element by the norm. Since it's done along dim=1, rows are contiguous (since x is [batch, dim], row-major, each row contiguous). Actually x shape is [batch_size, dim], so dim=1 dimension contiguous? In PyTorch default memory format, x is row-major: elements in a row contiguous. So dim=1 is contiguous. We can write a kernel where each thread block processes one row (batch element), and threads within a block perform a parallel reduction to compute sum of squares, then broadcast the final norm to divide all elements.
Given dim=65535, one row has 65535 elements. Each block can have up to 1024 threads (standard). For 65535 elements, each thread would handle ~64 elements. We can do a grid-stride loop within each row. For reduction, each thread maintains a local sum of squares, then do a block-level reduction (e.g., using __shared__ mem or warp shuffle). Since sum of squares of float32 might need careful accumulation but fp32 tolerance is 1e-4. Using Kahan? Probably not needed.
We want high memory bandwidth. We'll read input elements once, compute sum of squares, write output. Since norm is compute once per row.
Because H100 has 3.35 TB/s memory bandwidth, reading and writing float32 for 32768*65535 ~ 2.147 billion elements (since total <2^31 actually exactly 2,147,450,880? Wait 32768*65535=2,147,450,880). That's about 8GB read + 8GB write = 16GB. At 3.35TB/s, lower bound is ~4.8ms. Actually 16GB/3.35TB/s = 4.78 ms. Reference torch.norm + division might be somewhat slower due to kernel launch overhead and intermediate storage. We need to beat it.
We might want to leverage vectorized loads (float4) to improve bandwidth. Since row length is 65535, which is not a multiple of 4 (65535 % 4 = 3). We can handle tail with scalar loads.
Or we can use the fact that each thread block processes one row. Let's figure launch configuration.
Block size: 1024 threads per block. Grid: batch_size = 32768 blocks. So 32768 blocks. Each block does reduction over 65535 elements. Since 1024 threads each do ~64 elements. The block reduction adds some overhead but okay.
Alternatively, we can assign multiple blocks per row for larger parallelism, but then need atomic adds or second reduction pass. Since we have 32768 rows, total blocks 32768 is already enough to saturate GPU (H100 has 132 SMs? Actually H100 has 132 SMs? Wait H100 80GB has 132 SMs maybe 144? It has up to 132 SMs for H100 SXM? The spec mentions register file etc but not SM count. Typical H100 SXM5 has 132 SMs. 32768 blocks is enough.
So block-per-row is simplest. We can use warp shuffle reduction. For block reduction, we can accumulate per-thread sums, then reduce within warps using __shfl_down_sync, store warp sums in shared memory, then reduce across warps.
Since we have 1024 threads = 32 warps. Use __shared__ float warpSum[32].
Implementation steps for kernel:
__global__ void l2_norm_kernel(const float* __restrict__ input, float* __restrict__ output, int batch_size, int dim) {
int row = blockIdx.x;
if (row >= batch_size) return;
const float* row_in = input + row * dim;
float* row_out = output + row * dim;
int tid = threadIdx.x;
int lane = tid % 32;
int warpId = tid / 32;
float local_sq_sum = 0.0f;
// Grid-stride loop with vectorized loads (float4)
// Process 4 elements at a time.
int i = tid;
int n_vec = dim / 4;
// each thread handles indices i, i+blockDim.x, etc up to n_vec
for (int idx = i; idx < n_vec; idx += blockDim.x) {
float4 val = reinterpret_cast<const float4*>(row_in)[idx];
local_sq_sum += val.x * val.x;
local_sq_sum += val.y * val.y;
local_sq_sum += val.z * val.z;
local_sq_sum += val.w * val.w;
}
// Handle tail elements
int tail_start = n_vec * 4;
for (int idx = tail_start + tid; idx < dim; idx += blockDim.x) {
float val = row_in[idx];
local_sq_sum += val * val;
}
// Warp reduction
for (int offset = 16; offset > 0; offset /= 2) {
local_sq_sum += __shfl_down_sync(0xffffffff, local_sq_sum, offset);
}
__shared__ float warpSums[32]; // max 32 warps for 1024 threads
if (lane == 0) warpSums[warpId] = local_sq_sum;
__syncthreads();
// Reduce across warps
float sum = 0.0f;
if (tid < 32) {
sum = (tid < (blockDim.x + 31) / 32) ? warpSums[tid] : 0.0f;
for (int offset = 16; offset > 0; offset /= 2) {
sum += __shfl_down_sync(0xffffffff, sum, offset);
}
}
__syncthreads();
float norm = sqrtf(sum);
// To avoid division by zero, if norm == 0 set output = 0.
// However input is rand so norm > 0.
// Write output, vectorized if possible
// Use same loop pattern.
// Since we need to load again or we can compute from original input. We already have original row_in.
// To vectorize writes, we can write float4.
for (int idx = i; idx < n_vec; idx += blockDim.x) {
float4 val = reinterpret_cast<const float4*>(row_in)[idx];
val.x /= norm;
val.y /= norm;
val.z /= norm;
val.w /= norm;
reinterpret_cast<float4*>(row_out)[idx] = val;
}
for (int idx = tail_start + tid; idx < dim; idx += blockDim.x) {
float val = row_in[idx];
row_out[idx] = val / norm;
}
}
We need to be careful about alignment for float4. torch tensors are aligned to at least 256 bytes? Actually default allocator aligns to at least 512 bytes for large allocations. The row start might not be aligned to 16 bytes? Wait in a 2D tensor [batch, dim] with row-major, each row start is at offset row*dim. dim=65535. Since dim is not a multiple of 4, row offsets are row*65535 floats = row*262140 bytes. Is that divisible by 16? 262140 % 16 = 12? Let's compute: 65535*4=262140. 262140 / 16 = 16383.75. Not integer. Wait 65535 is odd, times 4 = 262140, which is divisible by 4 but not by 16? 262140 / 16 = 16383.75, actually 16*16383 = 262128, remainder 12. So row start alignment is not guaranteed to be 16-byte aligned for rows beyond the first. The tensor's base pointer is aligned to 512 bytes maybe, but offset row*dim might not be 16-byte aligned for all rows. For vectorized loads we need alignment. Since row*65535 floats offset mod 4 = row*65535 mod 4. Since 65535 mod 4 = 3, offset mod 4 cycles: row0 mod4=0, row1 mod4=3, row2 mod4=2, row3 mod4=1, row4 mod4=0, etc. So some rows are aligned to 16 bytes (when offset multiple of 4 floats = 16 bytes). In fact row%4==0 gives offset multiple of 4 floats. But others not. float4 loads require 16-byte alignment for each thread's address. Since each thread loads at index idx*4 (in floats), the address is row_start + idx*4*4 = row_start + idx*16 bytes. If row_start is not 16-byte aligned, then idx*16 is multiple of 16, so row_start + idx*16 alignment = row_start alignment. So if row_start is not 16-byte aligned, none of the vector loads will be aligned, causing misaligned address exception for ld.global.v4. We can either use unaligned vector loads? CUDA supports them? On modern GPUs, vector loads with .v4 require alignment to 16 bytes? I think load instructions with .v4 require 16-byte alignment. However if we use the float4 type and pointer cast, compiler may generate a single 128-bit load, but if address is not aligned to 16, it's UB? Actually CUDA compiler will generate a ld instruction with appropriate alignment. If we don't guarantee alignment, using __ldg with a pointer that isn't aligned may lead to misaligned access (not fatal on compute capability >= 7? Many sources say unaligned vector loads are safe but might be compiled as two 64-bit loads. Better to avoid vectorized loads if alignment uncertain.
Alternatively, we can use the built-in vector types but misalignment is generally okay for 128-bit loads on SM 70+? For safety, we can avoid vectorization and just use scalar loads with good coalescing. Since memory bus width is large, scalar loads within a warp will be coalesced automatically. The bandwidth difference may not be huge. But to maximize bandwidth, vectorization can help reduce instruction overhead. We could manually align the first few elements and then do vectorized loads for the bulk. But the bulk start may not be aligned. Another approach: we can transpose the operation? However dim is huge and contiguous.
Given the tolerance, perhaps scalar loads are enough. Let's see. H100 memory bandwidth 3.35 TB/s = 3350 GB/s. For 16 GB of data, 4.8ms. If we get ~90% of that, ~5.3ms. Reference in PyTorch might be around 7-8ms? Let's aim to be close to SOL.
Alternatively, we can process elements with each thread handling multiple contiguous elements (not strided) to increase per-thread work and reduce synchronization overhead. But with 1024 threads per block and dim=65535, each thread does about 64 elements. Strided access pattern (tid, tid+blockDim, ...) is fine, coalesced across threads (thread 0 reads index 0, thread1 reads 1, etc). Actually in the loop for (int idx = tid; idx < dim; idx += blockDim.x) each thread reads contiguous elements separated by blockDim.x. Within a warp, consecutive threads read consecutive elements (since idx = tid + k*blockDim). This is fully coalesced.
If we use vectorized loads, each thread reads 4 consecutive elements. But coalescing also works. So scalar is fine.
One nuance: we need to consider that the reference computes norm via torch.norm which likely uses a vectorized kernel that may be quite optimized. But reference does two passes: first compute norm (reduction), then division (elementwise). That means reading x twice from memory. Our fused kernel reads x once and writes output once. So total memory traffic: read x (8GB) + write output (8GB) = 16GB. Reference: read x for norm (~8GB) + write norm (small) + read x for division (~8GB) + write output (8GB) = ~24GB. So fused kernel should have ~1.5x advantage in memory bandwidth.
We should still aim for efficient implementation.
Let's think about using shared memory reduction vs warp shuffle. Warp shuffle is faster but we can also use shared memory for block reduction. We'll do per-thread accumulation, warp shuffle reduce to lane0, store in shared, then reduce warps. That is efficient.
Given dim=65535, blockDim.x can be e.g., 512 or 1024. Let's pick 512 or 1024. Using 1024 threads means 64 warps? Wait 1024/32=32 warps. 32 warps per block. Each warp handles 32 elements per iteration. In 64 iterations, 64*32 = 2048 elements per thread? Wait each thread iterates over dim with step blockDim.x. So each thread processes ceil(dim/blockDim.x) ~ ceil(65535/1024) = 64 elements. So each warp processes 32*64 = 2048 elements. That's fine.
But we need to be careful about blockDim.x > dim? Here dim is 65535, blockDim.x can be up to 1024, fine.
We should also consider using __restrict__ to allow compiler optimizations.
Potential optimization: use fast math sqrt? rsqrt maybe: reciprocal sqrt. We compute 1/sqrt(sum) and multiply each element. This is faster. We'll compute float inv_norm = rsqrtf(sum); then out = x * inv_norm. This avoids division per element. Since fp32 tolerance 1e-4, rsqrtf is accurate enough? rsqrtf is approximate but should be within tolerance for typical values? Since sum is O(dim), maybe ~65535*0.33 ~ 21600, rsqrtf accurate to ~1 ulp? Actually rsqrtf is approximate within 2 ulp? CUDA's rsqrtf is a fast approximate instruction that is accurate to ~1 ulp for normal inputs. It's used commonly. We could use rsqrtf for speed. Or we can compute norm = sqrtf(sum); and then multiply by 1.0f/norm. Since division per element is more expensive than multiplication. But using rsqrtf on sum and then multiply each element should be safe. Or compute inv_norm = 1.0f / sqrtf(sum). The division is done once per row. That's negligible.
Given we want the best performance, using multiplication by inv_norm is better.
Let's write kernel:
__global__ void l2_norm_forward(const float* __restrict__ input, float* __restrict__ output, int batch_size, int dim) {
int row = blockIdx.x;
if (row >= batch_size) return;
const float* row_in = input + (int64_t)row * dim;
float* row_out = output + (int64_t)row * dim;
int tid = threadIdx.x;
float local_sum = 0.0f;
#pragma unroll 4
for (int idx = tid; idx < dim; idx += blockDim.x) {
float val = row_in[idx];
local_sum += val * val;
}
// warp reduce
unsigned mask = 0xffffffff;
for (int offset = 16; offset > 0; offset >>= 1) {
local_sum += __shfl_down_sync(mask, local_sum, offset);
}
__shared__ float sdata[32];
int warpId = tid >> 5;
int lane = tid & 31;
if (lane == 0) sdata[warpId] = local_sum;
__syncthreads();
if (warpId == 0) {
local_sum = (tid < (blockDim.x >> 5)) ? sdata[tid] : 0.0f;
for (int offset = 16; offset > 0; offset >>= 1) {
local_sum += __shfl_down_sync(mask, local_sum, offset);
}
}
__syncthreads(); // ensure all threads see the result
float inv_norm = 0.0f;
if (tid == 0) {
float norm = sqrtf(local_sum);
inv_norm = (norm > 0.0f) ? (1.0f / norm) : 0.0f;
sdata[0] = inv_norm;
}
__syncthreads();
inv_norm = sdata[0];
// write output
#pragma unroll 4
for (int idx = tid; idx < dim; idx += blockDim.x) {
float val = row_in[idx];
row_out[idx] = val * inv_norm;
}
}
Wait, we need to handle warpId==0 reducing across warps. After __syncthreads, local_sum in thread 0 holds the total sum. Then compute inv_norm. We store in sdata[0] and broadcast.
The reduction code: after __syncthreads, thread 0 has total sum in its local_sum? Actually in the warp reduction code, only threads in warp 0 compute local_sum. Specifically, after storing per-warp sums in sdata, we have:
if (warpId == 0) {
local_sum = (tid < numWarps) ? sdata[tid] : 0.0f;
for offset... local_sum += __shfl_down_sync(...)
}
After the loop, lane 0 of warp 0 has local_sum = total. So we need to broadcast to all threads. We can do:
__shared__ float inv_norm_shared;
if (tid == 0) {
float norm = sqrtf(local_sum);
inv_norm_shared = (norm > 0.0f) ? 1.0f / norm : 0.0f;
}
__syncthreads();
float inv_norm = inv_norm_shared;
That is simple.
But note: inside `if (warpId == 0)` block, the __shfl_down_sync(...) may have divergence but all threads in warp 0 participate, so it's fine.
One more improvement: since dim is 65535, each thread does about 64 iterations in the loop. `#pragma unroll 4` might help but the loop count is not compile-time constant (dim is runtime). However blockDim.x is not known at compile time either. `#pragma unroll 4` will unroll the loop body by factor 4 and reduce loop overhead. Good.
But there is also register pressure. 1024 threads per block means up to 64K registers per SM. Max registers per thread 255. Our kernel likely uses ~32 registers, so occupancy fine.
Alternatively, we can use block size = 512. That gives 2 blocks per SM concurrently maybe? Occupancy might be higher. With 512 threads, each thread processes 128 elements. That may not matter much.
Let's try blockDim = 512 or 1024. Which gives best occupancy? H100 has 64K registers per SM. If we use 64 registers, 1024 threads per block = 64K registers, exactly one block per SM. If we use 512 threads, maybe two blocks per SM. Occupancy might be better with 512 threads if register count is high. But with 32 registers, 1024 threads uses 32K registers, leaving room for another block. Actually per thread registers are allocated in chunks. But we don't know exact count. The compiler decides.
I think 512 threads is a safe choice. But we have 32768 rows; 32768 blocks of 512 threads each = 16,777,216 threads. Total SMs maybe 132, so ~127k threads per SM. With 2048 threads per SM max, that's 65 blocks per SM if they fit. That's okay. With 1024 threads per block, 32k blocks => 256 blocks per SM max. That's also okay.
I'll start with block size 512.
Potential issue: Using `sqrtf` and division for each row once is fine. Use rsqrtf? The accuracy of rsqrtf might be slightly lower but within tolerance? Let's check: rsqrtf computes approximate reciprocal square root. For fp32, typical error is ~1-2 ulp. Sum of squares might be large. The final inv_norm error is tiny relative. Then we multiply each element, error maybe 1 ulp. FP32 tolerance is 1e-4. For typical values ~0...1 (rand), output values are around 1/sqrt(dim) ~ 0.0039. Relative error of 1e-4 absolute? Wait tolerance is atol=rtol=1e-4. So absolute tolerance 1e-4, relative tolerance 1e-4. If output ~0.0039, relative 1e-4 gives absolute 3.9e-7. That's tight. If we use rsqrtf approximate, could we exceed 3.9e-7? Possibly. Because the error on inv_norm might be ~1e-7 relative, times element ~0.0039 gives ~4e-10 absolute. That's fine. Actually let's compute: rsqrtf error relative about 1e-6? In CUDA, rsqrt.approx.f32 has maximum relative error of about 2^-14 (~6.1e-5)? Wait I recall rsqrt.approx has error ~1 ulp for normal inputs? Let me recall. The fast math rsqrtf has maximum relative error on the order of 2^-14 (around 6e-5) or maybe 2^-13. If relative error on inv_norm is 6e-5, then output error relative is same: 6e-5, which exceeds relative tolerance 1e-4? Wait 6e-5 < 1e-4, so it's okay. But absolute tolerance is 1e-4, so relative 6e-5 on output ~0.0039 is 2.3e-7, which is far below 1e-4. So rsqrtf is safe. However some elements could be close to 1 if one coordinate dominates. In worst case, output ~1, relative error 6e-5 gives absolute 6e-5 < 1e-4. So rsqrtf is safe for tolerance. But to be extra safe, using sqrtf + division once per row is fine; it's not a bottleneck.
Let's write the kernel in Python using torch.utils.cpp_extension.load_inline.
We need to define a host function:
torch::Tensor l2_norm_cuda(torch::Tensor x) {
auto output = torch::empty_like(x);
int batch_size = x.size(0);
int dim = x.size(1);
const int threads = 512;
// We could use 1024, but let's pick 512 or maybe 1024.
const dim3 blocks(batch_size);
l2_norm_kernel<<<blocks, threads>>>(x.data_ptr<float>(), output.data_ptr<float>(), batch_size, dim);
return output;
}
Wait, we need to pass a 1D launch config: dim3 blocks(batch_size). threads = some number.
Because x might not be contiguous? The reference input is contiguous and shape [batch, dim]. But forward may get non-contiguous? Not required, but we can enforce via x.contiguous() or check. Since we are replacing forward, we can call x.contiguous(). However making a copy costs time. But input is contiguous in reference. Let's assume contiguous for performance. Actually we can check and handle.
To be safe, we can do:
auto x_contig = x.contiguous();
auto output = torch::empty_like(x_contig);
...
return output;
But if x is already contiguous, contiguous does nothing. That's fine.
Wait, torch::empty_like(x_contig) returns a tensor with same shape and strides as x_contig (which is contiguous). Good.
Alternatively, we can compute batch_size = x.size(0) and dim = x.size(1) and assume 2D.
Now write the Python code.
One optimization: we can use `load_inline` with `cuda_sources` and `functions=['l2_norm_cuda']`.
Need to include `<torch/extension.h>` and `<cuda_runtime.h>`.
Let's draft the code.
Also, we should use `c10::cuda::CUDAStreamGuard` maybe? Not necessary if we just call the kernel. Kernels launch on the current CUDA stream.
The `ModelNew` class:
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x):
return l2_norm_cuda(x)
But we need to import the compiled extension at module load time.
So at top:
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cpp_source = '''
torch::Tensor l2_norm_cuda(torch::Tensor x);
'''
cuda_source = '''
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void l2_norm_kernel(const float* __restrict__ input,
float* __restrict__ output,
int batch_size,
int dim) {
int row = blockIdx.x;
if (row >= batch_size) return;
const float* row_in = input + static_cast<int64_t>(row) * dim;
float* row_out = output + static_cast<int64_t>(row) * dim;
int tid = threadIdx.x;
float local_sum = 0.0f;
#pragma unroll 4
for (int idx = tid; idx < dim; idx += blockDim.x) {
float val = row_in[idx];
local_sum += val * val;
}
// Warp reduction
unsigned mask = 0xffffffff;
for (int offset = 16; offset > 0; offset >>= 1) {
local_sum += __shfl_down_sync(mask, local_sum, offset);
}
__shared__ float sdata[32]; // max 32 warps
int warpId = tid >> 5;
int lane = tid & 31;
if (lane == 0) {
sdata[warpId] = local_sum;
}
__syncthreads();
// Reduce across warps
if (warpId == 0) {
local_sum = (lane < (blockDim.x >> 5)) ? sdata[lane] : 0.0f;
for (int offset = 16; offset > 0; offset >>= 1) {
local_sum += __shfl_down_sync(mask, local_sum, offset);
}
}
__shared__ float inv_norm;
if (tid == 0) {
float norm = sqrtf(local_sum);
inv_norm = (norm > 0.0f) ? (1.0f / norm) : 0.0f;
}
__syncthreads();
float scale = inv_norm;
#pragma unroll 4
for (int idx = tid; idx < dim; idx += blockDim.x) {
float val = row_in[idx];
row_out[idx] = val * scale;
}
}
torch::Tensor l2_norm_cuda(torch::Tensor x) {
auto x_contig = x.contiguous();
auto output = torch::empty_like(x_contig);
int batch_size = x_contig.size(0);
int dim = x_contig.size(1);
const int threads = 512;
const dim3 blocks(batch_size);
l2_norm_kernel<<<blocks, threads>>>(
x_contig.data_ptr<float>(),
output.data_ptr<float>(),
batch_size,
dim
);
return output;
}
'''
Then load:
l2_norm_ext = load_inline(
name='l2_norm_cuda_ext',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['l2_norm_cuda'],
extra_cuda_cflags=['-O3'],
verbose=False
)
class ModelNew(nn.Module):
def __init__(self):
super(ModelNew, self).__init__()
def forward(self, x):
return l2_norm_ext.l2_norm_cuda(x)
Wait, load_inline signature: `load_inline(name, cpp_sources, cuda_sources, functions, extra_cflags, extra_cuda_cflags, extra_ldflags, extra_include_paths, build_directory, ...)`. We should ensure we pass `extra_cuda_cflags=['-O3', '--use_fast_math']` maybe. `--use_fast_math` will make sqrtf faster, and maybe rsqrtf approximate, but we are using sqrtf. `--use_fast_math` is okay. We can also add `-maxrregcount=64` or something, but not needed.
Potential issues:
- `blockDim.x >> 5` in warp reduce: That's blockDim.x / 32. If blockDim.x is not multiple of 32? 512 is multiple of 32. It's fine.
- The sdata array size: 32. But if blockDim.x = 1024, warps = 32. If blockDim.x = 512, warps = 16. So size 32 is enough.
- We used `lane < (blockDim.x >> 5)`. Since we use `lane` as index, and the number of warps is blockDim.x/32. That's correct.
- In the warp reduction, we used `local_sum` for warp 0 only. At the end, tid 0 has total sum. Then we compute inv_norm. That's fine.
Now, compile and check correctness.
Need to output a complete Python file text.
But before final submission, we should test compile and correctness. Let's write a temporary kernel code string.
Wait, since we can only call tools, and each call takes a string of the full ModelNew kernel file, let's do that.
But we need to be careful: the `compile_kernel` and `run_correctness` take the full python source. So we will send the string.
One potential problem: load_inline compiles on the fly, may take some time. But it's okay. The evaluation environment likely caches or uses the same build directory? The tool might call it in a fresh environment each time, but compilation should succeed.
Also, the `cpp_sources` string with `torch::Tensor l2_norm_cuda(torch::Tensor x);` is used as the header for pybind. According to docs, `cpp_sources` can be just declarations or can include host wrapper functions. Actually `cpp_sources` is the C++ source that gets compiled together with the CUDA source. But `load_inline` auto-generates a pybind wrapper based on `functions`. The wrapper expects the declared function symbols. So we need to provide the host wrapper definitions in `cpp_sources` or `cuda_sources`. Since we have host wrapper in `cuda_sources`, that's fine because .cu files are compiled with nvcc and linked with the C++ code. But is the host wrapper in `cuda_sources` visible to the linker? Yes, because nvcc compiles it and exports the symbol. The `cpp_sources` can be empty or contain declarations. If we put the definition in `cuda_sources`, it's okay. Actually sometimes it is recommended to put everything in `cuda_sources`. Let's test.
Alternatively, we can put everything in `cpp_sources` as raw string with `__global__` etc? No, __global__ must be in .cu file compiled by nvcc. So we put kernel and host wrapper in `cuda_sources`, and leave `cpp_sources` empty or just include statements.
Wait, `load_inline` expects `cpp_sources` string for the C++ compilation unit and `cuda_sources` string for the CUDA compilation unit. Both are compiled and linked. The `functions` list is used to generate pybind11 bindings. It generates code like:
m.def("l2_norm_cuda", &l2_norm_cuda, ...);
So we need the symbol `l2_norm_cuda` to be defined in one of the sources. We'll define it in `cuda_sources`. That's fine.
But `cpp_sources` maybe needs to include `<torch/extension.h>` too? The auto-generated binding includes it, I think. But to be safe, we can leave `cpp_sources=''` or `'#include <torch/extension.h>'`. Actually the `cpp_sources` might just be empty. Let's see.
I know a typical usage:
```
import torch
from torch.utils.cpp_extension import load_inline
cuda_source = """
... kernel ...
torch::Tensor my_op(torch::Tensor x) { ... }
"""
cpp_source = "torch::Tensor my_op(torch::Tensor x);"
module = load_inline(
name="my_module",
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=["my_op"],
)
```
This works. Or even `cpp_sources=""`.
Let's do that.
But if we specify `cpp_sources=""` and `cuda_sources` includes the host wrapper, it should compile because pybind code will refer to the symbol. However the C++ compiler might need a declaration before the generated binding code. The generated binding code might include a forward declaration automatically based on auto-generation. I think `load_inline` parses the functions list and generates declarations. Actually not sure. Let's include a declaration in `cpp_sources` to be safe.
cpp_sources = '''
#include <torch/extension.h>
torch::Tensor l2_norm_cuda(torch::Tensor x);
'''
cuda_sources = '''...'''
Let's construct the code now and compile.
One more thing: we need to make sure `ModelNew` inherits `nn.Module` and `forward` signature matches.
Also for the `get_inputs` and `get_init_inputs` functions, they are provided in the evaluation harness, so we don't need to include them.
Wait, the problem says "Your submission is a complete Python file that defines a class called `ModelNew`." So we should not include `get_inputs` or `get_init_inputs`. But we can include them? Probably not needed. Let's just define `ModelNew`.
Potential optimization: Use a 2D grid to process each row with multiple blocks to increase parallelism and reduce loop iterations per thread. But total blocks already high. However with 65535 elements per row, one block of 512 threads does 128 loop iterations. That's fine. Memory access latency may be hidden by having many blocks per SM.
But reading 16GB at 3.35TB/s is 4.8ms. Each thread does 128 loads and 128 stores. 4.8ms is the target. Reference might do 7.2ms (24GB / 3.35TB/s). So we have a healthy margin.
But can we further reduce overhead? We can increase block size to 1024 threads, which halves the number of loop iterations per thread to 64, reducing instruction overhead. However block size 1024 may reduce occupancy if register usage is high. Let's estimate register usage: local_sum (1), row_in pointer (1), row_out pointer (1), tid (1), mask (1), etc. Maybe ~20 registers. So 1024 threads * 20 = 20K registers, well within 64K. So occupancy should still be high. But maximum threads per SM is 2048, so 1024 threads per block means max 2 blocks per SM concurrently. With 128 SMs, that's 256 blocks concurrently. Since we have 32768 blocks, that's enough. Actually each SM can have up to 32 blocks concurrently (max blocks per SM is 32). But limited by threads. With 1024 threads per block, max 2 blocks per SM = 256 blocks total. That's enough to hide latency? 256 * 1024 = 262k threads. GPU has many SMs. I think it's fine. With 512 threads, max 4 blocks per SM = 512 blocks concurrently. That may help occupancy.
I think 512 is a good default. But we can also let CUDA choose: threads = 256, 512, 1024. Let's test 512 first.
Potential issue: Using `#pragma unroll 4` on a loop when the iteration count is runtime might cause the compiler to unroll by 4, which is fine. But for the tail, there is none because we use dynamic loop condition. The compiler will generate a regular loop with some unrolling. Good.
One more optimization: Use `__ldg` for read-only loads. Since input is read-only (except we read again for output). Using `__ldg` can improve cache performance on read-only data. `const float* __restrict__ input` with `__ldg` might help. We can do:
float val = __ldg(row_in + idx);
But __ldg requires const float*? Actually __ldg takes a pointer. Yes.
However for the second read (when scaling), using __ldg again is fine.
I will write:
float val = __ldg(row_in + idx);
But note: __ldg might not be necessary and could potentially be slower if data is not in read-only cache? Typically helpful for streaming data. Since we read once for reduction and once for scaling, but each element is read by one thread only. __ldg is fine.
Let's incorporate __ldg.
Kernel:
__global__ void l2_norm_kernel(const float* __restrict__ input,
float* __restrict__ output,
int batch_size,
int dim) {
int row = blockIdx.x;
if (row >= batch_size) return;
const float* row_in = input + static_cast<int64_t>(row) * dim;
float* row_out = output + static_cast<int64_t>(row) * dim;
int tid = threadIdx.x;
float local_sum = 0.0f;
#pragma unroll 4
for (int idx = tid; idx < dim; idx += blockDim.x) {
float val = __ldg(row_in + idx);
local_sum += val * val;
}
unsigned mask = 0xffffffff;
for (int offset = 16; offset > 0; offset >>= 1) {
local_sum += __shfl_down_sync(mask, local_sum, offset);
}
__shared__ float sdata[32];
int warpId = tid >> 5;
int lane = tid & 31;
if (lane == 0) sdata[warpId] = local_sum;
__syncthreads();
if (warpId == 0) {
local_sum = (lane < (blockDim.x >> 5)) ? sdata[lane] : 0.0f;
for (int offset = 16; offset > 0; offset >>= 1) {
local_sum += __shfl_down_sync(mask, local_sum, offset);
}
}
__shared__ float inv_norm;
if (tid == 0) {
float norm = sqrtf(local_sum);
inv_norm = (norm > 0.0f) ? (1.0f / norm) : 0.0f;
}
__syncthreads();
float scale = inv_norm;
#pragma unroll 4
for (int idx = tid; idx < dim; idx += blockDim.x) {
float val = __ldg(row_in + idx);
row_out[idx] = val * scale;
}
}
Note: using __ldg for output write doesn't apply. Good.
Potential issue: The `sdata` array is declared inside kernel. Each block gets its own shared memory. Size 32 * 4 = 128 bytes. Very small.
Now, host wrapper:
torch::Tensor l2_norm_cuda(torch::Tensor x) {
TORCH_CHECK(x.is_cuda(), "input must be a CUDA tensor");
TORCH_CHECK(x.dim() == 2, "input must be 2D");
auto x_contig = x.contiguous();
auto output = torch::empty_like(x_contig);
int batch_size = x_contig.size(0);
int dim = x_contig.size(1);
const int threads = 512;
l2_norm_kernel<<<batch_size, threads>>>(
x_contig.data_ptr<float>(),
output.data_ptr<float>(),
batch_size,
dim
);
return output;
}
But wait: we need to include `<torch/extension.h>` and maybe `<c10/cuda/CUDAException.h>`? CUDA kernel launches don't synchronize by default, but PyTorch checks CUDA errors on subsequent calls. To be safe we can use `C10_CUDA_KERNEL_LAUNCH_CHECK();`. But need to include something. Actually `C10_CUDA_KERNEL_LAUNCH_CHECK` is defined in `c10/cuda/CUDAException.h`. But `torch/extension.h` might include it. Let's add it after kernel launch:
l2_norm_kernel<<<batch_size, threads>>>(...);
C10_CUDA_KERNEL_LAUNCH_CHECK();
Yes.
Now, compile.
But before finalizing, we should test if this compiles and passes. Let's write the full Python source and call `compile_kernel`.