perf: pure FP8 fwd/bwd and lean non-transposed GEMM
- drop the fused kernel; forward/backward are quantize + a pre-quantized GEMM - rename module fp8_mm -> fp8_ops (mm.cu -> ops.cu) - kernels/launchers fp8_gemm_kernel / launch_fp8_gemm; drop PqTraits/gather_trans/pack_fp8x4_vector - remove the in-kernel transposed-operand branches (TransA/TransB) - backward: quantize g once (amax_g here), explicit fp8 transposes, fast non-transposed GEMMs (dX = g@w^T, dW = g^T@x^T) - each pass uses a single FP8 format (E4M3 fwd / E5M2 bwd)
This commit is contained in:
+1
-1
@@ -17,7 +17,7 @@ CUDA_AVAIL = torch.cuda.is_available()
|
||||
KERNEL_AVAIL = CUDA_AVAIL and all(is_available(k) for k in KERNEL_NAMES)
|
||||
FP8_AVAIL = (
|
||||
CUDA_AVAIL
|
||||
and is_available("fp8_mm")
|
||||
and is_available("fp8_ops")
|
||||
and torch.cuda.get_device_capability() >= (8, 9)
|
||||
)
|
||||
skip_no_cuda = pytest.mark.skipif(not CUDA_AVAIL, reason="CUDA not available")
|
||||
|
||||
Reference in New Issue
Block a user