Move specialized gemm kernels and Triton allocator into NVIDIA backend; add iluvatar backend support

Author: AriesMTR

benchmark/test_gemm_perf.py

@@ -107,7 +107,7 @@ def gems_sgemm_wrapper(
     alpha_ptr,
     beta_ptr,
 ):
-    flag_blas.ops.sgemm(
+    flag_blas.sgemm(
         transa,
         transb,
         m,
@@ -197,7 +197,7 @@ def gems_hgemm_wrapper(
     alpha_ptr,
     beta_ptr,
 ):
-    flag_blas.ops.hgemm(
+    flag_blas.hgemm(
         transa,
         transb,
         m,
@@ -287,7 +287,7 @@ def gems_bfgemm_wrapper(
     alpha_ptr,
     beta_ptr,
 ):
-    flag_blas.ops.bfgemm(
+    flag_blas.bfgemm(
         transa,
         transb,
         m,
@@ -372,7 +372,7 @@ def gems_fp8gemm_wrapper(
     alpha_ptr,
     beta_ptr,
 ):
-    flag_blas.ops.fp8gemm(
+    flag_blas.fp8gemm(
         transa,
         transb,
         m,

src/flag_blas/__init__.py

@@ -2,20 +2,13 @@
 flag_blas - BLAS operations implemented with Triton
 """
 
-import torch
-import triton
-from packaging import version
-
-
-def _alloc_fn(size, alignment, stream):
-    return torch.empty(size, device="cuda", dtype=torch.int8)
+import warnings
 
-
-triton.set_allocator(_alloc_fn)
+import torch
 
 from flag_blas import runtime
-from flag_blas import testing
-from flag_blas.ops import *
+from flag_blas import testing  # noqa: F401
+from flag_blas.ops import *  # noqa: F401,F403
 from flag_blas.config import aten_patch_list, resolve_user_setting
 from flag_blas.runtime.register import Register