[attention-int8] Multiple build issues and numerical accuracy problem

[florianmattana]

Summary

I audited the attention-int8 kernel and found several issues preventing it from building and running correctly outside of a Codespaces environment. After fixing the build issues, the kernel runs but produces numerically inaccurate results.

Environment

• GPU: NVIDIA RTX 5070 (sm_120)
• CUDA: 12.9
• PyTorch: 2.9.1+cu128
• OS: Ubuntu (x86_64)

Issues Found

1. Makefile: hardcoded Codespaces path

The NIX_ENV variable in the Makefile points to /home/codespace/.nix-profile/etc/profile.d/nix.sh, which breaks on any non-Codespaces machine.

Fix:

# Before
NIX_ENV = . /home/codespace/.nix-profile/etc/profile.d/nix.sh

# After
NIX_ENV = . $(HOME)/.nix-profile/etc/profile.d/nix.sh

2. launch_int8_attention — inline prevents linking + type mismatch

In attention_int8.cu, the function launch_int8_attention is declared inline, which prevents it from being visible to the linker when torch_binding.cpp declares it as extern.

Additionally, the timestep parameter is int64_t in the .cu file but int in the extern declaration in torch_binding.cpp, causing a symbol mismatch.

Fix in attention_int8.cu:

// Before
inline cudaError_t launch_int8_attention(..., int64_t timestep, ...)

// After
cudaError_t launch_int8_attention(..., int timestep, ...)

3. static vs extern conflict in validation functions

Several functions in torch_binding.cpp are declared static but are also declared (non-static) in torch_binding.h. This causes a compilation error.

Fix: Remove static from validate_tensor, validate_shapes, validate_head_dim, validate_kv_constraint, validate_timestep_scales, and int8_attention_cuda in torch_binding.cpp.

4. example.py does not test the actual kernel

The example script creates a 1D tensor [1.0, 2.0, 3.0] and checks if the result is x + 1.0. This has nothing to do with an attention kernel that expects 4D tensors [B, H, N, D].

5. Numerical accuracy — results diverge from reference

After fixing all build issues, I compared the kernel output against torch.nn.functional.scaled_dot_product_attention:

Output shape: torch.Size([1, 8, 64, 64])
Mean difference: 0.168081
Max difference: 1.693359

For INT8 quantized attention, the expected mean error should be below ~0.01-0.02. A mean error of 0.168 suggests a problem in the quantization, the softmax, or the dequantization logic.

Test script used

import torch
torch.ops.load_library("./torch-ext/attention_int8/_attention_int8_cuda_dba582b_dirty.abi3.so")

Q = torch.randn(1, 8, 64, 64, dtype=torch.float16, device="cuda")
K = torch.randn(1, 8, 64, 64, dtype=torch.float16, device="cuda")
V = torch.randn(1, 8, 64, 64, dtype=torch.float16, device="cuda")

O = torch.ops.int8_attn.int8_attention_forward(Q, K, V)
ref = torch.nn.functional.scaled_dot_product_attention(Q, K, V)

diff_mean = (O.float() - ref.float()).abs().mean().item()
diff_max = (O.float() - ref.float()).abs().max().item()

print(f"Output shape: {O.shape}")
print(f"Mean difference: {diff_mean:.6f}")
print(f"Max difference: {diff_max:.6f}")

Suggestions

• Fix the build issues (1-3) so the kernel compiles out of the box
• Replace example.py with a real test
• Investigate the numerical accuracy issue — likely in the online softmax or INT8 quantization/dequantization path

[ParagEkbote]

Thank you for diagnosing the issues and writing about them in such depth 👍. After doing a bit of investigation, I have also seemed to uncover some possible reasons for the high amount of mean error in the kernel:

1. Quantization range:

For int8 quantization range, we assume a unsymmetric range of [-128,127] as seen in the helper code. Applying a range of [-127,127] i.e. a symmetric range could help reduce the error.

2. Incorrect max computation for the query values:

Similarly, in for q values, there seems to be incorrect maximum range for the kernel since it should be i ∈ [0, q_size * HEAD_DIM) but the max value is set at Q_head[q_start * HEAD_DIM + i]. This indexing needs to be corrected.

3. Time step scaling for Q quantization

For diffusion models, time step scaling is applied to improve the quality of images by adjusting the number of denoising steps. The quantization omits this, causing distortion in the scaling of issues.

4. Independent per-tile scaling of K

In the following code, scl_K is broken due to inconsistencies of attention logits being rescaled differently across tiles.

