Register AutocastCPU for the attention and DISCO custom ops - #246
Merged
Conversation
Both op families registered an autocast kernel only at the AutocastCUDA
dispatch key, so under torch.autocast("cpu", ...) nothing reconciled their
inputs. The two symptoms differ:
- Attention hard-fails. The kernels dispatch once on q's scalar type and then
reinterpret every activation pointer as that type, so they require k, v and q
to share a dtype and check it explicitly. Autocast casts some ops and not
others, so a module mixing projections with normalization can hand the op an
fp32 q next to an fp16 v, tripping the check:
RuntimeError: v dtype (Half) must match q dtype (Float)
- DISCO fails quietly. Its CPU kernel normalizes dtypes itself (disco_cpu_fwd.cpp
upcasts reduced-precision inp and vals to fp32), so nothing errors -- but the
op ran at whatever dtype it was handed rather than the autocast dtype, i.e.
CPU autocast simply had no effect on DISCO layers.
Both now register AutocastCPU alongside AutocastCUDA, via one closure per op so
the two keys cannot drift apart again.
Found on nvcr.io/nvidia/pytorch:24.12-py3 (torch ~2.6), where six CPU AMP rows of
test_custom_implementation failed. It was invisible on newer PyTorch, on CI and
on two other containers, because those happen to produce consistent dtypes for
this graph -- the missing registration was latent there, not absent.
The new test is therefore at op level rather than module level: it constructs the
mismatch directly (fp32 q, reduced-precision k/v) so the regression is detectable
on any torch version. Without autocast active the same inputs raise, which makes
the test discriminating rather than vacuous.
Suite: 596 passed, 15 skipped (attention + convolution).
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Both op families registered an autocast kernel only at the AutocastCUDA dispatch key, so under torch.autocast("cpu", ...) nothing reconciled their inputs. The two symptoms differ:
Attention hard-fails. The kernels dispatch once on q's scalar type and then reinterpret every activation pointer as that type, so they require k, v and q to share a dtype and check it explicitly. Autocast casts some ops and not others, so a module mixing projections with normalization can hand the op an fp32 q next to an fp16 v, tripping the check: RuntimeError: v dtype (Half) must match q dtype (Float)
DISCO fails quietly. Its CPU kernel normalizes dtypes itself (disco_cpu_fwd.cpp upcasts reduced-precision inp and vals to fp32), so nothing errors -- but the op ran at whatever dtype it was handed rather than the autocast dtype, i.e. CPU autocast simply had no effect on DISCO layers.
Both now register AutocastCPU alongside AutocastCUDA, via one closure per op so the two keys cannot drift apart again.
Found on nvcr.io/nvidia/pytorch:24.12-py3 (torch ~2.6), where six CPU AMP rows of test_custom_implementation failed. It was invisible on newer PyTorch, on CI and on two other containers, because those happen to produce consistent dtypes for this graph -- the missing registration was latent there, not absent.
The new test is therefore at op level rather than module level: it constructs the mismatch directly (fp32 q, reduced-precision k/v) so the regression is detectable on any torch version. Without autocast active the same inputs raise, which makes the test discriminating rather than vacuous.
Suite: 596 passed, 15 skipped (attention + convolution).