[PyTorch] Enable FP8 block scaling for fusible ops#3242
Conversation
Enables FP8 block scaling (FP8BS) in composable fusible-ops for both non-grouped and grouped linears. - Drop the blanket FP8BS NotImplementedError in ops/op.py; FP8BS now uses the generic recipe-state path like MXFP8. - Add `Float8BlockwiseQTensorStorage.view()` so quantized norm outputs reshape without dequantizing (mirrors `MXFP8TensorStorage.view()`). - Route grouped FP8BS through the Hopper graph-safe path (cuBLAS 13.4+). Fuse dbias in the rowwise pass when dgrad is required. - Exercise FP8BS across the fusible-ops test suite with 128-divisible sizes. Signed-off-by: Alp Dener <adener@nvidia.com>
Greptile SummaryThis PR enables FP8 block scaling (FP8BS) for TE/PyTorch fusible ops by removing the
Confidence Score: 5/5Safe to merge. All new code paths are well-guarded by hardware capability and cuBLAS version checks, and the view() implementation correctly preserves scale-inv invariants for both 1D and 2D block-scaling modes. The removal of the NotImplementedError is mechanical and routes FP8BS through the well-tested recipe-state path already used by MXFP8. The new view() precisely mirrors the MXFP8 reference implementation, accounting for the transposed columnwise storage layout. The grouped-linear changes correctly restrict the new graph-safe path to Hopper + cuBLAS 13.4+ and gate dbias fusion on whether a dgrad pass is actually needed. No incorrect computations, stale state, or dropped side-effects were found across the five changed files. No files require special attention. Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[FP8BS recipe active in fusible op] --> B{need_to_reset_recipe_state?}
B -- yes --> C[RecipeState.create for FP8BS]
C --> D[make_quantizers → Float8BlockQuantizer list]
B -- no --> E[reuse existing quantizers]
D & E --> F[Forward pass]
F --> G{GroupedLinear?}
G -- yes --> H{_use_grouped_compute?}
H -- yes --> I{CC >= 10.0 Blackwell?}
I -- yes --> J[Return False → split-quantize fallback]
I -- no Hopper --> K{cuBLAS >= 13.4?}
K -- yes --> L[graph-safe grouped GEMM path]
K -- no --> J
G -- no --> M[single linear path unchanged]
L --> N[Backward pass]
N --> O{input_requires_grad?}
O -- yes --> P[bgrad_group_quantize - fuse dbias in rowwise dgrad pass]
O -- no --> Q[group_quantize columnwise only]
Q --> R[compute_grouped_dbias separately]
P & R --> S[Return grads]
Reviews (2): Last reviewed commit: "[PyTorch] Enable FP8 block scaling acros..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if -1 in shape: | ||
| d_inferred = -math.prod(cur_shape) // math.prod(shape) | ||
| for i, d in enumerate(shape): | ||
| if d == -1: | ||
| shape[i] = d_inferred | ||
| break |
There was a problem hiding this comment.
-1 inference silently truncates for non-divisible shapes
d_inferred = -math.prod(cur_shape) // math.prod(shape) uses integer floor division. If the total element count isn't evenly divisible by the product of the explicit dims (e.g. cur_shape=(3, 128), shape=[2, -1, 128]), the inferred dim is silently truncated (d_inferred = 1 instead of fractional 1.5). The subsequent self._rowwise_data.view(*shape) will then raise a PyTorch error because 2 * 1 * 128 ≠ 3 * 128, but the inferred shape[i] value (1) is already wrong. This mirrors the MXFP8 implementation exactly — it is a pre-existing pattern — but worth documenting that -1 semantics only produce a meaningful result when the remaining dimensions divide evenly (same contract as torch.Tensor.view).
f517d02 to
639233e
Compare
|
/te-ci pytorch |
3 participants
Description
Enables FP8 block scaling (FP8BS) for TE/PyTorch fusible-ops with coverage for both non-grouped and grouped linears, following the MXFP8 paths.
IMPORTANT: FP8BS support for comm+GEMM overlap is deferred to a future PR.
Changes
NotImplementedErrorfor FP8BS, use the generic recipe-state path like MXFP8.Float8BlockwiseQTensorStorage.view()so quantized norm outputs reshape without dequantizing (mirrorsMXFP8TensorStorage.view()).Type of change
Checklist: