init

Author: cyx-6

flashinfer/fused_moe.py

@@ -39,6 +39,7 @@
 from .jit import JitSpec
 from .jit import env as jit_env
 from .jit import gen_jit_spec, setup_cubin_loader, sm100a_nvcc_flags
+from .jit.cubin_loader import get_cubin
 from .utils import _check_shape_dtype_device, register_custom_op, register_fake_op
 
 
@@ -773,6 +774,13 @@ def cutlass_fused_moe(
 
 
 def trtllm_gen_fused_moe_sm100_module() -> JitSpec:
+    hash = "f5deee96023f1d74b1ff71ac69f782a96741a053"
+    metainfo = get_cubin(
+        f"{hash}/batched_gemm-c603ed2-3fa89e1/include/KernelMetaInfo",
+        "d789c63aaeee1aa0a68ebf22fa693b6b82a7c2319bd933a00a10306ca08d9e0e",
+        ".h",
+    )
+    assert metainfo, "KernelMetaInfo.h not found"
     return gen_jit_spec(
         "fused_moe_sm100",
         [
@@ -790,6 +798,13 @@ def trtllm_gen_fused_moe_sm100_module() -> JitSpec:
             "-DENABLE_FP4",
         ]
         + sm100a_nvcc_flags,
+        extra_include_paths=[
+            jit_env.FLASHINFER_CACHE_DIR
+            / "cubins"
+            / hash
+            / "batched_gemm-c603ed2-3fa89e1"
+            / "include"
+        ],
         extra_ldflags=["-lcuda"],
     )
 

flashinfer/jit/attention/pytorch.py

@@ -22,6 +22,7 @@
 
 from .. import env as jit_env
 from ..core import JitSpec, gen_jit_spec, logger, sm90a_nvcc_flags, sm100a_nvcc_flags
+from ..cubin_loader import get_cubin
 from ..utils import (
     dtype_map,
     filename_safe_dtype_map,
@@ -1482,12 +1483,27 @@ def gen_fmha_cutlass_sm100a_module(
 
 
 def trtllm_gen_fmha_module():
+    hash = "5f2779e6df822bc0b26940b6d3b0059c86f0a6a1"
+    metainfo = get_cubin(
+        f"{hash}/fmha/trtllm-gen/include/flashInferMetaInfo",
+        "11f31dc81f996e39c3f1d85d773864c9113c5837619e21418a846befa4f8dddd",
+        ".h",
+    )
+    assert metainfo, "flashInferMetaInfo.h not found"
     return gen_jit_spec(
         "fmha_gen",
         [
             jit_env.FLASHINFER_CSRC_DIR / "trtllm_fmha_runner.cu",
             jit_env.FLASHINFER_CSRC_DIR / "trtllm_fmha_kernel_launcher.cu",
         ],
+        extra_include_paths=[
+            jit_env.FLASHINFER_CACHE_DIR
+            / "cubins"
+            / hash
+            / "fmha"
+            / "trtllm-gen"
+            / "include"
+        ],
         extra_ldflags=["-lcuda"],
     )
 

include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmInterface.h

@@ -645,7 +645,7 @@ int32_t BatchedGemmInterface::run(BatchedGemmConfig const& config, void* workspa
 
   auto fiModuleLoadData = [&](CUmodule* module) {
     const std::string sha256 = config.mHash ? config.mHash : "";
-    const std::string pipeline_hash = "991e7438224199de85ef08a2730ce18c12b4e0aa";
+    const std::string pipeline_hash = "f5deee96023f1d74b1ff71ac69f782a96741a053";
     const std::string cubin_path = pipeline_hash + "/" + std::string("batched_gemm-") +
                                    TLLM_GEN_COMMIT + "-" + TLLM_GEN_BATCHED_GEMM_CONFIG_HASH + "/";
     std::string fname_cubin = config.mFunctionName;