refactor: Improved metainfo for trtllm-gen kernels #1328

cyx-6 · 2025-07-25T05:58:20Z

📌 Description

Move trtllm-gen batched-gemm and gemm metainfo headers into cache directory and link when jit.

🔍 Related Issues

🚀 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.

🧪 Tests

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

Reviewer Notes

gemini-code-assist

Summary of Changes

Hello @cyx-6, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request refactors the handling of metainfo headers for trtllm-gen kernels within the FlashInfer library. It reverts a previous change, streamlining the JIT compilation process by moving metainfo headers to a cache directory and linking them directly during compilation, rather than relying on a runtime callback mechanism. This change aims to improve the efficiency and clarity of how kernel metadata is accessed and utilized.

Highlights

Metainfo Loading Refactor: The previous setup_metainfo_loader mechanism, which used a Python callback to fetch metainfo at runtime, has been removed from both Python and C++ components.
Direct Metainfo Inclusion for JIT: Metainfo headers for trtllm-gen kernels (both FMHA and Fused MoE) are now explicitly fetched from a cache directory using get_cubin and included directly in the JIT compilation process via extra_include_paths.
C++ Metainfo Integration: The C++ backend (fmhaKernels.cuh) no longer relies on a runtime getMetaInfo callback. Instead, kernel metadata is passed directly to the kernel factory, suggesting it's now statically defined or linked within the C++ binary.
Header File Restructuring: The kernelMetaInfo.h file has been deleted, indicating a shift in how kernel metadata is structured and accessed, moving towards a more integrated approach.
Pipeline Hash Update: A pipeline_hash in BatchedGemmInterface.h has been updated, likely reflecting changes in the build or generation process for batched GEMM kernels.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments or fill out our survey to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request refactors how kernel metainfo is handled for trtllm-gen kernels, moving from a dynamic callback-based loading mechanism to including header files at JIT compilation time. The changes are consistent across Python and C++ code. My review focuses on improving maintainability by replacing hardcoded strings with constants and suggesting modern C++ idioms.

gemini-code-assist · 2025-07-25T06:00:48Z

flashinfer/fused_moe.py

+    hash = "f5deee96023f1d74b1ff71ac69f782a96741a053"
+    metainfo = get_cubin(
+        f"{hash}/batched_gemm-c603ed2-3fa89e1/include/KernelMetaInfo",
+        "d789c63aaeee1aa0a68ebf22fa693b6b82a7c2319bd933a00a10306ca08d9e0e",
+        ".h",
+    )
+    assert metainfo, "KernelMetaInfo.h not found"


This function contains several hardcoded strings for hashes and path components. To improve maintainability and readability, define these as constants at the module level and reference them here. This also helps avoid duplicating the path construction logic.

_TRTLLM_FUSED_MOE_SM100_PIPELINE_HASH = "f5deee96023f1d74b1ff71ac69f782a96741a053" _TRTLLM_FUSED_MOE_SM100_CUBIN_PATH_COMPONENT = "batched_gemm-c603ed2-3fa89e1" _TRTLLM_FUSED_MOE_SM100_METAINFO_SHA256 = "d789c63aaeee1aa0a68ebf22fa693b6b82a7c2319bd933a00a10306ca08d9e0e"

and then use them inside the function. I've also de-duplicated the path construction in the suggestion.

pipeline_hash = "f5deee96023f1d74b1ff71ac69f782a96741a053" cubin_path_component = "batched_gemm-c603ed2-3fa89e1" metainfo_sha256 = "d789c63aaeee1aa0a68ebf22fa693b6b82a7c2319bd933a00a10306ca08d9e0e" cubin_path = f"{pipeline_hash}/{cubin_path_component}" metainfo = get_cubin( f"{cubin_path}/include/KernelMetaInfo", metainfo_sha256, ".h", )

Seems like a good suggestion to me!

gemini-code-assist · 2025-07-25T06:00:49Z

flashinfer/jit/attention/pytorch.py

+    hash = "5f2779e6df822bc0b26940b6d3b0059c86f0a6a1"
+    metainfo = get_cubin(
+        f"{hash}/fmha/trtllm-gen/include/flashInferMetaInfo",
+        "11f31dc81f996e39c3f1d85d773864c9113c5837619e21418a846befa4f8dddd",
+        ".h",
+    )
+    assert metainfo, "flashInferMetaInfo.h not found"


This function contains several hardcoded strings for hashes and path components. To improve maintainability and readability, it's better to define these as constants at the module level and reference them here. This also helps avoid duplicating the path construction logic.

For example, you could define:

_TRTLLM_FMHA_PIPELINE_HASH = "5f2779e6df822bc0b26940b6d3b0059c86f0a6a1" _TRTLLM_FMHA_CUBIN_PATH_COMPONENT = "fmha/trtllm-gen" _TRTLLM_FMHA_METAINFO_SHA256 = "11f31dc81f996e39c3f1d85d773864c9113c5837619e21418a846befa4f8dddd"

and then use them inside the function. I've also de-duplicated the path construction in the suggestion.

pipeline_hash = "5f2779e6df822bc0b26940b6d3b0059c86f0a6a1" cubin_path_component = "fmha/trtllm-gen" metainfo_sha256 = "11f31dc81f996e39c3f1d85d773864c9113c5837619e21418a846befa4f8dddd" cubin_path = f"{pipeline_hash}/{cubin_path_component}" metainfo = get_cubin( f"{cubin_path}/include/flashInferMetaInfo", metainfo_sha256, ".h", )

gemini-code-assist · 2025-07-25T06:00:49Z

