misc: jit: Import jit_env as a module

Author: abcdabcd987

flashinfer/cascade.py

@@ -20,7 +20,9 @@
 import torch
 
 from .decode import BatchDecodeWithPagedKVCacheWrapper
-from .jit import FLASHINFER_CSRC_DIR, JitSpec, gen_jit_spec, has_prebuilt_ops
+from .jit import JitSpec
+from .jit import env as jit_env
+from .jit import gen_jit_spec, has_prebuilt_ops
 from .prefill import BatchPrefillWithPagedKVCacheWrapper, single_prefill_with_kv_cache
 from .utils import register_custom_op, register_fake_op
 
@@ -31,8 +33,8 @@ def gen_cascade_module() -> JitSpec:
     return gen_jit_spec(
         "cascade",
         [
-            FLASHINFER_CSRC_DIR / "cascade.cu",
-            FLASHINFER_CSRC_DIR / "flashinfer_cascade_ops.cu",
+            jit_env.FLASHINFER_CSRC_DIR / "cascade.cu",
+            jit_env.FLASHINFER_CSRC_DIR / "flashinfer_cascade_ops.cu",
         ],
     )
 

flashinfer/custom_all_reduce.py

@@ -21,7 +21,9 @@
 
 import torch
 
-from .jit import FLASHINFER_CSRC_DIR, JitSpec, gen_jit_spec, has_prebuilt_ops
+from .jit import JitSpec
+from .jit import env as jit_env
+from .jit import gen_jit_spec, has_prebuilt_ops
 from .utils import register_custom_op
 
 _comm_module = None
@@ -31,8 +33,8 @@ def gen_comm_module() -> JitSpec:
     return gen_jit_spec(
         "comm",
         [
-            FLASHINFER_CSRC_DIR / "flashinfer_comm_ops.cu",
-            FLASHINFER_CSRC_DIR / "custom_all_reduce.cu",
+            jit_env.FLASHINFER_CSRC_DIR / "flashinfer_comm_ops.cu",
+            jit_env.FLASHINFER_CSRC_DIR / "custom_all_reduce.cu",
         ],
     )
 

flashinfer/gemm.py

@@ -21,14 +21,9 @@
 import torch
 import torch.nn.functional as F
 