You are also correct on the example.py not being updated with a real test. I'd like to add an example with FLUX.2-klein-4B that shows the usage of the kernel with get_kernel and show the perfomance improvements.

Would you like to test and fix the issues in your PR to get the kernel in working instead of a PoC? Is the build process with Nix usable enough or have you been testing with a different method?

cc: @florianmattana

[florianmattana]

Hi @ParagEkbote I will try to have a look but once again I dont when Ill get time. Glad it fixes the first problems.
Is there anything you'd like me to change on the PR before merging, or do you prefer to integrate the fixes yourself? Happy either way. Just in case you make any change, let me know to let me fetch the right branch

[florianmattana]

@ParagEkbote I spent some time tracing through the indexing for issue #2, here's where I'm at.

So if I understand correctly, Q_head points to the start of Q for the current batch and head. q_start is the first token of the tile, q_size is how many tokens are in it, and HEAD_DIM is the number of dimensions per token.

For the max computation, Q_head[q_start * HEAD_DIM + i] with i going from 0 to q_size * HEAD_DIM reads exactly the current tile's data.

Then the quantization loop does Q_head[(q_start + qi) * HEAD_DIM + di], which if you expand it gives q_start * HEAD_DIM + qi * HEAD_DIM + di, same range.

So I end up with both loops reading the same memory. I might be missing something obvious though, could you point me to a specific case where it breaks? Maybe something with the last tile when q_size < BQ?

[ParagEkbote]

Your reasoning about the memory range is correct.

Considering i ∈ [0, q_size * HEAD_DIM), then

Q_head[q_start * HEAD_DIM + i] does indeed cover the same contiguous region as Q_head[(q_start + qi) * HEAD_DIM + di]

But because i = qi * HEAD_DIM + di. So, both loops read the same tile of Q. Where the concern comes in is not the address range but the indexing model used in the two loops.

The max loop uses flat indexing:

Q_head[q_start * HEAD_DIM + i]

while the quantization loop uses structured (qi, di) indexing:

Q_head[(q_start + qi) * HEAD_DIM + di]

These are equivalent for the current layout but using different indexing models for two operations that conceptually iterate over the same tile can make the code more fragile (e.g., if vectorized loads or layout changes are introduced later).

A safer pattern is to use the same (qi, di) decomposition as the quantization loop:

for (int i = tid; i < q_size * HEAD_DIM; i += THREADS) {
    int qi = i / HEAD_DIM;
    int di = i % HEAD_DIM;

    float v = __half2float(Q_head[(q_start + qi) * HEAD_DIM + di]);
    lqmax = fmaxf(lqmax, fabsf(v));
}

Along with time step scaling for Q quantization, issue (2) where the Q scale is computed as:

inv_Q = 127.f / (abs_max_Q * ts)

but the quantization uses:

quantize_f32(Q, inv_Q)

This means the scale assumes the activations were multiplied by ts, but the quantization step never actually applies that factor. Effectively the reconstructed values become approximately Q / ts, which changes the magnitude of the logits before softmax. WDYT?

cc: @florianmattana

[florianmattana]

I see thanks for the calrification. Then it depends on your project objectives. Maybe we should first ignore ts to secure the computation first as we could put it back later for more granularity. Because changing on the nature of the data should be treaterd separatly I think
I will propose a fix to remove the timestep and let you judge if it is the right solution to go with

[florianmattana]

@ParagEkbote
For issue #4 (per-tile K scaling), I see a few possible approaches: computing a global K scale with an extra pass before the tile loop, or pre-quantizing K with a global scale. Each has performance trade-offs. What direction do you prefer?
As it is a design choice , I let the TL decide :)

[ParagEkbote]

For issue #4 (per-tile K scaling), I see a few possible approaches: computing a global K scale with an extra pass before the tile loop, or pre-quantizing K with a global scale. Each has performance trade-offs.

Well, I would like to think that my journey in tech has some distance to go before being a TL, but thank you for the endearing sentiment😃

As for the per-tile K scaling, does the global K scale approach lead us to the greater amount of reduction of mean error than the second approach? I would like to go with the approach that can reduce the mean error level upto a point you feel is acceptable for the kernel while keeping the code maintainable at a positive level.

cc: @florianmattana

[florianmattana]

Got it 😃 Thanks for the guidance. I'll go with the global K scale approach, adding a loop.
The implementation seems pretty straightforward with one extra pass over K to find the global max_abs before the tile loop, then use a single scl_K for all tiles.
If it works and we see some improvment with 1D mock data we can then try with regular 4d tensor