Skip to content

Float16 mul! falls through to GPUArrays' generic kernel (~65x slower than rocBLAS gemm_ex) #1048

Description

@gbaraldi

Float16 mul! currently reaches neither rocBLAS nor any tuned kernel — it falls through to GPUArrays' generic matmul. Measured on an MI300A (gfx942, Julia 1.12.6; identical numbers under ROCm 6.4.2, 7.2.4 and 7.14.0), a 4096³ f16 GEMM runs at 3.0 TFLOPS, when rocBLAS's tuned path delivers ~200 on the same machine.

Two stacked causes:

1. Dispatch fallthrough. The rocBLAS branch in LinearAlgebra.generic_matmatmul! (src/blas/highlevel.jl) requires T <: ROCBLASFloat, which excludes Float16 (and BFloat16). On Julia ≥ 1.12 the generic_matmatmul_wrapper! overload is declared for T <: ROCBLASFloatWithHalf — Float16 included — but it forwards to that same generic_matmatmul!, so Float16 falls to GPUArrays.generic_matmatmul! anyway. Julia 1.11 hits the same eltype branch, so the fallthrough exists there too.

2. The typed entry point would be the wrong target anyway. rocBLAS.gemm! maps Float16 to legacy rocblas_hgemm, which Tensile doesn't tune. The tuned path is rocblas_gemm_ex with f16 operands and an f32 compute type — which is also the numerically better route (f32 accumulation). AMDGPU.jl carries only the raw generated binding (rocblas_gemm_ex_64), with no high-level wrapper and nothing dispatching to it.

Measured ladder, 4096³ f16 (min of 20, MI300A):

path TFLOPS
mul! today (GPUArrays fallback) 3.0
rocBLAS.gemm!rocblas_hgemm 69.6–71.8
rocblas_gemm_ex, f16 in / f32 compute 189–224

For f32/f64, gemm_ex and the typed gemm! perform identically (50.2 vs 50.2 and 52.2 vs 52.7 TFLOPS respectively), so routing through gemm_ex carries no risk for the types that already work.

Repro:

using AMDGPU, LinearAlgebra
A, B = ROCArray(rand(Float16, 4096, 4096)), ROCArray(rand(Float16, 4096, 4096))
C = similar(A)
mul!(C, A, B)  # ~85 ms on MI300A; ~1.4 ms via rocblas_gemm_ex

PR incoming: a rocBLAS.gemm_ex! wrapper (mirroring CUBLAS.gemmEx!) plus dispatching Float16/BFloat16 mul! through it with Float32 accumulation — validated on MI300A at 190 (f16) / 208 (bf16) TFLOPS with the existing Level 3 tests passing.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions