Repository navigation
Conversation
tf.matmul of two fp8e4m3 operands returned fp8e4m3, so a block-scaled FP8 GEMM could not multiply its product by an f32 scale. No MMA sums in fp8e4m3: MatMul now types, costs and evaluates fp8e4m3 products in f32.
A slice bound written with dimension arithmetic, such as scale[:, gi * G:gi * G + G], stopped analyze with 'DimMul: no cost evaluator registered'. DimAdd, DimSub, DimMul, DimFloorDiv, DimMod, DimMin and DimMax now count as integer work.
SM90 had only the bf16/f16 Wgmma atom, so tiled_mma refused fp8e4m3 operands and schedule candidates offered no FP8 instruction. T.cuda.sm90.WgmmaFp8 declares the 64 x n x 32 e4m3 MMA with an f32 accumulator and K-major A and B in swizzled shared memory. The schedule evaluator sums its operands in f32, and facts lists each capability once.
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
1 participant
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.
Why
scheduled:
tf.matmulreturned fp8 products,analyzehad no cost fordimension arithmetic in slice bounds, and SM90 declared no FP8 WGMMA atom.
What
MatMultypes, costs and evaluates a product of twofp8e4m3operands inf32(matmul_result_dtype). Every other dtype keeps its own result dtype.DimAdd,DimSub,DimMul,DimFloorDiv,DimMod,DimMinandDimMaxcount as integer work in the compute cost.
T.cuda.sm90.WgmmaFp8(n=...):wgmma.mma_async64 x n x 32 for e4m3,f32 accumulator fragment as
Wgmma, A (64, 32) and B (32, n) K-major inshared memory, in swizzled rows of 16 to 128 bytes. It
joins the MMA family, so
schedule candidatesoffers it andfinalizelowersit. The schedule evaluator sums FP8 operands in f32.
schedule factslists each capability once (two atoms sharewgmma.mma_async).analysis case, an FP8 K-major WGMMA fixture with its TIR golden, facts
golden, smem/rmem budgets, and an exact-f32 evaluation check. The plain
candidates golden gains the two
WgmmaFp8refusal lines.Contract
docs/spec/hir.md: MatMul returnsf32forfp8e4m3operands. This changesthe result dtype of any program that already multiplied fp8 tensors; before,
such a program could not combine the product with an
f32value.docs/spec/tir.mdanddocs/spec/code-organization.mddeclareWgmmaFp8and its file.Risk
B must be K-major; an (N, K) row-major weight is written as its (K, N) view.
CUTLASS checkout); the schedule, analysis, evaluator and MatMul suites pass
(464 tests).