fix

Author: yzh119

csrc/flat_prefill_kernel_delta_rule_sm90_extern.inc

@@ -47,7 +47,7 @@ namespace flat {
 #define DECLARE_TEMPLATE_INSTANCE(is_gva, needs_beta, needs_alpha, init_state, ctype) \
 extern template void launch_delta_rule_prefill_kernel_gbai<is_gva, needs_beta, needs_alpha, init_state, cutlass::arch::Sm90, ctype, ctype, float>( \
     cudaStream_t, ctype*, float*, ctype const*, ctype const*, ctype const*, \
-    float const*, float const*, float const*, int64_t const*, int32_t, int32_t, \
+    float const*, float const*, float const*, int64_t const*, uint8_t*, int32_t, int32_t, \
     int32_t, int32_t, int32_t, int32_t, int64_t, float, int32_t);
 
 // Extern template declarations for half

csrc/gdn_prefill_sm90_kernel_inst.jinja

@@ -31,7 +31,7 @@ namespace flat {
 // Parameter types must exactly match the extern template declaration in prefill_kernel_delta_rule_sm90_extern.inc
 template void launch_delta_rule_prefill_kernel_gbai<{{ is_gva }}, {{ needs_beta }}, {{ needs_alpha }}, {{ init_state }}, cutlass::arch::Sm90, {{ dtype }}, {{ dtype }}, float>(
     cudaStream_t, {{ dtype }}*, float*, {{ dtype }} const*, {{ dtype }} const*, {{ dtype }} const*,
-    float const*, float const*, float const*, int64_t const*, int32_t, int32_t,
+    float const*, float const*, float const*, int64_t const*, uint8_t*, int32_t, int32_t,
     int32_t, int32_t, int32_t, int32_t, int64_t, float, int32_t);
 
 }  // namespace flat

flashinfer/aot.py

@@ -443,9 +443,6 @@ 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)