include/flashinfer/trtllm/fmha/fmhaKernels.cuh

+      sTllmGenFmhaKernelMetaInfos,
+      sizeof(sTllmGenFmhaKernelMetaInfos) / sizeof(sTllmGenFmhaKernelMetaInfos[0]), dtypeQ, dtypeKv,
+      dtypeOut, sm);


For better readability and to use a more modern C++ approach, consider using std::size to get the number of elements in the array. This requires including the <iterator> header at the top of the file.

The sizeof(array) / sizeof(array[0]) idiom is correct, but std::size is generally preferred in C++17 and later for type safety and clarity.

#include <iterator> return TllmFmhaKernelFactory::Get().getKernels( sTllmGenFmhaKernelMetaInfos, std::size(sTllmGenFmhaKernelMetaInfos), dtypeQ, dtypeKv, dtypeOut, sm);

flashinfer/fused_moe/core.py

…ainfo

flashinfer/gemm.py

Co-authored-by: Mehdi Amini <[email protected]>

yzh119 · 2025-08-13T07:57:19Z

flashinfer/fused_moe/core.py

@@ -846,6 +848,8 @@ def trtllm_gen_fused_moe_sm100_module() -> JitSpec:
        + sm100a_nvcc_flags,
        extra_ldflags=["-lcuda"],
        extra_include_paths=[
+            # link "include" sub-directory in cache
+            jit_env.FLASHINFER_CACHE_DIR / "cubins" / include_path,


This might be conflicting with #1462, considering using "FLASHINFER_CUBIN_DIR" instead

yzh119 · 2025-08-13T07:57:32Z

flashinfer/gemm.py

-            / "flashinfer/trtllm/gemm/trtllmGen_gemm_export",
-        ],
+        # link "include" sub-directory in cache
+        extra_include_paths=[jit_env.FLASHINFER_CACHE_DIR / "cubins" / include_path],


…ainfo

yzh119

LGTM, thanks @cyx-6 for the refactor work and @joker-eph for all the suggestions and discussions!

gemini-code-assist bot reviewed Jul 25, 2025

View reviewed changes

cyx-6 force-pushed the gemm-metainfo branch 2 times, most recently from e0eca6a to d98f090 Compare July 31, 2025 00:19

This was referenced Jul 31, 2025

feature: add fp4 mm using trtllm backend #1355

Merged

GPT-OSS Support: Add Blackwell MoE mxfp4 implementation from TRTLLM and Attention Sink #1389

Merged

rebase

c5a6ca1

cyx-6 force-pushed the gemm-metainfo branch from d98f090 to c5a6ca1 Compare August 9, 2025 06:24

upd

2f8d2e6

joker-eph reviewed Aug 9, 2025

View reviewed changes

flashinfer/fused_moe/core.py Show resolved Hide resolved

cyx-6 added 4 commits August 10, 2025 04:42

upd

11c7bf9

fix

d4138f5

Merge commit '7c4541228d1255311f472d882182b841ef8fd4f2' into gemm-met…

1219360

…ainfo

add documentation

c3945bf

joker-eph approved these changes Aug 12, 2025

View reviewed changes

joker-eph reviewed Aug 12, 2025

View reviewed changes

flashinfer/gemm.py Outdated Show resolved Hide resolved

cyx-6 and others added 3 commits August 12, 2025 13:21

Update flashinfer/gemm.py

d3356ac

Co-authored-by: Mehdi Amini <[email protected]>

upd

7aafe0d

upd hash

927a110

yzh119 reviewed Aug 13, 2025

View reviewed changes

cyx-6 added 3 commits August 13, 2025 00:57

Merge commit '30319e7680fa2640c54b28c0691f596df6e5cc54' into gemm-met…

349133f

…ainfo

use FLASHINFER_CUBIN_DIR

a368bb2

fix

cf3e0e0

yzh119 approved these changes Aug 13, 2025

View reviewed changes

yzh119 merged commit e0aa4e7 into flashinfer-ai:main Aug 13, 2025
1 check passed

This was referenced Aug 14, 2025

Restore llama4 fc2 required kernels #1417

Open

Potential header precedence issue with new cubin/artifact loading approach #1485

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

refactor: Improved metainfo for trtllm-gen kernels #1328

refactor: Improved metainfo for trtllm-gen kernels #1328

Uh oh!

cyx-6 commented Jul 25, 2025 •

edited

Loading

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Jul 25, 2025

Uh oh!

joker-eph Jul 25, 2025

Uh oh!

gemini-code-assist bot Jul 25, 2025

Uh oh!

gemini-code-assist bot Jul 25, 2025

Uh oh!

Uh oh!

Uh oh!

yzh119 Aug 13, 2025

Uh oh!

yzh119 Aug 13, 2025

Uh oh!

yzh119 left a comment

Uh oh!

Uh oh!

Uh oh!

refactor: Improved metainfo for trtllm-gen kernels #1328

refactor: Improved metainfo for trtllm-gen kernels #1328

Uh oh!

Conversation

cyx-6 commented Jul 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

📌 Description

🔍 Related Issues

🚀 Pull Request Checklist

✅ Pre-commit Checks

🧪 Tests

Reviewer Notes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Jul 25, 2025

Choose a reason for hiding this comment

Uh oh!

joker-eph Jul 25, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 25, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 25, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

yzh119 Aug 13, 2025

Choose a reason for hiding this comment

Uh oh!

yzh119 Aug 13, 2025

Choose a reason for hiding this comment

Uh oh!

yzh119 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

cyx-6 commented Jul 25, 2025 •

edited

Loading