Repository navigation
Update tl.dot precision handling - #3
Open
GMNGeoffrey wants to merge 6 commits into
Open
GMNGeoffrey wants to merge 6 commits into
GMNGeoffrey wants to merge 6 commits into
Conversation
The module exists only for the Triton kernels, so it lives with them. Its tests move to match. No change beyond the module path at the import sites for git rename tracking. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
I had a chance to redo this when updating OF3 to do something similar (aqlaboratory/openfold-3#457) and reworked some of the naming for clarity and better consistency with Triton and Torch: - `matmul_precision()` -> `current_dot_precision()` - `torch_precision()` -> `torch_fp32_precision()` - consistent usage of `dot_precision` for the resolved mode Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
As in the previous commit. I had a chance to redo this when updating OF3 to do something similar (aqlaboratory/openfold-3#457) and reworked some of the documentation. Bring those updates back over and remove unnecessary diffs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Left unset, `tl.dot` picks a reduced-precision mode for fp32 operands wherever the target has one. `tl_dot` makes the mode a required argument so a lint rule can ban bare `tl.dot` calls. The wrapper inlines completely during compilation. Triton's JIT does not support keyword-only parameters, so `input_precision` precedes the optional `acc` positionally. Tested with a minimal kernel including that ieee is more accurate than tf32. The no-setting test now stubs torch, so it no longer depends on the global precision other tests leave behind. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Mechanical: every kernel `tl.dot` becomes `tl_dot`, and the two call sites that passed `acc` positionally now pass it by keyword, since `tl_dot` takes `input_precision` before `acc`. Ruff's banned-api rule (TID251) rejects `tl.dot` anywhere else, so a call site can't leave off the precision. This replaces the gross test that parsed the package source. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Since bf16x3 is a software emulation, Triton reports it available on many targets where it is not faster than ieee. Following torch's tf32 setting there makes kernels slower and less accurate for no benefit, so the setting is now honored only on devices with a fast reduced-precision dot. `has_fast_reduced_precision_fp32_dot` mirrors torch.cuda.is_tf32_supported as of torch 2.8 (gfx94x/gfx95x on ROCm, sm >= 8.0 on CUDA). We reimplement it because torch's implementation varies by version: it's absent before 2.7, and False for all of ROCm in 2.7. Ours also checks that cuda is available. Tests that check how the setting is read and plumbed now mock the fast dot setting so they can work independent of hardware. Those comparing tf32 with ieee outputs skip mock the "is faster" piece so that they still exercise a (potentially slow) reduced precision dot where available. The shim's precision tests now mock the effective torch setting instead of setting it, so the don't need a GPU and remove the dependency on the main package testutil helpers (which would break tests run from an sdist). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
GMNGeoffrey
changed the title
This branch has not been deployed
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
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.
I had a chance to take another look at this when updating OF3 to do something similar (aqlaboratory/openfold-3#457) and reworked some of it. Bring those improvements back here. Some stuff is just wording I like better for function naming and docs. The material changes:
torch.cuda.is_tf32_supported, but we reimplement it because Torch's function is not available prior to 2.7 and in 2.7 it returns False on ROCm across the board. When we drop support for <=2.7, we could consider using Torch's setting instead, but should probably check that it's really answering the same question we're trying to ask. In the OF3 PR, I used Torch's function. I only realized the problem with it for older versions when creating this PR. One notable exclusion from the list is gfx1151. In my local tests on my laptop, bf16x3 is actually a speedup here, but for now we're deferring to Torch.tl.dotwith a wrapper function and Ruff ban for the API, replacing the gross test doing source inspection. I think this is cleaner, more robust, and provides quicker feedback.The commits are laid out for ease of review one-by-one as I separated out mechanical things to be quickly scanned.