reasoning
Let me analyze the reference model:
1. `alpha = torch.sigmoid(layout_logits).unsqueeze(-1)` - (B, T, 1)
2. `amax = x.detach().abs().max().clamp(min=1e-8)` - scalar
3. `scale = amax / 127.0` - scalar
4. `x_q = clamp(round(x / scale), -127, 127) * scale` - (B, T, D)
5. `mixed = alpha * x_q + (1.0 - alpha) * x` - (B, T, D)
6. `out = self.out_proj(mixed)` - linear (B, T, D) -> (B, T, D)
The GPU is H100 with 3.35 TB/s bandwidth and massive compute. For a small model like this (B=8, T=64, D=128), the key is to fuse operations to reduce memory bandwidth.
The most impactful fusion would be to fuse steps 2-6 into a single kernel. Since we also need a linear layer at the end, we can:
1. Fuse sigmoid, fake-quant, and mixing into one kernel
2. Or also fuse the linear layer if possible
Actually, looking more carefully:
- We have `out_proj` which is a `nn.Linear(dim, dim)` with bias
- We can't easily fuse the linear layer weights access with the previous operations unless we write a custom matmul
Given the small sizes, a fused kernel that does the elementwise ops (sigmoid, quant, mix) and then calls the linear layer might not be enough. But we CAN fuse the elementwise operations into one kernel, and for the linear layer, since D=128 is small, we could even consider fusing a small matrix multiplication.
However, the simplest and likely most effective approach is:
1. Fuse: sigmoid(layout_logits) -> compute alpha per element
2. Compute global amax and scale
3. Fuse: for each element of x, compute fake_quant and mixed output
4. Then apply the linear layer
Wait, since layout_logits is (B, T) and x is (B, T, D), we need to broadcast alpha across D. We can do this all in one kernel.
For the linear layer at the end, since D=128, we have a 128x128 matmul + bias. This is small but necessary. We can either use the built-in linear or write a custom matmul kernel.
Actually, let me think about this more. The reference does:
1. sigmoid on (B*T) elements
2. abs/max reduction on (B*T*D) elements
3. elementwise division, round, clamp, multiply on (B*T*D) elements
4. elementwise mix on (B*T*D) elements
5. linear on (B*T*D) -> (B*T*D)
The main wins are from fusing steps 1-4 into a single kernel to avoid intermediate memory traffic. Since step 2 needs a global reduction (amax), we either:
- Do a two-pass approach: first compute amax, then run the fused elementwise kernel
- Or do a cooperative kernel with shared memory reduction
Given the small problem size (8*64*128 = 65536 elements), a two-pass approach is fine. We can use a simple reduction kernel or even PyTorch's `.abs().max()`.
Let me write a fused kernel that:
1. Takes x (B, T, D), layout_logits (B, T), weight (D, D), bias (D)
2. Computes alpha = sigmoid(layout_logits) for each (b, t)
3. Computes global amax
4. For each element, computes mixed = alpha * fake_quant(x, scale) + (1-alpha) * x
5. Then computes out = mixed @ weight.T + bias
For step 5, since D=128 is relatively small and the batch dimension is 512 (8*64), we could write a small batched matmul. But using cuBLAS via PyTorch's linear might be competitive since it's well-optimized. However, if we can fuse the elementwise transformation directly into the matmul loading, that would be ideal.
Actually, I think the best approach is:
1. Write a fused kernel for steps 1-4 that outputs mixed (B, T, D)
2. Then call torch.nn.functional.linear on mixed
But we can do even better by fusing into the linear layer. Since D=128, a simple thread-per-output-element approach for the matmul might work, or we could tile it.
Let me think about the specific fusion. If we write a kernel where each thread block computes a tile of output elements, and each thread loads the mixed values on-the-fly, we save the mixed tensor write/read.
For D=128, if we have each thread compute one output element, we need 512*128 = 65536 threads. That's about 1024 blocks of 64 threads, which is fine.
Actually, let me consider a more structured approach:
- Each thread block handles a (TILE_B, TILE_D) region of the output
- Threads within the block cooperatively load x and weight tiles
- Each thread computes its output element by iterating over D
For D=128, if we use TILE_D = 32 and TILE_B = 16, then each block has 512 threads. We'd load 16*128 floats of x and 32*128 floats of weight per block. But the weights are transposed (out_proj.weight is (D_out, D_in), so we need D_in columns per output). Actually wait, linear is y = x @ W^T + b. So for each output element y[b,t,o], we sum over i: mixed[b,t,i] * W[o,i].
If we do this, each thread needs all D=128 inputs for each output it computes. With D=128, this is very manageable.
Let me design a kernel where:
1. Each block computes a 2D tile of (TILE_BT, TILE_DOUT) output elements
2. Threads load x values, compute fake_quant and mix on-the-fly
3. Threads accumulate the matmul result
4. Add bias at the end
For the layout_logits, we need alpha for each (b,t). Since layout_logits is (B, T), each thread block needs a subset of alphas. We can have threads load the alphas they need.
Let me think about tiling more carefully. The output is (B, T, DOUT=128).
- We can treat (B, T) as a single dimension M = 512
- N = 128, K = 128
This is a small GEMM: M=512, N=128, K=128. Using CUTLASS/cuBLAS should be very fast. The overhead of our custom kernel might not beat cuBLAS.
However, we save memory bandwidth by fusing the elementwise ops into the GEMM. cuBLAS would need to read x (65536 floats = 256KB), write mixed (256KB), then read mixed again for GEMM. The weights are 128*128 = 16K floats = 64KB. The total bandwidth for the unfused approach would be significant relative to compute.
Actually for this small size, cuBLAS GEMM time is probably dominated by overhead and memory bandwidth. If we fuse everything into one kernel, we:
- Read x once: 256KB
- Read layout_logits once: 2KB
- Read weights once: 64KB
- Read bias once: 512B
- Write output once: 256KB
Total: ~580KB read + 256KB write = ~836KB
Unfused with PyTorch:
- Read x for amax: 256KB
- Read x for quant: 256KB
- Write x_q: 256KB
- Read x_q and x for mix: 512KB
- Write mixed: 256KB
- Read mixed for linear: 256KB
- Read weights for linear: 64KB
- Read bias for linear: 0.5KB
- Write output: 256KB
Total: much more
So a fused kernel should definitely win. Let me design it.
Kernel design:
```cuda
__global__ void fused_quant_layout_linear_kernel(
const float* x, // (B, T, D)
const float* layout_logits, // (B, T)
const float* weight, // (D, D) - but linear uses W^T, so weight is (out_dim, in_dim)
const float* bias, // (D,)
float* out, // (B, T, D)
int B, int T, int D)
{
// Grid: blocks over (BT, D_out)
// Each block computes a tile of outputs
int tile_bt = blockIdx.x * TILE_BT + threadIdx.y;
int tile_dout = blockIdx.x * ... wait, I need 2D grid.
// Better: 2D grid
// blockIdx.y over BT tiles, blockIdx.x over DOUT tiles
}
```
Actually, let me simplify. I'll use a 1D grid of blocks where each block handles a contiguous chunk of (B*T) rows and all D output dimensions. Or each block handles a chunk of rows and a tile of output dimensions.
For D=128, let me just have each block compute 1 row (i.e., 1 (b,t) pair) and all 128 output dims. But that's only 128 threads. We'd need 512 blocks.
Or each block computes multiple rows. Let's say each block computes 8 rows and 128 output dims with 128 threads. Each thread handles 8 output elements (one per row). That's 1024 threads per block which is a lot.
Let me think about a better configuration. On H100, optimal is often 256 or 512 threads per block.
How about:
- Each block handles a (TILE_BT=8, TILE_DOUT=32) tile
- 256 threads arranged as (32, 8) or similar
- Actually, let me use 256 threads: threadIdx.x for DOUT (32), threadIdx.y for BT (8)
- Each thread computes 1 output element
Wait, 32*8 = 256 threads. Each block computes 8 rows and 32 output dims. We'd need:
- 512/8 = 64 blocks in BT dimension
- 128/32 = 4 blocks in DOUT dimension
- Total: 256 blocks
Each block:
- Loads 8 alphas (one per row)
- For K=128 iterations (in tiles):
- Load 8*K_TILE x values, compute fake_quant+mix
- Load 32*K_TILE weight values
- Compute partial sums
With K=128, we can process the entire K dimension in one go if we load it all into shared memory. But 8*128 = 1024 floats = 4KB for x, and 32*128 = 4096 floats = 16KB for weights. Total 20KB shared memory, which is fine.
Actually an even simpler approach: since K=128 is small, each thread can compute its output by directly loading from global memory without shared memory for x/weights. But that would cause uncoalesced access or redundant loads.
Wait, let me think about the weight access pattern. If each thread computes out[bt, dout], then thread (bt_i, dout_j) needs weights[dout_j, 0..K-1]. If threads in a warp have contiguous dout_j, then weight access is strided by K. For K=128 and warp_size=32, this could be coalesced if we transpose the weight storage.
Actually, in PyTorch, nn.Linear stores weight as (out_features, in_features). For y = x @ W^T, we need W^T which means accessing W[dout, k] where dout varies across threads. If threads in a warp have consecutive dout values, they access W[dout, k] with stride in_features=128. With 32 threads, that's 32*128*4B = 16KB stride which is bad for coalescing.
To fix this, we should transpose the weights before passing to CUDA, or access them in a transposed manner.
Hmm, actually for a custom kernel, I can just require the weights in a specific layout. Let me transpose W in the host wrapper so that in the kernel we access W_T[k, dout] with dout being the fastest varying dimension. Then threads with consecutive dout values access consecutive memory - perfect coalescing!
Wait, but I'm not writing the weight loading - I'll be reading from the PyTorch nn.Linear module. In the ModelNew constructor, I can copy and transpose the weights.
Let me reconsider. Maybe the simplest robust approach is:
1. Fused elementwise kernel: takes x, layout_logits, outputs mixed
2. Use torch.matmul for the linear layer
This avoids complex GEMM tuning and still gets most of the bandwidth savings. Let me calculate if this is likely to be fast enough.
For M=512, N=128, K=128, a cuBLAS GEMM is very fast. The elementwise part is bandwidth bound. By fusing the elementwise ops, we reduce that part from multiple memory passes to one.
Let me try this simpler approach first and see if it passes correctness, then optimize further if needed.
Actually, let me try to fuse the whole thing. For D=128, a naive per-thread matmul over K=128 is cheap. Let me write the fully fused kernel.
Approach: Each block computes a tile of output. For simplicity:
- Grid: 1D over output rows (B*T), with each block handling BLOCK_ROWS rows and all D output dims
- Block size: say 128 threads
- Each thread handles 1 output dimension for all rows in the block's chunk
Wait, let me do:
- Block size: 256 threads
- Each block handles OUTPUTS_PER_BLOCK output elements
- If each thread computes 1 output element, with 256 threads we compute 256 outputs per block
- Output has 512*128 = 65536 elements
- Need 65536/256 = 256 blocks
256 blocks on H100 is plenty. Each thread computes 1 output element out[bt, dout].
For out[bt, dout]:
out[bt, dout] = sum_k(mixed[bt, k] * W[dout, k]) + bias[dout]
where mixed[bt, k] = alpha[bt] * fake_quant(x[bt, k], scale) + (1-alpha[bt]) * x[bt, k]
In the kernel, each thread:
1. Computes bt and dout from thread/block indices
2. Loads alpha[bt]
3. For k = 0 to K-1:
a. Load x[bt, k]
b. Compute fake_quant(x[bt, k], scale) = clamp(round(x/scale), -127, 127) * scale
c. Compute mixed
d. Load weight[dout, k]
e. Accumulate mixed * weight
4. Add bias[dout]
5. Store out[bt, dout]
Wait, weight is W[dout, k]. If consecutive threads have consecutive dout, then weight access is strided by K=128. This is bad.
To fix this, in the host code I can transpose the weights to W_T[k, dout]. Then each thread loads W_T[k, dout] with consecutive dout accesses being coalesced.
Yes! I'll transpose the weight tensor in the constructor.
Similarly, x is stored as x[bt, k] with k being the fastest varying dimension. If threads in a warp have the same bt but different dout, they all load the same x[bt, k] each iteration - that's good! The x loads are broadcast within the warp.
Let me verify: threads with same bt, different dout. When loading x[bt, k], all threads in the warp load the same address. This is handled via L1 cache broadcast. Not as fast as coalesced, but okay for small K.
Actually, if we have 256 threads per block, and we organize them so that threads with consecutive threadIdx.x have the same bt... no wait, we want threadIdx.x to vary dout so weight access is coalesced.
Let's say:
- blockDim.x = 32 (warp size), blockDim.y = 8
- threadIdx.x varies dout, threadIdx.y varies bt within block
- Block handles 8 bt values and... wait, if we have 256 threads and we want each thread to compute 1 output, we need to cover some tile.
Actually, let's use a 2D grid and 2D thread blocks:
- tile_bt = 8, tile_dout = 32
- blockDim = (32, 8) = 256 threads
- Each block computes 8*32 = 256 outputs
Grid:
- gridDim.x = ceil(B*T / 8) = 64
- gridDim.y = ceil(D / 32) = 4
Total blocks: 256. Total threads: 65536.
For coalescing:
- threadIdx.x = dout within tile (0..31)
- threadIdx.y = bt within tile (0..7)
Consecutive threadIdx.x means consecutive dout. Accessing transposed_weight[k, dout] is coalesced. Good!
For x access: all threads with the same threadIdx.y (same bt) load the same x[bt, k]. This broadcasts. Threads with different bt load different x values.
Actually wait, in a warp of 32 threads, if threadIdx.x varies and threadIdx.y is constant within a warp... well, with blockDim=(32,8), a warp consists of 32 threads with threadIdx.x=0..31 and threadIdx.y fixed. So all threads in a warp have the same bt and different dout. They all load the same x[bt, k]. This is broadcast - efficient.
Then they load transposed_weight[k, dout] with consecutive dout - coalesced. Perfect!
Now for k loop, we'd iterate k=0 to 127. Each thread block would need to load 8*128 x values for its 8 bt positions, and 32*128 transposed_weight values for its 32 dout positions. We can use shared memory to reduce global memory traffic, or just rely on caching since we only have 256 blocks and the working set is small.
Actually, K=128 is small. x has 65536 floats = 256KB. transposed_weight has 16384 floats = 64KB. Total is 320KB which easily fits in L2 cache. So we don't need shared memory - just rely on L2.
But for even better performance, we can preload x for the block's bt rows into shared memory, and preload the weight tile into shared memory. But wait, with 320KB total data and L2 being much larger, cache should handle it well.
Hmm, but with 256 blocks, if each block loads 8*128*4B = 4KB of x and 32*128*4B = 16KB of weights, and some blocks run on the same SM, the L1 cache (probably 128KB per SM) might not hold everything. H100 has 256KB L1 per SM I think? Actually L1 is combined with shared memory, up to 228KB.
Anyway, let me just write a straightforward kernel and see. If it's not fast enough, I'll add shared memory tiling.
Also, I need to think about the amax/scale computation. The scalar amax needs to be computed before the kernel launches. I can do this with a simple PyTorch operation.
Let me write the kernel:
```cuda
__global__ void fused_kernel(
const float* __restrict__ x,
const float* __restrict__ layout_logits,
const float* __restrict__ weight_t,
const float* __restrict__ bias,
float* __restrict__ out,
int B, int T, int D,
float scale)
{
int tile_bt = blockIdx.x;
int tile_dout = blockIdx.y;
int local_bt = threadIdx.y;
int local_dout = threadIdx.x;
int bt = tile_bt * blockDim.y + local_bt;
int dout = tile_dout * blockDim.x + local_dout;
if (bt >= B * T || dout >= D) return;
float alpha = 1.0f / (1.0f + expf(-layout_logits[bt]));
// Actually need to be careful: layout_logits shape is (B, T)
// bt = b * T + t
float sum = 0.0f;
for (int k = 0; k < D; k++) {
float xv = x[bt * D + k];
float q = roundf(xv / scale);
q = fminf(fmaxf(q, -127.0f), 127.0f);
float xq = q * scale;
float mixed = alpha * xq + (1.0f - alpha) * xv;
float wv = weight_t[k * D + dout]; // transposed: weight_t[k, dout]
sum += mixed * wv;
}
out[bt * D + dout] = sum + bias[dout];
}
```
Wait, I need to double-check the layout_logits indexing. layout_logits is (B, T), so layout_logits[b, t] = layout_logits[b * T + t]. And bt = b * T + t. So layout_logits[bt] is correct.
Let me also think about numerical precision. The reference uses:
```python
q = torch.clamp(torch.round(x / scale), -127.0, 127.0)
return q * scale
```
In CUDA, `roundf` rounds to nearest even. PyTorch's `torch.round` also rounds to nearest even. So `roundf` should match.
But wait, there's a subtlety. The reference does `x.detach().abs().max().clamp(min=1e-8)`. The `.detach()` doesn't affect values, just gradients. So in forward we just need to compute `amax = x.abs().max().clamp(min=1e-8)`.
Also, the scale is computed with `amax / 127.0`. Then `x / scale` may have numerical differences between PyTorch and CUDA. But if we match the operations exactly, it should be fine.
One issue: `expf` for sigmoid. `torch.sigmoid` is `1 / (1 + exp(-x))`. For large negative x, `expf(x)` underflows to 0, and sigmoid becomes 1. For large positive x, `expf(-x)` underflows to 0, and sigmoid becomes 1/(1+0) = 1. For very large positive x, `expf(-x)` underflows to 0, which is fine. But for very large negative x, we compute `1.0f + expf(-x)` where `-x` is very large positive, so `expf(-x)` overflows to inf, and we get `1.0f + inf = inf`, then `1.0f / inf = 0`. That's correct.
Actually wait, for negative x: if x = -100, then `-x = 100`, `expf(100)` overflows to inf for float32 (max is about 88.7). Then `1 + inf = inf`, `1/inf = 0`. Correct.
A more stable sigmoid is:
```cuda
float alpha;
if (x >= 0) {
float z = expf(-x);
alpha = 1.0f / (1.0f + z);
} else {
float z = expf(x);
alpha = z / (1.0f + z);
}
```
But torch.sigmoid might use a different formula. Let me check if simple `1/(1+exp(-x))` matches torch.sigmoid well enough. For typical values of layout_logits (standard normal), x is within a few standard deviations, so exp(-x) won't overflow. I'll use the simple form.
Wait, actually PyTorch's sigmoid might have special handling. Let me think: for x = -100, torch.sigmoid(-100) should be very close to 0. In float32, `exp(100)` overflows, but PyTorch's sigmoid might handle this correctly too.
To be safe and match exactly, let me check if I should use a stable formulation. I'll use:
```cuda
float logistic(float x) {
if (x >= 0) {
float z = expf(-x);
return 1.0f / (1.0f + z);
} else {
float z = expf(x);
return z / (1.0f + z);
}
}
```
This is the standard numerically stable sigmoid.
Now, for the weight layout. PyTorch's `nn.Linear` stores weight as (out_features, in_features) = (D, D). I need to transpose it to (D, D) with layout [in_features, out_features] for coalesced access. I'll do this in ModelNew's constructor or in the forward pass.
Actually, I can just call `self.weight_t = self.out_proj.weight.t().contiguous()` in the constructor.
Let me think about whether to also fuse bias. Yes, bias is just adding at the end.
Now, one more optimization: since we're doing a full matmul, maybe we should use shared memory for tiling to reduce global memory bandwidth. But with K=128 and L2 cache, it might not be necessary. Let me start with the simple version and see.
Actually, let me add a shared memory optimization. We can load a tile of weight_t into shared memory. But wait, the weight is only 64KB. With 256 blocks, if they all load parts of weights, L2 should cache it well. But let me add shared memory for the K dimension.
For a thread block handling (tile_bt=8, tile_dout=32), in each K iteration:
- Load x[bt_0..7, k] into shared memory (8 floats)
- Load weight_t[k, dout_0..31] into shared memory (32 floats)
- Synchronize
- Compute partial products
With K=128 and tile size K_TILE=1, this is inefficient. If we use K_TILE=32:
- Load x[bt_0..7, k..k+31] into SM: 8*32 = 256 floats = 1KB
- Load weight_t[k..k+31, dout_0..31] into SM: 32*32 = 1024 floats = 4KB
- Total SM per block: 5KB
Then iterate K in steps of 32: 128/32 = 4 iterations. This is good.
But with only 256 blocks, the kernel is small. Let me try the simple version without shared memory first, as the data fits in L2 and the kernel is simple. If it's not fast enough, I'll add tiling.
Actually, given that this is H100 with huge compute and memory bandwidth, and the problem is small, a well-fused kernel without shared memory tiling should be quite fast. The main win comes from fusing the operations to reduce memory passes.
Let me also consider: should I use float16 or tensor cores? The reference uses fp32, and we need fp32 output. So we should stay in fp32.
OK let me write the kernel. I'll also add a reduction for amax. For amax, I can use PyTorch's `.abs().max()` or write a CUDA reduction. PyTorch's built-in reduction is likely well-optimized, so I'll use it.
Let me now write the complete ModelNew:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = '''
__global__ void fused_quant_layout_linear_kernel(
const float* __restrict__ x,
const float* __restrict__ layout_logits,
const float* __restrict__ weight_t,
const float* __restrict__ bias,
float* __restrict__ out,
int BT, int D,
float scale)
{
int tile_bt = blockIdx.x;
int tile_dout = blockIdx.y;
int local_bt = threadIdx.y;
int local_dout = threadIdx.x;
int bt = tile_bt * blockDim.y + local_bt;
int dout = tile_dout * blockDim.x + local_dout;
if (bt >= BT || dout >= D) return;
float logit = layout_logits[bt];
float alpha;
if (logit >= 0.0f) {
alpha = 1.0f / (1.0f + expf(-logit));
} else {
float z = expf(logit);
alpha = z / (1.0f + z);
}
float sum = 0.0f;
for (int k = 0; k < D; k++) {
float xv = x[bt * D + k];
float q = roundf(xv / scale);
q = fminf(fmaxf(q, -127.0f), 127.0f);
float xq = q * scale;
float mixed = alpha * xq + (1.0f - alpha) * xv;
float wv = weight_t[k * D + dout];
sum += mixed * wv;
}
out[bt * D + dout] = sum + bias[dout];
}
'''
cpp_source = '''
torch::Tensor fused_quant_layout_linear(
torch::Tensor x,
torch::Tensor layout_logits,
torch::Tensor weight_t,
torch::Tensor bias)
{
int B = x.size(0);
int T = x.size(1);
int D = x.size(2);
int BT = B * T;
auto out = torch::empty_like(x);
float amax = x.abs().max().item<float>();
amax = fmaxf(amax, 1e-8f);
float scale = amax / 127.0f;
const dim3 block(32, 8); // dout, bt
const dim3 grid((BT + block.y - 1) / block.y, (D + block.x - 1) / block.x);
AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), "fused_quant_layout_linear", [&] {
fused_quant_layout_linear_kernel<<<grid, block>>>(
x.data_ptr<float>(),
layout_logits.data_ptr<float>(),
weight_t.data_ptr<float>(),
bias.data_ptr<float>(),
out.data_ptr<float>(),
BT, D, scale);
});
return out;
}
'''
```
Wait, I should be careful. In C++ host code, `x.abs().max()` creates a computation graph. Actually no, `x` is a tensor and `.abs().max()` on it does the operation eagerly. But `item<float>()` synchronizes CPU with GPU which might be slow.
A better approach is to compute amax on GPU and pass it as a scalar to the kernel using `.item<float>()`. But for a small kernel, the synchronization overhead might be negligible compared to the kernel launch.
Alternatively, I can write a small reduction kernel to compute amax. But PyTorch's reduction is highly optimized. Let me stick with PyTorch for amax and accept the tiny sync overhead.
Actually, there's another approach: pass `amax` to the kernel by computing it in the C++ code and then `cudaMemcpyToSymbol` or just pass it as a kernel argument. Using `.item<float>()` is fine.
Wait, I need to think about this more carefully. `.item<float>()` does a device-to-host memcpy which synchronizes. For such a small kernel, this might add a few microseconds. Not a big deal.
Let me write the full code. I need to be careful about:
1. Using `__restrict__` for pointer aliasing
2. Block/grid dimensions
3. Proper error handling (not really needed for this task)
4. Matching PyTorch's numerical behavior exactly
Let me also think about `roundf`. In CUDA, `roundf` rounds half values to nearest even (banker's rounding). PyTorch's `torch.round` also uses banker's rounding. So they should match.
But there's a potential issue: `roundf` might not match exactly due to differences in how division is computed. The reference does `torch.round(x / scale)`. In CUDA, `xv / scale` is computed as float division. These should be very close but might differ by 1 ULP. For the tolerance of 1e-4, that should be fine.
Actually wait, if the rounding differs by 1, then `q * scale` could differ by `scale`, which for small amax could be significant relative to the tolerance. Let me think...
If amax is around 1.0, then scale ≈ 1/127 ≈ 0.00787. A difference of 1 in q means a difference of 0.00787 in x_q, which is within 1e-4 tolerance? No, 0.00787 > 1e-4. So we need the rounding to be numerically equivalent.
Hmm, but for values far from the rounding threshold, a 1 ULP difference in `x/scale` won't change the rounded value. Only values very close to a rounding boundary could be affected. The probability of this is low, but let me be careful.
PyTorch's division and CUDA's division might differ slightly. But both use IEEE 754 single-precision arithmetic. The division algorithm should produce the same result for most inputs. I'm not too worried about this.
One more thing: when scale is computed as `amax / 127.0`, in PyTorch this is done in fp32. In CUDA I do `scale = amax / 127.0f`. This should be the same.
Actually, in the C++ code, `amax` is a `float` from `item<float>()`. Then `scale = amax / 127.0f`. This is fine.
Let me write the complete code and compile it.
Wait, I should reconsider whether to write the fully fused kernel or just fuse the elementwise part. Given that D=128 and the matmul is small, a fully fused custom matmul might not be as fast as cuBLAS. But the bandwidth savings might make up for it.
Let me analyze: the custom kernel reads x once, layout_logits once, weights once, bias once, and writes output once. Total memory movement per output element: we read K weights and 1 x element per output element, plus 1 layout_logits per BT group, plus 1 bias per D.
Wait, each thread computing one output reads D x elements, D weight elements, 1 layout_logits element, and 1 bias element. Total reads per thread: 2*D + 2 = 258 floats for 1 output. That's 258:1 read-to-write ratio. Since we have 65536 outputs, total reads = 65536 * 258 floats ≈ 16.9M floats ≈ 67.6 MB. Total writes = 65536 floats ≈ 256 KB.
The arithmetic intensity is 2*D = 256 FLOPs per output (D mul-adds for matmul, plus some elementwise ops). Memory is ~67.6 MB. On H100 with 3.35 TB/s bandwidth, memory time would be ~20 microseconds. Compute time for 256 FLOPs * 65536 outputs = 16.7 MFLOPs is negligible.
A cuBLAS-based approach would:
1. PyTorch elementwise ops: multiple memory passes of x (B*T*D floats each)
- amax reduction: reads x once = 256KB
- fake_quant: reads x, writes x_q = 512KB
- mix: reads x_q, x, alpha, writes mixed = 768KB (but alpha is computed from sigmoid which reads layout_logits = 2KB and writes alpha = 256KB)
- linear: reads mixed, weights, bias, writes output = 256KB + 64KB + 0.5KB + 256KB = 576.5KB
Total: ~2400KB + some overhead
Actually wait, the PyTorch reference computes alpha as `torch.sigmoid(layout_logits).unsqueeze(-1)`. This produces a tensor of size (B, T, 1). The broadcasting in `alpha * x_q` reads both alpha and x_q. So memory for that op is: read alpha (256KB), read x_q (256KB), write mixed (256KB) = 768KB.
So total PyTorch reference memory:
- amax: read x = 256KB
- sigmoid: read layout_logits = 2KB, write alpha = 256KB
- fake_quant: read x = 256KB, write x_q = 256KB
- mix: read alpha = 256KB, read x_q = 256KB, read x = 256KB, write mixed = 256KB
- linear: read mixed = 256KB, read weights = 64KB, read bias = 0.5KB, write output = 256KB
Total reads: ~2100KB, writes: ~1000KB
My fused kernel: reads ~67600KB, writes ~256KB. Wait, that's MORE reads! Because each output element reads D x elements and D weight elements. In PyTorch's linear, the weights are read once total, not once per output.
Ah, I see the issue. In cuBLAS linear, the weights are read from memory once and heavily reused via shared memory/L1 cache. In my naive per-thread kernel, each thread reads D weight elements from global memory. Since there are BT*D outputs, weights are read BT*D*D = 512*128*128 = 8,388,608 times! But the actual weight data is only 128*128 = 16384 floats. So if the cache works well, most of these reads hit cache.
In fact, with L2 cache being much larger than 64KB, the weights should stay in L2 after the first few blocks read them. Similarly, x values for a given bt are read D times (once per output element), but again this should be L1 cached.
So the effective memory bandwidth is much lower than the naive calculation. The real question is: does the cache work well enough?
In my kernel, threads in a warp have the same bt and different dout. They all load the same x[bt, k] each iteration. This broadcasts well (all threads in warp read the same address). For weights, threads load weight_t[k, dout] with consecutive dout, which is coalesced and the entire warp reads 32 consecutive floats = 128 bytes.
For the K loop of 128 iterations:
- Each warp reads 128 broadcasts of x (same address each time within an iteration, but different addresses across iterations)
- Each warp reads 128 coalesced chunks of 32 floats from weights
If L1 cache is 256KB per SM, and we have ~4 warps per SM (假设), the x values might not all stay in L1. But the weight accesses for consecutive k values should hit L1 after the first load.
Actually, for each block of 256 threads (8 warps), the block computes 8 BT values and 32 DOUT values. For a fixed k, the block loads:
- x[bt0..7, k]: 8 floats (broadcast to all threads)
- weight_t[k, dout0..31]: 32 floats (coalesced for each warp)
The block iterates k=0..127. The 8 x values and 32 weight values for each k are different. Total x loaded by block: 8*128 = 1024 floats = 4KB. Total weights: 32*128 = 4096 floats = 16KB.
If L1 is say 128KB, both fit. But blocks might be scheduled on the same SM, so L1 cache pressure depends on how many blocks run concurrently per SM.
Actually, I think a more efficient approach would be to use shared memory for explicit caching. Load the 8 x values and 32 weight values for a chunk of K into shared memory, then compute from there.
Or better: use a standard GEMM approach where we tile in K. But given the tiny sizes, maybe it's overkill.
Alternatively, I could just fuse the elementwise ops and use cuBLAS for the linear layer. That would be:
1. Fused kernel: read x, layout_logits, write mixed
2. cuBLAS: read mixed, weights, bias, write output
The fused elementwise kernel memory:
- read x: 256KB
- read layout_logits: 2KB
- write mixed: 256KB
Then cuBLAS reads mixed, weights, bias and writes output. Since cuBLAS is highly optimized for 512x128x128 GEMM, and the memory for mixed is already written, it's not too bad.
But wait, the PyTorch reference already has all these ops in separate kernels. The fused elementwise kernel would save writing x_q and reading it + x + alpha in separate ops.
Hmm, let me reconsider. For B=8, T=64, D=128, M=512, N=128, K=128:
- The elementwise part is bounded by ~1-2 MB of memory traffic
- The GEMM part has compute ~2*M*N*K = 16.7 MFLOPs and reads mixed + weights = 320KB
At 3.35 TB/s, 2MB takes ~0.6us. At 67 TFLOPS, 16.7 MFLOPs takes ~0.25us. So both are very fast.
The overhead of kernel launches might dominate. PyTorch reference launches ~5 kernels. My fused approach would launch either 1 kernel (fully fused) or 2 kernels (elementwise fused + cuBLAS).
If fully fused is just one kernel, that's a clear win in terms of kernel launch overhead. Also, we avoid the intermediate write to mixed.
I think the fully fused approach is better. Let me optimize it with shared memory.
Shared memory tiling strategy:
- Block: (32, 8) threads -> handles 32 dout x 8 bt
- Tile K in chunks of, say, 32
- Shared memory:
- x_smem[8][32]: 8*32 = 256 floats = 1KB
- w_smem[32][32]: 32*32 = 1024 floats = 4KB
- Total: 5KB per block -> very small, allows many concurrent blocks
Wait, but x[bt, k] where bt varies across the block's 8 rows, and k is the K tile index. For K_TILE=32:
- x_smem holds x[bt0..7, k_start..k_start+31]
- w_smem holds weight_t[k_start..k_start+31, dout0..31]
Actually, weight_t[k, dout] where dout varies across 32 values and k varies across 32 values. So w_smem[32][32].
In each K tile iteration:
1. Load x_smem: each of the 8 bt rows needs 32 x values. We have 256 threads. We can have groups of threads load the x values.
- Thread (dout_i, bt_j) can load x[bt_j, k_start + dout_i] when dout_i < 32. But wait, K_TILE=32 and blockDim.x=32. So thread with local_dout=i can load x[local_bt, k_start + i]. This is coalesced (consecutive threads load consecutive k for same bt).
2. Load w_smem: weight_t[k_start + i, dout_j] for i in 0..31, j in 0..31 (but block only handles 32 dout values).
- We can have each thread load one weight value. Thread (dout_i, bt_j) loads weight_t[k_start + bt_j*4 + (dout_i/4), dout_i]? That's getting complicated.
Actually, let me think about this more carefully. We have 256 threads and need to load 32*32 = 1024 weight values. So each thread loads 4 weight values. Similarly for x, 8*32 = 256 values, each thread loads 1 x value.
For x loading:
- Thread (dout_i, bt_j) loads x[(tile_bt * 8 + bt_j) * D + (k_start + dout_i)]
- Threads with same bt_j and consecutive dout_i load consecutive memory: coalesced!
- We need a boundary check for dout_i >= K, but K_TILE=32 and K=128, so within tiles it's fine except possibly the last tile. But 128 is divisible by 32, so no problem.
For w loading:
- Transposed weight is [D, D] with layout weight_t[k, dout]
- We need w_smem[k_tile, dout_tile] for k_tile in 0..31, dout in 0..31
- Thread (dout_i, bt_j) needs to load 4 weight values: weight_t[k_start + bt_j*4 + 0..3, dout_i]
- Wait, bt_j ranges 0..7. If each thread loads 4 weights, and bt_j goes from 0..7, then bt_j*4 gives offsets 0, 4, 8, ..., 28. That's perfect for 32 k values per tile.
- But thread (dout_i, bt_j) loads weight_t[k_start + bt_j*4 + m, dout_i] for m=0,1,2,3
- This means threads with the same dout_i but different bt_j load weights at stride 4 in k dimension. Consecutive threads (same bt_j, consecutive dout_i) load weight_t[k, dout_i] with consecutive dout_i: coalesced!
Wait, if thread (dout_i, bt_j) loads weight_t[k_start + bt_j*4 + m, dout_i] for m=0,1,2,3:
- For m=0: address = &(weight_t[k_start + bt_j*4, dout_i])
- Threads with fixed bt_j and varying dout_i access consecutive memory: coalesced.
- But different bt_j threads access different k values.
Actually, let me organize the loading differently. We have 256 threads and need to load 1024 weight values. Thread tid (0..255) can load weights at positions:
k_idx = tid / 32, dout_idx = tid % 32... no wait.
Let me use a simpler approach. Since the data is small, I'll skip shared memory tiling. The L1/L2 cache on H100 is very effective. Let me write the naive version first and benchmark it.
Actually, I just realized something important. With blockDim=(32, 8):
- Each block computes 32*8 = 256 outputs
- Total outputs = 512*128 = 65536
- Total blocks = 256
- H100 has 132 SMs (I think for H100 SXM). Actually H100 has up to 144 SMs I believe? Let me check... H100 has 132 SMs for the chip I think. But even with 132 SMs, 256 blocks is enough to keep all SMs busy with 2 waves.
For each block, the K loop of 128 iterations:
- x loads: 8 different bt rows, 128 k values each. Total 1024 floats = 4KB
- weight loads: 32 different dout columns, 128 k values each. But weight_t[k, dout] with dout varying. Total 4096 floats = 16KB
If L1 is 256KB, and blocks don't share L1 (each block runs on one SM), then a block's working set of 20KB fits comfortably in L1. Multiple blocks can run concurrently on the same SM.
Wait, actually on H100, multiple thread blocks can run concurrently on the same SM, sharing the L1 cache. With 256KB L1 + shared memory, and each block using maybe 1KB of registers, there's plenty of room.
I think the naive version should work well due to caching. Let me write it.
But actually, for the weight loading: consecutive threads in a warp have consecutive dout values. They load weight_t[k, dout] with consecutive dout - this is coalesced. Good.
For x loading: all threads in a warp have the same bt value (since threadIdx.y is constant within a warp of 32 threads where threadIdx.x varies). They all load x[bt, k] for the same bt and k. This is a broadcast - all threads load the same address. In modern GPUs, this is handled efficiently via L1 cache broadcast.
So memory access patterns are good. Let me also add `__ldg` (load global through read-only cache) for the weights and x to hint to the compiler. Actually `const __restrict__` should already enable this.
OK, let me write the code now.
One more optimization: loop unrolling. We can use `#pragma unroll` for the K loop. With K=128, full unrolling might create too much code. I'll try `#pragma unroll 8` or just `#pragma unroll`.
Actually, for K=128, if the compiler unrolls it, that's a lot of instructions. Let me not force unrolling and see.
Hmm, actually with K=128 and each iteration having several instructions, a block of 256 threads doing 128 iterations is 32768 iterations total per block. At 256 blocks, that's a lot of instruction issue. Unrolling might help reduce loop overhead, but 128 is a lot. Let me use `#pragma unroll 4` or similar.
Wait, K is a runtime variable (though in practice it's 128). I can pass it as a template parameter if I want compile-time unrolling. But that requires generating code for each possible D. Since D=128 is fixed in this problem, I could template on 128. But to be general, let me just use the runtime value.
Actually, I think loop overhead is negligible compared to the arithmetic and memory in each iteration. Let me not worry about it.
Let me write the full code now.
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
CUDA_SOURCE = r'''
__global__ void fused_quant_layout_linear_kernel(
const float* __restrict__ x,
const float* __restrict__ layout_logits,
const float* __restrict__ weight_t,
const float* __restrict__ bias,
float* __restrict__ out,
int B, int T, int D,
float scale)
{
int tile_bt = blockIdx.x;
int tile_dout = blockIdx.y;
int local_bt = threadIdx.y;
int local_dout = threadIdx.x;
int bt = tile_bt * blockDim.y + local_bt;
int dout = tile_dout * blockDim.x + local_dout;
int BT = B * T;
if (bt >= BT || dout >= D) return;
float logit = layout_logits[bt];
float alpha;
if (logit >= 0.0f) {
alpha = 1.0f / (1.0f + expf(-logit));
} else {
float z = expf(logit);
alpha = z / (1.0f + z);
}
float one_minus_alpha = 1.0f - alpha;
float sum = 0.0f;
for (int k = 0; k < D; k++) {
float xv = x[bt * D + k];
float q = roundf(xv / scale);
q = fminf(fmaxf(q, -127.0f), 127.0f);
float xq = q * scale;
float mixed = alpha * xq + one_minus_alpha * xv;
float wv = weight_t[k * D + dout];
sum += mixed * wv;
}
out[bt * D + dout] = sum + bias[dout];
}
'''
CPP_SOURCE = r'''
#include <torch/extension.h>
torch::Tensor fused_quant_layout_linear(
torch::Tensor x,
torch::Tensor layout_logits,
torch::Tensor weight_t,
torch::Tensor bias)
{
int B = x.size(0);
int T = x.size(1);
int D = x.size(2);
int BT = B * T;
auto out = torch::empty_like(x);
float amax = x.abs().max().item<float>();
if (amax < 1e-8f) amax = 1e-8f;
float scale = amax / 127.0f;
const dim3 block(32, 8);
const dim3 grid((BT + block.y - 1) / block.y, (D + block.x - 1) / block.x);
fused_quant_layout_linear_kernel<<<grid, block>>>(
x.data_ptr<float>(),
layout_logits.data_ptr<float>(),
weight_t.data_ptr<float>(),
bias.data_ptr<float>(),
out.data_ptr<float>(),
B, T, D,
scale
);
return out;
}
'''
class ModelNew(nn.Module):
def __init__(self, dim: int):
super().__init__()
self.dim = dim
self.out_proj = nn.Linear(dim, dim, bias=True)
# Transpose weight for coalesced access in CUDA kernel
self.register_buffer('weight_t', self.out_proj.weight.t().contiguous())
def forward(self, x: torch.Tensor, layout_logits: torch.Tensor) -> torch.Tensor:
return fused_quant_layout_linear(x, layout_logits, self.weight_t, self.out_proj.bias)
```
Wait, I need to compile the CUDA code using `torch.utils.cpp_extension.load_inline`. Let me write the full module.
Also, I need to make sure `weight_t` stays synchronized with `out_proj.weight` if it changes during training. The problem says "forward computation", so we're only optimizing inference. For training, the weights might be updated by optimizer, but `weight_t` wouldn't be automatically updated. However, for the benchmark, we're probably only doing forward passes. I'll add a note about this.
Actually, looking at the reference, it just uses `self.out_proj = nn.Linear(dim, dim, bias=True)` and calls it. The weights would be created in `Model.__init__`. In `ModelNew`, I create my own `out_proj` so weights are initialized and I can transpose them.
But the evaluation code probably copies weights from `Model` to `ModelNew` or trains `ModelNew` from scratch. Either way, my approach should work for forward passes.
Let me also think about whether I need to handle the `weight_t` being a parameter or buffer. I'll make it a buffer so it's saved with the model state.
Actually, a simpler approach: in `forward`, I can transpose the weight on-the-fly. But that adds overhead. Better to pre-transpose in `__init__`.
Wait, if the weights are updated during training, the transposed copy becomes stale. For inference-only, it's fine. Since this is a kernel benchmark focused on forward pass, I'll assume inference or that weights don't change between forward calls.
Let me also think about whether `load_inline` works correctly. I need to specify `cuda_sources` and `functions`.
```python
fused_ops = load_inline(
name="fused_quant_layout",
cuda_sources=CUDA_SOURCE,
cpp_sources=CPP_SOURCE,
functions=["fused_quant_layout_linear"],
extra_cuda_cflags=["-O3", "--use_fast_math"],
)
```
Wait, `--use_fast_math` might change numerical precision. `expf` with fast math might have slightly different accuracy. For tolerance 1e-4, it should be fine. But to be safe, let me not use `--use_fast_math` initially. I can add it later if needed.
Actually, for sigmoid on standard normal random values (which are the inputs), `expf` differences of a few ULPs won't matter much. But I'll avoid `--use_fast_math` to be safe.
Let me also think about whether `torch::empty_like(x)` creates the right output shape. Yes, it creates a tensor of the same shape and dtype as x, on the same device.
One more thing: `x.abs().max()` returns a scalar tensor. `.item<float>()` extracts the value as float. This does a GPU->CPU sync. For such a small kernel, it's fine. But I can also write a tiny CUDA kernel to compute amax if I want to avoid the sync. Actually, let me see if there's an async way. PyTorch doesn't have a good async scalar extraction API. I'll stick with `.item<float>()`.
Wait, I could pass the amax computation into the CUDA kernel itself! Like, have the kernel do the reduction internally. But that makes the kernel much more complex. The sync overhead of `.item<float>()` is probably ~1-5 microseconds. For a kernel that takes ~20 microseconds, that's a 20% hit. Hmm.
Actually, for M=512, N=128, K=128, the custom matmul kernel might take only 5-10 microseconds on H100. A 5-microsecond sync would be significant.
Let me write a two-kernel approach: one kernel to compute amax (reduction), and another that reads the amax from device memory. The reduction kernel output stays on GPU, and the main kernel reads it. This avoids CPU sync!
I can just use `torch.max(torch.abs(x))` which returns a GPU tensor. Then in the C++ code, I can pass `amax_tensor` to the kernel and have each thread read `amax_tensor.data_ptr<float>()[0]`.
Wait, but the kernel will then read from device memory on every thread, which is slow. Better to read it once per thread block or use `__ldg`. Actually, the compiler might optimize this if I load it to a register at the start of the kernel.
Better yet: read it into a register in the host code. No wait, I can't read GPU memory in host code without sync.
Actually, I can compute amax on GPU using a simple kernel, store it in a 1-element tensor, and pass that tensor's pointer to the main kernel. Each thread block reads that value at the start. Since all thread blocks read the same address, it's cached via L2/L1 broadcast. This is efficient!
Let me write a simple amax reduction kernel, or just use PyTorch's `amax` which returns a GPU tensor:
```python
amax_tensor = x.detach().abs().amax() # or .max()
# clamp
amax_tensor = torch.clamp(amax_tensor, min=1e-8)
scale_tensor = amax_tensor / 127.0
```
Then pass `scale_tensor` to the CUDA kernel. Each thread reads `*scale_tensor.data_ptr<float>()` at the start. Wait, no - we want scale, not amax. Actually we can compute scale on CPU since it's just amax / 127. But that requires sync again.
Hmm, but if we pass `amax_tensor` to the kernel, each thread can compute `scale = max(amax, 1e-8) / 127.0` locally. This avoids the sync!
Actually wait, `amax_tensor` is still a GPU tensor. We can pass it to the kernel without syncing. The kernel reads it from GPU memory. Each thread computes `amax = *amax_ptr; amax = fmaxf(amax, 1e-8f); scale = amax / 127.0f;`. Since all threads read the same address, it's efficiently cached.
We can use PyTorch's built-in `torch.abs(x).max()` to get the GPU tensor. Then pass its data pointer to the CUDA kernel.
Wait, but `torch.max()` on a tensor returns a tensor with one element. Let me verify: `x.abs().max()` returns a scalar tensor. Yes.
Actually, `x.detach().abs().max()` creates a scalar tensor on GPU? Let me think. In PyTorch, `.max()` on a tensor returns a tensor. For a 1D tensor, it returns a 0D tensor. The computation happens on GPU, and the result is a 0D tensor stored on GPU. We can access `.data_ptr<float>()` on it, which gives a GPU pointer.
So the C++ code would be:
```cpp
torch::Tensor amax_tensor = x.detach().abs().max().clamp_min(1e-8f);
float* amax_ptr = amax_tensor.data_ptr<float>();
```
And in the kernel, each thread reads `*amax_ptr` and computes scale locally.
Wait, but is `amax_tensor` really on GPU? Let me verify. Yes, PyTorch reductions on CUDA tensors produce CUDA tensors. The result stays on the device.
This is great! No CPU sync needed.
But there's a subtle issue: `.max()` on a multi-dimensional tensor. For a tensor of shape (B, T, D), `.max()` without specifying dim returns the global maximum. Yes, that's what we want.
Wait, `.abs()` first, then `.max()`. In PyTorch, `x.abs().max()` is the same as `x.abs().amax()`. Actually, `.max()` without dim returns a scalar (0D) tensor. But you need `.item()` to get the Python float. However, in C++, we can access `.data_ptr<float>()` directly.
Wait, I need to double check. In PyTorch C++ API, `torch::max(x)` returns a `torch::Tensor` containing the max value. Let me verify: actually `torch::max(x)` in the C++ API returns a tuple of (max_values, max_indices). For a single tensor without dim, it returns a scalar tensor. Yes.
Hmm wait, no. In PyTorch C++ API:
- `torch::max(const Tensor& self)` returns `std::tuple<Tensor, Tensor>` (values and indices)
- So I'd need `std::get<0>(torch::max(x))` to get the value tensor.
Actually, for global max, `.max()` in Python returns a scalar tensor. In C++, `torch::max(x)` returns a tuple. The scalar tensor is the first element of the tuple.
Alternatively, I can use `x.abs().amax()` which returns just the values tensor. `amax` without dim returns a scalar tensor.
Actually, in PyTorch C++ API:
- `at::amax(const Tensor& self, IntArrayRef dim, bool keepdim=false)` - but without dim, does it work globally? Let me check... Actually `amax` without dim might throw an error. Hmm.
Actually, looking at PyTorch docs, `torch.max(input)` returns the maximum value of all elements in the input tensor. In C++: `torch::max(const Tensor& self) -> Tensor` actually might return just the values. Let me check the libtorch docs...
Actually, I think `torch::max(x)` in the C++ API does return a tuple for the global max (to be consistent with named tensors or something). But `torch::amax(x, {}, false)` might work.
Wait, let me just use `torch::abs(x).max()` and get the first element of the tuple. Or even simpler: I can write a quick custom reduction kernel in CUDA, or just use the CPU sync approach. Given that the target is probably not extremely sensitive to a few microseconds of sync overhead, and my kernel should be much faster than the reference anyway, let me just use the simple approach.
Actually, you know what, let me just compute amax and scale using PyTorch in the Python `forward` method, and then launch the kernel. The sync to get scale as a Python float is fine:
```python
def forward(self, x, layout_logits):
with torch.no_grad():
amax = x.detach().abs().max().clamp(min=1e-8)
scale = amax / 127.0
return fused_quant_layout_linear(x, layout_logits, self.weight_t, self.out_proj.bias, scale)
```
Wait, but if `amax` is a CUDA tensor, then `scale = amax / 127.0` is also a CUDA tensor. I can't pass it as a float argument to the C++ function without `.item()`, which syncs.
I can pass the scale tensor to the C++ function and have the kernel read from it:
```python
# Python
with torch.no_grad():
amax = x.detach().abs().max().clamp(min=1e-8)
scale = amax / 127.0
# C++ function takes scale tensor
torch::Tensor fused_quant_layout_linear(..., torch::Tensor scale_tensor)
```
In the kernel, read `float scale = *scale_tensor.data_ptr<float>();` and compute with it.
Is `scale_tensor` a 0D tensor (scalar tensor) on GPU? Yes! We can access its data_ptr and read in the kernel.
But wait, I said earlier that `.max()` returns a tuple in C++. In Python, `x.abs().max()` returns a tensor. Then `.clamp(min=1e-8)` returns a tensor. Then `/ 127.0` returns a tensor. So `scale` is a Python tensor variable pointing to GPU memory.
We can pass this tensor to the C++ function and the kernel reads from it. No CPU sync!
Actually, I realize I can compute amax in Python and pass the tensor. The C++ signature would be:
```cpp
torch::Tensor fused_quant_layout_linear(
torch::Tensor x,
torch::Tensor layout_logits,
torch::Tensor weight_t,
torch::Tensor bias,
torch::Tensor scale_tensor)
```
Then in Python:
```python
with torch.no_grad():
scale = x.detach().abs().max().clamp(min=1e-8) / 127.0
return fused_quant_layout_linear(x, layout_logits, self.weight_t, self.out_proj.bias, scale)
```
Wait, but `x.detach().abs().max()` - does `.max()` return a 0D tensor on GPU? Yes, in PyTorch Python API, `Tensor.max()` without arguments returns a tensor with the global max.
But actually, I just remembered: `Tensor.max()` in PyTorch returns a tensor, but for some older versions or edge cases... No, in modern PyTorch, `torch.max(x)` returns a `torch.return_types.max` namedtuple with `values` and `indices`. Actually no, for global max, I think it returns just a scalar tensor.
Wait, let me think again. In Python:
```python
>>> x = torch.randn(2, 3)
>>> x.max()
tensor(1.2345)
```
Yes, it returns a scalar tensor. So `x.abs().max().clamp(min=1e-8)` chain works and produces a scalar tensor.
But if I use it in the Python forward method, it will stay on GPU. Let me pass it to the C++ extension.
Actually, there's another concern: autograd. The `scale` tensor might participate in autograd if x requires grad. But `x.detach().abs().max()` means x is detached, so no grad. And the CUDA kernel won't produce gradients for scale (it's not backpropagated through in the reference anyway, since scale is computed with `torch.no_grad()`). So this is correct.
Let me now write the complete code with this approach.
Also, I should reconsider whether to write the fully fused kernel or just the elementwise fused + cuBLAS. The fully fused avoids the mixed tensor write/read. Given that D=128 and GEMM is small, custom GEMM should be competitive or better than cuBLAS due to the fusion.
Let me write the fully fused version.
One more optimization: in the kernel, we can reduce register pressure by recomputing instead of storing. But the compiler should handle this.
Also, let me check if I should use `fmaf` for the mixed computation. `mixed = fmaf(alpha, xq, one_minus_alpha * xv)` or similar. But plain multiply-add is fine.
Wait, I just realized: `one_minus_alpha * xv + alpha * xq` might be written more efficiently as `xv + alpha * (xq - xv)`. This saves one multiply:
```cuda
float diff = xq - xv;
float mixed = xv + alpha * diff;
```
This is one FMA instead of two multiplies and one add. Saves an instruction and might be slightly more numerically stable. I'll use this.
Also, `roundf(xv / scale)` - we can rewrite this if we want. But the division is necessary.
Let me now write the full code:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
CUDA_SOURCE = r'''
__global__ void fused_quant_layout_linear_kernel(
const float* __restrict__ x,
const float* __restrict__ layout_logits,
const float* __restrict__ weight_t,
const float* __restrict__ bias,
float* __restrict__ out,
int B, int T, int D,
const float* __restrict__ scale_ptr)
{
int tile_bt = blockIdx.x;
int tile_dout = blockIdx.y;
int local_bt = threadIdx.y;
int local_dout = threadIdx.x;
int bt = tile_bt * blockDim.y + local_bt;
int dout = tile_dout * blockDim.x + local_dout;
int BT = B * T;
if (bt >= BT || dout >= D) return;
float scale = *scale_ptr;
float logit = layout_logits[bt];
float alpha;
if (logit >= 0.0f) {
alpha = 1.0f / (1.0f + expf(-logit));
} else {
float z = expf(logit);
alpha = z / (1.0f + z);
}
float sum = 0.0f;
for (int k = 0; k < D; k++) {
float xv = x[bt * D + k];
float q = roundf(xv / scale);
q = fminf(fmaxf(q, -127.0f), 127.0f);
float xq = q * scale;
float mixed = xv + alpha * (xq - xv);
float wv = weight_t[k * D + dout];
sum = fmaf(mixed, wv, sum);
}
out[bt * D + dout] = sum + bias[dout];
}
'''
CPP_SOURCE = r'''
#include <torch/extension.h>
torch::Tensor fused_quant_layout_linear(
torch::Tensor x,
torch::Tensor layout_logits,
torch::Tensor weight_t,
torch::Tensor bias,
torch::Tensor scale)
{
int B = x.size(0);
int T = x.size(1);
int D = x.size(2);
int BT = B * T;
auto out = torch::empty_like(x);
const dim3 block(32, 8);
const dim3 grid((BT + block.y - 1) / block.y, (D + block.x - 1) / block.x);
fused_quant_layout_linear_kernel<<<grid, block>>>(
x.data_ptr<float>(),
layout_logits.data_ptr<float>(),
weight_t.data_ptr<float>(),
bias.data_ptr<float>(),
out.data_ptr<float>(),
B, T, D,
scale.data_ptr<float>()
);
return out;
}
'''
```
Wait, `torch::empty_like(x)` creates a tensor. But what device? It'll be on the same device as x. Good.
One issue: the `scale` tensor might not be contiguous, or it might be on a different device. I should probably add a check or call `.contiguous()`. But for a scalar tensor, it's always contiguous.
Let me also make sure to handle the shape of `layout_logits`. In the reference, it's (B, T). The kernel accesses it as 1D: `layout_logits[bt]` where `bt = b * T + t`. In PyTorch, a (B, T) tensor is row-major, so `layout_logits[b, t]` is at index `b * T + t`. This is correct.
But when passed to C++, the tensor might not be contiguous if it was created with stride. The reference doesn't create strided tensors, but to be safe, I can call `.contiguous().view(-1)` in Python. Or in C++, I can ensure it's contiguous.
Actually, let me just call `.contiguous().view(-1)` in Python forward:
```python
layout_logits_flat = layout_logits.contiguous().view(-1)
```
This is safe and has no cost if already contiguous.
Similarly, I should make sure `x` is contiguous. But `torch.randn` creates contiguous tensors.
Now let me think about whether I should even bother with the custom GEMM or just use PyTorch linear on the output of a fused elementwise kernel.
For comparison, a fused elementwise-only kernel would be:
```python
# Kernel: read x, layout_logits, write mixed
# Then: torch.nn.functional.linear(mixed, weight, bias)
```
This has the advantage of using cuBLAS for the GEMM, which is highly optimized. The disadvantage is writing mixed and reading it back.
For D=128, M=512, the GEMM in cuBLAS should be very fast, but there might be launch overhead. The memory traffic for mixed is 2 * 512 * 128 * 4 = 512KB (write + read). Not huge, but not zero.
Let me think about whether a custom GEMM can beat cuBLAS for this tiny size. For 512x128x128, cuBLAS is probably well-optimized and uses tensor cores or at least vectorized loads. My custom kernel might not use tensor cores, but it avoids the mixed memory traffic.
I think for this size, the custom fused kernel should be competitive or better. Let me try it.
Actually, I just realized I should consider using `torch.matmul` with the pre-transposed weights. In PyTorch, `torch.matmul(mixed, weight_t) + bias` where `weight_t` is the transpose of the original weight. This might use efficient GEMM and we avoid the transpose in the linear layer.
But the main win from fusion is still the elimination of intermediate mixed tensor.
Let me just write the fully fused kernel and test it.
Wait, I want to double check the sigmoid implementation. PyTorch's sigmoid is:
```python
def sigmoid(x):
return torch.where(x >= 0, 1 / (1 + (-x).exp()), x.exp() / (1 + x.exp()))
```
My CUDA implementation:
```cuda
if (logit >= 0.0f) {
alpha = 1.0f / (1.0f + expf(-logit));
} else {
float z = expf(logit);
alpha = z / (1.0f + z);
}
```
This should match PyTorch's sigmoid exactly (or very close). For values where `logit >= 0`, we compute `1 / (1 + exp(-logit))`. For negative, we compute `exp(logit) / (1 + exp(logit))`. This is the standard numerically stable formulation.
One more check: `torch.round` vs `roundf`. Both use round-half-to-even. Good.
Now, let me finalize the code and try to compile it.
Actually, wait. I need to be more careful about the `ModelNew` class structure. The evaluation code probably instantiates `ModelNew(dim)` and then calls `forward`. It might copy weights from a pre-trained `Model` instance.
If the eval code does something like:
```python
model = Model(dim)
model_new = ModelNew(dim)
model_new.load_state_dict(model.state_dict())
```
Then my `ModelNew` needs to have state dict keys that match or be robust to weight loading.
My `ModelNew` has:
- `self.out_proj = nn.Linear(dim, dim, bias=True)`
- `self.register_buffer('weight_t', ...)`
If `model.state_dict()` has keys like `out_proj.weight` and `out_proj.bias`, then `model_new.load_state_dict(model.state_dict())` will load `out_proj.weight` and `out_proj.bias`. But my `weight_t` buffer was initialized in `__init__` before loading. After loading, `out_proj.weight` changes but `weight_t` doesn't get updated!
This is a problem. To fix it, I can either:
1. Recompute `weight_t` after potential weight loading
2. Use a hook to update `weight_t` when `out_proj.weight` changes
3. Compute the transpose in `forward` on the fly (has overhead but simple)
For a benchmark, weights are typically loaded once. The simplest solution is to transpose in `forward`:
```python
def forward(self, x, layout_logits):
weight_t = self.out_proj.weight.t().contiguous()
...
```
But this adds a `.t().contiguous()` call on every forward, which involves memory allocation and copying 64KB. That's probably negligible compared to the kernel time, but it's extra work.
Another option: override `load_state_dict` or use a forward_pre_hook. But that's overkill.
For benchmark purposes, I can just transpose in forward. Or I can initialize `weight_t` and also handle the fact that for the benchmark, weights might not change after initialization. Actually, in most kernel benchmarks, the model weights are initialized randomly and not loaded from another model. So `weight_t` will be correct.
Let me check the benchmark code pattern. Usually it does:
1. Instantiate `ModelNew(dim)`
2. Call `forward` with inputs
3. Compare outputs with `Model(dim)` running the same inputs
If both use random initialization, they won't match unless weights are copied. The benchmark probably copies weights or runs its own comparison logic.
To be safe, I'll transpose in `forward` but use `view` if possible, or just transpose once and cache:
```python
def forward(self, x, layout_logits):
if not hasattr(self, '_weight_t') or self._weight_t is None:
self._weight_t = self.out_proj.weight.t().contiguous()
...
```
But if weights change, this is stale. For inference benchmark, weights don't change.
Actually, let me look at the problem statement again. It says "Replace the forward computation". The weights are the same. In typical KernelBench setup, they instantiate `ModelNew` and copy weights from `Model` if needed. But they can also just test with randomly initialized weights and compare against `Model` with same initialization.
Since `ModelNew` inherits from `nn.Module`, the `out_proj.weight` tensor will be initialized by PyTorch. My `weight_t` buffer is initialized in `__init__` at that time, so it matches. If weights are later loaded, it's a problem, but I think for this benchmark it's fine.
Still, to be completely safe, let me just do the transpose in forward with `.t().contiguous()`. The overhead of transposing a 128x128 matrix is tiny.
Actually, wait. If I do `self.out_proj.weight.t()`, I get a view (transpose) but not contiguous. Passing it to CUDA might cause strided access issues. I need `.contiguous()`.
Let me do this in forward:
```python
def forward(self, x, layout_logits):
W_t = self.out_proj.weight.t().contiguous()
with torch.no_grad():
scale = x.detach().abs().max().clamp(min=1e-8) / 127.0
return fused_ops.fused_quant_layout_linear(
x, layout_logits.contiguous().view(-1), W_t,
self.out_proj.bias, scale)
```
Hmm but `layout_logits` might be (B, T) and I want to flatten it to (B*T). Yes, `.contiguous().view(-1)` is good.
Wait, I need to think about whether `scale` tensor will be on the right device. `x.detach().abs().max().clamp(min=1e-8) / 127.0` produces a tensor on x's device. Yes.
But `amax_tensor` is a 0D tensor. In PyTorch, is it on CPU or GPU? For a CUDA tensor `x`, `x.max()` returns a CUDA scalar tensor. Yes.
However, I need to be careful: `x.detach().abs().max()` - `detach()` returns a new tensor that doesn't require grad. `abs()` returns a tensor. `max()` returns a scalar tensor. All on the same device. Good.
One more thing: the problem says `check_kernel=False with backend=cuda` for default ModelNew. This means the benchmark framework detects whether I have custom CUDA code. Once I provide it, it will run actual checks.
Let me also consider: should I use `torch.utils.cpp_extension.load_inline` at module level? Yes, that's the standard approach. It compiles the first time the module is imported.
I should make sure the compilation works. Let me write the complete file and test compilation.
Actually, I realize I should also consider whether I can write a warp-specialized kernel or use a different thread organization. But the blockDim=(32,8) with grid based on tiles should be fine.
Let me also check: `layout_logits[bt]` in the kernel where `bt = b * T + t`. If `layout_logits` was flattened with `.view(-1)` in Python, then it's a 1D tensor of size B*T, and `layout_logits[bt]` is correct. If it wasn't flattened, and the stride is not T, it would be wrong. So I should definitely flatten it.
Wait, actually in C++, if I pass a (B, T) tensor, `data_ptr<float>()` gives the base pointer. The element at (b, t) is at `base + b * stride[0] + t * stride[1]`. If the tensor is contiguous, stride[0] = T and stride[1] = 1, so `base + b * T + t` is correct. But if it's not contiguous (e.g., from slicing), the stride might differ.
To be safe, I'll pass it as `.contiguous().view(-1)` from Python. Then it's guaranteed 1D contiguous.
OK, I think I've thought about this enough. Let me write the complete code and compile it.
One more consideration: since `x` is fp32 and we do `roundf(xv / scale)`, for very large or very small values this should be fine. The tolerance is 1e-4.
Actually, double-checking: `scale = amax / 127.0`. If `amax` is large, say 1000, then `scale ≈ 7.87`. Then `x / scale` for x around 1000 gives around 127, which rounds to 127. Then `xq = 127 * scale ≈ 999.49`. So `xq` approximates `x`. For x close to amax, this is fine.
For x=0, `q = roundf(0) = 0`, `xq = 0`. Good.
The mixed output is `alpha * xq + (1-alpha) * x`. For alpha=0.5, it averages. For alpha=1, it's fully quantized. For alpha=0, it's identity.
One potential issue: the reference computes `alpha = torch.sigmoid(layout_logits).unsqueeze(-1)`. This gives alpha of shape (B, T, 1). But for the kernel, we just need the scalar alpha per bt. Since `alpha` is the same for all D dimensions of a given (b, t), we compute it once per bt in the kernel. This is correct.
Now, let me finalize the code. I'll also add `extra_include_paths` or other compilation flags if needed. The simplest compilation flags should work.
Wait, I also need to think about whether `fminf` and `fmaxf` are the right functions. Yes, `fminf(a, b)` returns the smaller of two floats, and `fmaxf(a, b)` returns the larger. So `fminf(fmaxf(q, -127.0f), 127.0f)` clamps q to [-127, 127]. This matches PyTorch's `torch.clamp(q, -127.0, 127.0)`.
There's also `min` and `max` for floats, but `fminf` and `fmaxf` are more explicit.
Also, I should check if I need `#include <cmath>` or similar. In CUDA, `roundf`, `expf`, `fminf`, `fmaxf` are available without explicit includes as they're built-in. In the C++ host code, I need `torch/extension.h`.
Let me write the code now.
Actually, one more thing. I should think about whether to use `torch::Tensor` or `at::Tensor`. `torch::Tensor` is a typedef for `at::Tensor`. `torch/extension.h` includes everything needed.
Also, I notice I don't need to pass `B, T, D` separately if I can compute `BT = B * T`. But I'll pass all of them.
Let me also check grid dimensions. With `BT = 512` and `block.y = 8`: `(512 + 8 - 1) / 8 = 64`. With `D = 128` and `block.x = 32`: `(128 + 32 - 1) / 32 = 4`. So grid = (64, 4) = 256 blocks. Total threads launched: 256 * 256 = 65536. Each thread computes exactly 1 output element. Good.
Wait, `(D + block.x - 1) / block.x = (128 + 31) / 32 = 159 / 32 = 4`. Yes.
For the last blocks in each dimension, some threads might be out of bounds. We handle this with the `if (bt >= BT || dout >= D) return;` check. Good.
One potential issue: `layout_logits[bt]`. If `layout_logits` is a 1D tensor of size B*T, then `bt` ranges 0..B*T-1, which is correct.
But what if B*T is not divisible by block.y=8? Then the last tile has some valid threads and some out-of-bounds threads. The out-of-bounds threads return early. The valid threads still compute correctly. Good.
Wait, but the block.sync is not needed since there's no shared memory. So threads returning early is fine.
If I add shared memory later, I'd need to be more careful. But without SM, it's fine.
Let me write the code now and compile it.
Actually, I want to add one optimization. Currently each thread loads the scale pointer from global memory. I can load it once at the start:
```cuda
float scale = *scale_ptr;
```
Then use it in the loop. Since scale is the same for all threads, and we read it once, it's fine.
Also for bias:
```cuda
out[bt * D + dout] = sum + bias[dout];
```
Each thread reads `bias[dout]`. Threads with consecutive dout read consecutive bias values. This is coalesced. Good.
Now let me write the final code.
Wait, I need to think about `load_inline` usage. The function `load_inline` needs `name`, `cuda_sources`, `cpp_sources`, `functions`, and optionally `extra_cuda_cflags`.
Let me also think about error handling. If compilation fails, I need to debug. I'll catch compilation errors.
Here's the plan:
1. Write complete ModelNew code
2. Compile with `compile_kernel`
3. Fix any compilation errors
4. Run correctness check
5. If issues, debug
6. Submit
Let me write the code now.
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
CUDA_SOURCE = """
__global__ void fused_quant_layout_linear_kernel(
const float* __restrict__ x,
const float* __restrict__ layout_logits,
const float* __restrict__ weight_t,
const float* __restrict__ bias,
float* __restrict__ out,
int B, int T, int D,
const float* __restrict__ scale_ptr)
{
int tile_bt = blockIdx.x;
int tile_dout = blockIdx.y;
int local_bt = threadIdx.y;
int local_dout = threadIdx.x;
int bt = tile_bt * blockDim.y + local_bt;
int dout = tile_dout * blockDim.x + local_dout;
int BT = B * T;
if (bt >= BT || dout >= D) return;
float scale = *scale_ptr;
float logit = layout_logits[bt];
float alpha;
if (logit >= 0.0f) {
alpha = 1.0f / (1.0f + expf(-logit));
} else {
float z = expf(logit);
alpha = z / (1.0f + z);
}
float sum = 0.0f;
for (int k = 0; k < D; ++k) {
float xv = x[bt * D + k];
float q = roundf(xv / scale);
q = fminf(fmaxf(q, -127.0f), 127.0f);
float xq = q * scale;
float mixed = xv + alpha * (xq - xv);
float wv = weight_t[k * D + dout];
sum = fmaf(mixed, wv, sum);
}
out[bt * D + dout] = sum + bias[dout];
}
"""
CPP_SOURCE = """
#include <torch/extension.h>
torch::Tensor fused_quant_layout_linear(
torch::Tensor x,
torch::Tensor layout_logits,
torch::Tensor weight_t,
torch::Tensor bias,
torch::Tensor scale)
{
int B = x.size(0);
int T = x.size(1);
int D = x.size(2);
int BT = B * T;
auto out = torch::empty_like(x);
const dim3 block(32, 8);
const dim3 grid((BT + block.y - 1) / block.y, (D + block.x - 1) / block.x);
AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), "fused_quant_layout_linear", [&] {
fused_quant_layout_linear_kernel<<<grid, block>>>(
x.data_ptr<scalar_t>(),
layout_logits.data_ptr<scalar_t>(),
weight_t.data_ptr<scalar_t>(),
bias.data_ptr<scalar_t>(),
out.data_ptr<scalar_t>(),
B, T, D,
scale.data_ptr<scalar_t>()
);
});
return out;
}
"""
fused_ops = load_inline(
name="fused_quant_layout_linear",
cuda_sources=CUDA_SOURCE,
cpp_sources=CPP_SOURCE,
functions=["fused_quant_layout_linear"],
extra_cuda_cflags=["-O3"],
)
class ModelNew(nn.Module):
def __init__(self, dim: int):
super().__init__()
self.dim = dim
self.out_proj = nn.Linear(dim, dim, bias=True)
def forward(self, x: torch.Tensor, layout_logits: torch.Tensor) -> torch.Tensor:
W_t = self.out_proj.weight.t().contiguous()
with torch.no_grad():
amax = x.detach().abs().max().clamp(min=1e-8)
scale = amax / 127.0
return fused_ops.fused_quant_layout_linear(
x, layout_logits.contiguous().view(-1), W_t,
self.out_proj.bias, scale)
```
Wait, I used `AT_DISPATCH_FLOATING_TYPES` which is good for fp32 and fp64. But the kernel uses `float` specifically. If `scalar_t` is `double`, then `float*` wouldn't match. But the problem says "Produce numerically equivalent outputs at fp32 precision". So inputs are fp32. `AT_DISPATCH_FLOATING_TYPES_AND2` or just using `float` directly would work.
Actually, to keep it simple and since we know it's fp32, I can just cast to float in the kernel call. But `AT_DISPATCH_FLOATING_TYPES` is safer and good practice.
But in the kernel, all pointers are `float*`. If `scalar_t` is `double`, we'd have a type mismatch. Since we only expect fp32, I can either:
1. Use `AT_DISPATCH_FLOATING_TYPES` but the kernel takes `scalar_t*` instead of `float*`
2. Just use `x.data_ptr<float>()` directly and skip dispatch
I'll use option 2 and remove the dispatch macro, since we know it's fp32.
Actually, using `AT_DISPATCH_FLOATING_TYPES` is the PyTorch way. Let me change the kernel to use `scalar_t`:
```cuda
template <typename scalar_t>
__global__ void fused_quant_layout_linear_kernel(
const scalar_t* __restrict__ x,
...
```
Wait, but `roundf`, `expf`, `fmaf` are float functions. If `scalar_t` is `double`, we'd need `round`, `exp`, `fma`. But again, we're only handling fp32. To keep it simple, I'll just use `float` and `data_ptr<float>()`.
Let me simplify:
```cpp
torch::Tensor fused_quant_layout_linear(...) {
// ...
fused_quant_layout_linear_kernel<<<grid, block>>>(
x.data_ptr<float>(),
layout_logits.data_ptr<float>(),
weight_t.data_ptr<float>(),
bias.data_ptr<float>(),
out.data_ptr<float>(),
B, T, D,
scale.data_ptr<float>()
);
return out;
}
```
This assumes all tensors are float32. Good.
One more thing: `layout_logits.contiguous().view(-1)` in Python returns a contiguous flattened tensor. Good.
Actually, I want to check something about `scale`. It's a 0D tensor. When we do `scale = amax / 127.0` in Python, what type is it? `amax` is a float tensor, `127.0` is a Python float. The division returns a float tensor. Good.
When we pass it to C++, `scale.data_ptr<float>()` should work. We dereference it in the kernel with `*scale_ptr`. This gives the float value.
One potential issue: `scale` might be on CPU. Let me think. `x` is on CUDA (since inputs are CUDA tensors). `x.detach().abs().max()` returns a CUDA scalar tensor. `clamp(min=1e-8)` is on CUDA. `/ 127.0` is on CUDA. So `scale` is a CUDA scalar tensor. Good.
Let me now compile this and test. I'll write the full code string.
Actually, I want to double check: `nn.Linear` weight is stored as `(out_features, in_features)`. So for `nn.Linear(128, 128)`, weight is (128, 128). The linear operation is `y = x @ W^T + b`. So `W^T` is (128, 128) with the original W's columns becoming rows.
In my code, `W_t = self.out_proj.weight.t().contiguous()`. This gives the transpose of W, so `W_t` has shape (128, 128). In the kernel, `weight_t[k * D + dout]` accesses `W_t[k, dout]`, which corresponds to `W[dout, k]`. And the matmul is `mixed[bt, k] * W[dout, k]` summed over k. This is `mixed @ W^T`. Correct!
Wait, let me verify once more:
- `mixed` is (BT, D) with `mixed[bt, k]`
- `W` is (D, D) from `nn.Linear`
- Linear computes `output[bt, dout] = sum_k mixed[bt, k] * W[dout, k]`
- `W_t = W.t()` so `W_t[k, dout] = W[dout, k]`
- Kernel computes `sum += mixed[bt, k] * W_t[k, dout]`
- This equals `sum_k mixed[bt, k] * W[dout, k]`. Correct!
Good.
One more optimization: the kernel does `expf` on each thread. For threads in the same warp with the same `bt`, they all compute the same `alpha`. But wait, in my block layout, threads in the same warp have the same `bt` because `threadIdx.y` is constant across a warp of 32. So all threads in a warp compute the same `alpha`. The compiler might not optimize this automatically, but it's fine - `expf` is only computed once per thread, and warp divergence isn't an issue since all threads in the warp take the same branch.
Actually, wait. In a warp of 32 threads, `local_bt` is constant (since the warp spans `threadIdx.x = 0..31` and a fixed `threadIdx.y`). So `bt = tile_bt * 8 + local_bt` is the same for all threads in the warp. Thus `layout_logits[bt]` is the same, and `alpha` is the same. The `expf` computation is done redundantly by all 32 threads, but since it's the same value, it could ideally be computed once per warp. But compilers don't typically do this optimization. However, `expf` is only called once per thread per kernel launch, and we have 65536 threads. The total number of `expf` calls is 65536, but since each warp of 32 threads computes the same value, we're doing 32x redundant work. But `expf` is relatively cheap, and the kernel is likely memory-bound anyway. So it's fine.
If I wanted to optimize further, I could have the first thread in each bt row compute alpha and store it in shared memory for the other threads. But that complicates the kernel. I'll skip it for now.
OK, let me write the complete code and test compilation.
Actually, one more thing about `load_inline`. The `cuda_sources` should not contain the `__global__` function in the `cpp_sources`. That's what I have. But I need to make sure the `cpp_sources` doesn't contain `__global__` or device code.
My `CPP_SOURCE` only contains the host wrapper. `CUDA_SOURCE` contains the kernel. Good.
Also, `load_inline` inserts the CUDA code into a `.cu` file and the C++ code into a `.cpp` file. It should handle compilation correctly.
Let me write the final code.
Actually, I just noticed: in my kernel, I use `xv + alpha * (xq - xv)` for the mix. The reference does `alpha * x_q + (1.0 - alpha) * x`. These are mathematically equivalent but might differ slightly in floating point.
`alpha * xq + (1-alpha) * xv` vs `xv + alpha * (xq - xv)`.
Second form: `xv + alpha * xq - alpha * xv`. If expanded: `(1-alpha) * xv + alpha * xq`. In exact arithmetic, they're the same. In floating point:
- Form 1: compute `t1 = alpha * xq`, `t2 = (1-alpha) * xv`, `result = t1 + t2`
- Form 2: compute `diff = xq - xv`, `prod = alpha * diff`, `result = xv + prod`
Both have 2 multiplies and 2 adds. The rounding might differ slightly. For tolerance 1e-4, it should be fine. But to be safe and match the reference more closely, let me use the exact same expression:
```cuda
float mixed = alpha * xq + (1.0f - alpha) * xv;
```
Wait, but I earlier optimized to `xv + alpha * (xq - xv)` to save storing `one_minus_alpha`. Let me just use the reference form. I can compute `one_minus_alpha` once and reuse it.
```cuda
float one_minus_alpha = 1.0f - alpha;
// ...
for (...) {
// ...
float mixed = alpha * xq + one_minus_alpha * xv;
}
```
This is the reference's exact computation order (though `alpha * xq` and `(1-alpha) * xv` might be computed in a different order than PyTorch's broadcasted multiplication). Again, for tolerance 1e-4 it should be fine.
Actually, PyTorch does element-wise operations in a vectorized manner, which might change the order of operations slightly. But as long as it matches within 1e-4, we're good.
Let me use:
```cuda
float mixed = alpha * xq + one_minus_alpha * xv;
```
And with `fmaf`:
```cuda
float mixed = fmaf(one_minus_alpha, xv, alpha * xq);
```
Hmm, `fmaf(one_minus_alpha, xv, alpha * xq)` computes `one_minus_alpha * xv + alpha * xq` with a single rounding for the final add. This is more accurate than two separate multiplies and an add. In fact, it might be MORE accurate than PyTorch, which could cause issues if PyTorch's result is at the boundary of the tolerance.
But wait, `fmaf` might not be used by PyTorch. If PyTorch computes `alpha * xq` first, rounds to float, then computes `one_minus_alpha * xv`, rounds to float, then adds them, rounding again. `fmaf` does `one_minus_alpha * xv` in extended precision (or exact), adds `alpha * xq`, then rounds once. The result could differ by 1-2 ULP.
For tolerance 1e-4, 1-2 ULP of fp32 is about 1e-7 for values near 1. So the difference is negligible. But to be safe and match exactly, I'll avoid `fmaf` for the mixed computation and just use regular arithmetic. I'll only use `fmaf` for the accumulation `sum += mixed * wv` which is not directly in the reference (since the reference uses PyTorch matmul which might use tensor cores or other methods).
Actually for the accumulation, PyTorch's matmul uses whatever cuBLAS does, which might use tensor cores or FMAs in various orders. The tolerance is 1e-4, so small differences in the matmul are fine.
So I'll just use regular operations without `fmaf` to keep it simple and closer to the reference's behavior.
Actually, let me reconsider. For very large tensors, the matmul order matters. But with D=128 and the kernel doing a simple dot product, the difference between FMA and separate mul-add might be tiny.
OK, I'll use regular code without explicit `fmaf`. The compiler might still emit FMA instructions if `-O3` and `--use_fast_math` are enabled, but without `--use_fast_math`, the compiler follows IEEE more strictly.
Actually, I'll just write it naturally:
```cuda
sum += mixed * wv;
```
OK, final code:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
CUDA_SOURCE = """
__global__ void fused_quant_layout_linear_kernel(
const float* __restrict__ x,
const float* __restrict__ layout_logits,
const float* __restrict__ weight_t,
const float* __restrict__ bias,
float* __restrict__ out,
int B, int T, int D,
const float* __restrict__ scale_ptr)
{
int tile_bt = blockIdx.x;
int tile_dout = blockIdx.y;
int local_bt = threadIdx.y;
int local_dout = threadIdx.x;
int bt = tile_bt * blockDim.y + local_bt;
int dout = tile_dout * blockDim.x + local_dout;
int BT = B * T;
if (bt >= BT || dout >= D) return;
float scale = *scale_ptr;
float logit = layout_logits[bt];
float alpha;
if (logit >= 0.0f) {
alpha = 1.0f / (1.0f + expf(-logit));
} else {
float z = expf(logit);
alpha = z / (1.0f + z);
}
float one_minus_alpha = 1.0f - alpha;
float sum = 0.0f;
for (int k = 0; k < D; ++k) {
float xv = x[bt * D + k];
float q = roundf(xv / scale);
q = fminf(fmaxf(q, -127.0f), 127.0f);
float xq = q * scale;
float mixed = alpha * xq + one_minus_alpha * xv;
float wv = weight_t[k * D + dout];
sum += mixed * wv;
}
out[bt * D + dout] = sum + bias[dout];
}
"""
CPP_SOURCE = """
#include <torch/extension.h>
torch::Tensor fused_quant_layout_linear(
torch::Tensor x,
torch::Tensor layout_logits,
torch::Tensor weight_t,
torch::Tensor bias,
torch::Tensor scale)
{
int B = x.size(0);
int T = x.size(1);
int D = x.size(2);
int BT = B * T;
auto out = torch::empty_like(x);
const dim3 block(32, 8);
const dim3 grid((BT + block.y - 1) / block.y, (D + block.x - 1) / block.x);
fused_quant_layout_linear_kernel<<<grid, block>>>(
x.data_ptr<float>(),
layout_logits.data_ptr<float>(),
weight_t.data_ptr<float>(),
bias.data_ptr<float>(),
out.data_ptr<float>(),
B, T, D,
scale.data_ptr<float>()
);
return out;
}
"""
fused_ops = load_inline(
name="fused_quant_layout",
cuda_sources=CUDA_SOURCE,
cpp_sources=CPP_SOURCE,
functions=["fused_quant_layout_linear"],
extra_cuda_cflags=["-O3"],
)
class ModelNew(nn.Module):
def __init__(self, dim: int):
super().__init__()
self.dim = dim
self.out_proj = nn.Linear(dim, dim, bias=True)
def forward(self, x: torch.Tensor, layout_logits: torch.Tensor) -> torch.Tensor:
W_t = self.out_proj.weight.t().contiguous()
with torch.no_grad():
amax = x.detach().abs().max().clamp(min=1e-8)
scale = amax / 127.0
return fused_ops.fused_quant_layout_linear(
x, layout_logits.contiguous().view(-1), W_t,
self.out_proj.bias, scale)
```
Wait, I need to check if `torch.utils.cpp_extension.load_inline` is called at module load time. Yes, it compiles and caches the extension. This is the standard pattern.
One issue: if the extension fails to compile, the whole script fails. That's expected.
Let me also think about whether to add `extra_cuda_cflags=["-O3", "--ptxas-options=-v"]` to see register usage. But I don't need that for submission.
Actually, let me add `verbose=False` or `verbose=True` if I want to debug. I'll leave it as default.
Now let me write this as the complete Python file and compile it.
Actually, I realize there's one more issue. The evaluation framework might run the script in a sandbox where the compiled extension cache is persistent. `load_inline` uses a hash of the source code to name the extension, so changes to the source will recompile.
OK, let me compile it now. I'll create the full module code.
Hmm, wait. I just realized that `load_inline` might fail if there are issues with the `name` being too long or containing invalid characters. `"fused_quant_layout"` should be fine.
Also, `torch.utils.cpp_extension.load_inline` is imported from `torch.utils.cpp_extension`. I need to make sure this is available.
Let me write the code as a string for `compile_kernel`.
Actually, wait. I need to think about whether the kernel will be called with inputs on GPU or CPU. The problem states `backend=cuda`, so inputs are CUDA tensors.
One more subtle issue: `x.detach().abs().max().clamp(min=1e-8)` returns a tensor. `/ 127.0` returns a tensor. But `.clamp(min=1e-8)` might return a tensor with `min` as a Python float. PyTorch handles this by creating a scalar tensor. Yes, it works.
But what about `amax / 127.0`? If `amax` is a GPU tensor and `127.0` is a Python float, the result is a GPU tensor. Good.
Let me now compile it.
Actually, I just want to make sure: `x.detach()` creates a new tensor that doesn't require grad. `abs()` creates a new tensor. `max()` returns a scalar tensor. `clamp(min=1e-8)` returns a scalar tensor. All on GPU. Good.
And then `scale = amax / 127.0`. Since `amax` is a scalar tensor, `scale` is also a scalar tensor. Passing it to the C++ function as `torch::Tensor scale` works.
I just want to verify that `scale.data_ptr<float>()` works on a 0D tensor. Yes, 0D tensors have a valid data pointer to their single element.
What about `layout_logits.contiguous().view(-1)`? If `layout_logits` is (B, T), then `.view(-1)` gives (B*T,). `data_ptr<float>()` works.
OK, I'll now compile.
Wait, let me also think about a potential numerical issue with `roundf`. In CUDA, `roundf` rounds to nearest, ties to even. In PyTorch, `torch.round` also rounds to nearest, ties to even. They should match.
However, `xv / scale` might differ by a tiny bit between PyTorch and CUDA. The division algorithms could round differently. For values very close to `n + 0.5`, this could cause a difference in the rounded integer.
The probability of a value being exactly on a rounding boundary is extremely low for random floats. For the benchmark, it shouldn't matter. If it does, I could try to match PyTorch's division exactly, but that's very difficult without knowing its exact implementation.
For tolerance 1e-4, a difference of 1 in `q` means a difference of `scale` in `xq`. Since `scale = amax / 127`, and `amax` is typically around 3-4 for standard normal inputs, `scale` is around 0.02-0.03. This is much larger than 1e-4. A single off-by-1 in rounding could cause a difference of ~0.03, which exceeds the tolerance!
So I do need the rounding to match exactly. Let me think about how to ensure this.
The issue is whether `xv / scale` in CUDA produces exactly the same result as in PyTorch. Both use IEEE 754 single-precision division. PyTorch's division is implemented in its ATen library, typically using standard C++ division for CPU and CUDA's `/` operator for GPU. On GPU, PyTorch's `torch.div` or `/` operator typically maps directly to CUDA's `__fdiv_rn` or similar intrinsic.
Actually, let me verify: in PyTorch, `x / scale` where both are CUDA float32 tensors uses CUDA's float division. In my kernel, `xv / scale` also uses CUDA's float division. So they should be exactly the same, since both run on the same GPU hardware with the same rounding mode.
Wait, but PyTorch might use tensor cores or some other path? For elementwise division of two tensors, PyTorch typically uses a CUDA kernel with standard division. For scalar division, it might use `__fdiv_rn` or multiply by the reciprocal. But in any case, the GPU hardware division should be the same.
Hmm, but there's a subtlety. PyTorch computes `x / scale` where `scale` is a CPU scalar (Python float). Actually, in the reference:
```python
scale = amax / 127.0
x_q = _fake_quant_symmetric_int8(x, scale)
```
Here `scale` is a Python float (CPU scalar). When PyTorch does `x / scale`, it creates a CUDA scalar from the Python float and then does elementwise division. The actual computation is `xv / scale_cuda` where `scale_cuda` is the same float value as the Python float.
In my kernel, `scale` is read from a CUDA tensor that was created from `amax / 127.0` in Python. But `amax` is a CUDA tensor, `amax / 127.0` is computed on GPU. The result in the CUDA tensor should be the same as what PyTorch would use if it computed on GPU.
Wait, in PyTorch, the reference computes:
```python
with torch.no_grad():
amax = x.detach().abs().max().clamp(min=1e-8)
scale = amax / 127.0
```
Here `amax` is a scalar tensor. `amax / 127.0` is a tensor operation. It computes on GPU. So `scale` is a GPU scalar tensor.
Then `_fake_quant_symmetric_int8(x, scale)` does `x / scale`. Here `scale` is a CPU scalar? No, `scale` is a tensor. But wait, when passed to a function, if `scale` is a 0D tensor, then `x / scale` is tensor division. PyTorch might promote it to a CUDA scalar and do elementwise division.
Actually, in the Python function:
```python
def _fake_quant_symmetric_int8(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
q = torch.clamp(torch.round(x / scale), -127.0, 127.0)
return q * scale
```
Wait, `scale` is explicitly typed as `torch.Tensor`. But in the reference's `forward`:
```python
scale = amax / 127.0
x_q = _fake_quant_symmetric_int8(x, scale)
```
Here `scale` is a Python float, not a tensor! Because `amax` is obtained via `.item()`? No, `amax = x.detach().abs().max().clamp(min=1e-8)` is a tensor. Then `scale = amax / 127.0` is a tensor. Then passed to the function which expects a Tensor but receives... a scalar tensor.
Actually wait, `amax = x.detach().abs().max()` - does `.max()` return a tensor or a scalar? In PyTorch, `Tensor.max()` without dim returns a tensor. But then `.clamp(min=1e-8)` on a scalar tensor returns a scalar tensor. And `amax / 127.0` is tensor-scalar division, returning a tensor.
So `scale` is a 0D tensor. Then `x / scale` in the fake_quant function is elementwise division of x by the scalar tensor scale. PyTorch broadcasts scale to x's shape and does elementwise division on GPU.
In my kernel, I compute `scale` as a 0D tensor on GPU and read it in the kernel. Then each thread does `xv / scale`. This should match PyTorch's behavior exactly.
But wait, let me double check: in the PyTorch reference, when scale is a 0D tensor, does PyTorch convert it to a scalar CUDA value and pass it to a kernel? Yes, the elementwise kernel does `xv / scale_value` for each element. This is exactly what my kernel does.
So the division should match. Good.
However, there's another subtlety. The reference does:
```python
amax = x.detach().abs().max().clamp(min=1e-8)
```
But wait, `x.detach().abs().max()` returns a 0D tensor. Then `.clamp(min=1e-8)` clamps it. Then `scale = amax / 127.0`.
In my code, I do:
```python
with torch.no_grad():
amax = x.detach().abs().max().clamp(min=1e-8)
scale = amax / 127.0
```
This matches exactly. Good.
One more issue: PyTorch's `round` behavior for `.5` cases. PyTorch rounds to nearest even. CUDA's `roundf` also rounds to nearest even. They match.
But wait, there's `nearbyint` in CUDA which also rounds to nearest even. `roundf` rounds half values away from zero by default in some implementations? Let me check.
Actually, I need to be very careful here. In C/C++, `roundf` rounds half away from zero (i.e., 0.5 -> 1.0, -0.5 -> -1.0). But PyTorch's `torch.round` uses round-to-nearest-even (i.e., 0.5 -> 0.0, 1.5 -> 2.0, 2.5 -> 2.0).
This is a critical difference! If the value after division is exactly `n + 0.5`, PyTorch rounds to the nearest even integer, while `roundf` rounds away from zero.
For example:
- `torch.round(torch.tensor(0.5))` = 0
- `roundf(0.5f)` = 1.0f
This could cause differences!
In CUDA, the function that implements round-to-nearest-even is `nearbyintf` or `rintf`. Let me verify:
- `nearbyintf(0.5f)` = 0.0f (round to nearest even)
- `rintf(0.5f)` = 0.0f (round to nearest even, using current rounding mode which is typically round-to-nearest-even)
Yes! I should use `nearbyintf` instead of `roundf` to match PyTorch.
Let me verify with another example:
- `torch.round(torch.tensor(1.5))` = 2
- `nearbyintf(1.5f)` = 2.0f
- `roundf(1.5f)` = 2.0f (same for this case)
- `torch.round(torch.tensor(2.5))` = 2
- `nearbyintf(2.5f)` = 2.0f
- `roundf(2.5f)` = 3.0f (different!)
Yes, I must use `nearbyintf` to match PyTorch exactly.
Actually, let me double check PyTorch's `torch.round` documentation. It says "rounds each element to the nearest integer, with halfway cases rounded to the nearest even integer." This is banker's rounding.
So I'll change `roundf` to `nearbyintf` in my kernel.
Let me also check if there are other rounding differences. The division `xv / scale` followed by `nearbyintf` should match `torch.round(x / scale)`.
OK, let me update the kernel to use `nearbyintf`.
Wait, let me also consider what happens for negative numbers:
- `torch.round(torch.tensor(-0.5))` = 0 (nearest even)
- `nearbyintf(-0.5f)` = -0.0f (which is 0.0f)
- `roundf(-0.5f)` = -1.0f
Different! But `nearbyintf` matches PyTorch.
- `torch.round(torch.tensor(-1.5))` = -2
- `nearbyintf(-1.5f)` = -2.0f
- `roundf(-1.5f)` = -2.0f
Same for odd integer + 0.5.
- `torch.round(torch.tensor(-2.5))` = -2
- `nearbyintf(-2.5f)` = -2.0f
- `roundf(-2.5f)` = -3.0f
Different! `nearbyintf` matches PyTorch.
Great catch! I'll use `nearbyintf`.
Let me also verify that `nearbyintf` is available in CUDA. Yes, it's part of the standard math library.
Now let me write the updated code.
Actually, one more thought. The reference computes `amax = x.detach().abs().max().clamp(min=1e-8)`. Then `scale = amax / 127.0`. Then `q = torch.round(x / scale)`.
For large `x`, `x / scale` is large, and `nearbyintf` returns a float. For values above 2^24, `nearbyintf` might not change the value since float32 can't represent all integers at that magnitude. But `x / scale` is at most about `127 * 128 / 1 = 16256` for the worst case? Wait, no. `x` values can be anything (standard normal, so around ±3). `amax` is around 3-4. `scale = amax / 127 ≈ 0.02-0.03`. `x / scale` is around `±3 / 0.03 = ±100`. So `q` is around ±100, well within the range where float32 represents all integers exactly. No issue.
OK, let me finalize the code and compile.