Besides Apache TVM mentioned in #19 we can also consider using the MLIR / OpenXLA ecosystem that powers JAX (and PyTorch can export to it).
See:
Note that matmul / tensor contraction is named dot_general: https://openxla.org/stablehlo/spec#dot_general
It's unsure if that uses tensor cores for types like FP8 or if it can implement fast NVFP4 attention https://github.com/openxla/stablehlo/pull/2096/changes.
Besides Apache TVM mentioned in #19 we can also consider using the MLIR / OpenXLA ecosystem that powers JAX (and PyTorch can export to it).
See:
Note that matmul / tensor contraction is named dot_general: https://openxla.org/stablehlo/spec#dot_general
It's unsure if that uses tensor cores for types like FP8 or if it can implement fast NVFP4 attention https://github.com/openxla/stablehlo/pull/2096/changes.