refactor: Improved metainfo for trtllm-gen fmha (#1292)

Refactor the metainfo for trtllm gen fmha.

<!-- .github/pull_request_template.md -->

## 📌 Description

<!-- What does this PR do? Briefly describe the changes and why they’re
needed. -->

## 🔍 Related Issues

<!-- Link any related issues here -->

## 🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull
request, please make sure the following items are complete.

### ✅ Pre-commit Checks

- [ ] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [ ] I have installed the hooks with `pre-commit install`.
- [ ] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.

> If you are unsure about how to set up `pre-commit`, see [the
pre-commit documentation](https://pre-commit.com/).

## 🧪 Tests

- [ ] Tests have been added or updated as needed.
- [ ] All tests are passing (`unittest`, etc.).

## Reviewer Notes

<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->

---------

Co-authored-by: Yaxing Cai <[email protected]>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Yingyi Huang <[email protected]>

Author: 3 people

flashinfer/decode.py

@@ -33,6 +33,7 @@
     get_batch_prefill_uri,
     get_single_decode_uri,
     setup_cubin_loader,
+    setup_metainfo_loader,
     trtllm_fmha_gen_module,
     trtllm_mla_gen_module,
 )
@@ -302,6 +303,7 @@ def get_trtllm_fmha_gen_module():
     mod = trtllm_fmha_gen_module()
     op = mod.build_and_load()
     setup_cubin_loader(mod.get_library_path())
+    setup_metainfo_loader(mod.get_library_path())
     return op
 
 
@@ -310,6 +312,7 @@ def get_trtllm_mla_gen_module():
     mod = trtllm_mla_gen_module()
     op = mod.build_and_load()
     setup_cubin_loader(mod.get_library_path())
+    setup_metainfo_loader(mod.get_library_path())
     return op
 
 

flashinfer/jit/__init__.py

@@ -69,7 +69,7 @@
 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 .cubin_loader import setup_cubin_loader
+from .cubin_loader import setup_cubin_loader, setup_metainfo_loader
 
 
 @functools.cache

flashinfer/jit/cubin_loader.py

@@ -188,3 +188,32 @@ def get_cubin_callback(name, sha256):
     dll_cubin_handlers[dll_path] = cb
 
     _LIB.FlashInferSetCubinCallback(cb)
+
+
+dll_metainfo_handlers = {}
+
+
+def setup_metainfo_loader(dll_path: str):
+    if dll_path in dll_metainfo_handlers:
+        return
+
+    _LIB = ctypes.CDLL(dll_path)
+
+    # Define the correct callback type
+    CALLBACK_TYPE = ctypes.CFUNCTYPE(
+        None, ctypes.c_char_p, ctypes.c_char_p, ctypes.c_char_p
+    )
+
+    def get_metainfo_callback(name, sha256, extension):
+        metainfo = get_cubin(
+            name.decode("utf-8"), sha256.decode("utf-8"), extension.decode("utf-8")
+        )
+        _LIB.FlashInferSetCurrentMetaInfo(
+            convert_to_ctypes_char_p(metainfo), ctypes.c_int(len(metainfo))
+        )
+
+    # Create the callback and keep a reference to prevent GC
+    cb = CALLBACK_TYPE(get_metainfo_callback)
+    dll_metainfo_handlers[dll_path] = cb
+
+    _LIB.FlashInferSetMetaInfoCallback(cb)

include/flashinfer/cubin_loader.h

@@ -56,3 +56,31 @@ std::string getCubin(const std::string& name, const std::string& sha256) {
   callbackGetCubin(name.c_str(), sha256.c_str());
   return current_cubin;
 }
+
+void (*callbackGetMetaInfo)(const char* path, const char* sha256, const char* extension) = nullptr;
+
+// Set the python callback, called by the python code using ctypes.
+extern "C" void FlashInferSetMetaInfoCallback(void (*callback)(const char* path, const char* sha256,
+                                                               const char* extension)) {
+  callbackGetMetaInfo = callback;
+}
+
+// Thread-local variable that stores the current metainfo.
+// It is reset on every call to `getMetaInfo()`.
+thread_local std::string raw_metainfo;
+
+// Called by the callback to set the current metainfo.
+extern "C" void FlashInferSetCurrentMetaInfo(const char* binary, int size) {
+  raw_metainfo = std::string(binary, size);
+}
+
+// Get the metainfo from the python callback.
+// This is the API for the native library to use.
+std::string getMetaInfo(const std::string& name, const std::string& sha256,
+                        const std::string& extension) {
+  if (!callbackGetMetaInfo) {
+    throw std::runtime_error("FlashInferSetMetaInfoCallback not set");
+  }
+  callbackGetMetaInfo(name.c_str(), sha256.c_str(), extension.c_str());
+  return raw_metainfo;
+}