-from .jit import (
-    FLASHINFER_CSRC_DIR,
-    JitSpec,
-    gen_jit_spec,
-    has_prebuilt_ops,
-    sm90a_nvcc_flags,
-    sm100a_nvcc_flags,
-)
+from .jit import JitSpec
+from .jit import env as jit_env
+from .jit import gen_jit_spec, has_prebuilt_ops, sm90a_nvcc_flags, sm100a_nvcc_flags
 from .utils import (
     _get_cache_buf,
     determine_gemm_backend,
@@ -46,9 +41,9 @@ def gen_gemm_module() -> JitSpec:
     return gen_jit_spec(
         "gemm",
         [
-            FLASHINFER_CSRC_DIR / "bmm_fp8.cu",
-            FLASHINFER_CSRC_DIR / "group_gemm.cu",
-            FLASHINFER_CSRC_DIR / "flashinfer_gemm_ops.cu",
+            jit_env.FLASHINFER_CSRC_DIR / "bmm_fp8.cu",
+            jit_env.FLASHINFER_CSRC_DIR / "group_gemm.cu",
+            jit_env.FLASHINFER_CSRC_DIR / "flashinfer_gemm_ops.cu",
         ],
         extra_ldflags=["-lcublas", "-lcublasLt"],
     )
@@ -157,10 +152,10 @@ def gen_gemm_sm100_module() -> JitSpec:
     return gen_jit_spec(
         "gemm_sm100",
         [
-            FLASHINFER_CSRC_DIR / "gemm_groupwise_sm100.cu",
-            FLASHINFER_CSRC_DIR / "group_gemm_groupwise_sm100.cu",
-            FLASHINFER_CSRC_DIR / "gemm_sm100_pybind.cu",
-            FLASHINFER_CSRC_DIR / "group_gemm_sm100_pybind.cu",
+            jit_env.FLASHINFER_CSRC_DIR / "gemm_groupwise_sm100.cu",
+            jit_env.FLASHINFER_CSRC_DIR / "group_gemm_groupwise_sm100.cu",
+            jit_env.FLASHINFER_CSRC_DIR / "gemm_sm100_pybind.cu",
+            jit_env.FLASHINFER_CSRC_DIR / "group_gemm_sm100_pybind.cu",
         ],
         extra_cuda_cflags=sm100a_nvcc_flags,
     )
@@ -181,14 +176,14 @@ def gen_gemm_sm90_module() -> JitSpec:
     return gen_jit_spec(
         "gemm_sm90",
         [
-            FLASHINFER_CSRC_DIR / "group_gemm_sm90.cu",
-            FLASHINFER_CSRC_DIR / "flashinfer_gemm_sm90_ops.cu",
-            FLASHINFER_CSRC_DIR / "group_gemm_f16_f16_sm90.cu",
-            FLASHINFER_CSRC_DIR / "group_gemm_bf16_bf16_sm90.cu",
-            FLASHINFER_CSRC_DIR / "group_gemm_e4m3_f16_sm90.cu",
-            FLASHINFER_CSRC_DIR / "group_gemm_e5m2_f16_sm90.cu",
-            FLASHINFER_CSRC_DIR / "group_gemm_e4m3_bf16_sm90.cu",
-            FLASHINFER_CSRC_DIR / "group_gemm_e5m2_bf16_sm90.cu",
+            jit_env.FLASHINFER_CSRC_DIR / "group_gemm_sm90.cu",
+            jit_env.FLASHINFER_CSRC_DIR / "flashinfer_gemm_sm90_ops.cu",
+            jit_env.FLASHINFER_CSRC_DIR / "group_gemm_f16_f16_sm90.cu",
+            jit_env.FLASHINFER_CSRC_DIR / "group_gemm_bf16_bf16_sm90.cu",
+            jit_env.FLASHINFER_CSRC_DIR / "group_gemm_e4m3_f16_sm90.cu",
+            jit_env.FLASHINFER_CSRC_DIR / "group_gemm_e5m2_f16_sm90.cu",
+            jit_env.FLASHINFER_CSRC_DIR / "group_gemm_e4m3_bf16_sm90.cu",
+            jit_env.FLASHINFER_CSRC_DIR / "group_gemm_e5m2_bf16_sm90.cu",
         ],
         extra_cuda_cflags=sm90a_nvcc_flags,
     )

flashinfer/jit/__init__.py

@@ -18,6 +18,7 @@
 import os
 
 # Re-export
+from . import env as env
 from .activation import gen_act_and_mul_module as gen_act_and_mul_module
 from .activation import get_act_and_mul_cu_str as get_act_and_mul_cu_str
 from .attention import gen_batch_decode_mla_module as gen_batch_decode_mla_module
@@ -61,7 +62,6 @@
 from .core import gen_jit_spec as gen_jit_spec
 from .core import sm90a_nvcc_flags as sm90a_nvcc_flags
 from .core import sm100a_nvcc_flags as sm100a_nvcc_flags
-from .env import *
 
 cuda_lib_path = os.environ.get(
     "CUDA_LIB_PATH", "/usr/local/cuda/targets/x86_64-linux/lib/"

flashinfer/jit/activation.py

@@ -18,8 +18,8 @@
 
 import jinja2
 
+from . import env as jit_env
 from .core import JitSpec, gen_jit_spec
-from .env import FLASHINFER_GEN_SRC_DIR
 from .utils import write_if_different
 
 activation_templ = r"""
@@ -77,7 +77,7 @@ def get_act_and_mul_cu_str(act_func_name: str, act_func_def: str) -> str:
 
 
 def gen_act_and_mul_module(act_func_name: str, act_func_def: str) -> JitSpec:
-    gen_directory = FLASHINFER_GEN_SRC_DIR
+    gen_directory = jit_env.FLASHINFER_GEN_SRC_DIR
     os.makedirs(gen_directory, exist_ok=True)
     sources = [gen_directory / f"{act_func_name}_and_mul.cu"]
     write_if_different(