refactor: download trtllm gemm metadata from server (#1378)

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

## 📌 Description

refactor: download trtllm gemm metadata from server

## 🔍 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

- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] 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.
-->

Author: ttyio

flashinfer/artifacts.py

@@ -53,7 +53,7 @@ class ArtifactPath:
         "991e7438224199de85ef08a2730ce18c12b4e0aa/batched_gemm-c603ed2-2dc78d9/"
     )
     TRTLLM_GEN_GEMM: str = (
-        "fffd607babb0844f24225997409747ca38229333/gemm-c603ed2-f2b0c24/"
+        "07a5f242a649533ff6885f87c42b2476a9e46233/gemm-c603ed2-434a6e1/"
     )
     CUDNN_SDPA: str = "4c623163877c8fef5751c9c7a59940cd2baae02e/fmha/cudnn/"
     DEEPGEMM: str = "d25901733420c7cddc1adf799b0d4639ed1e162f/deep-gemm/"
@@ -64,12 +64,18 @@ class MetaInfoHash:
         "8c5630020c0452fb1cd1ea7e3b8fdbb7bf94f71bd899ed5b704a490bdb4f7368"
     )
     DEEPGEMM: str = "69aa277b7f3663ed929e73f9c57301792b8c594dac15a465b44a5d151b6a1d50"
+    TRTLLM_GEN_GEMM: str = (
+        "50c5627324003c822efbdd1d368b1e569f4f67f4bb0a2fbb7397cd56c6d14c2a"
+    )
 
 
 def download_artifacts() -> bool:
     env_backup = os.environ.get("FLASHINFER_CUBIN_CHECKSUM_DISABLED", None)
     os.environ["FLASHINFER_CUBIN_CHECKSUM_DISABLED"] = "1"
-    cubin_files = [(ArtifactPath.TRTLLM_GEN_FMHA + "flashInferMetaInfo", ".h")]
+    cubin_files = [
+        (ArtifactPath.TRTLLM_GEN_FMHA + "flashInferMetaInfo", ".h"),
+        (ArtifactPath.TRTLLM_GEN_GEMM + "KernelMetaInfo", ".h"),
+    ]
     for kernel in [
         ArtifactPath.TRTLLM_GEN_FMHA,
         ArtifactPath.TRTLLM_GEN_BMM,

flashinfer/gemm.py

@@ -25,7 +25,7 @@
 import torch
 import torch.nn.functional as F
 
-from .artifacts import ArtifactPath
+from .artifacts import ArtifactPath, MetaInfoHash
 from .autotuner import (
     AutoTuner,
     ConstraintSpec,
@@ -38,6 +38,7 @@
     get_last_power_of_2_num_tokens_buckets,
     last_positive_power_of_2,
 )
+from .jit.cubin_loader import get_cubin
 
 CUDNN_AVAILABLE = False
 try:
@@ -302,6 +303,13 @@ def get_gemm_sm100_module():
 
 
 def trtllm_gemm_gen_module() -> JitSpec:
+    header_name = "KernelMetaInfo"
+    metainfo = get_cubin(
+        f"{ArtifactPath.TRTLLM_GEN_GEMM}/{header_name}",
+        MetaInfoHash.TRTLLM_GEN_GEMM,
+        ".h",
+    )
+    assert metainfo, f"{header_name}.h not found"
     return gen_jit_spec(
         "trtllm_gemm",
         [
@@ -313,6 +321,11 @@ def trtllm_gemm_gen_module() -> JitSpec:
             f'-DTLLM_GEN_GEMM_CUBIN_PATH=\\"{ArtifactPath.TRTLLM_GEN_GEMM}\\"',
         ]
         + sm100a_nvcc_flags,
+        extra_include_paths=[
+            jit_env.FLASHINFER_CACHE_DIR / "cubins" / ArtifactPath.TRTLLM_GEN_GEMM,
+            jit_env.FLASHINFER_INCLUDE_DIR
+            / "flashinfer/trtllm/gemm/trtllmGen_gemm_export",
+        ],
         extra_ldflags=["-lcuda"],
     )