bugfix

Author: yzh119

csrc/flat/prefill/prefill_kernel_delta_rule_sm90_extern.inc

@@ -18,6 +18,10 @@
 
 #pragma once
 
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include "cutlass/arch/arch.h"
+
 namespace flat {
 
 // clang-format off

flashinfer/aot.py

@@ -443,6 +443,9 @@ def gen_all_modules(
     add_misc: bool,
     add_xqa: bool,
 ) -> List[JitSpec]:
+    # TEMPORARY: Only compile gdn_prefill_sm90 for testing
+    return [gen_gdn_prefill_sm90_module()]
+
     jit_specs: List[JitSpec] = []
     jit_specs.append(gen_spdlog_module())
     has_sm90 = sm_capabilities.get("sm90", False)