Simon Layton cb6e4d7d82 User-passed alpha to scaled_gemm (#165563) ...

Summary:

Add optional user-passed `alpha` argument to
`at::cuda::blas::scaled_gemm`, necessary for two-level-scaled NVFP4 gemm
calls (where the global de-scales are folded into the `alpha` argument.

Global de-scales are naturally device tensors, but using cublas'
device-pointer mode for `alpha`/`beta` has an interesting lifetime
implication - the `alpha` tensor must be valid & correct until the end
of the matmul call, *not* just the launch (as for host values). To
enable this, I added device-constant memory for `one` and `zero`, along
with a statically-held single-fp32-value tensor, which is valid from the
first passed-`alpha` invocation of `scaled_gemm` to the end of the
program. User-passed values are copied into this perpetual buffer to
ensure lifetime requirements are met.

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
Signed-off-by: Simon Layton <simonlayton@meta.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/165563
Approved by: https://github.com/drisspg, https://github.com/eqy

2025-10-17 09:42:33 +00:00

..

conda

PyTorch -> C++17 (#98209) (#100557)

2023-05-19 00:49:08 +00:00

src

User-passed alpha to scaled_gemm (#165563)

2025-10-17 09:42:33 +00:00

tools

Adds Issue#153109 as a test for CUDAPluggableAllocator (#163575)

2025-10-01 09:07:48 +00:00

CMakeLists.txt

Revert "Use official CUDAToolkit module in CMake (#154595)"

2025-06-23 21:15:31 +00:00