[PyTorch] Custom kernel to compute reciprocal of a single float

[timmoon10]

Description

FP8 training is frequently bottlenecked by CPU overheads and a non-trivial fraction of CPU overhead comes from small PyTorch operations. For example, when I benchmark the forward pass of small Linear modules on an L40, I estimate ~20% of runtime is spent in handling the FP8 scaling factors (mainly in reciprocal and clone operations). This PR attempts to mitigate these overheads by adding a scalar_reciprocal kernel that operates on a single float, bringing the kernel launch cost down from ~20 us to ~10 us. In my benchmark of Linear forwards, I see a 8% reduction in runtime.

Alternative approaches:

• Modify the cast and cast-transpose kernels to perform the scale-inv update, similar to how they perform the amax update. Logically, the scale is part of the FP8 recipe and the scale-inv is part of the data.
• Use torch.compile to fuse FP8 scale operations. We would require significant refactoring to avoid incurring extra overhead from graph breaks, especially in how we deal with the FP8TensorMeta class.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refractor

Changes

• Custom kernel to compute reciprocal of a single float

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[timmoon10]

/te-ci pytorch

[ksivaman]

Why?

[timmoon10]

fwd_scale_inverses is created with a clone operation (~20 us), while input_fp8_scale_inv is created with the scalar reciprocal kernel (~10 us).

[ksivaman]

This is only set for float8tensor case

[timmoon10]

We handle three cases:

• Non-FP8: scale-inv is None
  

  TransformerEngine/transformer_engine/pytorch/module/linear.py

  Line 120 in 23258bb

  inputmat_fp8_scale_inv = None

• FP8, Float8Tensor input: scale-inv is taken from Float8Tensor
  

  TransformerEngine/transformer_engine/pytorch/module/linear.py

  Line 125 in 23258bb

  inputmat_fp8_scale_inv = inputmat._scale_inv

• FP8, non-Float8Tensor input: scale-inv is computed with fast kernel
  

  TransformerEngine/transformer_engine/pytorch/module/linear.py

  Line 158 in 23258bb

  inputmat_fp8_scale_inv = tex.scalar_reciprocal(

[timmoon10]

Closed by #1083