Refactor device file structures (#182)

* Refactor SM100 files

* Make SM90 work

* Minor fix

* Lint

* Minor fix

---------

Co-authored-by: Zhean Xu <xza@deepseek.com>

Author: LyricZhao

csrc/apis/attention.hpp

@@ -5,7 +5,7 @@
 #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE
 #include "../jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp"
 #include "../jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp"
-#include "../jit_kernels/impls/sm100_fp8_gemm_1d1d.hpp"
+#include "../jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp"
 #include "../jit_kernels/impls/smxx_fp8_mqa_logits.hpp"
 #include "../jit_kernels/impls/smxx_fp8_paged_mqa_logits.hpp"
 #include "../jit_kernels/impls/smxx_clean_logits.hpp"
@@ -64,7 +64,7 @@ static void fp8_gemm_nt_skip_head_mid(const std::pair<torch::Tensor, torch::Tens
 
     // Dispatch into different implements
     const auto arch_major = device_runtime->get_arch_major();
-    const auto epilogue_type = fmt::format("EpilogueHeadSplits<{}, {}, {}>", left, mid, right);
+    const auto epilogue_type = fmt::format("epilogue::transform::EpilogueHeadSplits<{}, {}, {}>", left, mid, right);
     if (arch_major == 9 and sfa.scalar_type() == torch::kFloat and std::get<1>(recipe.value()) != 1) {
         const auto major_sfb = get_major_type_ab(sfb);
         sm90_fp8_gemm_1d2d(a.first, sfa, b.first, sfb, std::nullopt, d, m, n, k, major_a, major_b, major_sfb, compiled_dims, epilogue_type);

csrc/apis/gemm.hpp

@@ -6,7 +6,7 @@
 #include "../jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp"
 #include "../jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp"
 #include "../jit_kernels/impls/sm90_bf16_gemm.hpp"
-#include "../jit_kernels/impls/sm100_fp8_gemm_1d1d.hpp"
+#include "../jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp"
 #include "../jit_kernels/impls/sm100_bf16_gemm.hpp"
 #endif 
 

csrc/indexing/main.cu

@@ -3,12 +3,14 @@
 #include <deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh>
 #include <deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh>
 #include <deep_gemm/impls/sm100_bf16_gemm.cuh>
-#include <deep_gemm/impls/sm100_fp8_gemm_1d1d.cuh>
+#include <deep_gemm/impls/sm100_fp8_fp4_gemm_1d1d.cuh>
 
 // Attention kernels
 #include <deep_gemm/impls/sm90_fp8_mqa_logits.cuh>
 #include <deep_gemm/impls/sm90_fp8_paged_mqa_logits.cuh>
+#include <deep_gemm/impls/sm100_fp4_mqa_logits.cuh>
 #include <deep_gemm/impls/sm100_fp8_mqa_logits.cuh>
+#include <deep_gemm/impls/sm100_fp4_paged_mqa_logits.cuh>
 #include <deep_gemm/impls/sm100_fp8_paged_mqa_logits.cuh>
 /* oss-ignore-begin */
 #include <deep_gemm/impls/sm100_sparse_mqa_logits_bwd.cuh>

csrc/jit_kernels/heuristics/common.hpp

@@ -1,7 +1,7 @@
 #pragma once
 
 #include <unordered_set>
-#include <deep_gemm/common/types.hpp>
+#include <deep_gemm/common/types.cuh>
 
 #include "config.hpp"
 #include "runtime.hpp"

csrc/jit_kernels/heuristics/config.hpp

@@ -2,7 +2,7 @@
 
 #include <cute/arch/mma_sm100_desc.hpp>
 #include <c10/core/ScalarType.h>
-#include <deep_gemm/common/types.hpp>
+#include <deep_gemm/common/types.cuh>
 
 #include "../../utils/math.hpp"
 

csrc/jit_kernels/heuristics/sm100.hpp

@@ -2,7 +2,7 @@
 
 #include <cute/arch/mma_sm100_desc.hpp>
 // Reuse some types in the JIT modules
-#include <deep_gemm/common/types.hpp>
+#include <deep_gemm/common/types.cuh>
 
 #include "common.hpp"
 #include "runtime.hpp"

csrc/jit_kernels/heuristics/sm90.hpp

@@ -2,7 +2,7 @@
 
 #include <cute/arch/mma_sm100_desc.hpp>
 // Reuse some types in the JIT modules
-#include <deep_gemm/common/types.hpp>
+#include <deep_gemm/common/types.cuh>
 
 #include "common.hpp"
 #include "utils.hpp"

csrc/jit_kernels/heuristics/utils.hpp

@@ -2,7 +2,7 @@
 
 #include <cute/arch/mma_sm100_desc.hpp>
 // Reuse some types in the JIT modules
-#include <deep_gemm/common/types.hpp>
+#include <deep_gemm/common/types.cuh>
 
 #include "common.hpp"
 #include "../../utils/exception.hpp"

csrc/jit_kernels/impls/epilogue.hpp

@@ -6,7 +6,7 @@
 namespace deep_gemm {
 
 static std::string get_default_epilogue_type(const std::optional<std::string>& epilogue_type) {
-    return epilogue_type.value_or("EpilogueIdentity");
+    return epilogue_type.value_or("epilogue::transform::EpilogueIdentity");
 }
 
 } // namespace deep_gemm

…it_kernels/impls/sm100_fp8_gemm_1d1d.hpp …ernels/impls/sm100_fp8_fp4_gemm_1d1d.hppcsrc/jit_kernels/impls/sm100_fp8_gemm_1d1d.hpp renamed to csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp csrc/jit_kernels/impls/sm100_fp8_gemm_1d1d.hpp renamed to csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp

@@ -38,12 +38,12 @@ class SM100FP8FP4Gemm1D1DRuntime final: public LaunchRuntime<SM100FP8FP4Gemm1D1D
     static std::string generate_impl(const Args& args) {
         // TODO: rename files
         return fmt::format(R"(
-#include <deep_gemm/impls/sm100_fp8_gemm_1d1d.cuh>
+#include <deep_gemm/impls/sm100_fp8_fp4_gemm_1d1d.cuh>
 
 using namespace deep_gemm;
 
 static void __instantiate_kernel() {{
-    auto ptr = reinterpret_cast<void*>(&sm100_fp8_gemm_1d1d_impl<
+    auto ptr = reinterpret_cast<void*>(&sm100_fp8_fp4_gemm_1d1d_impl<
         {}, {},
         {}, {},
         {}, {}, {},