Remove star imports for the top-level module

Author: tengqm

src/flag_gems/__init__.py

@@ -242,10 +242,12 @@
     "add",
     "add_",
     "addcdiv",
+    "addcmul",
     "addmm",
     "addmm_out",
     "addmv",
     "addmv_out",
+    "addr",
     "all",
     "all_dim",
     "all_dims",
@@ -423,6 +425,7 @@
     "ones",
     "ones_like",
     "pad",
+    "per_token_group_quant_fp8",
     "polar",
     "pow_scalar",
     "pow_tensor_scalar",
@@ -441,7 +444,6 @@
     "reciprocal_",
     "relu",
     "relu_",
-    "addcmul",
     "softplus",
     "remainder",
     "remainder_",
@@ -506,10 +508,8 @@
     "upsample_nearest2d",
     "var_mean",
     "vdot",
-    "addr",
     "vector_norm",
     "vstack",
-    "per_token_group_quant_fp8",
     "SUPPORTED_FP8_DTYPE",
     "weight_norm_interface",
     "weight_norm_interface_backward",

src/flag_gems/ops/__init__.py

@@ -6,6 +6,7 @@
 import torch
 
 import flag_gems
+from flag_gems import testing
 
 from .conftest import QUICK_MODE, TO_CPU
 
@@ -205,14 +206,14 @@ def to_cpu(res, ref):
 
 def gems_assert_close(res, ref, dtype, equal_nan=False, reduce_dim=1, atol=1e-4):
     res = to_cpu(res, ref)
