From bcda0c9ee6a6e687e53ac933f3541dd5c5a1fe9b Mon Sep 17 00:00:00 2001
From: Yineng Zhang <me@zhyncs.com>
Date: Wed, 22 Jan 2025 20:33:13 +0800
Subject: [PATCH] sync the upstream updates of flashinfer (#3051)

---
 .github/workflows/pr-test-sgl-kernel.yml | 1 +
 sgl-kernel/3rdparty/flashinfer           | 2 +-
 sgl-kernel/setup.py                      | 6 ++++++
 3 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/pr-test-sgl-kernel.yml b/.github/workflows/pr-test-sgl-kernel.yml
index 3d9802658..794a73f36 100644
--- a/.github/workflows/pr-test-sgl-kernel.yml
+++ b/.github/workflows/pr-test-sgl-kernel.yml
@@ -40,6 +40,7 @@ jobs:
         run: |
           pip3 install torch==2.5.1 && pip3 install pytest && pip3 install vllm
           pip3 uninstall sgl-kernel -y || true
+          find . -name index.lock -delete
           cd sgl-kernel
           git submodule deinit --all --force && git submodule sync --recursive && git submodule update --init --force --recursive
           pip3 install .
diff --git a/sgl-kernel/3rdparty/flashinfer b/sgl-kernel/3rdparty/flashinfer
index a0e99a3a8..4e8eb1879 160000
--- a/sgl-kernel/3rdparty/flashinfer
+++ b/sgl-kernel/3rdparty/flashinfer
@@ -1 +1 @@
-Subproject commit a0e99a3a820109763d9a757138a5cdf7bbcd1f85
+Subproject commit 4e8eb1879f9c3ba6d75511e5893183bf8f289a62
diff --git a/sgl-kernel/setup.py b/sgl-kernel/setup.py
index 1197611d6..b9324c355 100644
--- a/sgl-kernel/setup.py
+++ b/sgl-kernel/setup.py
@@ -47,6 +47,7 @@ include_dirs = [
     cutlass.resolve() / "tools" / "util" / "include",
     root / "src" / "sgl-kernel" / "csrc",
     flashinfer.resolve() / "include",
+    flashinfer.resolve() / "include" / "gemm",
     flashinfer.resolve() / "csrc",
 ]
 nvcc_flags = [
@@ -91,7 +92,12 @@ ext_modules = [
             "src/sgl-kernel/csrc/sampling_scaling_penalties.cu",
             "src/sgl-kernel/csrc/sgl_kernel_ops.cu",
             "src/sgl-kernel/csrc/rotary_embedding.cu",
+            "3rdparty/flashinfer/csrc/activation.cu",
+            "3rdparty/flashinfer/csrc/bmm_fp8.cu",
+            "3rdparty/flashinfer/csrc/group_gemm.cu",
+            "3rdparty/flashinfer/csrc/group_gemm_sm90.cu",
             "3rdparty/flashinfer/csrc/norm.cu",
+            "3rdparty/flashinfer/csrc/sampling.cu",
         ],
         include_dirs=include_dirs,
         extra_compile_args={