Skip to content

Fix degeneracy in attention kernel backward tests#258

Open
GMNGeoffrey wants to merge 2 commits into
aqlaboratory:mainfrom
GMNGeoffrey:kernel-tests
Open

Fix degeneracy in attention kernel backward tests#258
GMNGeoffrey wants to merge 2 commits into
aqlaboratory:mainfrom
GMNGeoffrey:kernel-tests

Conversation

@GMNGeoffrey

@GMNGeoffrey GMNGeoffrey commented Jun 12, 2026

Copy link
Copy Markdown
Contributor

Summary
The attention kernel backward tests were degenerate because the gradients they were comparing were smaller than the tolerances used for comparison (loss = mean(out) -> gradients < 1e-6). Reworked these so they're actually testing something and made some other changes to avoid future issues. I made similar changes to the forward kernel tests, but didn't touch the e2e tests for now. I can update those to match as well, but wanted to get feedback first in case you dislike this approach.

Changes

  • Use a sum instead of mean for the loss
  • Multiply the output by a random cotangent to get the loss, keeping gradients at full scale, but allowing some cancellation.
  • Use randn (normal mean 0, std 1) rather than rand (uniform [0, 1)) for inputs. I think this is more realistic and the mix of negative and positive values again avoids everything blowing up.
  • Scale down the pair bias so it doesn't dominate the actual inputs.
  • Use fp64 for the reference implementation. The reference implementation has its own floating point errors in lower precision dtypes, which forces very loose tolerances. I validated the tolerances used against the error of the lower-precision reference to ensure we're not using something too loose. The inputs are still generated in the target dtype and upcast, so they're exactly representable, avoiding introducing additional error just from casting the input.
  • Only check the input gradients (q, kv, bias). The weight gradients are a downstream autograd contraction that adds no kernel coverage while forcing a much looser tolerance.

To avoid future issues, I added a custom assert assert_close_nondegenerate that tries to guard against this sort of degeneracy. In addition to checking tensors are close with torch's assert_close, it checks:

  • All values are finite (no infs or nans)
  • The expected output is not all close to zero (within the same atol and rtol used for the closeness check). This would've caught the issue here.
  • The expected and actual values are not suspiciously close to a very tight tolerance. This guards against accidentally testing the reference against itself, or something like that.

I also seeded the RNG for all kernel tests via the seeded_rng fixture for determinism.

Testing
I ran the Triton kernel tests on Radeon 8060S and the cueq kernel tests on RTX 3090. I wasn't able to run the deepspeed kernels because they appear to not be supported on RTX 3090. Hopefully CI coverage there is sufficient.

The backward tests used loss = mean(out), which divided every input
gradient by ~230k elements down to ~1e-6, far below the absolute
tolerance they were compared against. So the tests passed regardless of
whether the kernel backward was correct.

To mitigate this, I reworked the comparisons in a few ways:
- Use a sum instead of mean
- Multiply the output by a random cotangent to get the loss, keeping
  gradients at full scale, but allowing some cancellation.
- Use randn (normal mean 0, std 1) rather than rand (uniform [0, 1)) for
  inputs. I think this is more realistic and the mix of negative and
  positive values again avoids everything blowing up.
- Scale down the pair bias so it doesn't dominate the actual inputs.
- Use fp64 for the reference implementation. The reference
  implementation has its own floating point errors in lower precision
  dtypes, which forces very loose tolerances. I validated the tolerances
  used against the error of the lower-precision reference to ensure
  we're not using something too loose. The inputs are still generated in
  the target dtype and upcast, so they're exactly representable,
  avoiding introducing additional error just from casting the input.
- Only check the input gradients (q, kv, bias). The weight gradients are
  a downstream autograd contraction that adds no kernel coverage while
  forcing a much looser tolerance.

To avoid future issues, I added a custom assert
`assert_close_nondegenerate` that tries to guard against this sort of
degeneracy. In addition to checking tensors are close with torch's
assert_close, it checks:
- All values are finite (no infs or nans)
- The expected output is not all close to zero (within the same atol and
  rtol used for the closeness check). This would've caught the issue
  here.
- The expected and actual values are not suspiciously close to a very
  tight tolerance. This guards against accidentally testing the
  reference against itself, or something like that.

I also seeded the RNG for all kernel tests via the seeded_rng fixture
for determinism.
I ran cuequivariance on RTX 3090 and it need the same looser tolerances.
I wasn't able to run Deepspeed, but we explicitly cast the inputs there
to bf16, so I assume it will apply as well. So I just set the tolerances
default to that for the half type and got rid of the more complicated
stuff.
@christinaflo christinaflo added the safe-to-test Internal only label used to indicate PRs that are ready for automated CI testing. label Jun 20, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

safe-to-test Internal only label used to indicate PRs that are ready for automated CI testing.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants