refactor

Author: yzh119

csrc/flat/prefill/prefill_kernel_delta_rule_sm90.cu

@@ -15,11 +15,11 @@
  */
 #include <cuda_bf16.h>
 
-#include "prefill_kernel_delta_rule_sm90.cuh"
+#include "flashinfer/flat/prefill/prefill_kernel_delta_rule_sm90.cuh"
 
 // Extern template declarations prevent implicit instantiation here.
 // Explicit instantiations are in separate generated files for parallel compilation.
-#include "prefill_kernel_delta_rule_sm90_extern.inc"
+#include "flat_prefill_kernel_delta_rule_sm90_extern.inc"
 
 namespace flat {
 

…refill_kernel_delta_rule_sm90_extern.inc …refill_kernel_delta_rule_sm90_extern.inccsrc/flat/prefill/prefill_kernel_delta_rule_sm90_extern.inc renamed to csrc/flat_prefill_kernel_delta_rule_sm90_extern.inc csrc/flat/prefill/prefill_kernel_delta_rule_sm90_extern.inc renamed to csrc/flat_prefill_kernel_delta_rule_sm90_extern.inc

@@ -25,7 +25,7 @@
 #include <iostream>
 #include <sstream>
 
-#include "flat/prefill/prefill_kernel.hpp"
+#include "flashinfer/flat/prefill/prefill_kernel.hpp"
 
 using tvm::ffi::Optional;
 using tvm::ffi::TensorView;

csrc/gdn_prefill_launcher.cu

@@ -23,7 +23,7 @@
 
 // Include the header which defines the function template
 // The header includes all necessary CUTLASS type definitions
-#include "prefill_kernel_delta_rule_sm90.cuh"
+#include "flashinfer/flat/prefill/prefill_kernel_delta_rule_sm90.cuh"
 
 namespace flat {
 

csrc/gdn_prefill_sm90_kernel_inst.jinja

@@ -65,7 +65,8 @@ def gen_gdn_prefill_sm90_module() -> JitSpec:
             )
             write_if_different(dest_path, source)
 
-    # Copy source files to gen_directory (like POD module does)
+    # Copy source files to gen_directory for compilation
+    # Headers are now in include/flashinfer/flat/ and accessible via standard include paths
     for filename in [
         "gdn_prefill_launcher.cu",
         "flat/prefill/prefill_kernel_delta_rule_sm90.cu",
@@ -75,17 +76,8 @@ def gen_gdn_prefill_sm90_module() -> JitSpec:
         source_paths.append(dest_path)
         write_if_different(dest_path, src_path.read_text())
 
-    # Copy header files so relative includes work
-    for filename in [
-        "flat/prefill/prefill_kernel_delta_rule_sm90.cuh",
-        "flat/prefill/prefill_kernel_delta_rule_sm90_extern.inc",
-    ]:
-        src_path = jit_env.FLASHINFER_CSRC_DIR / filename
-        write_if_different(gen_directory / src_path.name, src_path.read_text())
-
     return gen_jit_spec(
         uri,
         source_paths,
         extra_cuda_cflags=sm90a_nvcc_flags + ["-DFLAT_SM90A_ENABLED", "-std=c++20"],
-        extra_include_paths=[gen_directory, jit_env.FLASHINFER_CSRC_DIR],
     )

flashinfer/jit/gdn.py

@@ -18,7 +18,7 @@
 #include "cute/tensor.hpp"
 #include "cutlass/arch/barrier.h"
 #include "cutlass/cutlass.h"
-#include "flat/cute_ext.hpp"
+#include "flashinfer/flat/cute_ext.hpp"
 
 namespace flat::collective {
 

…e/collective/flat_collective_inverse.hpp …e/collective/flat_collective_inverse.hppcsrc/flat/ampere/collective/flat_collective_inverse.hpp renamed to include/flashinfer/flat/ampere/collective/flat_collective_inverse.hpp csrc/flat/ampere/collective/flat_collective_inverse.hpp renamed to include/flashinfer/flat/ampere/collective/flat_collective_inverse.hpp

@@ -18,7 +18,7 @@
 #include "cute/tensor.hpp"
 #include "cutlass/cutlass.h"
 #include "cutlass/pipeline/sm90_pipeline.hpp"
-#include "flat/unused.hpp"
+#include "flashinfer/flat/unused.hpp"
 
 namespace flat::collective {
 

…pere/collective/flat_collective_load.hpp …pere/collective/flat_collective_load.hppcsrc/flat/ampere/collective/flat_collective_load.hpp renamed to include/flashinfer/flat/ampere/collective/flat_collective_load.hpp csrc/flat/ampere/collective/flat_collective_load.hpp renamed to include/flashinfer/flat/ampere/collective/flat_collective_load.hpp

@@ -19,7 +19,7 @@
 #include <stdexcept>
 #include <string>
 
-#include "debug.hpp"
+#include "flashinfer/flat/debug.hpp"
 
 #define FLAT_UNUSED_PARAMETER(x) (void)x