-    flag_gems.testing.assert_close(
+    testing.assert_close(
         res, ref, dtype, equal_nan=equal_nan, reduce_dim=reduce_dim, atol=atol
     )
 
 
 def gems_assert_equal(res, ref, equal_nan=False):
     res = to_cpu(res, ref)
-    flag_gems.testing.assert_equal(res, ref, equal_nan=equal_nan)
+    testing.assert_equal(res, ref, equal_nan=equal_nan)
 
 
 def unsqueeze_tuple(t, max_len):

tests/accuracy_utils.py

@@ -9,6 +9,7 @@
 import triton
 
 import flag_gems
+from flag_gems import fused
 
 try:
     import vllm.vllm_flash_attn.flash_attn_interface as _vllm  # noqa: F401
@@ -1186,7 +1187,7 @@ def test_concat_and_cache_mla(
         else:
             ref_kv_cache = to_reference(ref_temp)
         with flag_gems.use_gems():
-            flag_gems.concat_and_cache_mla(
+            fused.concat_and_cache_mla(
                 kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale
             )
 
@@ -1295,7 +1296,7 @@ def test_reshape_and_cache(
         cloned_value_cache = value_cache.clone()
 
         # Call the reshape_and_cache kernel.
-        flag_gems.reshape_and_cache(
+        fused.reshape_and_cache(
             key,
             value,
             key_cache,
@@ -1444,7 +1445,7 @@ def test_reshape_and_cache_flash(
         cloned_value_cache = value_cache.clone()
 
         # Call the reshape_and_cache kernel.
-        flag_gems.reshape_and_cache_flash(
+        fused.reshape_and_cache_flash(
             key,
             value,
             key_cache,
@@ -1580,7 +1581,7 @@ def ref_mla(
         dv,
         causal,
     )
-    res_out = flag_gems.flash_mla(
+    res_out = fused.flash_mla(
         q,
         block_table,
         blocked_k,

tests/test_attention_ops.py

@@ -7,6 +7,7 @@
 import torch
 
 import flag_gems
+from flag_gems import fused
 
 from .accuracy_utils import (
     ALL_FLOAT_DTYPES,
@@ -977,7 +978,7 @@ def test_accuracy_gelu_and_mul(shape, approximate, dtype):
         torch.nn.functional.gelu(ref_inp1, approximate=approximate), ref_inp2
     )
     with flag_gems.use_gems():
-        res_out = flag_gems.gelu_and_mul(inp1, inp2, approximate)
+        res_out = fused.gelu_and_mul(inp1, inp2, approximate)
 
     out_grad = torch.randn_like(res_out)
     ref_grad = to_reference(out_grad, True)
@@ -1389,7 +1390,7 @@ def test_accuracy_silu_and_mul(shape, dtype):
 
     ref_out = torch.mul(torch.nn.functional.silu(ref_inp1), ref_inp2)
     with flag_gems.use_gems():
-        res_out = flag_gems.silu_and_mul(inp1, inp2)
+        res_out = fused.silu_and_mul(inp1, inp2)
 
     out_grad = torch.randn_like(res_out)
     ref_grad = to_reference(out_grad, True)

tests/test_binary_pointwise_ops.py

@@ -6,6 +6,7 @@
 import torch
 
 import flag_gems
+from flag_gems import fused, ops
 
 from .accuracy_utils import (
     FLOAT_DTYPES,
@@ -143,7 +144,7 @@ def test_accuracy_baddbmm(M, N, K, scalar, dtype):
     alpha = beta = scalar
 
     ref_out = torch.baddbmm(ref_bias, ref_mat1, ref_mat2, alpha=alpha, beta=beta)
-    res_out = flag_gems.baddbmm(bias, mat1, mat2, alpha=alpha, beta=beta)
+    res_out = ops.baddbmm(bias, mat1, mat2, alpha=alpha, beta=beta)
 
     gems_assert_close(res_out, ref_out, dtype, reduce_dim=K)
 
@@ -174,7 +175,7 @@ def test_accuracy_baddbmm_backward(M, N, K, scalar, dtype):
     alpha = beta = scalar
 
     ref_out = torch.baddbmm(ref_bias, ref_mat1, ref_mat2, alpha=alpha, beta=beta)
-    res_out = flag_gems.baddbmm(bias, mat1, mat2, alpha=alpha, beta=beta)
+    res_out = ops.baddbmm(bias, mat1, mat2, alpha=alpha, beta=beta)
 
     out_grad = torch.randn_like(res_out)
     ref_grad = to_reference(out_grad, True)
@@ -303,7 +304,7 @@ def test_accuracy_outer(M, N, dtype):
     ref_inp2 = to_reference(inp2, True)
 
     ref_out = torch.outer(ref_inp1, ref_inp2)
-    res_out = flag_gems.outer(inp1, inp2)
+    res_out = fused.outer(inp1, inp2)
     gems_assert_close(res_out, ref_out, dtype)
 
     out_grad = torch.randn_like(res_out)

tests/test_blas_ops.py

@@ -6,6 +6,7 @@
 import torch
 
 import flag_gems
+from flag_gems import fused, ops
 
 from .accuracy_utils import (
     FLOAT_DTYPES,
@@ -366,7 +367,7 @@ def test_accuracy_instancenorm(
         eps=eps,
     )
 
-    res_out = flag_gems.instance_norm(
+    res_out = fused.instance_norm(
         inp,
         weight=weight,
         bias=bias,
@@ -544,7 +545,7 @@ def _torch_rms_norm(x, weight, eps):
         return weight * hidden_states
 
     ref_out = _torch_rms_norm(ref_inp, weight=ref_weight, eps=eps)
-    res_out = flag_gems.rms_norm(inp, list(layer_shape), weight=weight, eps=eps)
+    res_out = ops.rms_norm(inp, list(layer_shape), weight=weight, eps=eps)
 
     res_grad = torch.tensor(
         np_grad, dtype=dtype, device=flag_gems.device, requires_grad=True
@@ -587,7 +588,7 @@ def test_accuracy_skip_layernorm(shape, dtype):
         bias=ref_bias,
         eps=eps,
     )
-    res_out = flag_gems.skip_layer_norm(
+    res_out = fused.skip_layer_norm(
         inp, residual, list(layer_shape), weight=weight, bias=bias, eps=eps
     )
 
@@ -624,7 +625,7 @@ def _torch_fused_add_rms_norm(x, residual, weight, eps):
         eps=eps,
     )
 
-    res_out, res_new_residual = flag_gems.fused_add_rms_norm(
+    res_out, res_new_residual = fused.fused_add_rms_norm(
         inp, residual, list(layer_shape), weight=weight, eps=eps
     )