diff --git a/csrc/fp4_bindings.cpp b/csrc/fp4_bindings.cpp index 6a0194f7..d4207c5d 100644 --- a/csrc/fp4_bindings.cpp +++ b/csrc/fp4_bindings.cpp @@ -950,6 +950,86 @@ The downstream GEMM keeps its original K = N_il/2 weight. R"pbdoc( Skinny-M half-width fused GeGLU GEMM on the decoder tile (128x64x256); same contract as cutlass_fp4_gemm_geglu_il_hw. +)pbdoc"); + + m.def("cutlass_fp4_gemm_geglu_il_hw_nod", + [](uintptr_t A_packed, uintptr_t SFA, + uintptr_t B_packed, uintptr_t SFB, + uintptr_t D_dummy, uintptr_t compact_packed, uintptr_t compact_sfa, + int M, int N_il, int K, uintptr_t stream) -> int { + const auto shape = fp4_kernel_shape({{"M", M}, {"N_il", N_il}, {"K", K}}); + require_fp4_ptrs("cutlass_fp4_gemm_geglu_il_hw_nod", + {{"A_packed", A_packed}, {"SFA", SFA}, + {"B_packed", B_packed}, {"SFB", SFB}, + {"D_dummy", D_dummy}, + {"compact_packed", compact_packed}, + {"compact_sfa", compact_sfa}}, shape); + require_fp4(M > 0 && N_il > 0 && K > 0 && (N_il % 32) == 0 && + (K % 16) == 0, + "cutlass_fp4_gemm_geglu_il_hw_nod", + "M must be positive, N_il a positive multiple of 32 " + "and K a positive multiple of 16", + shape); + return flash_rt::fp4::cutlass_fp4_gemm_geglu_il_hw_nod( + reinterpret_cast(A_packed), + reinterpret_cast(SFA), + reinterpret_cast(B_packed), + reinterpret_cast(SFB), + reinterpret_cast(D_dummy), + reinterpret_cast(compact_packed), + reinterpret_cast(compact_sfa), + M, N_il, K, + reinterpret_cast(stream)); + }, + py::arg("A_packed"), py::arg("SFA"), + py::arg("B_packed"), py::arg("SFB"), + py::arg("D_dummy"), py::arg("compact_packed"), py::arg("compact_sfa"), + py::arg("M"), py::arg("N_il"), py::arg("K"), + py::arg("stream") = 0, + R"pbdoc( +Half-width fused GeGLU GEMM with the collective's own D store elided: +compact_packed/compact_sfa are the only outputs and D_dummy is never +written (still validated; the host-side TMA descriptor needs a real +pointer). Same contract as cutlass_fp4_gemm_geglu_il_hw otherwise. +)pbdoc"); + + m.def("cutlass_fp4_gemm_geglu_il_hw_nod_v10", + [](uintptr_t A_packed, uintptr_t SFA, + uintptr_t B_packed, uintptr_t SFB, + uintptr_t D_dummy, uintptr_t compact_packed, uintptr_t compact_sfa, + int M, int N_il, int K, uintptr_t stream) -> int { + const auto shape = fp4_kernel_shape({{"M", M}, {"N_il", N_il}, {"K", K}}); + require_fp4_ptrs("cutlass_fp4_gemm_geglu_il_hw_nod_v10", + {{"A_packed", A_packed}, {"SFA", SFA}, + {"B_packed", B_packed}, {"SFB", SFB}, + {"D_dummy", D_dummy}, + {"compact_packed", compact_packed}, + {"compact_sfa", compact_sfa}}, shape); + require_fp4(M > 0 && N_il > 0 && K > 0 && (N_il % 32) == 0 && + (K % 16) == 0, + "cutlass_fp4_gemm_geglu_il_hw_nod_v10", + "M must be positive, N_il a positive multiple of 32 " + "and K a positive multiple of 16", + shape); + return flash_rt::fp4::cutlass_fp4_gemm_geglu_il_hw_nod_v10( + reinterpret_cast(A_packed), + reinterpret_cast(SFA), + reinterpret_cast(B_packed), + reinterpret_cast(SFB), + reinterpret_cast(D_dummy), + reinterpret_cast(compact_packed), + reinterpret_cast(compact_sfa), + M, N_il, K, + reinterpret_cast(stream)); + }, + py::arg("A_packed"), py::arg("SFA"), + py::arg("B_packed"), py::arg("SFB"), + py::arg("D_dummy"), py::arg("compact_packed"), py::arg("compact_sfa"), + py::arg("M"), py::arg("N_il"), py::arg("K"), + py::arg("stream") = 0, + R"pbdoc( +Skinny-M no-D-store fused GeGLU GEMM on the decoder tile (128x64x256); +same contract as cutlass_fp4_gemm_geglu_il_hw_nod. )pbdoc"); #ifdef FLASHRT_HAVE_COSMOS3_EDGE diff --git a/csrc/gemm/fp4/cutlass_fp4_gemm_bias_f32b_f16out_sm100.cu b/csrc/gemm/fp4/cutlass_fp4_gemm_bias_f32b_f16out_sm100.cu new file mode 100644 index 00000000..74b20a1c --- /dev/null +++ b/csrc/gemm/fp4/cutlass_fp4_gemm_bias_f32b_f16out_sm100.cu @@ -0,0 +1,153 @@ +// ============================================================================ +// FlashRT — NVFP4 GEMM with fp32 per-column bias and fp16 output +// (SM100/SM110). See header for the contract. +// ============================================================================ + +#include "gemm/fp4/cutlass_fp4_gemm_bias_f32b_f16out_sm100.cuh" + +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/fusion/operations.hpp" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/util/packed_stride.hpp" +#include "cutlass/detail/sm100_blockscaled_layout.hpp" +#include "cute/tensor.hpp" + +#include +#include + +namespace flash_rt { +namespace fp4 { + +namespace bias_f16out { + +using namespace cute; + +using ElementA = cutlass::nv_float4_t; +using LayoutATag = cutlass::layout::RowMajor; +constexpr int AlignmentA = 32; + +using ElementB = cutlass::nv_float4_t; +using LayoutBTag = cutlass::layout::ColumnMajor; +constexpr int AlignmentB = 32; + +using ElementAccumulator = float; +using ElementCompute = float; +using ArchTag = cutlass::arch::Sm100; +using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; + +using ElementD = cutlass::half_t; +using ElementC = cutlass::half_t; +constexpr int AlignmentCD = 8; + +using MmaTileShape = Shape<_128, _128, _256>; +using ClusterShape = Shape<_1, _1, _1>; + +// per-shape CUTLASS workspace cache (capture-safe: growth happens during +// the uncaptured warmup evaluation) +struct ws_key { + int M, N, K; + bool operator==(const ws_key & o) const { return M == o.M && N == o.N && K == o.K; } +}; +struct ws_key_hash { + size_t operator()(const ws_key & k) const noexcept { + return (size_t) k.M * 1315423911u ^ (size_t) k.N * 2654435761u ^ (size_t) k.K; + } +}; +inline void * get_ws(int M, int N, int K, size_t needed) { + static std::unordered_map, ws_key_hash> cache; + static std::mutex mu; + std::lock_guard lk(mu); + auto & e = cache[ws_key{M, N, K}]; + if (e.second < needed) { + if (e.first) { cudaFree(e.first); } + cudaMalloc(&e.first, needed); + e.second = needed; + } + return e.first; +} + +using FusionOperation = cutlass::epilogue::fusion::LinCombPerColBias< + ElementD, ElementCompute, float, ElementC, ElementCompute>; + +using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementAccumulator, + ElementC, cutlass::layout::RowMajor, AlignmentCD, + ElementD, cutlass::layout::RowMajor, AlignmentCD, + cutlass::epilogue::collective::EpilogueScheduleAuto, + FusionOperation>::CollectiveOp; + +using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, + ElementA, LayoutATag, AlignmentA, + ElementB, LayoutBTag, AlignmentB, + ElementAccumulator, MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp; + +using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, CollectiveMainloop, CollectiveEpilogue, void>; +using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + +} // namespace bias_f16out + +int gemm_bias_f16out( + const void * A_packed, const void * SFA, + const void * B_packed, const void * SFB, + const void * bias_f32, void * D_f16, + int M, int N, int K, + cudaStream_t stream) { + using namespace bias_f16out; + + auto stride_A = cutlass::make_cute_packed_stride( + typename Gemm::GemmKernel::StrideA{}, {M, K, 1}); + auto stride_B = cutlass::make_cute_packed_stride( + typename Gemm::GemmKernel::StrideB{}, {N, K, 1}); + auto stride_C = cutlass::make_cute_packed_stride( + typename Gemm::GemmKernel::StrideC{}, {M, N, 1}); + auto stride_D = cutlass::make_cute_packed_stride( + typename Gemm::GemmKernel::StrideD{}, {M, N, 1}); + using Cfg = + typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + auto layout_SFA = Cfg::tile_atom_to_shape_SFA(make_shape(M, N, K, 1)); + auto layout_SFB = Cfg::tile_atom_to_shape_SFB(make_shape(M, N, K, 1)); + + using EA = typename ElementA::DataType; + using SA = typename ElementA::ScaleFactorType; + + typename Gemm::Arguments args{ + cutlass::gemm::GemmUniversalMode::kGemm, {M, N, K, 1}, + {reinterpret_cast(A_packed), stride_A, + reinterpret_cast(B_packed), stride_B, + reinterpret_cast(SFA), layout_SFA, + reinterpret_cast(SFB), layout_SFB}, + {{}, + reinterpret_cast(D_f16), stride_C, + reinterpret_cast(D_f16), stride_D}}; + args.epilogue.thread.alpha = 1.0f; + args.epilogue.thread.beta = 0.0f; + args.epilogue.thread.bias_ptr = reinterpret_cast(bias_f32); + + Gemm gemm; + auto st = gemm.can_implement(args); + if (st != cutlass::Status::kSuccess) return static_cast(st) | 0x10000; + size_t ws_sz = Gemm::get_workspace_size(args); + void* ws = ws_sz > 0 ? get_ws(M, N, K, ws_sz) : nullptr; + st = gemm.initialize(args, ws, stream); + if (st != cutlass::Status::kSuccess) return static_cast(st) | 0x20000; + st = gemm.run(stream); + return (st == cutlass::Status::kSuccess) ? 0 + : (static_cast(st) | 0x30000); +} + +} // namespace fp4 +} // namespace flash_rt diff --git a/csrc/gemm/fp4/cutlass_fp4_gemm_bias_f32b_f16out_sm100.cuh b/csrc/gemm/fp4/cutlass_fp4_gemm_bias_f32b_f16out_sm100.cuh new file mode 100644 index 00000000..fd1cab03 --- /dev/null +++ b/csrc/gemm/fp4/cutlass_fp4_gemm_bias_f32b_f16out_sm100.cuh @@ -0,0 +1,30 @@ +// ============================================================================ +// FlashRT — NVFP4 GEMM with fp32 per-column bias and fp16 output +// (SM100/SM110). +// +// D_f16[M, N] = A @ B^T + bias[N]. For hosts that keep biases in fp32 but +// consume the projection in fp16 (e.g. attention inputs cast for flash +// attention): the fp16 conversion happens once in the epilogue from the +// fp32 accumulator, matching an fp32-output GEMM followed by an fp16 cast +// bit for bit. +// ============================================================================ +#pragma once + +#include + +namespace flash_rt { +namespace fp4 { + +// A: [M, K] NVFP4 packed row-major + SFA (tile-interleaved). +// B: [N, K] NVFP4 packed column-major + SFB. +// bias_f32: [N] fp32, broadcast over rows. D_f16: [M, N] fp16 row-major. +// Returns 0 on success; CUTLASS status | stage flag otherwise. +int gemm_bias_f16out( + const void * A_packed, const void * SFA, + const void * B_packed, const void * SFB, + const void * bias_f32, void * D_f16, + int M, int N, int K, + cudaStream_t stream); + +} // namespace fp4 +} // namespace flash_rt diff --git a/csrc/gemm/fp4/cutlass_fp4_gemm_geglu_il_sm100.cu b/csrc/gemm/fp4/cutlass_fp4_gemm_geglu_il_sm100.cu index 16f9a332..76ea721e 100644 --- a/csrc/gemm/fp4/cutlass_fp4_gemm_geglu_il_sm100.cu +++ b/csrc/gemm/fp4/cutlass_fp4_gemm_geglu_il_sm100.cu @@ -24,6 +24,7 @@ #include "cute/tensor.hpp" #include "gemm/fp4/sm100_gelu_mul_blockscale_visitor.hpp" +#include "gemm/fp4/sm100_epilogue_nod.hpp" namespace flash_rt { namespace fp4 { @@ -173,6 +174,48 @@ using GemmKernelHwV10 = cutlass::gemm::kernel::GemmUniversal< using GemmHwV10 = cutlass::gemm::device::GemmUniversalAdapter; +// ── No-D-store instantiations ── +// The builder's collective still stages the unread D tile through smem and +// TMA-stores it (row-aliased, but the smem->L2 traffic and store instructions +// remain). Rebind the built epilogue onto CollectiveEpilogueNoD +// (sm100_epilogue_nod.hpp), which elides that store entirely; the compact +// store node is the only writer. SharedStorage is unchanged, so the +// mainloop carveout from the builder output stays valid. +template +struct MakeNoD; + +template +struct MakeNoD, + Rest...>> { + using type = cutlass::epilogue::collective::CollectiveEpilogueNoD< + StagesC, StagesD, FragmentSize, ReuseSmemC, DelayTmaStore, Rest...>; +}; + +using CollectiveEpilogueHwNoD = typename MakeNoD::type; +static_assert(sizeof(typename CollectiveEpilogueHwNoD::SharedStorage) == + sizeof(typename CollectiveEpilogueHw::SharedStorage), + "NoD epilogue must keep the builder's smem footprint"); + +using GemmKernelHwNoD = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloopHw, CollectiveEpilogueHwNoD, void>; + +using GemmHwNoD = cutlass::gemm::device::GemmUniversalAdapter; + +using CollectiveEpilogueHwNoDV10 = typename MakeNoD::type; +static_assert(sizeof(typename CollectiveEpilogueHwNoDV10::SharedStorage) == + sizeof(typename CollectiveEpilogueHwV10::SharedStorage), + "NoD epilogue must keep the builder's smem footprint"); + +using GemmKernelHwNoDV10 = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloopHwV10, CollectiveEpilogueHwNoDV10, void>; + +using GemmHwNoDV10 = cutlass::gemm::device::GemmUniversalAdapter; + } // namespace geglu_il int cutlass_fp4_gemm_geglu_il( @@ -253,7 +296,12 @@ static int run_geglu_il_hw( auto stride_A = cutlass::make_cute_packed_stride(StrideAT{}, {M, K, 1}); auto stride_B = cutlass::make_cute_packed_stride(StrideBT{}, {N_il, K, 1}); auto stride_C = cutlass::make_cute_packed_stride(StrideCT{}, {M, N_il, 1}); + // The compact store node writes the real output itself; the collective's + // own D store lands in a buffer nothing reads. Aliasing every row onto + // row 0 shrinks that write from M*N_il to a single row, which at encoder + // shape is 12.8 MB of pure waste per call. auto stride_D = cutlass::make_cute_packed_stride(StrideDT{}, {M, N_il, 1}); + cute::get<0>(stride_D) = 0; auto layout_SFA = CfgT::tile_atom_to_shape_SFA(make_shape(M, N_il, K, 1)); auto layout_SFB = CfgT::tile_atom_to_shape_SFB(make_shape(M, N_il, K, 1)); @@ -319,5 +367,31 @@ int cutlass_fp4_gemm_geglu_il_hw_v10( M, N_il, K, stream); } +int cutlass_fp4_gemm_geglu_il_hw_nod( + void const* A_packed, void const* SFA, + void const* B_packed, void const* SFB, + void* D_dummy, + void* compact_packed, + void* compact_sfa, + int M, int N_il, int K, + cudaStream_t stream) { + return geglu_il::run_geglu_il_hw( + A_packed, SFA, B_packed, SFB, D_dummy, compact_packed, compact_sfa, + M, N_il, K, stream); +} + +int cutlass_fp4_gemm_geglu_il_hw_nod_v10( + void const* A_packed, void const* SFA, + void const* B_packed, void const* SFB, + void* D_dummy, + void* compact_packed, + void* compact_sfa, + int M, int N_il, int K, + cudaStream_t stream) { + return geglu_il::run_geglu_il_hw( + A_packed, SFA, B_packed, SFB, D_dummy, compact_packed, compact_sfa, + M, N_il, K, stream); +} + } // namespace fp4 } // namespace flash_rt diff --git a/csrc/gemm/fp4/cutlass_fp4_gemm_geglu_il_sm100.cuh b/csrc/gemm/fp4/cutlass_fp4_gemm_geglu_il_sm100.cuh index 34e1bd4c..590822bc 100644 --- a/csrc/gemm/fp4/cutlass_fp4_gemm_geglu_il_sm100.cuh +++ b/csrc/gemm/fp4/cutlass_fp4_gemm_geglu_il_sm100.cuh @@ -54,5 +54,26 @@ int cutlass_fp4_gemm_geglu_il_hw_v10( int M, int N_il, int K, cudaStream_t stream); +// No-D-store variants: identical contract, but the collective's own D store +// is elided (D_dummy is never written; the pointer is still required for the +// host-side TMA descriptor and may be the same small shared buffer). +int cutlass_fp4_gemm_geglu_il_hw_nod( + void const* A_packed, void const* SFA, + void const* B_packed, void const* SFB, + void* D_dummy, + void* compact_packed, + void* compact_sfa, + int M, int N_il, int K, + cudaStream_t stream); + +int cutlass_fp4_gemm_geglu_il_hw_nod_v10( + void const* A_packed, void const* SFA, + void const* B_packed, void const* SFB, + void* D_dummy, + void* compact_packed, + void* compact_sfa, + int M, int N_il, int K, + cudaStream_t stream); + } // namespace fp4 } // namespace flash_rt diff --git a/csrc/gemm/fp4/cutlass_fp4_gemm_siglip_ffn_f32out_sm100.cu b/csrc/gemm/fp4/cutlass_fp4_gemm_siglip_ffn_f32out_sm100.cu new file mode 100644 index 00000000..323c5606 --- /dev/null +++ b/csrc/gemm/fp4/cutlass_fp4_gemm_siglip_ffn_f32out_sm100.cu @@ -0,0 +1,282 @@ +// ============================================================================ +// FlashRT — NVFP4 GEMM pair for a SigLIP-style vision-tower FFN with fp32 +// bias/residual boundaries (SM100/SM110). See header for the contract. +// +// Up: D_fp4[M, N] = blockscale( gelu_tanh(A @ B^T + bias[N]) ) +// Down: D_f32[M, N] = A @ B^T + bias[N] + beta * C_f32[M, N] +// ============================================================================ + +#include "gemm/fp4/cutlass_fp4_gemm_siglip_ffn_f32out_sm100.cuh" + +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/thread/activation.h" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/fusion/operations.hpp" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/util/packed_stride.hpp" +#include "cutlass/detail/sm100_blockscaled_layout.hpp" +#include "cute/tensor.hpp" + +#include +#include + +namespace flash_rt { +namespace fp4 { + +namespace siglip_ffn { + +using namespace cute; + +using ElementA = cutlass::nv_float4_t; +using LayoutATag = cutlass::layout::RowMajor; +constexpr int AlignmentA = 32; + +using ElementB = cutlass::nv_float4_t; +using LayoutBTag = cutlass::layout::ColumnMajor; +constexpr int AlignmentB = 32; + +using ElementAccumulator = float; +using ElementCompute = float; +using ArchTag = cutlass::arch::Sm100; +using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; +constexpr int SFVecSize = 16; + +using UpTileShape = Shape<_128, _256, _256>; +using DownTileShape = Shape<_128, _128, _256>; +using ClusterShape = Shape<_1, _1, _1>; + +// per-shape CUTLASS workspace cache (capture-safe: growth happens during +// the uncaptured warmup evaluation) +struct ws_key { + int which, M, N, K; + bool operator==(const ws_key & o) const { return which == o.which && M == o.M && N == o.N && K == o.K; } +}; +struct ws_key_hash { + size_t operator()(const ws_key & k) const noexcept { + return (size_t) k.which * 40503u ^ (size_t) k.M * 1315423911u ^ (size_t) k.N * 2654435761u ^ (size_t) k.K; + } +}; +inline void * get_ws(int which, int M, int N, int K, size_t needed) { + static std::unordered_map, ws_key_hash> cache; + static std::mutex mu; + std::lock_guard lk(mu); + auto & e = cache[ws_key{which, M, N, K}]; + if (e.second < needed) { + if (e.first) { cudaFree(e.first); } + cudaMalloc(&e.first, needed); + e.second = needed; + } + return e.first; +} + +// ── Up: bias + tanh-GELU + fp4/SFA output ────────────────────────────────── +namespace up { + +using ElementD = cutlass::float_e2m1_t; +using ElementC = ElementD; +using ElementSFD = cutlass::float_ue4m3_t; +constexpr int AlignmentD = 32; + +using MmaTileShape = UpTileShape; + +using FusionOperation = + cutlass::epilogue::fusion::LinCombPerColBiasEltActBlockScaleFactor< + cutlass::epilogue::thread::GELU_taylor, SFVecSize, + ElementD, ElementCompute, ElementSFD, cutlass::layout::RowMajor, + float, ElementC, ElementCompute>; + +using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementAccumulator, + ElementC, cutlass::layout::RowMajor, AlignmentD, + ElementD, cutlass::layout::RowMajor, AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto, + FusionOperation>::CollectiveOp; + +using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, + ElementA, LayoutATag, AlignmentA, + ElementB, LayoutBTag, AlignmentB, + ElementAccumulator, MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp; + +using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, CollectiveMainloop, CollectiveEpilogue, void>; +using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + +} // namespace up + +// ── Down: bias + residual source, fp32 output ────────────────────────────── +namespace down { + +using ElementD = float; +using ElementC = float; +constexpr int AlignmentCD = 4; + +using MmaTileShape = DownTileShape; + +using FusionOperation = cutlass::epilogue::fusion::LinCombPerColBias< + ElementD, ElementCompute, float, ElementC, ElementCompute>; + +using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementAccumulator, + ElementC, cutlass::layout::RowMajor, AlignmentCD, + ElementD, cutlass::layout::RowMajor, AlignmentCD, + cutlass::epilogue::collective::EpilogueScheduleAuto, + FusionOperation>::CollectiveOp; + +using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, + ElementA, LayoutATag, AlignmentA, + ElementB, LayoutBTag, AlignmentB, + ElementAccumulator, MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp; + +using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, CollectiveMainloop, CollectiveEpilogue, void>; +using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + +} // namespace down + +} // namespace siglip_ffn + +int siglip_ffn_up_gelu_fp4out( + const void * A_packed, const void * SFA, + const void * B_packed, const void * SFB, + const void * bias_f32, + void * D_packed, void * D_SFD, + int M, int N, int K, + cudaStream_t stream) { + using namespace siglip_ffn; + using Gemm = up::Gemm; + + auto stride_A = cutlass::make_cute_packed_stride( + typename Gemm::GemmKernel::StrideA{}, {M, K, 1}); + auto stride_B = cutlass::make_cute_packed_stride( + typename Gemm::GemmKernel::StrideB{}, {N, K, 1}); + auto stride_C = cutlass::make_cute_packed_stride( + typename Gemm::GemmKernel::StrideC{}, {M, N, 1}); + auto stride_D = cutlass::make_cute_packed_stride( + typename Gemm::GemmKernel::StrideD{}, {M, N, 1}); + using Cfg = + typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + auto layout_SFA = Cfg::tile_atom_to_shape_SFA(make_shape(M, N, K, 1)); + auto layout_SFB = Cfg::tile_atom_to_shape_SFB(make_shape(M, N, K, 1)); + + using EA = typename ElementA::DataType; + using SA = typename ElementA::ScaleFactorType; + + typename Gemm::Arguments args{ + cutlass::gemm::GemmUniversalMode::kGemm, {M, N, K, 1}, + {reinterpret_cast(A_packed), stride_A, + reinterpret_cast(B_packed), stride_B, + reinterpret_cast(SFA), layout_SFA, + reinterpret_cast(SFB), layout_SFB}, + {{}, + reinterpret_cast(D_packed), stride_C, + reinterpret_cast(D_packed), stride_D}}; + args.epilogue.thread.alpha = 1.0f; + args.epilogue.thread.beta = 0.0f; + args.epilogue.thread.bias_ptr = reinterpret_cast(bias_f32); + static float* d_norm = nullptr; + if (!d_norm) { + if (cudaMalloc(&d_norm, sizeof(float)) != cudaSuccess) return -1; + float h = 1.0f; + cudaMemcpyAsync(d_norm, &h, sizeof(float), cudaMemcpyHostToDevice, + stream); + } + args.epilogue.thread.block_scale_factor_ptr = + reinterpret_cast(D_SFD); + args.epilogue.thread.norm_constant_ptr = d_norm; + + Gemm gemm; + auto st = gemm.can_implement(args); + if (st != cutlass::Status::kSuccess) return static_cast(st) | 0x10000; + size_t ws_sz = Gemm::get_workspace_size(args); + void* ws = ws_sz > 0 ? get_ws(0, M, N, K, ws_sz) : nullptr; + st = gemm.initialize(args, ws, stream); + if (st != cutlass::Status::kSuccess) return static_cast(st) | 0x20000; + st = gemm.run(stream); + return (st == cutlass::Status::kSuccess) ? 0 + : (static_cast(st) | 0x30000); +} + +int siglip_ffn_down_bias_res_f32( + const void * A_packed, const void * SFA, + const void * B_packed, const void * SFB, + const void * bias_f32, + const void * C_f32, void * D_f32, + int M, int N, int K, + cudaStream_t stream, float beta) { + using namespace siglip_ffn; + using Gemm = down::Gemm; + + auto stride_A = cutlass::make_cute_packed_stride( + typename Gemm::GemmKernel::StrideA{}, {M, K, 1}); + auto stride_B = cutlass::make_cute_packed_stride( + typename Gemm::GemmKernel::StrideB{}, {N, K, 1}); + auto stride_C = cutlass::make_cute_packed_stride( + typename Gemm::GemmKernel::StrideC{}, {M, N, 1}); + auto stride_D = cutlass::make_cute_packed_stride( + typename Gemm::GemmKernel::StrideD{}, {M, N, 1}); + using Cfg = + typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + auto layout_SFA = Cfg::tile_atom_to_shape_SFA(make_shape(M, N, K, 1)); + auto layout_SFB = Cfg::tile_atom_to_shape_SFB(make_shape(M, N, K, 1)); + + using EA = typename ElementA::DataType; + using SA = typename ElementA::ScaleFactorType; + + typename Gemm::Arguments args{ + cutlass::gemm::GemmUniversalMode::kGemm, {M, N, K, 1}, + {reinterpret_cast(A_packed), stride_A, + reinterpret_cast(B_packed), stride_B, + reinterpret_cast(SFA), layout_SFA, + reinterpret_cast(SFB), layout_SFB}, + {{}, + reinterpret_cast(C_f32), stride_C, + reinterpret_cast(D_f32), stride_D}}; + args.epilogue.thread.alpha = 1.0f; + args.epilogue.thread.beta = beta; + args.epilogue.thread.bias_ptr = reinterpret_cast(bias_f32); + + Gemm gemm; + auto st = gemm.can_implement(args); + if (st != cutlass::Status::kSuccess) return static_cast(st) | 0x10000; + size_t ws_sz = Gemm::get_workspace_size(args); + void* ws = ws_sz > 0 ? get_ws(1, M, N, K, ws_sz) : nullptr; + st = gemm.initialize(args, ws, stream); + if (st != cutlass::Status::kSuccess) return static_cast(st) | 0x20000; + st = gemm.run(stream); + return (st == cutlass::Status::kSuccess) ? 0 + : (static_cast(st) | 0x30000); +} + +int gemm_bias_f32out( + const void * A_packed, const void * SFA, + const void * B_packed, const void * SFB, + const void * bias_f32, void * D_f32, + int M, int N, int K, + cudaStream_t stream) { + // the Down configuration with beta = 0: D = A@B + bias + return siglip_ffn_down_bias_res_f32(A_packed, SFA, B_packed, SFB, bias_f32, + /*C=*/D_f32, D_f32, M, N, K, stream, /*beta=*/0.0f); +} + +} // namespace fp4 +} // namespace flash_rt diff --git a/csrc/gemm/fp4/cutlass_fp4_gemm_siglip_ffn_f32out_sm100.cuh b/csrc/gemm/fp4/cutlass_fp4_gemm_siglip_ffn_f32out_sm100.cuh new file mode 100644 index 00000000..98da67e2 --- /dev/null +++ b/csrc/gemm/fp4/cutlass_fp4_gemm_siglip_ffn_f32out_sm100.cuh @@ -0,0 +1,53 @@ +// ============================================================================ +// FlashRT — NVFP4 GEMM pair for a SigLIP-style vision-tower FFN with fp32 +// bias/residual boundaries (SM100/SM110). +// +// Variant of cutlass_fp4_gemm_siglip_ffn_sm100 for hosts that keep the FFN +// bias and residual tensors in fp32 (rather than fp16): the Up projection +// fuses bias + tanh-GELU and emits FP4 (e2m1) packed output + SFD, and the +// Down projection fuses bias + fp32 residual add with fp32 output. The +// CUTLASS workspace is cached per shape instead of allocated per call, so +// steady-state calls are graph-capture safe. +// +// Up: D_fp4[M, N] = blockscale( gelu_tanh(A @ B^T + bias[N]) ) +// Down: D_f32[M, N] = A @ B^T + bias[N] + beta * C_f32[M, N] +// ============================================================================ +#pragma once + +#include + +namespace flash_rt { +namespace fp4 { + +// A: [M, K] NVFP4 packed row-major + SFA (tile-interleaved). +// B: [N, K] NVFP4 packed column-major + SFB. +// bias_f32: [N] fp32, broadcast over rows. +// D_packed: [M, N] NVFP4 packed row-major; D_SFD: SFD tile-interleaved. +// Returns 0 on success; CUTLASS status | stage flag otherwise. +int siglip_ffn_up_gelu_fp4out( + const void * A_packed, const void * SFA, + const void * B_packed, const void * SFB, + const void * bias_f32, + void * D_packed, void * D_SFD, + int M, int N, int K, + cudaStream_t stream); + +// C_f32/D_f32: [M, N] fp32 row-major (may alias). beta scales C. +int siglip_ffn_down_bias_res_f32( + const void * A_packed, const void * SFA, + const void * B_packed, const void * SFB, + const void * bias_f32, + const void * C_f32, void * D_f32, + int M, int N, int K, + cudaStream_t stream, float beta); + +// Down configuration with beta = 0: D = A@B + bias (no residual read). +int gemm_bias_f32out( + const void * A_packed, const void * SFA, + const void * B_packed, const void * SFB, + const void * bias_f32, void * D_f32, + int M, int N, int K, + cudaStream_t stream); + +} // namespace fp4 +} // namespace flash_rt diff --git a/csrc/gemm/fp4/sm100_epilogue_nod.hpp b/csrc/gemm/fp4/sm100_epilogue_nod.hpp new file mode 100644 index 00000000..46f8d47c --- /dev/null +++ b/csrc/gemm/fp4/sm100_epilogue_nod.hpp @@ -0,0 +1,1304 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Functor performing elementwise operations used by epilogues. + + FlashRT fork of cutlass/epilogue/collective/sm100_epilogue_tma_warpspecialized.hpp + (CUTLASS v4.4.2): CollectiveEpilogueNoD is the Sm100TmaWarpSpecialized + collective with the D store elided. It exists for fusion operations whose + store node writes the real output itself (e.g. the compact GeGLU store), + leaving the collective's own smem->TMA D path writing a buffer nothing + reads. Sm90 and the Sm100 ptr-array epilogues already support this via + is_destination_supported / ElementD=void; this collective never picked + that up, so the same guards are applied here. Semantic changes are the + three `if constexpr (StoreD)` sites; everything else is verbatim. +*/ + + + +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/arch/barrier.h" +#include "cutlass/conv/convnd_problem_shape.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/detail.hpp" +#include "cutlass/epilogue/thread/scale_type.h" +#include "cutlass/epilogue/fusion/callbacks.hpp" +#include "cutlass/epilogue/fusion/sm100_callbacks_tma_warpspecialized.hpp" +#include "cutlass/detail/layout.hpp" +#include "cutlass/detail/helper_macros.hpp" +#include "cutlass/trace.h" + +#include "cutlass/conv/detail.hpp" +#include "cute/tensor.hpp" +#include "cutlass/cuda_host_adapter.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::epilogue::collective { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + int StagesC_, + int StagesD_, + int FragmentSize_, + bool ReuseSmemC_, + bool DelayTmaStore_, + class CtaTileShape_, // (CTA_M,CTA_N,CTA_K, optional: Tile_L) + class EpilogueTile_, // (EPI_TILE_M, EPI_TILE_N) + class ElementC_, + class StrideC_, + class ElementD_, + class StrideD_, + class FusionCallbacks_, + class CopyOpT2R_, + class CopyOpG2S_, + class SmemLayoutAtomC_, + class CopyOpS2R_, + class CopyOpS2G_, + class SmemLayoutAtomD_, + class CopyOpR2S_, + class CopyOpR2R_ +> +class CollectiveEpilogueNoD { +public: + // + // Type Aliases + // + using DispatchPolicy = Sm100TmaWarpSpecialized; + using CtaTileShape = CtaTileShape_; + using EpilogueTile = EpilogueTile_; + using FusionCallbacks = FusionCallbacks_; + using ElementC = ElementC_; + using StrideC = StrideC_; + using ElementD = ElementD_; + using StrideD = StrideD_; + using CopyOpT2R = CopyOpT2R_; + using CopyOpG2S = CopyOpG2S_; + using SmemLayoutAtomC = SmemLayoutAtomC_; + using CopyOpS2R = CopyOpS2R_; + using CopyOpS2G = CopyOpS2G_; + using SmemLayoutAtomD = SmemLayoutAtomD_; + using CopyOpR2S = CopyOpR2S_; + using CopyOpR2R = CopyOpR2R_; + + using ThreadEpilogueOp = typename epilogue::fusion::FusionCallbacksTraits::Operation; + using GmemTiledCopyC = CopyOpG2S; + using GmemTiledCopyD = CopyOpS2G; + + // The whole point of this fork: the fusion op's store node writes the real + // output, so the collective's own D store never issues. ReuseSmemC couples + // the C load pipeline release to D store completion and must stay off. + constexpr static bool StoreD = false; + static_assert(not ReuseSmemC_, "CollectiveEpilogueNoD requires ReuseSmemC == false"); + + constexpr static int ThreadCount = 128; + + static_assert(!is_layout::value && is_tuple::value, "EpilogueTile must be a cute::Tile or cute::Shape"); + static_assert(rank(EpilogueTile{}) == 2, "EpilogueTile must be rank-2: [EPI_TILE_M, EPI_TILE_N]"); + +private: + using GmemElementD = ElementD; + using GmemElementC = cute::conditional_t,ElementD,ElementC>; // prevents void ref breakages + using SmemElementD = typename cutlass::detail::get_unpacked_element_type::type; + using SmemElementC = typename cutlass::detail::get_unpacked_element_type::type; + constexpr static int StagesC = StagesC_; + constexpr static int StagesD = StagesD_; + static_assert(StagesC >= 1, "StagesC must be >= 1"); + static_assert(StagesD >= 1, "StagesD must be >= 1"); + + constexpr static bool ReuseSmemC = ReuseSmemC_; + constexpr static bool is_source_supported = not cute::is_void_v; + + constexpr static bool is_m_major_C = detail::is_m_major(); + constexpr static bool is_m_major_D = detail::is_m_major(); + + constexpr static bool is_im2col_C = cute::is_same_v; + constexpr static bool is_im2col_D = cute::is_same_v; + + using SmemLayoutStageC = decltype(tile_to_shape(SmemLayoutAtomC{}, product_each(shape(EpilogueTile{})), + cute::conditional_t, Step<_1,_2>>{} )); + using SmemLayoutStageD = decltype(tile_to_shape(SmemLayoutAtomD{}, product_each(shape(EpilogueTile{})), + cute::conditional_t, Step<_1,_2>>{} )); + + constexpr static int StageCBits = cosize_v * sizeof_bits_v; + constexpr static int StageDBits = cosize_v * sizeof_bits_v; + constexpr static int MaxStageBits = cute::max(StageCBits, StageDBits); + constexpr static int StrideStageC = (ReuseSmemC ? MaxStageBits : StageCBits) / sizeof_bits_v; + constexpr static int StrideStageD = (ReuseSmemC ? MaxStageBits : StageDBits) / sizeof_bits_v; + + using SmemLayoutC = decltype(cute::append<3>(SmemLayoutStageC{}, Layout, Int>{})); + using SmemLayoutD = decltype(cute::append<3>(SmemLayoutStageD{}, Layout, Int>{})); + + constexpr static bool support_smem_reuse = is_source_supported && StagesD <= StagesC + && MaxStageBits % sizeof_bits_v == 0 + && MaxStageBits % sizeof_bits_v == 0; + static_assert(not (ReuseSmemC && not support_smem_reuse), "Smem reuse requirements not met"); + + constexpr static size_t SmemAlignmentC = cutlass::detail::alignment_for_swizzle(SmemLayoutC{}); + constexpr static size_t SmemAlignmentD = cutlass::detail::alignment_for_swizzle(SmemLayoutD{}); + constexpr static size_t MaxSmemAlignment = cute::max(SmemAlignmentC, SmemAlignmentD); + + // Not unroll epi subtile loop when the activation op is heavy to reduce instruction size and register pressure. + constexpr static bool UnrollEpiLoop = + not cutlass::epilogue::thread::kIsHeavy_member_or_false::value; + // TMA store delay only benefits with loop unrolling + constexpr static bool DelayTmaStore = DelayTmaStore_ and UnrollEpiLoop; + + struct CollectiveStorageWithC { + alignas(SmemAlignmentC) ArrayEngine> smem_C; + alignas(SmemAlignmentD) ArrayEngine> smem_D; + }; + + union CollectiveStorageWithoutC { + cute::array smem_C; + alignas(SmemAlignmentD) ArrayEngine> smem_D; + }; + + union CollectiveStorageReuseC { + alignas(MaxSmemAlignment) ArrayEngine> smem_C; + alignas(MaxSmemAlignment) ArrayEngine> smem_D; + }; + +public: + // TMA pipeline for loading C + using LoadPipeline = cutlass::PipelineTransactionAsync; + using LoadPipelineState = cutlass::PipelineState; + constexpr static uint32_t TmaTransactionBytes = StageCBits / 8; + + // TMA pipeline for storing D + using StorePipeline = cute::conditional_t, + cutlass::PipelineTmaStore>; + using StorePipelineState = cutlass::PipelineState; + + struct SharedStorage { + struct TensorStorage { + using CollectiveStorage = cute::conditional_t>; + CollectiveStorage collective; + + using FusionStorage = typename FusionCallbacks::SharedStorage; + FusionStorage thread; + } tensors; + + using PipelineStorage = typename LoadPipeline::SharedStorage; + PipelineStorage pipeline; + }; + using TensorStorage = typename SharedStorage::TensorStorage; + using PipelineStorage = typename SharedStorage::PipelineStorage; + + // Planar complex kernels have two accumulator copies for the real and imaginary tensors. + constexpr static int NumAccumulatorMtxs = 1; + + // Host side epilogue arguments + struct Arguments { + typename FusionCallbacks::Arguments thread{}; + ElementC const* ptr_C = nullptr; + StrideC dC{}; + ElementD* ptr_D = nullptr; + StrideD dD{}; + }; + +private: + static constexpr auto + get_tma_epi_tile() { + return cute::transform_apply(EpilogueTile{}, seq<0,1>{}, + [] (auto epi_tiler, auto mode) { + auto cta_tiler_shape = get(CtaTileShape{}); + // Use a dynamic stride to prevent mode coalescing + auto cta_tiler_stride = repeat_like(cta_tiler_shape, 0); + auto cta_tiler = make_layout(cta_tiler_shape, cta_tiler_stride); + // This is a multimodal CTA tiler, transform before returning + if constexpr (depth(cta_tiler) > 0) { + // This is an implicit multimodal tiler, match profile and return + if constexpr (tuple_size_v == 1) { + return make_tile(epi_tiler); + } + // This is an explicit multimodal tiler, compose out epi tiler + else { + return shape(composition(cta_tiler, epi_tiler)); + } + } + // This is a flat CTA tiler, no need for transformation + else { + return epi_tiler; + } + }, + [] (auto... epi_tilers) { + return make_tile(epi_tilers...); + } + ); + } + + using TmaEpilogueTile = decltype(get_tma_epi_tile()); + + template + static constexpr auto + get_tma_load_c(ProblemShapeMNL const& problem_shape_mnl, Arguments const& args) { + Tensor tensor_c = make_tensor(make_gmem_ptr(args.ptr_C), + make_layout(problem_shape_mnl, append<3>(args.dC, _0{}))); + return make_tma_copy(CopyOpG2S{}, tensor_c, SmemLayoutStageC{}, TmaEpilogueTile{}, _1{}); + } + + template + static constexpr auto + get_tma_store_d(ProblemShapeMNL const& problem_shape_mnl, Arguments const& args) { + Tensor tensor_d = make_tensor(make_gmem_ptr(args.ptr_D), + make_layout(problem_shape_mnl, append<3>(args.dD, _0{}))); + return make_tma_copy(CopyOpS2G{}, tensor_d, SmemLayoutStageD{}, TmaEpilogueTile{}, _1{}); + } + +public: + // Device side epilogue params + struct Params { + using TMA_C = decltype(get_tma_load_c (repeat_like(append<3>(StrideC{},_1{}), int32_t(0)), Arguments{})); + using TMA_D = decltype(get_tma_store_d(repeat_like(append<3>(StrideD{},_1{}), int32_t(0)), Arguments{})); + + typename FusionCallbacks::Params thread{}; + TMA_C tma_load_c; + TMA_D tma_store_d; + }; + + // + // Gemm Host Functions + // + + template + static constexpr Params + to_underlying_arguments( + ProblemShape const& problem_shape, + Arguments const& args, + [[maybe_unused]] void* workspace) { + // Optionally append 1s until problem shape is rank-4 in case its is only rank-3 (MNK) + auto problem_shape_mnl = select<0,1,3>(append<4>(problem_shape, 1)); + typename Params::TMA_C tma_load_c{}; + if constexpr (is_source_supported) { + tma_load_c = get_tma_load_c(problem_shape_mnl, args); + } + + typename Params::TMA_D tma_store_d = get_tma_store_d(problem_shape_mnl, args); + + return { + FusionCallbacks::to_underlying_arguments(problem_shape, args.thread, workspace), + tma_load_c, + tma_store_d + }; + } + + template + static size_t + get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) { + return FusionCallbacks::get_workspace_size(problem_shape, args.thread); + } + + template + static cutlass::Status + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, + CudaHostAdapter* cuda_adapter = nullptr) { + return FusionCallbacks::initialize_workspace(problem_shape, args.thread, workspace, stream, cuda_adapter); + } + + template + static bool + can_implement( + ProblemShape const& problem_shape, + [[maybe_unused]] Arguments const& args) { + constexpr int tma_alignment_bits_d = cutlass::detail::get_output_alignment_bits(); + auto problem_shape_MNKL = append<4>(problem_shape, 1); + auto [M,N,K,L] = problem_shape_MNKL; + auto shape = cute::make_shape(M,N,L); + + bool implementable = true; + constexpr int min_tma_aligned_elements_D = tma_alignment_bits_d / cutlass::sizeof_bits::value; + if constexpr (cute::is_same_v) { // ignore L stride for implicit gemm + implementable = implementable && cutlass::detail::check_alignment(take<0,2>(shape), take<0,2>(StrideD{})); + } + else { + implementable = implementable && cutlass::detail::check_alignment(shape, StrideD{}); + } + + if constexpr (is_source_supported) { + constexpr int tma_alignment_bits_c = cutlass::detail::get_output_alignment_bits(); + constexpr int min_tma_aligned_elements_C = tma_alignment_bits_c / cutlass::sizeof_bits::value; + if constexpr (cute::is_same_v) { // ignore L stride for implicit gemm + implementable = implementable && cutlass::detail::check_alignment(take<0,2>(shape), take<0,2>(StrideC{})); + } + else { + implementable = implementable && cutlass::detail::check_alignment(shape, StrideC{}); + } + } + + if (!implementable) { + CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n"); + } + + bool fusion_implementable = FusionCallbacks::can_implement(problem_shape, args.thread); + + if (!fusion_implementable) { + CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum requirements for FusionCallbacks.\n"); + } + + return implementable && fusion_implementable; + } + + // + // Conv Host Functions + // + + template + static constexpr Params + to_underlying_arguments(cutlass::conv::ConvProblemShape const& problem_shape, Arguments const& args, void* workspace) { + return to_underlying_arguments(cutlass::conv::detail::get_transformed_problem_shape_MNKL(problem_shape), args, workspace); + } + + template + static size_t + get_workspace_size(cutlass::conv::ConvProblemShape const& problem_shape, Arguments const& args) { + return get_workspace_size(cutlass::conv::detail::get_transformed_problem_shape_MNKL(problem_shape), args); + } + + template + static cutlass::Status + initialize_workspace(cutlass::conv::ConvProblemShape const& problem_shape, Arguments const& args, + void* workspace, cudaStream_t stream, CudaHostAdapter* cuda_adapter = nullptr) { + return initialize_workspace(cutlass::conv::detail::get_transformed_problem_shape_MNKL(problem_shape), args, workspace, stream, cuda_adapter); + } + + template + static bool + can_implement(cutlass::conv::ConvProblemShape const& problem_shape, Arguments const& args) { + return can_implement(cutlass::conv::detail::get_transformed_problem_shape_MNKL(problem_shape), args); + } + + // + // Static Device Functions + // + + template + CUTLASS_DEVICE + static constexpr int + get_load_pipe_increment(CtaTileMNK const& cta_tile_mnk) { + // Compute number of epilogue subtiles + return size<1>(zipped_divide(make_layout(take<0,2>(cta_tile_mnk)), EpilogueTile{})); + } + + template + CUTLASS_DEVICE + static constexpr int + get_store_pipe_increment(CtaTileMNK const& cta_tile_mnk) { + return get_load_pipe_increment(cta_tile_mnk); + } + + /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance + CUTLASS_DEVICE static void + prefetch_tma_descriptors(Params const& epilogue_params) { + cute::prefetch_tma_descriptor(epilogue_params.tma_load_c.get_tma_descriptor()); + cute::prefetch_tma_descriptor(epilogue_params.tma_store_d.get_tma_descriptor()); + } + + // + // Constructor and Data Members + // + CUTLASS_DEVICE + CollectiveEpilogueNoD(Params const& params_, TensorStorage& shared_tensors) + : params(params_), fusion_callbacks(params_.thread, shared_tensors.thread) {} + +private: + Params const& params; + FusionCallbacks fusion_callbacks; + + // + // Non-static Device Functions + // +public: + CUTLASS_DEVICE bool + is_producer_load_needed() const { + return fusion_callbacks.is_producer_load_needed(); + } + + template< + bool ReuseTmem = false, + class ProblemShapeMNKL, + class CtaTileMNK, + class CtaCoordMNKL, + class MmaTileMNK, + class TiledMma + > + CUTLASS_DEVICE auto + load( + LoadPipeline load_pipeline, + LoadPipelineState load_pipe_producer_state, + ProblemShapeMNKL problem_shape_mnkl, + CtaTileMNK cta_tile_mnk, + CtaCoordMNKL cta_coord_mnkl, + MmaTileMNK mma_tile_mnk, + TiledMma tiled_mma, + TensorStorage& shared_tensors, + bool reverse_epi_n = false) { + using namespace cute; + + int lane_idx = canonical_lane_idx(); + auto [M, N, K, L] = problem_shape_mnkl; + auto [m_coord, n_coord, k_coord, l_coord] = cta_coord_mnkl; + + // The tma tensor C under im2col mode only has two modes (M, N) which + // should be local tiled with only (m_coord, n_coord). + auto coord_shape = + conditional_return(make_coord(m_coord, n_coord), make_coord(m_coord, n_coord, l_coord)); + + // Represent the full source tensor, slice to get the tile this CTA is currently responsible for + Tensor mC_mn = params.tma_load_c.get_tma_tensor(make_shape(M,N,L)); // (M,N,L) + Tensor mC = coalesce(mC_mn, take<0,2>(cta_tile_mnk)); + Tensor gC = local_tile(mC, take<0,2>(cta_tile_mnk), coord_shape); // (CTA_M,CTA_N) + + // Apply epilogue subtile, get matching smem tensor + auto ptr_sC = shared_tensors.collective.smem_C.begin(); + Tensor gC_epi = flat_divide(gC, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + Tensor sC_epi = make_tensor(make_smem_ptr(ptr_sC), SmemLayoutC{}); // (EPI_TILE_M,EPI_TILE_N,PIPE_C) + + // Prepare the thread(b)lock's (G)mem to (S)mem TMA tiled copy (bGS_) + ThrCopy thrblk_g2s = params.tma_load_c.get_slice(Int<0>{}); + Tensor bGS_gC = thrblk_g2s.partition_S(gC_epi); // (TMA,TMA_M,TMA_N,EPI_M,EPI_N) + Tensor bGS_sC = thrblk_g2s.partition_D(sC_epi); // (TMA,TMA_M,TMA_N,PIPE_C) + + // Get the fusion callbacks for the producer load warp + auto pld_args = cutlass::epilogue::fusion::detail::ProducerLoadArgs{ + problem_shape_mnkl, + cta_tile_mnk, + cta_coord_mnkl, + tiled_mma, + EpilogueTile{}, + lane_idx + }; + auto pld_callbacks = fusion_callbacks.get_producer_load_callbacks(pld_args); + bool is_C_load_needed = is_source_supported && fusion_callbacks.is_C_load_needed(); + + // Predication for TMA load (one thread issues TMA load) + bool issue_tma_load = cute::elect_one_sync(); + + // Pre-loop fusion callback entry point + pld_callbacks.begin(); + + CUTLASS_PRAGMA_UNROLL + for (int iter_n = 0; iter_n < size<3>(gC_epi); ++iter_n) { + CUTLASS_PRAGMA_UNROLL + for (int iter_m = 0; iter_m < size<2>(gC_epi); ++iter_m) { + int epi_m = iter_m, epi_n = iter_n; + if constexpr (ReuseTmem) { + if (reverse_epi_n) { + epi_n = size<3>(gC_epi) - 1 - iter_n; + } + } + // Acquire the lock for this stage + constexpr uint16_t mcast_mask = 0; + uint64_t* tma_barrier = load_pipeline.producer_get_barrier(load_pipe_producer_state); + load_pipeline.producer_acquire(load_pipe_producer_state); + + // Execute the TMA load for C if needed + if (issue_tma_load && is_C_load_needed) { + copy(params.tma_load_c.with(*tma_barrier, mcast_mask), + bGS_gC(_,_,_,epi_m,epi_n), bGS_sC(_,_,_,load_pipe_producer_state.index())); + load_pipeline.producer_expect_transaction(load_pipe_producer_state); + } + + // Loop fusion callback entry point + pld_callbacks.step(tma_barrier, epi_m, epi_n, load_pipe_producer_state.count(), issue_tma_load); + + // Commit TMA loads for this stage and release the lock + load_pipeline.producer_commit(load_pipe_producer_state); + ++load_pipe_producer_state; + } + } + + // Post-loop fusion callback entry point + pld_callbacks.end(); + + return load_pipe_producer_state; + } + + CUTLASS_DEVICE void + load_tail( + LoadPipeline load_pipeline, + LoadPipelineState load_pipe_producer_state, + [[maybe_unused]] StorePipeline store_pipeline, + [[maybe_unused]] StorePipelineState store_pipe_producer_state) { + load_pipeline.producer_tail(load_pipe_producer_state); + } + + template< + bool ReuseTmem = false, + class AccumulatorPipeline, + class AccumulatorPipelineState, + class ProblemShapeMNKL, + class CtaTileMNK, + class CtaCoordMNKL, + class MmaTileMNK, + class TiledMma, + class AccEngine, + class AccLayout + > + CUTLASS_DEVICE auto + store( + LoadPipeline load_pipeline, + LoadPipelineState load_pipe_consumer_state, + StorePipeline store_pipeline, + StorePipelineState store_pipe_producer_state, + AccumulatorPipeline acc_pipeline, + AccumulatorPipelineState acc_pipe_consumer_state, + ProblemShapeMNKL problem_shape_mnkl, + CtaTileMNK cta_tile_mnk, + CtaCoordMNKL cta_coord_mnkl, + MmaTileMNK mma_tile_mnk, + TiledMma tiled_mma, + cute::Tensor accumulators, + TensorStorage& shared_tensors + ) { + using namespace cute; + using ElementAccumulator = typename AccEngine::value_type; + using ElementCompute_ = typename epilogue::fusion::FusionCallbacksTraits::ElementCompute; + using ElementCompute = cute::conditional_t,ElementAccumulator,ElementCompute_>; + + static_assert(is_tmem::value, "Accumulator must be TMEM resident."); + static_assert(rank(accumulators) == 3, "Accumulators must be MMA-partitioned: [MMA, MMA_M, MMA_N]"); + static_assert(size<1>(accumulators) == 1 && size<2>(accumulators) == 1, "TiledMMA must match partitioned ShapeMN"); + static_assert(rank(ProblemShapeMNKL{}) == 4, "ProblemShapeMNKL must be rank 4"); + static_assert(rank(CtaCoordMNKL{}) == 4, "CoordMNKL must be rank 4"); + + // Indexing variables + auto [M, N, K, L] = problem_shape_mnkl; + auto [m_coord, n_coord, k_coord, l_coord] = cta_coord_mnkl; + int thread_idx = threadIdx.x % ThreadCount; + int warp_idx = thread_idx / NumThreadsPerWarp; + [[maybe_unused]] int lane_idx = thread_idx % NumThreadsPerWarp; + + // The tma tensor D under im2col mode only has two modes (M, N) which + // should be local tiled with only (m_coord, n_coord). + auto coord_shape = + conditional_return(make_coord(m_coord, n_coord), make_coord(m_coord, n_coord, l_coord)); + + // Represent the full output tensor, slice to get the tile this CTA is responsible for + Tensor mD_mn = params.tma_store_d.get_tma_tensor(make_shape(M,N,L)); // (M,N,L) + Tensor mD = coalesce(mD_mn, take<0,2>(cta_tile_mnk)); + Tensor gD = local_tile(mD, take<0,2>(cta_tile_mnk), coord_shape); // (CTA_M,CTA_N) + + Tensor tAcc = accumulators(make_coord(_,_),_0{},_0{}); // (CTA_M,CTA_N) + + // Apply epilogue subtiling + Tensor tAcc_epi = flat_divide(tAcc, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + Tensor gD_epi = flat_divide( gD, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + + // Construct the corresponding pipelined smem tensors + auto ptr_sC = shared_tensors.collective.smem_C.begin(); + auto ptr_sD = shared_tensors.collective.smem_D.begin(); + Tensor sC_epi = cute::as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(ptr_sC), SmemLayoutC{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_C) + Tensor sD_epi = cute::as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(ptr_sD), SmemLayoutD{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_D) + + // (t)hread-partition for (t)mem to (r)egister copy (tTR_) + TiledCopy tiled_t2r = make_tmem_copy(CopyOpT2R{}, tAcc_epi(_,_,_0{},_0{})); + ThrCopy thread_t2r = tiled_t2r.get_slice(thread_idx); + Tensor tTR_tAcc = thread_t2r.partition_S(tAcc_epi); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N) + Tensor tTR_sD = thread_t2r.partition_D(sD_epi(_,_,_0{})); // (T2R,T2R_M,T2R_N) + + // Allocate D and accumulator registers + // Does directly store the visitor into smem. + constexpr bool IsDirectR2S = cute::is_same_v>; + using RegisterElementD = cute::conditional_t; + Tensor tTR_rAcc = make_tensor(shape(tTR_sD)); // (T2R,T2R_M,T2R_N) + Tensor tTR_rD = make_tensor(shape(tTR_sD)); // (T2R,T2R_M,T2R_N) + + // Vectorized fragment view + constexpr int FragmentSize = DispatchPolicy::FragmentSize; + Tensor tTR_rAcc_frg = recast>(coalesce(tTR_rAcc)); // (EPI_V) + Tensor tTR_rD_frg = recast>(coalesce(tTR_rD)); // (EPI_V) + CUTE_STATIC_ASSERT(size(tTR_rAcc) % DispatchPolicy::FragmentSize == 0, "Fragment size does not vectorize properly"); + + // (t)hread-partition for (s)mem to (r)egister copy (tSR_) + TiledCopy tiled_s2r = make_tiled_copy_D(Copy_Atom{}, tiled_t2r); + ThrCopy thread_s2r = tiled_s2r.get_slice(thread_idx); + Tensor tSR_sC = thread_s2r.partition_S(sC_epi); // (S2R,S2R_M,S2R_N,PIPE_C) + Layout tSR_rC_layout = thread_s2r.retile_D(tTR_rD).layout(); // (S2R,S2R_M,S2R_N) + + // Allocate C registers + // If C smem load is a non-vectorized dst(i) = src(i) then we can allocate C registers directly in the compute type + // to eliminate some redundant pack+unpack instruction sequences for sub-word types + constexpr bool IsDirectS2R = cute::is_same_v> + && decltype(max_common_vector(tSR_rC_layout, tSR_sC.layout()))::value <= 1; + using RegisterElementC = cute::conditional_t; + Tensor tTR_rC = make_tensor(shape(tTR_sD)); // (T2R,T2R_M,T2R_N) + Tensor tSR_rC = thread_s2r.retile_D(tTR_rC); // (S2R,S2R_M,S2R_N) + + // (t)hread-partition for (r)egister to (r)egister copy (tRR_) + TiledCopy tiled_r2r = make_tiled_copy_D(Copy_Atom{}, tiled_t2r); + ThrCopy thread_r2r = tiled_r2r.get_slice(thread_idx); + Tensor tRR_rD_src = thread_r2r.retile_S(tTR_rD); // (R2R,R2R_M,R2R_N,EPI_M,EPI_N) + Tensor tRR_rD_dst = thread_r2r.retile_D(tTR_rD); // (R2R,R2R_M,R2R_N,EPI_M,EPI_N) + + // (t)hread-partition for (r)egister to (s)mem copy (tRS_) + TiledCopy tiled_r2s = make_tiled_copy_D(Copy_Atom{}, tiled_r2r); + ThrCopy thread_r2s = tiled_r2s.get_slice(thread_idx); + Tensor tRS_sD = thread_r2s.partition_D(sD_epi); // (R2S,R2S_M,R2S_N,PIPE_D) + Tensor tRS_rD = [&]() CUTLASS_LAMBDA_FUNC_INLINE { + if constexpr (!IsDirectR2S) { + return make_tensor(shape(tRS_sD(_,_,_,_0{}))); + } + else{ + return thread_r2s.retile_S(tTR_rD); // (R2S,R2S_M,R2S_N) + } + }(); + + Tensor tRR_rD_dst_frg = recast>(coalesce(tRR_rD_dst)); + Tensor tRS_rD_frg = recast>(coalesce(tRS_rD)); + + // thread(b)lock-partition for (s)mem to (g)mem copy (bSG_) + ThrCopy thrblk_s2g = params.tma_store_d.get_slice(Int<0>{}); + Tensor bSG_sD = thrblk_s2g.partition_S(sD_epi); // (S2G,S2G_M,S2G_N,PIPE_D) + Tensor bSG_gD = thrblk_s2g.partition_D(gD_epi); // (S2G,S2G_M,S2G_N,EPI_M,EPI_N) + + // OOB predication for tile quantization "residue" + // Absolute coordinate tensors (dynamic) + Tensor mD_crd = make_identity_tensor(make_shape(M,N)); // (M,N) + Tensor cD_mn = local_tile(mD_crd, take<0,2>(cta_tile_mnk), make_coord(m_coord, n_coord)); // (CTA_M,CTA_N) + Tensor tTR_cD_mn = thread_t2r.partition_D(flat_divide(cD_mn, EpilogueTile{})); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N) + // Relative coordinate tensors (static) + Tensor cD = make_coord_tensor(cD_mn.layout()); // (CTA_M,CTA_N) + Tensor tTR_cD = make_coord_tensor(tTR_cD_mn.layout()); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N) + // Subtract the global "bottom right" corner from the local "top left" corner to get the max relative coordinate + auto residue_cD = make_coord(M,N) - cD_mn(_0{}); // (m,n) + auto residue_tTR_cD = make_coord(M,N) - tTR_cD_mn(_0{}); // (m,n) + + // Arguments for the fusion callbacks for the consumer store warps + constexpr bool RefSrc = false; // Register tensors reference T2R copy dst layout + auto cst_args = cutlass::epilogue::fusion::detail::ConsumerStoreArgs{ + problem_shape_mnkl, + cta_tile_mnk, + cta_coord_mnkl, + tiled_mma, + EpilogueTile{}, + tiled_t2r, + cD, + residue_cD, + tTR_cD, + residue_tTR_cD, + tTR_rC, + thread_idx + }; + + // Thread synchronizer for previously issued waits or fences + // to ensure visibility of smem reads/writes to threads or TMA unit + auto synchronize = [] () { cutlass::arch::NamedBarrier::sync(ThreadCount, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); }; + + // Predication for sub-128 thread T2R tiled copy + Layout tmem_warp_layout = typename decltype(make_tmem_warp_partitioner(tAcc_epi(_,_,0,0)))::TiledLayout_TV{}; + constexpr bool predicate_tmem_load = size(tmem_warp_layout) != cosize(tmem_warp_layout); + bool issue_tmem_load = true; + + // If tmem doesn't have enough capacity to support double buffering, a portion of tmem (a column of epilogue tiles) + // is overlapped between 2 pseudo-buffers. The shared tmem portion corresponds to the last epilogue tile column of + // tmem accumulator buffer 0, and the first epilogue tile column of tmem accumulator 1. + // Thus, whenever we are processing tmem accumulator buffer 0, we process the epilogue tiles with reversed column order. + // Once the last epilogue tile column is loaded from tmem, the acc_pipeline is released. + // Then, the next accumulation stage for buffer 1 can start. + [[maybe_unused]] bool reverse_epi_n = ReuseTmem && acc_pipe_consumer_state.phase() == 0; + static_assert(not (ReuseTmem && AccumulatorPipeline::Stages != 1), "Tmem reuse requires 1 accumulator stage"); + + // Predication for TMA store (one warp issues TMA store) + bool issue_tma_store = warp_idx == 0; + + // In the reuse smem configuration we have StagesC smem buffers and at most StagesD committed TMA stores in flight. + // The TMA store pipeline producer acquire returns when at most StagesD-1 committed stores are in-flight, so we can + // only guarantee store completion after StagesD iterations, then we can begin issuing releases on the smem buffer locks. + // store_pipe_producer_state tracks the acquire and load_pipe_consumer_state tracks the release, in circular buffer fashion. + // If TMA store supported async transaction mbarriers we would not need this synchronous release behavior. + LoadPipelineState load_wait_state = load_pipe_consumer_state; + if constexpr (ReuseSmemC) { + load_wait_state = store_pipe_producer_state; + load_wait_state.phase_ ^= 1; + } + + // We can delay issue of TMA store by one iteration to achieve better interleaving of non-TMA instructions + // Sync requirements of smem reuse may preclude this optimization + // Delayed stores cause delayed stage releases which causes deadlock when StagesC == StagesD + [[maybe_unused]] int epi_m_prev = 0; + [[maybe_unused]] int epi_n_prev = 0; + static_assert(not (DelayTmaStore and ReuseSmemC and StagesC <= StagesD), "This TMA epilogue configuration will deadlock"); + + // The Epilogue Loop + auto epi_loop_fn = [&] (auto& cst_callbacks) CUTLASS_LAMBDA_FUNC_INLINE { + bool is_producer_load_needed = fusion_callbacks.is_producer_load_needed(); + bool is_C_load_needed = is_source_supported && fusion_callbacks.is_C_load_needed(); + + // The TMA store sequence for one epilogue loop iteration + auto tma_store_fn = [&] (int epi_m, int epi_n) CUTLASS_LAMBDA_FUNC_INLINE { + // Write the tile from smem to gmem with TMA + cutlass::arch::fence_view_async_shared(); // ensure smem writes are visible to TMA + synchronize(); // ensure all threads have issued their async fence + if (issue_tma_store) { + if constexpr (StoreD) { + copy(params.tma_store_d, bSG_sD(_,_,_,store_pipe_producer_state.index()), bSG_gD(_,_,_,epi_m,epi_n)); + } + } + + // Post async fence, pre TMA commit callback entry point + cst_callbacks.tma_store(epi_m, epi_n, store_pipe_producer_state.count(), issue_tma_store); + + // Commit the TMA stores for this stage + if (issue_tma_store) { + store_pipeline.producer_commit(store_pipe_producer_state); + } + ++store_pipe_producer_state; + + // Wait for the next smem buffer to be available + if (issue_tma_store) { + store_pipeline.producer_acquire(store_pipe_producer_state); + } + synchronize(); + + if constexpr (ReuseSmemC) { + // producer_acquire returns when at most StagesD-1 committed stores are pending + bool store_finished = store_pipe_producer_state.count() > StorePipeline::UnacquiredStages; + // Let dma warp know earliest smem buffer is consumed and empty after StagesD producer commits + if (store_finished) { + if (is_producer_load_needed) { + load_pipeline.consumer_release(load_pipe_consumer_state); + } + ++load_pipe_consumer_state; + } + } + }; // tma_store_fn + + cst_callbacks.begin(); + if (cst_callbacks.begin_sync_needed()) { + synchronize(); + } + + // Begin the wait for the producer load results + ConsumerToken load_wait_token{BarrierStatus::WaitDone}; + if (is_producer_load_needed) { + load_wait_token = load_pipeline.consumer_try_wait(load_wait_state); + } + // Begin the wait for the accumulator results + ConsumerToken acc_wait_token = acc_pipeline.consumer_try_wait(acc_pipe_consumer_state); + + // For each epilogue subtile within the CTA tile + constexpr int NumEpiSubtilesN = CUTE_STATIC_V(size<3>(gD_epi)); + constexpr int NumEpiSubtilesM = CUTE_STATIC_V(size<2>(gD_epi)); + #pragma unroll(UnrollEpiLoop ? NumEpiSubtilesN : 1) + for (int iter_n = 0; iter_n < NumEpiSubtilesN; ++iter_n) { + #pragma unroll(UnrollEpiLoop ? NumEpiSubtilesM : 1) + for (int iter_m = 0; iter_m < NumEpiSubtilesM; ++iter_m) { + int epi_m = iter_m, epi_n = iter_n; + bool is_first_iteration = iter_m == 0 && iter_n == 0; + bool is_last_iteration = iter_m == size<2>(gD_epi)-1 && iter_n == size<3>(gD_epi)-1; + bool do_acc_release = is_last_iteration; + + // Reverse subtile order for tmem reuse if necessary + if constexpr (ReuseTmem) { + if (reverse_epi_n) { + epi_n = size<3>(gD_epi) - 1 - iter_n; + } + do_acc_release = iter_m == size<2>(gD_epi)-1 && iter_n == 0; + } + + cst_callbacks.begin_loop(epi_m, epi_n); + + if (is_producer_load_needed) { + // Wait for the producer load to fill smem + load_pipeline.consumer_wait(load_wait_state, load_wait_token); + + if (is_C_load_needed) { + // Copy source tile from smem to register + copy(tiled_s2r, tSR_sC(_,_,_,load_wait_state.index()), tSR_rC); + // Ensure smem loads are complete before reusing smem for mixed types/layouts + if constexpr (ReuseSmemC && not (SmemLayoutC{} == SmemLayoutD{})) { + synchronize(); + } + } + } + + // First loop fusion callback entry point + cst_callbacks.previsit(epi_m, epi_n, load_wait_state.count(), is_producer_load_needed); + + if (is_producer_load_needed) { + // Let producer load warp know smem buffers are consumed and empty + if constexpr (not ReuseSmemC) { + cutlass::arch::fence_view_async_shared(); + load_pipeline.consumer_release(load_pipe_consumer_state); + ++load_pipe_consumer_state; + } + ++load_wait_state; + } + + if (is_first_iteration) { + // Wait for mma warp to fill tmem buffer with accumulator results + acc_pipeline.consumer_wait(acc_pipe_consumer_state, acc_wait_token); + } + + // The current tile in tmem + Tensor tTR_tAcc_mn = tTR_tAcc(_,_,_,epi_m,epi_n); + + // Compute tmem load predication if necessary + if constexpr (predicate_tmem_load) { + // Issue tmem load if this tile's tmem subpartition is accessible by this warp + int subpart_idx = (tTR_tAcc_mn.data().dp_ / 32) % 4; + issue_tmem_load = warp_idx == subpart_idx; + } + bool issue_smem_store = issue_tmem_load; + + // Copy accumulator tile from tmem to register + if (issue_tmem_load) { + copy(tiled_t2r, tTR_tAcc_mn, tTR_rAcc); + } + + // After the last tmem load, signal that tmem buffer is consumed and empty + if (do_acc_release) { + cutlass::arch::fence_view_async_tmem_load(); + acc_pipeline.consumer_release(acc_pipe_consumer_state); + ++acc_pipe_consumer_state; + } + + // Vectorized fragment loop with visitor callback entry point + CUTLASS_PRAGMA_UNROLL + for (int epi_v = 0; epi_v < size(tTR_rD_frg); ++epi_v) { + tTR_rD_frg(epi_v) = cst_callbacks.visit(tTR_rAcc_frg(epi_v), epi_v, epi_m, epi_n); + } + + // The latest we can delay the TMA store is right before the smem store of the next iteration + // since the current TMA store needs to be committed before we can acquire the next smem buffer + if constexpr (DelayTmaStore) { + // Issue TMA stores for the previous subtile + if (not is_first_iteration) { + tma_store_fn(epi_m_prev, epi_n_prev); + } + epi_m_prev = epi_m; + epi_n_prev = epi_n; + } + + if constexpr (StoreD) { + if constexpr (!IsDirectR2S) { + // At present, only FP4 col output with scalefactor generation fusion would go into these branch + copy(tiled_r2r, tRR_rD_src, tRR_rD_dst); + } + tRS_rD_frg(_0{}) = cutlass::NumericArrayConverter{}(tRR_rD_dst_frg(_0{})); + } + + // Smem reduction callback entry point using current store buffer for workspace + Tensor reduction_buffer = make_tensor(raw_pointer_cast(sD_epi(_,_,store_pipe_producer_state.index()).data()), + make_layout(stride<2>(get_nonswizzle_portion(SmemLayoutD{})), _1{})); + cst_callbacks.reduce(reduction_buffer, synchronize, epi_m, epi_n, is_last_iteration, tRS_rD_frg); + + // Copy output tile from register to smem + if (issue_smem_store) { + if constexpr (StoreD) { + copy(tiled_r2s, tRS_rD, tRS_sD(_,_,_,store_pipe_producer_state.index())); + } + } + + // Post reduction, pre TMA store callback entry point + cst_callbacks.postreduce(epi_m, epi_n, store_pipe_producer_state.count(), issue_smem_store); + + if constexpr (not DelayTmaStore) { + // Issue TMA stores for this subtile + tma_store_fn(epi_m, epi_n); + } + + cst_callbacks.end_loop(epi_m, epi_n); + + if (is_producer_load_needed) { + // Begin the wait for the next subtile producer load + load_wait_token = load_pipeline.consumer_try_wait(load_wait_state, is_last_iteration); + } + } // for epi_m + } // for epi_n + + if constexpr (DelayTmaStore) { + // Issue TMA stores for the last subtile + tma_store_fn(epi_m_prev, epi_n_prev); + } + + cst_callbacks.end(); + }; // epi_loop_fn + + // + // BEGIN EPILOGUE + // + auto cst_callbacks = fusion_callbacks.template get_consumer_store_callbacks(cst_args); + epi_loop_fn(cst_callbacks); + return cute::make_tuple(load_pipe_consumer_state, store_pipe_producer_state, acc_pipe_consumer_state); + } + + // API with Global Accumulator in registers for FastFP32 (emulated MMA) kernels. + // The accumulator in TMEM periodically loaded into the registers so that the MMA can clear out the TMEM accumulator + // values for better accuracy. This epilogue accepts the accumulator in registers and take TiledCopy for the + // TMEM->Reg as a parameter to be used in partitioning GMEM tensors C and D. + template< + class ProblemShapeMNKL, + class CtaTileMNK, + class CtaCoordMNKL, + class MmaTileMNK, + class TiledMma, + class AccEngine, + class AccLayout, + class TiledCopyT2R + > + CUTLASS_DEVICE auto + store( + LoadPipeline load_pipeline, + LoadPipelineState load_pipe_consumer_state, + StorePipeline store_pipeline, + StorePipelineState store_pipe_producer_state, + ProblemShapeMNKL problem_shape_mnkl, + CtaTileMNK cta_tile_mnk, + CtaCoordMNKL cta_coord_mnkl, + MmaTileMNK mma_tile_mnk, + TiledMma tiled_mma, + cute::Tensor& tTR_rAcc, // (T2R,T2R_M,T2R_N,EPI_M,EPI_N) + TensorStorage& shared_tensors, + TiledCopyT2R tiled_t2r + ) { + using namespace cute; + using ElementAccumulator = typename AccEngine::value_type; + using ElementCompute_ = typename epilogue::fusion::FusionCallbacksTraits::ElementCompute; + using ElementCompute = cute::conditional_t,ElementAccumulator,ElementCompute_>; + + static_assert(is_rmem::value, "Accumulator must be Register resident."); + static_assert(rank(AccLayout{}) == 5, "Accumulators must be copy-partitioned: (T2R,T2R_M,T2R_N,EPI_M,EPI_N)"); + static_assert(rank(ProblemShapeMNKL{}) == 4, "ProblemShapeMNKL must be rank 4"); + static_assert(rank(CtaCoordMNKL{}) == 4, "CoordMNKL must be rank 4"); + + // Indexing variables + auto [M, N, K, L] = problem_shape_mnkl; + auto [m_coord, n_coord, k_coord, l_coord] = cta_coord_mnkl; + int thread_idx = threadIdx.x % ThreadCount; + int warp_idx = thread_idx / NumThreadsPerWarp; + [[maybe_unused]] int lane_idx = thread_idx % NumThreadsPerWarp; + + // The tma tensor D under im2col mode only has two modes (M, N) which + // should be local tiled with only (m_coord, n_coord). + auto coord_shape = + conditional_return(make_coord(m_coord, n_coord), make_coord(m_coord, n_coord, l_coord)); + + // Represent the full output tensor, slice to get the tile this CTA is responsible for + Tensor mD_mn = params.tma_store_d.get_tma_tensor(make_shape(M,N,L)); // (M,N,L) + Tensor mD = coalesce(mD_mn, take<0,2>(cta_tile_mnk)); + Tensor gD = local_tile(mD, take<0,2>(cta_tile_mnk), coord_shape); // (CTA_M,CTA_N) + + // Apply epilogue subtiling + Tensor gD_epi = flat_divide( gD, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + + // Construct the corresponding pipelined smem tensors + auto ptr_sC = shared_tensors.collective.smem_C.begin(); + auto ptr_sD = shared_tensors.collective.smem_D.begin(); + Tensor sC_epi = cute::as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(ptr_sC), SmemLayoutC{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_C) + Tensor sD_epi = cute::as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(ptr_sD), SmemLayoutD{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_D) + + // (t)hread-partition for (t)mem to (r)egister copy (tTR_) + ThrCopy thread_t2r = tiled_t2r.get_slice(thread_idx); + Tensor tTR_sD = thread_t2r.partition_D(sD_epi(_,_,_0{})); // (T2R,T2R_M,T2R_N) + + // Allocate D and accumulator registers + Tensor tTR_rD = make_tensor(shape(tTR_sD)); // (T2R,T2R_M,T2R_N) + + // Vectorized fragment view + constexpr int FragmentSize = DispatchPolicy::FragmentSize; + Tensor tTR_rD_frg = recast>(coalesce(tTR_rD)); // (EPI_V) + + // (t)hread-partition for (s)mem to (r)egister copy (tSR_) + TiledCopy tiled_s2r = make_tiled_copy_D(Copy_Atom{}, tiled_t2r); + ThrCopy thread_s2r = tiled_s2r.get_slice(thread_idx); + Tensor tSR_sC = thread_s2r.partition_S(sC_epi); // (S2R,S2R_M,S2R_N,PIPE_C) + Layout tSR_rC_layout = thread_s2r.retile_D(tTR_rD).layout(); // (S2R,S2R_M,S2R_N) + + // Allocate C registers + // If C smem load is a non-vectorized dst(i) = src(i) then we can allocate C registers directly in the compute type + // to eliminate some redundant pack+unpack instruction sequences for sub-word types + constexpr bool IsDirectS2R = cute::is_same_v> + && decltype(max_common_vector(tSR_rC_layout, tSR_sC.layout()))::value <= 1; + using RegisterElementC = cute::conditional_t; + Tensor tTR_rC = make_tensor(shape(tTR_sD)); // (T2R,T2R_M,T2R_N) + Tensor tSR_rC = thread_s2r.retile_D(tTR_rC); // (S2R,S2R_M,S2R_N) + + // (t)hread-partition for (r)egister to (s)mem copy (tRS_) + TiledCopy tiled_r2s = make_tiled_copy_D(Copy_Atom{}, tiled_t2r); + ThrCopy thread_r2s = tiled_r2s.get_slice(thread_idx); + Tensor tRS_rD = thread_r2s.retile_S(tTR_rD); // (R2S,R2S_M,R2S_N) + Tensor tRS_sD = thread_r2s.partition_D(sD_epi); // (R2S,R2S_M,R2S_N,PIPE_D) + + // thread(b)lock-partition for (s)mem to (g)mem copy (bSG_) + ThrCopy thrblk_s2g = params.tma_store_d.get_slice(Int<0>{}); + Tensor bSG_sD = thrblk_s2g.partition_S(sD_epi); // (S2G,S2G_M,S2G_N,PIPE_D) + Tensor bSG_gD = thrblk_s2g.partition_D(gD_epi); // (S2G,S2G_M,S2G_N,EPI_M,EPI_N) + + // OOB predication for tile quantization "residue" + // Absolute coordinate tensors (dynamic) + Tensor mD_crd = make_identity_tensor(make_shape(M,N)); // (M,N) + Tensor cD_mn = local_tile(mD_crd, take<0,2>(cta_tile_mnk), make_coord(m_coord, n_coord)); // (CTA_M,CTA_N) + Tensor tTR_cD_mn = thread_t2r.partition_D(flat_divide(cD_mn, EpilogueTile{})); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N) + // Relative coordinate tensors (static) + Tensor cD = make_coord_tensor(cD_mn.layout()); // (CTA_M,CTA_N) + Tensor tTR_cD = make_coord_tensor(tTR_cD_mn.layout()); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N) + // Subtract the global "bottom right" corner from the local "top left" corner to get the max relative coordinate + auto residue_cD = make_coord(M,N) - cD_mn(_0{}); // (m,n) + auto residue_tTR_cD = make_coord(M,N) - tTR_cD_mn(_0{}); // (m,n) + + // Get the fusion callbacks for the consumer store warps + constexpr bool RefSrc = false; // Register tensors reference T2R copy dst layout + auto cst_args = cutlass::epilogue::fusion::detail::ConsumerStoreArgs{ + problem_shape_mnkl, + cta_tile_mnk, + cta_coord_mnkl, + tiled_mma, + EpilogueTile{}, + tiled_t2r, + cD, + residue_cD, + tTR_cD, + residue_tTR_cD, + tTR_rC, + thread_idx + }; + + auto cst_callbacks = fusion_callbacks.template get_consumer_store_callbacks(cst_args); + bool is_producer_load_needed = fusion_callbacks.is_producer_load_needed(); + bool is_C_load_needed = is_source_supported && fusion_callbacks.is_C_load_needed(); + + // Thread synchronizer for previously issued waits or fences + // to ensure visibility of smem reads/writes to threads or TMA unit + auto synchronize = [] () { cutlass::arch::NamedBarrier::sync(ThreadCount, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); }; + + // Predication for TMA store (one warp issues TMA store) + bool issue_tma_store = warp_idx == 0; + + // In the reuse smem configuration we have StagesC smem buffers and at most StagesD committed TMA stores in flight. + // The TMA store pipeline producer acquire returns when at most StagesD-1 committed stores are in-flight, so we can + // only guarantee store completion after StagesD iterations, then we can begin issuing releases on the smem buffer locks. + // store_pipe_producer_state tracks the acquire and load_pipe_consumer_state tracks the release, in circular buffer fashion. + // If TMA store supported async transaction mbarriers we would not need this synchronous release behavior. + LoadPipelineState load_wait_state = load_pipe_consumer_state; + if constexpr (ReuseSmemC) { + load_wait_state = store_pipe_producer_state; + load_wait_state.phase_ ^= 1; + } + + // We can delay issue of TMA store by one iteration to achieve better interleaving of non-TMA instructions + // Sync requirements of smem reuse may preclude this optimization + // Delayed stores cause delayed stage releases which causes deadlock when StagesC == StagesD + int epi_m_prev = 0, epi_n_prev = 0; + static_assert(not (DelayTmaStore and ReuseSmemC and StagesC <= StagesD), "This TMA epilogue configuration will deadlock"); + + // The TMA store sequence for one subtile iteration + auto tma_store_fn = [&] (int epi_m, int epi_n) CUTLASS_LAMBDA_FUNC_INLINE { + // Write the tile from smem to gmem with TMA + cutlass::arch::fence_view_async_shared(); // ensure smem writes are visible to TMA + synchronize(); // ensure all threads have issued their async fence + if (issue_tma_store) { + copy(params.tma_store_d, bSG_sD(_,_,_,store_pipe_producer_state.index()), bSG_gD(_,_,_,epi_m,epi_n)); + } + + // Post async fence, pre TMA commit callback entry point + cst_callbacks.tma_store(epi_m, epi_n, store_pipe_producer_state.count(), issue_tma_store); + + // Commit the TMA stores for this stage + if (issue_tma_store) { + store_pipeline.producer_commit(store_pipe_producer_state); + } + ++store_pipe_producer_state; + + // Wait for the next smem buffer to be available + if (issue_tma_store) { + store_pipeline.producer_acquire(store_pipe_producer_state); + } + synchronize(); + + if constexpr (ReuseSmemC) { + // producer_acquire returns when at most StagesD-1 committed stores are pending + bool store_finished = store_pipe_producer_state.count() > StorePipeline::UnacquiredStages; + // Let dma warp know earliest smem buffer is consumed and empty after StagesD producer commits + if (store_finished) { + if (is_producer_load_needed) { + load_pipeline.consumer_release(load_pipe_consumer_state); + } + ++load_pipe_consumer_state; + } + } + }; + + // + // BEGIN EPILOGUE + // + + cst_callbacks.begin(); + if (cst_callbacks.begin_sync_needed()) { + synchronize(); + } + + // Begin the wait for the producer load results + ConsumerToken load_wait_token{BarrierStatus::WaitDone}; + if (is_producer_load_needed) { + load_wait_token = load_pipeline.consumer_try_wait(load_wait_state); + } + + // For each epilogue subtile within the CTA tile + constexpr int NumEpiSubtilesN = CUTE_STATIC_V(size<3>(gD_epi)); + constexpr int NumEpiSubtilesM = CUTE_STATIC_V(size<2>(gD_epi)); + #pragma unroll(UnrollEpiLoop ? NumEpiSubtilesN : 1) + for (int iter_n = 0; iter_n < NumEpiSubtilesN; ++iter_n) { + #pragma unroll(UnrollEpiLoop ? NumEpiSubtilesM : 1) + for (int iter_m = 0; iter_m < NumEpiSubtilesM; ++iter_m) { + int epi_m = iter_m, epi_n = iter_n; + bool is_first_iteration = iter_m == 0 && iter_n == 0; + bool is_last_iteration = iter_m == size<2>(gD_epi)-1 && iter_n == size<3>(gD_epi)-1; + + cst_callbacks.begin_loop(epi_m, epi_n); + + if (is_producer_load_needed) { + // Wait for the producer load to fill smem + load_pipeline.consumer_wait(load_wait_state, load_wait_token); + + if (is_C_load_needed) { + // Copy source tile from smem to register + copy(tiled_s2r, tSR_sC(_,_,_,load_wait_state.index()), tSR_rC); + // Ensure smem loads are complete before reusing smem for mixed types/layouts + if constexpr (ReuseSmemC && not (SmemLayoutC{} == SmemLayoutD{})) { + synchronize(); + } + } + } + + // First loop fusion callback entry point + cst_callbacks.previsit(epi_m, epi_n, load_wait_state.count(), is_producer_load_needed); + + if (is_producer_load_needed) { + // Let producer load warp know smem buffers are consumed and empty + if constexpr (not ReuseSmemC) { + cutlass::arch::fence_view_async_shared(); + load_pipeline.consumer_release(load_pipe_consumer_state); + ++load_pipe_consumer_state; + } + ++load_wait_state; + } + + Tensor tTR_rAcc_epi_tile = tTR_rAcc(_,_,_,epi_m,epi_n); + Tensor tTR_rAcc_frg = recast>(coalesce(tTR_rAcc_epi_tile)); // (EPI_V) + + // Vectorized fragment loop with visitor callback entry point + CUTLASS_PRAGMA_UNROLL + for (int epi_v = 0; epi_v < size(tTR_rD_frg); ++epi_v) { + tTR_rD_frg(epi_v) = cst_callbacks.visit(tTR_rAcc_frg(epi_v), epi_v, epi_m, epi_n); + } + + // The latest we can delay the TMA store is right before the smem store of the next iteration + // since the current TMA store needs to be committed before we can acquire the next smem buffer + if constexpr (DelayTmaStore) { + // Issue TMA stores for the previous subtile + if (not is_first_iteration) { + tma_store_fn(epi_m_prev, epi_n_prev); + } + epi_m_prev = epi_m; + epi_n_prev = epi_n; + } + + // Smem reduction callback entry point using current store buffer for workspace + Tensor reduction_buffer = make_tensor(raw_pointer_cast(sD_epi(_,_,store_pipe_producer_state.index()).data()), + make_layout(stride<2>(get_nonswizzle_portion(SmemLayoutD{})), _1{})); + cst_callbacks.reduce(reduction_buffer, synchronize, epi_m, epi_n, is_last_iteration, tTR_rD_frg); + + // Copy output tile from register to smem + bool issue_smem_store = true; + if (issue_smem_store) { + copy(tiled_r2s, tRS_rD, tRS_sD(_,_,_,store_pipe_producer_state.index())); + } + + // Post reduction, pre TMA store callback entry point + cst_callbacks.postreduce(epi_m, epi_n, store_pipe_producer_state.count(), issue_smem_store); + + if constexpr (not DelayTmaStore) { + // Issue TMA stores for this subtile + tma_store_fn(epi_m, epi_n); + } + + cst_callbacks.end_loop(epi_m, epi_n); + + if (is_producer_load_needed) { + // Begin the wait for the next subtile producer load + load_wait_token = load_pipeline.consumer_try_wait(load_wait_state, is_last_iteration); + } + } // for epi_m + } // for epi_n + + if constexpr (DelayTmaStore) { + // Issue TMA stores for the last subtile + tma_store_fn(epi_m_prev, epi_n_prev); + } + + cst_callbacks.end(); + + return cute::make_tuple(load_pipe_consumer_state, store_pipe_producer_state); + } + + template + CUTLASS_DEVICE void + store_tail( + LoadPipeline load_pipeline, + LoadPipelineState load_pipe_consumer_state, + StorePipeline store_pipeline, + StorePipelineState store_pipe_producer_state, + CtaTileMNK cta_tile_mnk) { + if constexpr (ReuseSmemC) { + if (fusion_callbacks.is_producer_load_needed()) { + // wait for all TMA stores to complete + store_pipeline.producer_tail(store_pipe_producer_state); + + // Issue releases on up to StagesD-1 previously issued TMA stores + constexpr int release_stages = cute::min(StorePipeline::UnacquiredStages, get_load_pipe_increment(cta_tile_mnk)); + CUTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < release_stages; ++stage) { + load_pipeline.consumer_release(load_pipe_consumer_state); + ++load_pipe_consumer_state; + } + } + } + } +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::epilogue::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/kernels/fp4_w4a4_mma_warpsplit_mrows_f32out_sm120.cu b/csrc/kernels/fp4_w4a4_mma_warpsplit_mrows_f32out_sm120.cu new file mode 100644 index 00000000..e0c28e46 --- /dev/null +++ b/csrc/kernels/fp4_w4a4_mma_warpsplit_mrows_f32out_sm120.cu @@ -0,0 +1,200 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// See header. Kernel body proven in the llama.cpp SM120 adapter (bit-exact +// duplicated-token replay across M variants, perplexity-neutral). + +#include "fp4_w4a4_mma_warpsplit_mrows_f32out_sm120.cuh" + +#include + +#include "cute/arch/mma_sm120.hpp" +#include "cutlass/numeric_types.h" + +namespace flash_rt { +namespace gemm { + +namespace { + +#if defined(__CUDA_ARCH_FEAT_SM120_ALL) || !defined(__CUDA_ARCH__) +#define FR_WS_SM120A_OK 1 +#endif + +__device__ __forceinline__ void pdl_sync() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + cudaGridDependencySynchronize(); +#endif +} +__device__ __forceinline__ void pdl_lc() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +using AtomType = cute::SM120::BLOCKSCALED::SM120_16x8x64_TN_VS< + cutlass::float_e2m1_t, cutlass::float_e2m1_t, float, + cutlass::float_ue4m3_t, 16>; + +__device__ __forceinline__ uint32_t fa(const uint8_t * s, int t0, int t1, int r) { + int ro = ((r & 1) ? (t1 + 8) : t1) * 32; + return *reinterpret_cast(s + ro + t0 * 4 + ((r >> 1) & 1) * 16); +} +__device__ __forceinline__ uint32_t fb(const uint8_t * s, int t0, int t1, int r) { + return *reinterpret_cast(s + t1 * 32 + t0 * 4 + r * 16); +} +__device__ __forceinline__ uint32_t fsa(const uint8_t * p, int u) { + return *reinterpret_cast(p + u * 4); +} +__device__ __forceinline__ void cpa(uint8_t * d, const uint8_t * s) { + uint32_t i = __cvta_generic_to_shared(d); + asm volatile("cp.async.ca.shared.global.L2::128B [%0], [%1], 4;\n" :: "r"(i), "l"(s)); +} +__device__ __forceinline__ void commit() { asm volatile("cp.async.commit_group;\n" ::); } +template __device__ __forceinline__ void waitg() { + asm volatile("cp.async.wait_group %0;\n" :: "n"(N)); +} + +template +__global__ void warpsplit_kernel_f32out( + const uint8_t * __restrict__ A, const uint8_t * __restrict__ B, + const uint8_t * __restrict__ SFA, const uint8_t * __restrict__ SFB, + float * __restrict__ D, float alpha, int N, int K) { + constexpr int M = MT; +#if defined(FR_WS_SM120A_OK) + pdl_lc(); pdl_sync(); + __shared__ uint8_t sA[WARPS][STAGES][16 * 32]; + __shared__ uint8_t sSFA[WARPS][STAGES][16 * 4]; + __shared__ uint8_t sB[WARPS][STAGES][8 * 32]; + __shared__ uint8_t sSFB[WARPS][STAGES][8 * 4]; + __shared__ float s_red[WARPS][4 * 8]; + + int tid = threadIdx.x, warp = tid >> 5, lane = tid & 31; + int my_n = blockIdx.x * 8; + const int KI = K / 64, KIw = KI / WARPS; + const int kt0 = warp * KIw; + const int KH = K / 2, ncs = (K / 16 + 3) / 4; + int t0 = lane & 3, t1 = lane >> 2, sau = (lane & 1) * 8 + (lane >> 2), sbu = lane >> 2; + float c0 = 0, c1 = 0, c2 = 0, c3 = 0; + + uint8_t (*mA)[16 * 32] = sA[warp]; + uint8_t (*mSFA)[16 * 4] = sSFA[warp]; + uint8_t (*mB)[8 * 32] = sB[warp]; + uint8_t (*mSFB)[8 * 4] = sSFB[warp]; + + if (lane >= 1 && lane < 16) { +#pragma unroll + for (int st = 0; st < STAGES; ++st) { + int4 * av = reinterpret_cast(mA[st]); int4 z{0, 0, 0, 0}; + av[lane * 2] = z; av[lane * 2 + 1] = z; + } + if (lane < 4) for (int st = 0; st < STAGES; ++st) + for (int i = 4 + lane; i < 64; i += 4) mSFA[st][i] = 0; + } + __syncwarp(); // M>1: row-1 cp.async below must not race the zero-init + auto ld = [&](int bf, int kt) { + int bo = kt * 32; + if (lane < 8) cpa(mA[bf] + lane * 4, A + bo + lane * 4); + if (lane == 0) cpa(mSFA[bf], SFA + kt * 512); +#pragma unroll + for (int rr = 1; rr < MT; ++rr) { // extra token rows: act tile + atom-layout scales (row r -> +r*16) + if (lane < 8) cpa(mA[bf] + rr * 32 + lane * 4, A + (size_t) rr * KH + bo + lane * 4); + if (lane == 0) cpa(mSFA[bf] + rr * 4, SFA + kt * 512 + rr * 16); + } + for (int c = 0; c < 2; ++c) { int ch = lane + c * 32, col = ch >> 3, off = ch & 7; + cpa(mB[bf] + ch * 4, B + (size_t)(my_n + col) * KH + bo + off * 4); } + if (lane < 8) { int col = my_n + lane, rb = col >> 7, ri = col & 127; + int si = rb * ncs + kt, ib = (ri & 31) * 16 + ((ri >> 5) & 3) * 4; + cpa(mSFB[bf] + lane * 4, SFB + (size_t)si * 512 + ib); } + }; +#pragma unroll + for (int st = 0; st < STAGES - 1; ++st) { if (st < KIw) ld(st, kt0 + st); commit(); } + for (int j = 0; j < KIw; ++j) { + int cb = j % STAGES, jp = j + STAGES - 1; + if (jp < KIw) ld(jp % STAGES, kt0 + jp); + commit(); waitg(); __syncwarp(); + uint32_t a0 = fa(mA[cb], t0, t1, 0), a1 = fa(mA[cb], t0, t1, 1); + uint32_t a2 = fa(mA[cb], t0, t1, 2), a3 = fa(mA[cb], t0, t1, 3); + uint32_t b0 = fb(mB[cb], t0, t1, 0), b1 = fb(mB[cb], t0, t1, 1); + uint32_t sfa_v = fsa(mSFA[cb], sau), sfb_v = fsa(mSFB[cb], sbu); + float d0, d1, d2, d3; + AtomType::fma(d0, d1, d2, d3, a0, a1, a2, a3, b0, b1, c0, c1, c2, c3, sfa_v, sfb_v); + c0 = d0; c1 = d1; c2 = d2; c3 = d3; + } + // m16n8 C fragment: {c0,c1} hold row (lane>>2) -> token t lives in lanes 4t..4t+3. + int q = lane >> 2, r = lane & 3; + if (q < M) { s_red[warp][q * 8 + r * 2] = c0; s_red[warp][q * 8 + r * 2 + 1] = c1; } + __syncthreads(); + if (warp == 0 && lane < 8) { + int col = my_n + lane; + if (col < N) { +#pragma unroll + for (int t = 0; t < MT; ++t) { + float acc = 0.f; +#pragma unroll + for (int w = 0; w < WARPS; ++w) acc += s_red[w][t * 8 + lane]; + D[(size_t) t * N + col] = acc * alpha; + } + } + } +#endif // FR_WS_SM120A_OK +} + +template +int launch(const uint8_t * A, const uint8_t * B, const uint8_t * SFA, + const uint8_t * SFB, float * D, float alpha, int N, int K, + bool pdl, cudaStream_t stream) { + const dim3 grid(N / 8), block(WARPS * 32); + if (pdl) { + cudaLaunchAttribute attr{}; + attr.id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr.val.programmaticStreamSerializationAllowed = 1; + cudaLaunchConfig_t cfg{}; + cfg.gridDim = grid; cfg.blockDim = block; cfg.dynamicSmemBytes = 0; + cfg.stream = stream; cfg.attrs = &attr; cfg.numAttrs = 1; + return (int) cudaLaunchKernelEx(&cfg, warpsplit_kernel_f32out, + A, B, SFA, SFB, D, alpha, N, K); + } + warpsplit_kernel_f32out<<>>( + A, B, SFA, SFB, D, alpha, N, K); + return (int) cudaGetLastError(); +} + +template +int launch_m(const uint8_t * A, const uint8_t * B, const uint8_t * SFA, + const uint8_t * SFB, float * D, float alpha, int M, int N, int K, + bool pdl, cudaStream_t stream) { + switch (M) { + case 1: return launch(A, B, SFA, SFB, D, alpha, N, K, pdl, stream); + case 2: return launch(A, B, SFA, SFB, D, alpha, N, K, pdl, stream); + case 3: return launch(A, B, SFA, SFB, D, alpha, N, K, pdl, stream); + default: return launch(A, B, SFA, SFB, D, alpha, N, K, pdl, stream); + } +} + +} // namespace + +int fp4_w4a4_mma_sm120_warpsplit_mrows_f32out( + const void * A_packed, const void * B_packed, float * D, int M, int N, + int K, const void * SFA, const void * SFB, float alpha, int warps, + int stages, bool pdl, cudaStream_t stream) { + if (M < 1 || M > 4 || N % 8 != 0 || K % 64 != 0) return -1; + if ((K / 64) % warps != 0) return -1; + const uint8_t * A = (const uint8_t *) A_packed; + const uint8_t * B = (const uint8_t *) B_packed; + const uint8_t * sfa = (const uint8_t *) SFA; + const uint8_t * sfb = (const uint8_t *) SFB; + const int cfg = stages * 10 + warps; + switch (cfg) { + case 34: return launch_m<3, 4>(A, B, sfa, sfb, D, alpha, M, N, K, pdl, stream); + case 44: return launch_m<4, 4>(A, B, sfa, sfb, D, alpha, M, N, K, pdl, stream); + case 64: return launch_m<6, 4>(A, B, sfa, sfb, D, alpha, M, N, K, pdl, stream); + case 38: return launch_m<3, 8>(A, B, sfa, sfb, D, alpha, M, N, K, pdl, stream); + case 48: return launch_m<4, 8>(A, B, sfa, sfb, D, alpha, M, N, K, pdl, stream); + case 32: return launch_m<3, 2>(A, B, sfa, sfb, D, alpha, M, N, K, pdl, stream); + case 62: return launch_m<6, 2>(A, B, sfa, sfb, D, alpha, M, N, K, pdl, stream); + case 42: return launch_m<4, 2>(A, B, sfa, sfb, D, alpha, M, N, K, pdl, stream); + default: return -1; + } +} + +} // namespace gemm +} // namespace flash_rt diff --git a/csrc/kernels/fp4_w4a4_mma_warpsplit_mrows_f32out_sm120.cuh b/csrc/kernels/fp4_w4a4_mma_warpsplit_mrows_f32out_sm120.cuh new file mode 100644 index 00000000..92c34698 --- /dev/null +++ b/csrc/kernels/fp4_w4a4_mma_warpsplit_mrows_f32out_sm120.cuh @@ -0,0 +1,30 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Small-M (M = 1..4) warp-split-K NVFP4 W4A4 GEMV/GEMM for sm_120, f32 +// output. Next generation of fp4_w4a4_mma_warpsplit_mrows_sm120 (additive: +// that entry is unchanged): the row count is a template parameter instead +// of a kernel argument (a runtime M in the MMA hot loop costs measurable +// time even at M=1), the epilogue writes f32 directly, and the launch can +// join the caller's programmatic-dependent-launch (PDL) chain — hosts that +// overlap every launch lose ground to any kernel that breaks the chain. +// The SM120_16x8x64 blockscaled MMA atom computes a 16-row tile, so rows +// 2..M ride the same weight HBM traffic as M=1: token t occupies A-tile +// row t (smem +t*32B), SFA atom-layout row t (+t*16B), and C-fragment +// lanes 4t..4t+3. +#pragma once +#include + +namespace flash_rt { +namespace gemm { + +// A_packed (M, K/2) row-major e2m1 pairs, B_packed (N, K/2), D f32 (M, N). +// SFA: atom-layout scales for rows 0..M-1 of problem (M, K); SFB: atom +// layout for (N, K). M in 1..4, warps in {2,4,8}, stages in {3,4,6}, +// N % 8 == 0, K % 64 == 0, (K/64) % warps == 0. Returns 0 on success. +int fp4_w4a4_mma_sm120_warpsplit_mrows_f32out( + const void * A_packed, const void * B_packed, float * D, int M, int N, + int K, const void * SFA, const void * SFB, float alpha, int warps, + int stages, bool pdl, cudaStream_t stream); + +} // namespace gemm +} // namespace flash_rt diff --git a/csrc/quantize/f32_act_to_nvfp4_swizzled_mrows_sm120.cu b/csrc/quantize/f32_act_to_nvfp4_swizzled_mrows_sm120.cu new file mode 100644 index 00000000..f3ff9ab4 --- /dev/null +++ b/csrc/quantize/f32_act_to_nvfp4_swizzled_mrows_sm120.cu @@ -0,0 +1,72 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// See header. The kernel is compile-time specialized on the row count and +// optionally joins the caller's programmatic-dependent-launch (PDL) chain: +// hosts that overlap every launch (llama.cpp CUDA backend) lose measurable +// time to any kernel that breaks the chain. + +#include "f32_act_to_nvfp4_swizzled_mrows_sm120.cuh" + +namespace flash_rt { +namespace quantize { + +namespace { + +__device__ __forceinline__ void pdl_sync() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + cudaGridDependencySynchronize(); +#endif +} +__device__ __forceinline__ void pdl_lc() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +template +__global__ void f32_act_to_nvfp4_kernel( + const float * __restrict__ x, + uint2 * __restrict__ dst_packed, + uint8_t * __restrict__ dst_sfa, + int D, long long x_srow) { + pdl_lc(); pdl_sync(); +#pragma unroll + for (int r = blockIdx.x; r < MT; r += gridDim.x) // launch with grid = MT + f32_act_to_nvfp4_row(x + (size_t) r * x_srow, dst_packed, dst_sfa, D, r); +} + +template +int launch(const float * x, void * dst_packed, void * dst_sfa, + int D, long long x_srow, bool pdl, cudaStream_t stream) { + const dim3 grid(MT), block(256); + if (pdl) { + cudaLaunchAttribute attr{}; + attr.id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr.val.programmaticStreamSerializationAllowed = 1; + cudaLaunchConfig_t cfg{}; + cfg.gridDim = grid; cfg.blockDim = block; cfg.dynamicSmemBytes = 0; + cfg.stream = stream; cfg.attrs = &attr; cfg.numAttrs = 1; + return (int) cudaLaunchKernelEx(&cfg, f32_act_to_nvfp4_kernel<256, MT>, + x, (uint2 *) dst_packed, (uint8_t *) dst_sfa, D, x_srow); + } + f32_act_to_nvfp4_kernel<256, MT><<>>( + x, (uint2 *) dst_packed, (uint8_t *) dst_sfa, D, x_srow); + return (int) cudaGetLastError(); +} + +} // namespace + +int f32_act_to_nvfp4_swizzled_mrows( + const float * x, void * dst_packed, void * dst_sfa, + int D, int M, long long x_srow, bool pdl, cudaStream_t stream) { + if (D % 16 != 0 || M < 1 || M > 4) return -1; + switch (M) { + case 1: return launch<1>(x, dst_packed, dst_sfa, D, x_srow, pdl, stream); + case 2: return launch<2>(x, dst_packed, dst_sfa, D, x_srow, pdl, stream); + case 3: return launch<3>(x, dst_packed, dst_sfa, D, x_srow, pdl, stream); + default: return launch<4>(x, dst_packed, dst_sfa, D, x_srow, pdl, stream); + } +} + +} // namespace quantize +} // namespace flash_rt diff --git a/csrc/quantize/f32_act_to_nvfp4_swizzled_mrows_sm120.cuh b/csrc/quantize/f32_act_to_nvfp4_swizzled_mrows_sm120.cuh new file mode 100644 index 00000000..21ad81f3 --- /dev/null +++ b/csrc/quantize/f32_act_to_nvfp4_swizzled_mrows_sm120.cuh @@ -0,0 +1,96 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// F32 activation rows -> NVFP4 packed + swizzled SFA (Sm1xx atom layout), +// M rows in one launch (M = 1..4, speculative-decode verify batches). +// Framework-free: raw pointers + cudaStream_t; the device body is exposed +// so host-adapter kernels can inline the quantization into fused producers. +// Additive: new file + new entry point (the single-row bf16 weight/act +// quantizers are separate, older entries). +#pragma once +#include +#include +#include + +namespace flash_rt { +namespace quantize { + +// ---- device body (shared with fused producers) ---------------------------- + +__device__ __forceinline__ int nvfp4_sfa_offset_128x64(int row, int k, int dim) { + const int row_block = row >> 7; + const int row_in_block = row & 127; + const int k_block = k >> 6; + const int k_in_block = k & 63; + const int k_blocks = (dim + 63) >> 6; + return row_block * k_blocks * 512 + k_block * 512 + + (row_in_block & 31) * 16 + (row_in_block >> 5) * 4 + + (k_in_block >> 4); +} + +// activation-quant boundary convention (<=; the weight packers use strict <) +__device__ __forceinline__ uint8_t nvfp4_act_f32_to_e2m1(float x) { + uint8_t sign = (x < 0.f) ? 0x8u : 0x0u; + float ax = fabsf(x); + uint8_t mant; + if (ax <= 0.25f) mant = 0u; + else if (ax <= 0.75f) mant = 1u; + else if (ax <= 1.25f) mant = 2u; + else if (ax <= 1.75f) mant = 3u; + else if (ax <= 2.5f) mant = 4u; + else if (ax <= 3.5f) mant = 5u; + else if (ax <= 5.0f) mant = 6u; + else mant = 7u; + return sign | mant; +} + +// quantize one f32 row of length D into packed e2m1 + SFA; `row` selects the +// SFA atom-layout row and the packed output row (row-major, D/16 uint2). +// Callable from any single participating block of THREADS threads. +template +__device__ __forceinline__ void f32_act_to_nvfp4_row( + const float * __restrict__ x, + uint2 * __restrict__ dst_packed, + uint8_t * __restrict__ dst_sfa, + int D, int row = 0) { + const int n_blocks = D / 16; + uint2 * dst_row = dst_packed + (size_t) row * n_blocks; + for (int b = threadIdx.x; b < n_blocks; b += THREADS) { + float vals[16]; + float amax = 0.f; +#pragma unroll + for (int i = 0; i < 16; ++i) { + vals[i] = x[b * 16 + i]; + const float a = fabsf(vals[i]); + if (a > amax) amax = a; + } + float desired = amax / 6.f; + if (desired < 1e-12f) desired = 1e-12f; + __nv_fp8_e4m3 bs_q = __nv_fp8_e4m3(fmaxf(desired, 0.f)); + const float bs_dq = static_cast(bs_q); + dst_sfa[nvfp4_sfa_offset_128x64(row, b * 16, D)] = *reinterpret_cast(&bs_q); + const float inv_bs = 1.f / bs_dq; + uint2 out; + uint8_t * ob = reinterpret_cast(&out); +#pragma unroll + for (int p = 0; p < 8; ++p) { + const uint8_t lo = nvfp4_act_f32_to_e2m1(vals[2 * p] * inv_bs); + const uint8_t hi = nvfp4_act_f32_to_e2m1(vals[2 * p + 1] * inv_bs); + ob[p] = static_cast(lo | (hi << 4)); + } + dst_row[b] = out; + } +} + +// ---- host entry ------------------------------------------------------------ + +// Quantize M f32 rows (row t at x + t*x_srow) into packed [M, D/2] + SFA in +// the atom layout for problem rows 0..M-1. M in 1..4 (compile-time +// specialized; a runtime M in the hot loop costs measurable time). +// pdl: join the caller's programmatic-dependent-launch chain. +// Returns 0 on success. +int f32_act_to_nvfp4_swizzled_mrows( + const float * x, void * dst_packed, void * dst_sfa, + int D, int M, long long x_srow, bool pdl, cudaStream_t stream); + +} // namespace quantize +} // namespace flash_rt diff --git a/docs/pi05_thor.md b/docs/pi05_thor.md new file mode 100644 index 00000000..c9de95cf --- /dev/null +++ b/docs/pi05_thor.md @@ -0,0 +1,432 @@ +# Pi0.5 on Jetson AGX Thor — Usage and Performance + +Canonical guide for running Pi0.5 (PaliGemma vision-language encoder + +Gemma action-expert decoder) on Jetson AGX Thor (SM110). It covers the +four supported precision tiers, how to select them, what each one costs +and delivers, and how to reproduce every number. + +`pi05_thor_decoder_fp4_e2e.md` is the chronological engineering record +behind these results — individual sections there describe intermediate +states and are superseded by this document. + +--- + +## 1. What runs on the device + +| stage | shape | precision (default tier) | +|---|---|---| +| SigLIP vision tower, 27 layers | `num_views × 256` tokens, d=1152 | FFN NVFP4 + AWQ, attention FA4 | +| Encoder (PaliGemma), 18 layers | `num_views × 256 + prompt` tokens, d=2048, H=16384, GQA 8/1, head_dim 256 | 17 live FFNs NVFP4 + AWQ, attention-O NVFP4, QKV FP8, attention FA4 | +| Decoder (Gemma action expert), 18 layers × 10 denoise steps | 10 action tokens, d=1024, H=4096 | all four projections NVFP4, attention cuBLAS FP16 | + +One `infer()` call runs the vision tower, the encoder prefill, and ten +denoise steps of the action expert, and returns a `(10, 7)` action chunk. +The whole path is CUDA-graph captured; there is no Python in the hot loop. + +--- + +## 2. Requirements + +- Jetson AGX Thor, compute capability `(11, 0)`, MAXN power mode +- CUDA 13, PyTorch ≥ 2.10 +- CUTLASS v4.4.2 at `third_party/cutlass` (vendored or symlinked) +- FA4 (`flashrt_fa4.cute`) importable — required for the SigLIP and + encoder attention path used by every number below + +```bash +cmake -B build -S . -DGPU_ARCH=110 +cmake --build build -j8 # .so files land in flash_rt/ +``` + +Locked clocks are required for reproducible timing: + +```bash +sudo nvpmodel -m 0 && sudo jetson_clocks +``` + +--- + +## 3. Quick start + +```python +from flash_rt.frontends.torch.pi05_thor_fp4 import Pi05TorchFrontendThorFP4 + +pipe = Pi05TorchFrontendThorFP4( + checkpoint_dir, + num_views=3, + use_fa4=True, + # ---- production NVFP4 tier ---- + use_fp4_encoder_ffn=True, fp4_layers=tuple(range(17)), + use_awq=True, awq_alpha=0.8, use_p1_split_gu=True, + use_fp4_encoder_attn=True, # attention-O projections + use_fp4_siglip_ffn=True, # all 27 SigLIP FFNs + use_fp4_decoder=True, # all four decoder projections +) + +pipe.set_prompt("pick up the black bowl and place it on the plate") +pipe.calibrate(observations, percentile=99.9) # 8 real observations +actions = pipe.infer(observation)["actions"] # (10, 7) +``` + +`observation` is a dict with `image` (`(224, 224, 3)` uint8/float16), +`state`, and — depending on `num_views` — `wrist_image` and +`wrist_image_right`. + +The FP8 baseline is the same call on the base class: + +```python +from flash_rt.frontends.torch.pi05_thor import Pi05TorchFrontendThor +pipe = Pi05TorchFrontendThor(checkpoint_dir, num_views=3, use_fa4=True) +``` + +**Calibration is not optional.** The FP4 tiers derive activation scales +and AWQ per-channel weight scales from real observations; calibrating on +synthetic data or skipping it degrades accuracy well past the gates below. + +--- + +## 4. The four precision tiers + +All four share the same encoder and SigLIP configuration; they differ in +how the **decoder** stores weights and activations. + +| tier | selection | decoder weights | decoder activations | +|---|---|---|---| +| **FP8** | `Pi05TorchFrontendThor` (no FP4 flags) | E4M3, per-tensor static scales | E4M3 | +| **NVFP4** (default) | `use_fp4_decoder=True` | E2M1 + per-16 UE4M3 block scale, per-block MSE scale search | dynamic E2M1 + per-16 scale | +| **INT4** | `decoder_weight_format="e0m3"`, `decoder_act_format="e0m3"` | uniform INT4 (E0M3, ±0..7), scale = amax/7 | uniform INT4 | +| **INT4+RHT** | above plus `decoder_rht=True` | INT4 after a per-16 orthonormal Hadamard rotation | INT4, rotation fused into the quantize kernel | + +**NVFP4** uses NVIDIA's block-scaled FP4: a 4-bit E2M1 element with a +UE4M3 scale per 16 elements. Weight scales come from a per-block MSE +search; activation scales are computed at run time in the fused +norm/quantize kernels. + +**INT4 (E0M3)** exploits the SM110 tcgen05 block-scaled MMA: the +instruction descriptor's 3-bit element-format field selects +sign-magnitude uniform INT4 at *run time*, so the same tensor-core path +serves both formats with no binary patching. The packed and scale-factor +layouts are identical to NVFP4, so buffers are interchangeable. + +**INT4+RHT** rotates every 16-element block by an orthonormal Hadamard +matrix (H16/4, symmetric) — weights offline, activations fused into the +quantize kernel. The rotation is mathematically inert for the GEMM +(verified at kernel level and end to end) but gaussianizes the +per-block distribution, which suits a uniform grid. It buys accuracy at +roughly +0.35 ms. + +An unquantized FP16 path also exists (`use_fp8=False`). It is far too +slow for deployment but serves as the common accuracy reference in §5.3, +where FP8 itself is measured against it. + +--- + +## 5. Performance and accuracy + +### 5.1 How these numbers were taken + +Thor exhibits **day-scale whole-machine drift** (thermal / EMC state): +the same binary measured hours apart can differ by ~2.5 ms, and the FP8 +reference moves with it. Therefore: + +- absolute milliseconds are only comparable **within one measurement + batch**; +- **speedup against the same-run FP8 reference is the stable metric**; +- every table below comes from one batch (the FP8 references land within + 0.2 ms of each other at each view count). + +Each run is a separate process pair — FP8 child then FP4 child — over +eight real LIBERO observations with matched noise seeds, 20 warmup +iterations and 100 timed iterations, with locked clocks verified from +`/sys` at start. Within-run spread is small (p95 − p50 ≤ 0.09 ms). + +### 5.2 Latency and accuracy by tier + +Cosines are against the FP8 reference, per sample over the eight +observations. `raw` is the pre-unnormalization action tensor; `act` is +the final action chunk. + +**3 views** + +| tier | p50 (ms) | speedup | raw cos | raw min | act cos | act min | gates | +|---|---|---|---|---|---|---|---| +| FP8 (reference) | 46.62 | 1.000 | — | — | — | — | — | +| **NVFP4 (default)** | **31.98** | **1.458** | 0.99904 | 0.99766 | 0.99974 | 0.99944 | PASS | +| INT4 | 32.54 | 1.461 | 0.99838 | 0.99512 | 0.99961 | 0.99939 | PASS | +| INT4+RHT | 32.60 | 1.452 | 0.99918 | 0.99742 | 0.99983 | 0.99970 | PASS | + +**2 views** + +| tier | p50 (ms) | speedup | raw cos | raw min | act cos | act min | gates | +|---|---|---|---|---|---|---|---| +| FP8 (reference) | 38.50 | 1.000 | — | — | — | — | — | +| **NVFP4 (default)** | **27.25** | **1.413** | 0.99921 | 0.99803 | 0.99972 | 0.99916 | PASS | +| INT4 | 27.72 | 1.387 | 0.99879 | 0.99751 | 0.99965 | 0.99928 | PASS | +| INT4+RHT | 27.88 | 1.381 | 0.99941 | 0.99828 | 0.99977 | 0.99922 | PASS | + +**1 view** — the accuracy gates do not pass at any fully-quantized tier, +for reasons that are not implementation defects (§6). + +| tier | p50 (ms) | speedup | raw cos | raw min | act cos | act min | gates | +|---|---|---|---|---|---|---|---| +| FP8 (reference) | 32.65 | 1.000 | — | — | — | — | — | +| NVFP4 (default) | **22.98** | **1.421** | 0.99137 | 0.96502 | 0.99347 | 0.97112 | accuracy FAIL | +| INT4 | 23.65 | 1.381 | 0.99101 | 0.96112 | 0.99332 | 0.96827 | accuracy FAIL | +| INT4+RHT | 23.75 | 1.375 | 0.99222 | 0.96385 | 0.99406 | 0.97021 | accuracy FAIL | +| FP8 encoder + INT4+RHT decoder | 29.06 | 1.122 | 0.99974 | 0.99954 | 0.99990 | 0.99973 | accuracy PASS | + +Gates: `raw cos ≥ 0.995`, worst-sample `raw cos ≥ 0.995`, +`action cos ≥ 0.999`, worst-sample `action cos ≥ 0.995`, plus a latency +gate (3-view p50 ≤ 40 ms, 2-view p95 ≤ 40 ms). + +### 5.3 Cosine against a common FP16 reference + +The tables above measure each quantized tier against FP8, which leaves +FP8's own error unmeasured. Running the same protocol against the FP16 +path (`use_fp8=False`) puts every tier on one yardstick: + +| views | tier | raw cos | raw min | act cos | act min | +|---|---|---|---|---|---| +| 3 | FP8 | 0.99994 | 0.99992 | 0.99997 | 0.99995 | +| 3 | NVFP4 | 0.99913 | 0.99812 | 0.99976 | 0.99948 | +| 3 | INT4 | 0.99848 | 0.99518 | 0.99963 | 0.99936 | +| 3 | INT4+RHT | 0.99928 | 0.99742 | 0.99985 | 0.99972 | +| 2 | FP8 | 0.99995 | 0.99994 | 0.99998 | 0.99997 | +| 2 | NVFP4 | 0.99929 | 0.99830 | 0.99976 | 0.99931 | +| 2 | INT4 | 0.99887 | 0.99775 | 0.99969 | 0.99943 | +| 2 | INT4+RHT | 0.99950 | 0.99854 | 0.99982 | 0.99937 | +| 1 | FP8 | 0.99876 | **0.99421** | 0.99905 | 0.99529 | +| 1 | NVFP4 | 0.99187 | 0.96338 | 0.99380 | 0.96939 | +| 1 | INT4 | 0.99118 | 0.95940 | 0.99342 | 0.96644 | +| 1 | INT4+RHT | 0.99204 | 0.96243 | 0.99381 | 0.96856 | + +Two things follow. + +**FP8 is essentially exact at two and three views** (0.9999+ on every +metric), so using it as the reference in §5.2 costs nothing — those +numbers are within 1e-4 of the same measurement against FP16. + +The same harness also times each tier, which places the unquantized path +on the scale. One locked-clock batch at three views: + +| tier | p50 (ms) | p95 (ms) | vs FP16 | +|---|---|---|---| +| FP16 | 80.230 | 80.605 | 1.000 | +| FP8 | 47.325 | 49.418 | 1.695 | +| NVFP4 | **31.842** | 31.893 | **2.520** | +| INT4 | 32.646 | 32.718 | 2.458 | +| INT4+RHT | 32.545 | 32.596 | 2.465 | + +**At one view even FP8 loses its worst sample**, to 0.99421 — below the +0.995 gate that the quantized tiers also miss. FP8 differs from FP16 by +a very small perturbation, so a sample that moves this much under it is +not being broken by 4-bit quantization; it is sitting somewhere that any +perturbation moves it. See §6. + +Reproduce with `tests/bench_pi05_precision_vs_fp16.py` (one subprocess +per tier, same prompt / observations / seeds as the strict suite). + +### 5.4 Choosing a tier + +- **Default to NVFP4.** It is both the fastest tier and, on 3 views, has + the best worst-sample raw cosine. The decoder FFN fusion (§8) only + applies to NVFP4 weights, which is why the INT4 tiers now sit ~0.7 ms + behind. +- **INT4+RHT when accuracy matters most.** It leads on aggregate cosine + and on worst-sample action cosine at both view counts, for ~0.9 ms. +- **Plain INT4 has no niche today** — it is slower than NVFP4 and less + accurate than INT4+RHT. It exists because it is the base the rotation + is applied to, and because it demonstrates the runtime-descriptor path. +- **FP8** remains the reference for correctness comparisons and for any + deployment that cannot calibrate on real observations. + +--- + +## 6. One-view accuracy + +At one view the per-sample cosine gates fail for every quantized tier. +This is characterized, not open: + +- **FP8 fails the same gate at one view** (worst-sample raw cosine 0.99421 + against FP16, §5.3) while being exact to 0.9999 at two and three views. + FP8 is a far smaller perturbation than 4-bit quantization, so whatever + moves that sample is not a property of the FP4 kernels. +- Ablations flip **different** samples under different quantization + configurations (full FP4 flips sample 0; FP8-encoder + FP4-decoder + flips sample 3 instead). +- The failure mode is a whole-trajectory direction change in the worst + sample's dominant motion component, not a magnitude error. The gripper + dimension is exact (cos = 1.0). +- Decoder accuracy improvements move it monotonically: INT4+RHT lifts the + worst sample from 0.843 to 0.930. + +The reading: with one view the observation underdetermines some samples, +which sit near a decision boundary of the flow-matching velocity field, +and any small perturbation — including FP8's — selects the other branch. +Such a sample produces *a different valid action candidate*, which a +per-sample cosine gate cannot distinguish from an error. Task success +rate is the meaningful judge; that evaluation is out of scope here. + +A configuration that passes every **accuracy** gate at one view — FP8 +encoder with an INT4+RHT decoder — at 29.06 ms (8/8 samples, worst-sample +raw cosine 0.99954, worst-sample action cosine 0.99973): + +```bash +--num-views 1 \ +--encoder-fp4-layer-count 0 --siglip-ffn-fp4 0 --encoder-attn-o-fp4 0 \ +--decoder-weight-format e0m3 --decoder-act-format e0m3 --decoder-rht 1 +``` + +It still trips the suite's published-SOTA latency gate (which wants +≤ 28.5 ms at one view), so `result.json` reports `passed: false` with +every accuracy gate green. Keeping the encoder in FP8 is what buys the +fidelity: it costs 6.1 ms against the fully quantized tier. + +--- + +## 7. Knobs + +Constructor keyword / bench flag pairs. Defaults are the production tier. + +| knob | default | effect | +|---|---|---| +| `decoder_weight_format` / `--decoder-weight-format` | `nvfp4` | `nvfp4` or `e0m3` | +| `decoder_act_format` / `--decoder-act-format` | `nvfp4` | `e0m3` requires `e0m3` weights | +| `decoder_rht` / `--decoder-rht` | `False` | per-16 Hadamard rotation; requires `e0m3` activations | +| `decoder_fused_geglu` / `--decoder-fused-geglu` | `True` | fuse the decoder GeGLU into the gate_up GEMM epilogue (NVFP4 weights only) | +| `encoder_p1_combiner` / `--encoder-p1-combiner` | `epilogue_hw_nod` | `epilogue_hw_nod` (fused, compact store, collective D store elided), `epilogue_hw` (fused, compact store), `epilogue` (fused, full width — parity with the old path), `lut_native` (separate GEMMs + combiner kernel) | +| `use_fp4_encoder_attn_qkv` / `--encoder-attn-qkv-fp4` | `False` | implemented and passing, but that GEMM is not weight-bandwidth-bound, so FP4 only matches FP8 while costing an extra quantize step | +| `decoder_fused_attn` / `--decoder-fused-attn` | `False` | folds the seqused mask into softmax (bit-identical, one fewer launch). Only the fixed-shape state-prompt path takes the seqused kernels, which this suite does not exercise | +| `awq_alpha` / `--awq-alpha` | `0.8` | AWQ per-channel scale exponent | +| `encoder_down_variant`, `decoder_*_variant` | `7`, `10` | GEMM tile selection | + +**Tile selection warning.** Cluster-launch GEMM variants invert between +isolated and in-pipeline benchmarks on Thor: the isolated-best tile for +one projection cost +2.2 ms end to end, and larger clusters +11–14 ms. +Always A/B tiles inside the pipeline. + +--- + +## 8. Where the time goes + +Single-frame kernel trace, 3 views, NVFP4 tier +(`nsys --cuda-graph-trace=node` around one `cudaProfilerStart/Stop` +window). Frame 31.85 ms under the profiler; 99% is GPU kernel time. + +| component | ms/frame | share | +|---|---|---| +| block-scaled FP4 GEMM (incl. fused GeGLU epilogues) | 21.3 | 67% | +| normalization / AdaRMS (incl. 350 decoder AdaRMS calls) | 3.3 | 10% | +| FA4 attention (SigLIP + encoder) | 2.1 | 7% | +| decoder attention cuBLAS chain (QK^T / softmax / AV) | 1.9 | 6% | +| encoder FP8 GEMM (attention QKV / O) | 1.2 | 4% | +| activation quantize | 1.1 | 3% | +| RoPE / QKV split | 0.6 | 2% | +| other | 0.4 | 1% | + +Headroom is thin and mostly hard floors: + +1. The decoder's 9.7 ms of GEMM sits against a ~7.3–7.8 ms weight-bandwidth + floor; the gap is fixed overhead across 720 GEMM launches, which would + need a bespoke persistent mainloop to recover. +2. Decoder AdaRMS (350 × 2.92 µs) and RoPE (180 × 1.58 µs) are at the + kernel-launch floor. +3. The SigLIP attention projections were characterized with L2-sector + counters and ruled out: at M=768 they are L2-bandwidth/compute bound + (32.4 MB of L2 reads per qkv call at ~1.16 TB/s, with the measured + time bracketed between the weight-only and all-DRAM rooflines), so + FP4 conversion has no bandwidth dividend and the extra quantize step + makes it a net loss. AWQ for the SigLIP up-projection remains the one + accuracy-gated candidate. + +### Approaches measured and rejected + +- **Single-kernel decoder attention.** Implemented and numerically + validated, then measured at 5–7× the existing chain across three + schedule designs. The skinny attention shape leaves the GEMM work + tensor-core-bound (the two cuBLAS calls are ~1 µs of tensor-core math), + and per-(head, row) grids multiply KV re-reads past the L2 budget. + FlashAttention-4 at this shape measures 24.6 µs (head_dim 256 has no + KV-split path). Fuse the glue *between* GEMMs, not the GEMMs. +- **Full-width fused GeGLU epilogue.** The combiner kernel it removes is + exactly cancelled by the doubled weight streaming of the K-expanded + down projection. The half-width compact store is the form that wins; + with its unread D store elided outright (`epilogue_hw_nod`, a fork of + the SM100 epilogue collective with the same `is_destination_supported` + guards the SM90 and SM100 ptr-array collectives already carry) it is + the default: a five-leg alternating sandwich measures −0.57 ms + (32.513/32.550/32.547 vs 31.956/31.967, drift ≤ 0.037 ms). Two + measurement caveats worth keeping: the isolated kernel benchmark scored + the elision as a regression (one more entry for the tile-selection + warning above), and under `nsys --cuda-graph-trace=node` the two + variants converge entirely — the win only exists unprofiled, so + per-kernel traces cannot attribute it. +- **No-D-store decoder GeGLU.** The same elision applied to the decoder + tile is a wash (the dummy store there is 0.04 MB), so it ships opt-in + (`--decoder-fused-geglu-nod`) and stays off by default. + +--- + +## 9. Reproducing + +```bash +# kernel numerical contracts +pytest tests/test_pi05_fp4_fusion_kernels.py tests/test_pi05_decoder_fp4_kernels.py + +# every tier against a common FP16 reference (§5.3) +python tests/bench_pi05_precision_vs_fp16.py \ + --checkpoint \ + --fixture /libero_obs3v_n8.npz \ + --num-views 3 --output-dir + +# strict end-to-end suite (requires a clean tracked worktree) +python tests/bench_pi05_decoder_fp4_e2e.py \ + --checkpoint \ + --num-views 3 \ + --fixture /libero_obs3v_n8.npz \ + --output-dir + +# tier switches + --decoder-weight-format e0m3 --decoder-act-format e0m3 # INT4 + --decoder-weight-format e0m3 --decoder-act-format e0m3 --decoder-rht 1 # INT4+RHT +``` + +The fixture is an npz of eight real LIBERO observations +(keys `n`, `img_i`, `state_i`, `wrist_i`, `wrist_right_i`). The suite +writes `result.json` with per-iteration timings, the verified clock +state, per-gate verdicts, and `.so` SHA256s, plus the FP4 and FP8 action +tensors. + +Measurement discipline: + +- compare back-to-back within one batch; across batches use speedup; +- warm up at least 20 iterations before timing; +- one process, exclusive GPU, no concurrent load. + +--- + +## 10. Implementation map + +| area | files | +|---|---| +| frontend, tier selection, weight prep | `flash_rt/frontends/torch/pi05_thor_fp4.py` | +| benchmarks | `tests/bench_pi05_decoder_fp4_e2e.py` (strict E2E), `tests/bench_pi05_precision_vs_fp16.py` (common-reference accuracy) | +| decoder pipeline | `flash_rt/models/pi05/pipeline_thor.py` | +| SigLIP / encoder pipeline | `flash_rt/hardware/thor/shared_primitives_fp4.py` | +| attention dispatch (FA4, cuBLAS, seqused) | `flash_rt/hardware/thor/attn_backend.py` | +| NVFP4 / INT4 GEMM runners | `csrc/gemm/fp4/` | +| fused GeGLU store epilogue | `csrc/gemm/fp4/sm100_gelu_mul_blockscale_visitor.hpp` | +| fused norm / quantize / activation kernels | `csrc/fused_fp4/`, `csrc/quantize/` | +| E0M3 quantizer and activation kernels | `csrc/quantize/quantize_e0m3_sfa.cu`, `csrc/fused_fp4/pi05_e0m3_act.cu` | + +--- + +## 11. Known limitations + +- One-view accuracy gates fail at every quantized tier (§6); a passing + configuration exists at reduced speed. +- Task-level (rollout) validation is out of scope for this document. +- The benchmark requires `--checkpoint` and `--fixture` explicitly; there + are no default dataset paths. +- Numbers here are Thor-specific. The SM110 runtime-descriptor INT4 path + in particular has no equivalent on other architectures. diff --git a/flash_rt/api.py b/flash_rt/api.py index 9caa96fa..e6715c60 100644 --- a/flash_rt/api.py +++ b/flash_rt/api.py @@ -460,9 +460,10 @@ def load_model(checkpoint, framework="torch", num_views=2, autotune=3, awq_alpha: AWQ activation scaling exponent. ``None`` selects 0.8 for the encoder-FP4 + decoder-FP4 preset and 0.5 otherwise. encoder_p1_combiner: FP4 encoder split-GU combiner. ``None`` resolves - to the preset: ``"epilogue_hw"`` (fused GeGLU epilogue) with - ``use_fp4_decoder=True``, ``"lut_native"`` otherwise. - ``"direct"``, ``"lut"`` and ``"epilogue"`` remain available for + to the preset: ``"epilogue_hw_nod"`` (fused GeGLU epilogue with + the collective's D store elided) with ``use_fp4_decoder=True``, + ``"lut_native"`` otherwise. ``"direct"``, ``"lut"``, + ``"epilogue"`` and ``"epilogue_hw"`` remain available for explicit A/B runs. encoder_down_variant: Cutlass NVFP4 encoder Down GEMM variant (production default ``7``). @@ -653,7 +654,7 @@ def load_model(checkpoint, framework="torch", num_views=2, autotune=3, if use_fp4_siglip_ffn is None: use_fp4_siglip_ffn = bool(use_fp4_decoder) if encoder_p1_combiner is None: - encoder_p1_combiner = ("epilogue_hw" if use_fp4_decoder + encoder_p1_combiner = ("epilogue_hw_nod" if use_fp4_decoder else "lut_native") else: if fp4_layers is None: diff --git a/flash_rt/catalog/__init__.py b/flash_rt/catalog/__init__.py index af00f582..ec8a6fee 100644 --- a/flash_rt/catalog/__init__.py +++ b/flash_rt/catalog/__init__.py @@ -10,7 +10,7 @@ no kernels, no host adapters and no runtime machinery, and is readable without torch. Its consumers are: -- the native ggml host adapter (``hosts/ggml``), whose qualification +- the native ggml host adapter (FlashRT-llama.cpp), whose qualification gates check a binding against the catalog; - the runtime exporters, which serialise ``BindingSpec.manifest()``; - the FlashRT-Structures package, which attaches structures onto an diff --git a/flash_rt/catalog/bindings/jetson_pi_edge_pi05.yaml b/flash_rt/catalog/bindings/jetson_pi_edge_pi05.yaml new file mode 100644 index 00000000..05c37ad2 --- /dev/null +++ b/flash_rt/catalog/bindings/jetson_pi_edge_pi05.yaml @@ -0,0 +1,97 @@ +binding: jetson_pi_edge_pi05 +structure: vla_tick_pipeline + +# Host: ggml/llama.cpp server (PKU-SEC-Lab/Jetson-PI-Edge fork) running pi0.5 +# on Jetson AGX Thor (SM110) with the native ggml adapter +# (the FlashRT-llama.cpp overlay). Structure regions execute as fused +# subgraph windows matched inside ggml-cuda's graph evaluation; the mapping +# below is window <-> catalog structure. + +stages: + obs_encode: + seam: "clip batched ViT encode (SigLIP graph) + llama_encode prefix" + capture: cuda_graphs_keyed # llama.cpp keyed multi-graph CUDA graphs + outputs: + cond_features: "prefix KV cache + device-resident image embeddings" + action_denoise: + seam: "llama_decode denoise loop over the pi0 action-expert stack" + loop_steps: 10 + noise_window: "seeded host-side noise tensor (PI0 seed window)" + capture: cuda_graphs_keyed + +cadences: + observation: obs_encode + tick: [obs_encode, action_denoise] + replan: action_denoise + +coverage: + contract: complete_hot_path + hot_path: + - observation_inputs + - vision_encoder + - vision_projection + - image_embd_residency + - prefix_prefill + - prefix_kv + - denoise_control + - timestep_conditioning + - action_expert_transformer + - euler_update + - action_readout + segments: + - name: observation_inputs + stage: obs_encode + classification: host_stage + seam: "mtmd bitmap load, patchify, and state ingestion" + - name: vision_encoder + stage: obs_encode + classification: structure + seam: "SigLIP tower graph (flash-attention path, padded heads baked into repacked weights)" + structures: [vision_ffn, qkv_pack, attention_core, linear_proj] + - name: vision_projection + stage: obs_encode + classification: structure + seam: "multimodal projector head" + structures: [linear_proj] + - name: image_embd_residency + stage: obs_encode + classification: state_region + seam: "device-resident image embeddings (D2D into a persistent buffer, no host round-trip)" + structures: [cadence_static] + - name: prefix_prefill + stage: obs_encode + classification: structure + seam: "llama_encode over the prefix backbone layers" + structures: [decoder_ffn, qkv_pack, attention_core, linear_proj] + - name: prefix_kv + stage: obs_encode + classification: state_region + seam: "persistent prefix KV cache reused across denoise steps" + structures: [cadence_static] + - name: denoise_control + stage: action_denoise + classification: control + seam: "per-step row-index upload and loop synchronization" + - name: timestep_conditioning + stage: action_denoise + classification: state_region + seam: "precomputed per-step modulation table, selected by row indices in-graph" + structures: [adaln_producer] + - name: action_expert_transformer + stage: action_denoise + classification: structure + seam: "action-expert decoder layers (modulated norms, fused QKV+RoPE, GeGLU FFN, gated residuals)" + structures: [modnorm_qkv_chain, qkv_pack, attention_core, decoder_ffn, + linear_proj, norm_fused, adaln_producer] + - name: euler_update + stage: action_denoise + classification: host_stage + seam: "host-side Euler integration between denoise steps" + - name: action_readout + stage: action_denoise + classification: host_stage + seam: "action chunk fetch and denormalization" + +hosts: + jetson_pi_edge: + versions: "feat/flashrt-thor-kernels" diff --git a/flash_rt/catalog/bindings/llamacpp_qwen36_35b_sm120.yaml b/flash_rt/catalog/bindings/llamacpp_qwen36_35b_sm120.yaml new file mode 100644 index 00000000..86bf4814 --- /dev/null +++ b/flash_rt/catalog/bindings/llamacpp_qwen36_35b_sm120.yaml @@ -0,0 +1,169 @@ +binding: llamacpp_qwen36_35b_sm120 +structure: autoregressive_decode_pipeline + +# Host: stock llama.cpp (upstream master lineage) running Qwen3.6-35B-A3B +# UD-Q4_K_M on RTX 5090 (SM120) with the native ggml adapter +# (FlashRT-llama.cpp: fr_win_qwen36_sm120.cu). Structure +# regions execute as fused subgraph windows matched inside ggml-cuda's +# graph evaluation. Two quality tiers share one binary: the safe tier keeps +# the target head and all matched math bit-comparable to stock (24-chunk +# PPL receipt), the full tier adds the NVFP4 lm-head swap. The speculative +# form runs the pipeline's spec_draft stage through the model's trained MTP +# head (single head, 3 chained drafts/step, verify batches M = 4). + +stages: + prefill: + seam: "llama_decode prompt processing (stock kernels; adapter windows are decode-shaped and stay out)" + capture: none + decode: + seam: "llama_decode token generation, M in [1,4] (speculative verify batches)" + capture: cuda_graphs_keyed + token_select: + seam: "server sampler chain; greedy or speculative acceptance (backend sampling via -bs)" + capture: host_dependent + spec_draft: + seam: "MTP draft chain evals + M=4 h-extraction eval (common/speculative draft-mtp driver)" + capture: cuda_graphs_keyed + +cadences: + token: [decode, token_select] + spec_step: [decode, token_select, spec_draft] + +# Machine-readable model facts consumed by the ggml adapter (build-time +# header generation: tools/gen_binding_header.py in FlashRT-llama.cpp). This is the +# single source for every model-specific constant in the window code; the +# window logic itself is family-level. +host_params: + arch: qwen35moe + n_layer: 40 + layer_scan_max: 64 # matcher bound for blk.%d parsing + d_model: 2048 + regions: # qkv_pack instances served by the fused NVFP4 GEMV + - kind: 0 # GDN in_proj pack + layers: 30 + k: 2048 + members: # [tensor suffix, rows, row offset in the pack] + - {name: attn_qkv.weight, rows: 8192, row_off: 0, leader: true} + - {name: attn_gate.weight, rows: 4096, row_off: 8192} + - {name: ssm_alpha.weight, rows: 32, row_off: 12288} + - {name: ssm_beta.weight, rows: 32, row_off: 12320} + - kind: 1 # full-attention qkv pack + layers: 10 + k: 2048 + members: + - {name: attn_q.weight, rows: 8192, row_off: 0, leader: true} + - {name: attn_k.weight, rows: 512, row_off: 8192} + - {name: attn_v.weight, rows: 512, row_off: 8704} + gdn: # gated_delta_core instance dims + qkvz_rows: 12288 # qkv 8192 + z 4096 (in_proj GEMV span served to the cell) + qkv_rows: 8192 + z_rows: 4096 + conv_channels: 8192 + conv_window: 3 + conv_cache_row: 24576 # conv_channels * conv_window + state_size: 524288 # out_dim * head_dim (32 heads x 128 x 128) + out_dim: 4096 + head_dim: 128 + moe: # moe_expert_ffn instance (dims read from tensors at run time) + gate_exps: ffn_gate_exps + up_exps: ffn_up_exps + down_exps: ffn_down_exps + gate_shexp: ffn_gate_shexp.weight + up_shexp: ffn_up_shexp.weight + down_shexp: ffn_down_shexp.weight + gate_inp_shexp: ffn_gate_inp_shexp.weight + n_expert: 256 + draft_prefix: mtp_ # MTP draft graph node-name prefix + out_names: [ffn_moe_out, ffn_out] + resid_names: [l_out, mtp_post_ffn] + out_proj: # linear_proj instances (Q8_0 GEMV + residual) + names: [ssm_out.weight, attn_output.weight] + k: 4096 + head: # lm head linear_proj instance + name: output.weight + n_vocab: 248320 + draft_copy_type: q8_0 # the spec draft's head copy is stored Q8_0; + # the target head is not — this discriminates them + +coverage: + contract: complete_hot_path + hot_path: + - token_embed + - gdn_in_proj_pack + - gdn_cell_span + - attn_qkv_pack + - attention_full + - norms + - moe_expert_span + - out_proj + - recurrent_state + - lm_head + - token_select + segments: + - name: token_embed + stage: decode + classification: host_stage + seam: "embedding row gather (stock get_rows)" + - name: gdn_in_proj_pack + stage: decode + classification: structure + seam: "qkv|gate|alpha|beta sibling pack served from one NVFP4 W4A4 GEMV staging pass (30 GDN layers); qkv_pack leaf form, fp4 block-scaled input (variant value proposed, see pins)" + structures: [qkv_pack] + - name: gdn_cell_span + stage: decode + classification: structure + seam: "conv + gated-delta recurrence + gated norm + epilogue as one 4-launch region; decode_recurrent, l2 in kernel; M>1 verify batches under snapshot_per_token + replay_in_region" + structures: [gated_delta_core] + - name: attn_qkv_pack + stage: decode + classification: structure + seam: "attention q|k|v sibling pack, same GEMV (10 attention layers)" + structures: [qkv_pack] + - name: attention_full + stage: decode + classification: host_stage + seam: "flash-attention core runs on the host's own kernels (not taken over)" + - name: norms + stage: decode + classification: host_stage + seam: "rms/l2 norm chain outside the fused regions runs on stock kernels" + - name: moe_expert_span + stage: decode + classification: structure + seam: "K0 quant+meta / K1 gate|up|GLU / K2 down+combine over ggml-native K-quant blocks (format_native consumption); shared expert folded; routing fused for M>=2" + structures: [moe_expert_ffn] + - name: out_proj + stage: decode + classification: structure + seam: "ssm_out/attn_output Q8_0 GEMV + residual epilogue, consuming the GDN epilogue's q8 activation handoff" + structures: [linear_proj] + - name: recurrent_state + stage: decode + classification: state_region + seam: "conv/state caches updated in place at M=1, per-token snapshots + checkpoint replay under speculative verify" + - name: lm_head + stage: decode + classification: structure + seam: "full tier: NVFP4 W4A4 head (M<=4 one pass, weights from the BF16-checkpoint pack); safe tier: stock Q6_K mmvq" + structures: [linear_proj] + - name: token_select + stage: token_select + classification: host_stage + seam: "server sampler chain; backend sampling (-bs) keeps greedy/acceptance on device" + - name: prefill_stock + stage: prefill + classification: host_stage + seam: "prompt processing runs entirely on stock kernels (adapter windows are decode-shaped)" + - name: spec_draft_serving + stage: spec_draft + classification: structure + seam: "draft-side lm-head serving (the draft GGUF's Q8_0 head copy from the FP4 pack — acceptance-only by spec math) + the draft MTP layer's Q8_0 experts through the same MoE span" + structures: [linear_proj, moe_expert_ffn] + - name: host_graph_slots + stage: spec_draft + classification: host_stage + seam: "host-tree change (llama-context): per-shape graph slots keyed (n_tokens, gtype, has_embd), each owning its scheduler" + +hosts: + llamacpp_sm120: + versions: "upstream 95b8e33 + FlashRT hook guards (GGML_CUDA_FLASHRT_SM120)" diff --git a/flash_rt/catalog/structures/README.md b/flash_rt/catalog/structures/README.md new file mode 100644 index 00000000..60011398 --- /dev/null +++ b/flash_rt/catalog/structures/README.md @@ -0,0 +1,40 @@ +# Structure catalog + +## What this is + +A structure entry is a **local boundary expression**: what a computation +region is — its symbolic dimensions, inputs/outputs, weight slots, variant +semantics — plus an executable reference implementation that defines what +"the same computation" means. The catalog exists for two purposes: + +1. **Context alignment against native pipelines.** When an adapter or a + host integration is being built (or read), the catalog is the map that + says which region of the host graph corresponds to which structure and + under which variant — so N host implementations of one boundary can be + compared, ported, and reasoned about as one thing. +2. **Distribution boundary management on the torch side** — the frontend, + discovery, and swap machinery consume these boundaries to decide what + can be handed to an implementation and at what seam. + +## What this is not + +**The catalog adjudicates nothing.** No performance claims, no expected +wins, no negative results, no campaign case histories, no tuning guidance. +Every judgment of that kind is conditional on a model, a hardware +generation, a driver, and a host version, and it expires the moment any of +those move — the only arbiter of whether an implementation is correct or +faster is a test run against the live system, never a statement recorded +here. Dated results belong to campaign records and per-binding +qualification gates, which are re-established by re-running them, not by +being quoted. + +Practically, for an entry in this directory: + +- **Belongs here**: boundary math, dimension/stride contracts, variant + *semantics* (including correctness-critical ones such as state snapshot + or rollback semantics — properties of the computation itself), the + executable reference, version numbers. +- **Does not belong here**: throughput or latency numbers, "X was judged + negative/positive", hardware-specific observations, host-specific war + stories, anything phrased as a verdict. If it can go stale without this + file changing, it goes elsewhere. diff --git a/flash_rt/catalog/structures/autoregressive_decode_pipeline/structure.yaml b/flash_rt/catalog/structures/autoregressive_decode_pipeline/structure.yaml index 465b7adc..c0e9b50b 100644 --- a/flash_rt/catalog/structures/autoregressive_decode_pipeline/structure.yaml +++ b/flash_rt/catalog/structures/autoregressive_decode_pipeline/structure.yaml @@ -1,7 +1,7 @@ structure: autoregressive_decode_pipeline kind: stage_pipeline family: autoregressive_decode -version: 2 +version: 3 description: > Schedule-layer structure for autoregressive generation: optional host input preparation and modality encoding feed a causal prefill that @@ -56,13 +56,31 @@ stages: - {name: token_id, window: swap} - {name: stop, window: host} capture: host_dependent + - name: spec_draft + optional: true + cadence: token + description: > + Speculative drafting between decode steps: a draft model (or trained + MTP head chain) proposes n_draft tokens; the next decode step becomes + a verify batch over 1 + n_draft rows whose logits feed acceptance. + Draft state (KV / recurrent / conv) advances tentatively and is + rolled back to the accepted position. + inputs: + - {name: token_id, window: swap} + - {name: target_hidden, buffer: mutable, optional: true} # MTP h rows + - {name: draft_state, buffer: mutable} + outputs: + - {name: draft_tokens, window: swap} + - {name: draft_state, buffer: mutable} + capture: bucketed_or_eager -embedded_regions: [decoder_ffn, vision_ffn, qkv_pack, qk_norm_rope, qkv_rope, - attention_core, gated_delta_core, linear_proj, - patch_projection] +embedded_regions: [decoder_ffn, vision_ffn, moe_expert_ffn, qkv_pack, + qk_norm_rope, qkv_rope, attention_core, gated_delta_core, + linear_proj, patch_projection] state_regions: - {name: kv_state, writers: [prefill, decode], reader: decode} + - {name: draft_state, writers: [spec_draft, decode], reader: spec_draft, optional: true} - {name: graph_buckets, cadence: shape_or_position} conformance: @@ -71,6 +89,12 @@ conformance: - token_selection_is_outside_decoder_blocks - modality_features_have_request_cadence - graph_bucket_fallback_is_reported + # speculative (spec_draft stage present): + - spec_output_distribution_matches_target # emitted tokens equal target-only + # decoding; the draft affects + # speed, never output + - draft_state_rollback_is_explicit # verify batches leave every + # tentative state rollback-capable gates: parity: diff --git a/flash_rt/catalog/structures/gated_delta_core/structure.yaml b/flash_rt/catalog/structures/gated_delta_core/structure.yaml index 68b69ac4..f05fbced 100644 --- a/flash_rt/catalog/structures/gated_delta_core/structure.yaml +++ b/flash_rt/catalog/structures/gated_delta_core/structure.yaml @@ -1,5 +1,5 @@ structure: gated_delta_core -version: 1 +version: 2 description: > Stateful Gated DeltaNet recurrence over Q/K/V, log-decay and update strength. The structure owns the recurrent state transition and exposes the @@ -27,8 +27,16 @@ weights: [] variants: phase: [decode_recurrent, sequence_scan, wy_chunk] - state_update: [in_place, explicit_output] + state_update: [in_place, explicit_output, snapshot_per_token] + # snapshot_per_token: each token of a batch writes its own + # state snapshot and the source state slot stays pristine, + # so the host can roll back to any position qk_norm: [l2_in_kernel, host] + checkpoint: [none, replay_in_region] + # replay_in_region: checkpoint saves that the host graph + # places inside the region's span are part of the boundary + # and must be reproduced by the region (or the region must + # decline the whole span) calibration: points: [] diff --git a/flash_rt/catalog/structures/moe_expert_ffn/reference.py b/flash_rt/catalog/structures/moe_expert_ffn/reference.py new file mode 100644 index 00000000..acb8602e --- /dev/null +++ b/flash_rt/catalog/structures/moe_expert_ffn/reference.py @@ -0,0 +1,84 @@ +"""Ground-truth reference for the sparse MoE feed-forward block. + +Plainest possible PyTorch, never executed on a serving hot path. Routing +replicates the Qwen3.x convention exactly: softmax over all experts, +iterative top-k with lower-index tie-break, clamp, renormalize. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + + +def _route_softmax_topk_clamp_renorm( + logits: torch.Tensor, k_top: int, clamp_min: float = 1e-20 +) -> tuple[torch.Tensor, torch.Tensor]: + """Softmax over E -> iterative top-k (lower index wins ties) -> renorm.""" + probs = torch.softmax(logits.float(), dim=-1) + m, _ = probs.shape + ids = torch.empty(m, k_top, dtype=torch.int64, device=logits.device) + vals = torch.empty(m, k_top, dtype=torch.float32, device=logits.device) + work = probs.clone() + for j in range(k_top): + # argmax returns the first (lowest-index) maximum, matching the + # host convention this structure binds to. + idx = work.argmax(dim=-1) + ids[:, j] = idx + vals[:, j] = work.gather(-1, idx[:, None]).squeeze(-1) + work.scatter_(-1, idx[:, None], float("-inf")) + vals = vals.clamp_min(clamp_min) + vals = vals / vals.sum(dim=-1, keepdim=True) + return ids, vals + + +def moe_expert_ffn_ref( + x: torch.Tensor, + w_gate_exps: torch.Tensor, + w_up_exps: torch.Tensor, + w_down_exps: torch.Tensor, + *, + w_router: torch.Tensor | None = None, + expert_ids: torch.Tensor | None = None, + expert_weights: torch.Tensor | None = None, + k_top: int = 8, + w_gate_shexp: torch.Tensor | None = None, + w_up_shexp: torch.Tensor | None = None, + w_down_shexp: torch.Tensor | None = None, + w_gate_inp_shexp: torch.Tensor | None = None, + residual: torch.Tensor | None = None, +) -> torch.Tensor: + """Routed expert GLU-FFN with optional sigmoid-gated shared expert. + + Either ``w_router`` (fused routing) or ``expert_ids``/``expert_weights`` + (external routing) must be provided. + """ + if expert_ids is None: + assert w_router is not None, "fused routing needs w_router" + logits = x.float() @ w_router.float() + expert_ids, expert_weights = _route_softmax_topk_clamp_renorm(logits, k_top) + assert expert_weights is not None + + m = x.shape[0] + y = torch.zeros(m, w_down_exps.shape[-1], dtype=torch.float32, device=x.device) + for t in range(m): + xt = x[t].float() + for j in range(expert_ids.shape[1]): + e = int(expert_ids[t, j]) + g = xt @ w_gate_exps[e].float() + u = xt @ w_up_exps[e].float() + h = F.silu(g) * u + y[t] += float(expert_weights[t, j]) * (h @ w_down_exps[e].float()) + + if w_gate_shexp is not None: + assert w_up_shexp is not None and w_down_shexp is not None + assert w_gate_inp_shexp is not None + for t in range(m): + xt = x[t].float() + sig = torch.sigmoid(xt @ w_gate_inp_shexp.float()) + h = F.silu(xt @ w_gate_shexp.float()) * (xt @ w_up_shexp.float()) + y[t] += sig * (h @ w_down_shexp.float()) + + if residual is not None: + y = y + residual.float() + return y.to(x.dtype) diff --git a/flash_rt/catalog/structures/moe_expert_ffn/structure.yaml b/flash_rt/catalog/structures/moe_expert_ffn/structure.yaml new file mode 100644 index 00000000..eb38bb04 --- /dev/null +++ b/flash_rt/catalog/structures/moe_expert_ffn/structure.yaml @@ -0,0 +1,67 @@ +structure: moe_expert_ffn +version: 1 +description: > + Sparse mixture-of-experts feed-forward block: route x[M,D] over E experts, + gather the top-k expert gate/up/down triples, activation(gate) * up per + expert, down projection, routing-weighted combine, plus an optional + always-on shared expert whose output is sigmoid-gated by its own scalar + projection. The routing decision (softmax/top-k/renorm) is part of the + boundary under the fused variant; under the external variant the block + consumes precomputed ids/weights. Per-token expert sets are independent, + so token batches (M > 1) are part of the boundary, not a special case. + +reference: + module: moe_expert_ffn.reference + entrypoint: moe_expert_ffn_ref + +boundary: + symbolic_dims: [M, D, F, E, K_top] + inputs: + - {name: x, dims: [M, D], dtype: "@binding"} + - {name: expert_ids, dims: [M, K_top], dtype: int32, optional: true} # external routing + - {name: expert_weights, dims: [M, K_top], dtype: "@binding", optional: true} + - {name: residual, dims: [M, D], dtype: "@binding", optional: true} + outputs: + - {name: y, dims: [M, D], dtype: "@binding"} + +weights: + - {slot: w_router, dims: [D, E], optional: true} # fused routing only + - {slot: w_gate_exps, dims: [E, D, F]} + - {slot: w_up_exps, dims: [E, D, F]} + - {slot: w_down_exps, dims: [E, F, D]} + - {slot: w_gate_shexp, dims: [D, F_sh], optional: true} + - {slot: w_up_shexp, dims: [D, F_sh], optional: true} + - {slot: w_down_shexp, dims: [F_sh, D], optional: true} + - {slot: w_gate_inp_shexp, dims: [D], optional: true} # sigmoid gate scalar + +variants: + routing: [external, fused] # fused = router GEMM + exact top-k in-boundary + topk_norm: [softmax_topk_clamp_renorm] # softmax over E, iterative top-k with + # lower-index tie-break, clamp, renormalize + shared_expert: [none, sigmoid_gated] + activation: [silu] + weight_consumption: [dequant_gemm, format_native] + # format_native = expert weights are + # consumed in the host's own storage + # format, activations quantized to the + # host's own convention — no repack, no + # dequant materialization + batch: [m1, m_le_4] + +calibration: + points: [x] + +gates: + parity: + metrics: [cosine, max_abs, p99_abs] + data: real_distribution + +boundary_notes: + - "expert_ids/expert_weights rows may be strided views of a sort output; + the row stride is part of the boundary, never assumed dense" + - "the routing decision under the fused variant is part of the boundary, + including the tie-break and normalization semantics" + - "per-token expert sets are independent: batch variants of the same token + denote the same computation row for row" + - "the combine (weighted sum + shared expert + residual) is part of the + boundary" diff --git a/flash_rt/frontends/torch/pi05_thor_fp4.py b/flash_rt/frontends/torch/pi05_thor_fp4.py index 930dae9d..28958849 100644 --- a/flash_rt/frontends/torch/pi05_thor_fp4.py +++ b/flash_rt/frontends/torch/pi05_thor_fp4.py @@ -84,7 +84,7 @@ def __init__(self, checkpoint_dir, num_views: int = 2, awq_alpha: float = 0.5, awq_calib_iters: int = 8, use_p1_split_gu: bool = False, - encoder_p1_combiner: str = "epilogue_hw", + encoder_p1_combiner: str = "epilogue_hw_nod", encoder_down_variant: int = 7, encoder_down_x_variant: int = 6, decoder_qkv_variant: int = 10, @@ -95,6 +95,7 @@ def __init__(self, checkpoint_dir, num_views: int = 2, decoder_act_format: str = "nvfp4", decoder_fused_attn: bool = False, decoder_fused_geglu: bool = True, + decoder_fused_geglu_nod: bool = False, decoder_rht: bool = False, use_fp8: bool = True, state_prompt_mode: str = "exact", @@ -147,11 +148,12 @@ def __init__(self, checkpoint_dir, num_views: int = 2, "fp4_layers contains non-live encoder FFN layers " f"{invalid_layers}; valid layers are [0, {self.Le - 2}]") if encoder_p1_combiner not in ("direct", "lut", "lut_native", - "epilogue", "epilogue_hw"): + "epilogue", "epilogue_hw", + "epilogue_hw_nod"): raise ValueError( "encoder_p1_combiner must be 'direct', 'lut', " - "'lut_native', 'epilogue', or 'epilogue_hw', got " - f"{encoder_p1_combiner!r}") + "'lut_native', 'epilogue', 'epilogue_hw', or " + f"'epilogue_hw_nod', got {encoder_p1_combiner!r}") self.encoder_p1_combiner = encoder_p1_combiner self.encoder_down_variant = int(encoder_down_variant) # Tile variant for the K-expanded Down GEMM of the 'epilogue' P1 @@ -189,6 +191,10 @@ def __init__(self, checkpoint_dir, num_views: int = 2, # the gate_up GEMM + GeGLU-quantize kernel (nvfp4 weights only). self.decoder_fused_geglu = (bool(decoder_fused_geglu) and decoder_weight_format == 'nvfp4') + # No-D-store variant of the fused decoder GeGLU (same contract, + # the collective's own D store elided). + self.decoder_fused_geglu_nod = (bool(decoder_fused_geglu_nod) + and self.decoder_fused_geglu) if self._fp4_layers: if not _HAS_FP4: @@ -639,7 +645,8 @@ def get(k): g_fp16 = gu_fp16[:He, :].contiguous() u_fp16 = gu_fp16[He:, :].contiguous() if self.encoder_p1_combiner in ('epilogue', - 'epilogue_hw'): + 'epilogue_hw', + 'epilogue_hw_nod'): # Fused-epilogue paths: single interleaved gate/up # weight (down AWQ inv_s folded into the up rows). # Full-width 'epilogue' additionally needs the @@ -1119,9 +1126,12 @@ def _alloc_fp4_scratch_for_Se(self, Se: int): # Fused-epilogue path: one full-width [Se, 2H] FP4 buffer # between the interleaved GEMM and the K-expanded Down. self._fp4_p1_il = FP4Buffer(Se, 2 * He, device='cuda') - elif self.encoder_p1_combiner == 'epilogue_hw': + elif self.encoder_p1_combiner in ('epilogue_hw', + 'epilogue_hw_nod'): # Half-width path writes Down's stock input scratch; the - # collective's D output lands in this reusable dummy. + # collective's D output lands in this reusable dummy + # ('epilogue_hw_nod' never writes it but the host-side TMA + # descriptor still needs a real pointer). self._fp4_p1_dummy = torch.zeros( Se, He, dtype=torch.uint8, device='cuda') else: @@ -1150,7 +1160,8 @@ def _alloc_fp4_scratch_for_Se(self, Se: int): self._fp4_scratch_dict['p1_il_p4'] = self._fp4_p1_il.packed.data_ptr() self._fp4_scratch_dict['p1_il_sfa'] = self._fp4_p1_il.sfa.data_ptr() self._fp4_scratch_dict['variant_dn_x'] = self.encoder_down_x_variant - elif self.encoder_p1_combiner == 'epilogue_hw': + elif self.encoder_p1_combiner in ('epilogue_hw', + 'epilogue_hw_nod'): self._fp4_scratch_dict['p1_dummy'] = self._fp4_p1_dummy.data_ptr() else: self._fp4_scratch_dict['p1_gate_p4'] = self._fp4_p1_gate.packed.data_ptr() @@ -1402,6 +1413,7 @@ def _capture_enc_ae_graph(self): ae_dims['act_format'] = self.decoder_act_format ae_dims['rht'] = self.decoder_rht ae_dims['fused_geglu'] = self.decoder_fused_geglu + ae_dims['fused_geglu_nod'] = self.decoder_fused_geglu_nod if self._attn is not None: # Fold the decoder seqused mask into the softmax kernel. self._attn.use_fused_softmax = self.decoder_fused_attn diff --git a/flash_rt/hardware/thor/shared_primitives_fp4.py b/flash_rt/hardware/thor/shared_primitives_fp4.py index b7caf4ce..bec72e67 100644 --- a/flash_rt/hardware/thor/shared_primitives_fp4.py +++ b/flash_rt/hardware/thor/shared_primitives_fp4.py @@ -119,7 +119,7 @@ def encoder_forward_with_fp4_subset(gemm, fvk, fvk_fp4, bufs, weights, dims, p1_il_p4 = fp4_scratch['p1_il_p4'] p1_il_sfa = fp4_scratch['p1_il_sfa'] variant_dn_x = fp4_scratch['variant_dn_x'] - elif p1_combiner == 'epilogue_hw': + elif p1_combiner in ('epilogue_hw', 'epilogue_hw_nod'): p1_dummy = fp4_scratch['p1_dummy'] else: p1_gate_p4 = fp4_scratch['p1_gate_p4'] @@ -251,8 +251,14 @@ def encoder_forward_with_fp4_subset(gemm, fvk, fvk_fp4, bufs, weights, dims, # 'epilogue_hw' quantizes at compact granularity in the # epilogue and writes Down's stock input buffer directly. w_il = fp4_weights[l]['gu_il'] - if p1_combiner == 'epilogue_hw': - _check(fvk_fp4.cutlass_fp4_gemm_geglu_il_hw( + if p1_combiner in ('epilogue_hw', 'epilogue_hw_nod'): + # 'epilogue_hw_nod' elides the collective's own D + # store (p1_dummy is never written); otherwise the + # two are identical. + hw_gemm = (fvk_fp4.cutlass_fp4_gemm_geglu_il_hw_nod + if p1_combiner == 'epilogue_hw_nod' + else fvk_fp4.cutlass_fp4_gemm_geglu_il_hw) + _check(hw_gemm( sc_gu.packed.data_ptr(), sc_gu.sfa.data_ptr(), w_il['packed'].data_ptr(), w_il['sfb'].data_ptr(), p1_dummy, diff --git a/flash_rt/models/pi05/pipeline_thor.py b/flash_rt/models/pi05/pipeline_thor.py index 181815c8..3e0bc5a0 100644 --- a/flash_rt/models/pi05/pipeline_thor.py +++ b/flash_rt/models/pi05/pipeline_thor.py @@ -287,6 +287,7 @@ def decoder_forward_fp4(ctx, fvk, fvk_fp4, bufs, weights, dims, stream=0, *, act_format = str(dims.get('act_format', 'nvfp4')) rht = 1 if dims.get('rht') else 0 fused_geglu = bool(dims.get('fused_geglu')) and weight_format == 'nvfp4' + fused_geglu_nod = bool(dims.get('fused_geglu_nod')) and fused_geglu act_e0m3 = act_format == 'e0m3' if act_format not in ('nvfp4', 'e0m3'): raise ValueError( @@ -429,7 +430,10 @@ def dec_gemm(variant, *args): # gelu(gate)*up per column pair and writes the Down input # (hid_fp4/hid_sfa) directly; the gate_up fp16 output and # the GeGLU-quantize kernel disappear. - rc = fvk_fp4.cutlass_fp4_gemm_geglu_il_hw_v10( + geglu_v10 = (fvk_fp4.cutlass_fp4_gemm_geglu_il_hw_nod_v10 + if fused_geglu_nod + else fvk_fp4.cutlass_fp4_gemm_geglu_il_hw_v10) + rc = geglu_v10( xn_fp4, xn_sfa, weights['gwil_fp4'][l], weights['gwil_sfb'][l], weights['gu_dummy'], hid_fp4, hid_sfa, diff --git a/tests/bench_pi05_decoder_fp4_e2e.py b/tests/bench_pi05_decoder_fp4_e2e.py index d2ef3419..fc6b4f12 100644 --- a/tests/bench_pi05_decoder_fp4_e2e.py +++ b/tests/bench_pi05_decoder_fp4_e2e.py @@ -47,7 +47,7 @@ # configuration the public API reproduces. PUBLIC_API_PRESET = { "encoder_gu_mode": "p1", - "encoder_p1_combiner": "epilogue_hw", + "encoder_p1_combiner": "epilogue_hw_nod", "encoder_down_variant": 7, "encoder_down_x_variant": 6, "decoder_gate_up_variant": 10, @@ -56,6 +56,7 @@ "decoder_rht": 0, "decoder_fused_attn": 0, "decoder_fused_geglu": 1, + "decoder_fused_geglu_nod": 0, "awq_alpha": 0.8, "encoder_attn_o_fp4": 1, "encoder_attn_qkv_fp4": 0, @@ -146,8 +147,9 @@ def main() -> int: "--encoder-gu-mode", choices=("p1", "merged"), default="p1") parser.add_argument( "--encoder-p1-combiner", - choices=("direct", "lut", "lut_native", "epilogue", "epilogue_hw"), - default="epilogue_hw") + choices=("direct", "lut", "lut_native", "epilogue", "epilogue_hw", + "epilogue_hw_nod"), + default="epilogue_hw_nod") parser.add_argument("--encoder-down-variant", type=int, default=7) parser.add_argument("--encoder-down-x-variant", type=int, default=6) parser.add_argument("--decoder-gate-up-variant", type=int, default=10) @@ -175,6 +177,10 @@ def main() -> int: "--decoder-fused-geglu", type=int, choices=(0, 1), default=1, help="Fused-epilogue decoder FFN: one interleaved GeGLU GEMM " "replaces the gate_up GEMM + GeGLU kernel (nvfp4 weights only)") + parser.add_argument( + "--decoder-fused-geglu-nod", type=int, choices=(0, 1), default=0, + help="No-D-store variant of the fused decoder GeGLU (requires " + "--decoder-fused-geglu 1)") parser.add_argument("--awq-alpha", type=float, default=0.8) parser.add_argument( "--encoder-attn-o-fp4", type=int, choices=(0, 1), default=1, @@ -291,6 +297,7 @@ def main() -> int: decoder_rht=bool(args.decoder_rht), decoder_fused_attn=bool(args.decoder_fused_attn), decoder_fused_geglu=bool(args.decoder_fused_geglu), + decoder_fused_geglu_nod=bool(args.decoder_fused_geglu_nod), use_fp4_decoder=True, use_fa4=True, use_fp4_encoder_attn=bool(args.encoder_attn_o_fp4), @@ -397,6 +404,8 @@ def main() -> int: "decoder_rht": bool(args.decoder_rht), "decoder_fused_attn": bool(args.decoder_fused_attn), "decoder_fused_geglu": bool(args.decoder_fused_geglu), + "decoder_fused_geglu_nod": bool( + args.decoder_fused_geglu_nod), "decoder_gate_up_variant": args.decoder_gate_up_variant, "attention": "fa4_siglip_encoder", } @@ -463,6 +472,7 @@ def main() -> int: "--decoder-rht", str(args.decoder_rht), "--decoder-fused-attn", str(args.decoder_fused_attn), "--decoder-fused-geglu", str(args.decoder_fused_geglu), + "--decoder-fused-geglu-nod", str(args.decoder_fused_geglu_nod), "--awq-alpha", str(args.awq_alpha), "--encoder-fp4-layer-count", str(args.encoder_fp4_layer_count), "--encoder-attn-o-fp4", str(args.encoder_attn_o_fp4), diff --git a/tests/bench_pi05_precision_vs_fp16.py b/tests/bench_pi05_precision_vs_fp16.py new file mode 100644 index 00000000..4a2bf2cf --- /dev/null +++ b/tests/bench_pi05_precision_vs_fp16.py @@ -0,0 +1,240 @@ +"""Measure every Pi0.5 precision tier against a common FP16 reference. + +The strict E2E suite (``bench_pi05_decoder_fp4_e2e.py``) reports each +quantized tier's cosine against FP8, which leaves FP8's own deviation +unmeasured. This harness runs FP16, FP8 and the three quantized tiers +through the identical observation / prompt / seed protocol and reports +every tier against FP16, so the tiers sit on one yardstick. + +One model per subprocess: Thor cannot hold two of these pipelines in a +single process. + + python tests/bench_pi05_precision_vs_fp16.py \ + --checkpoint \ + --fixture /libero_obs3v_n8.npz \ + --num-views 3 --output-dir +""" +import argparse +import json +import os +from pathlib import Path +import statistics +import subprocess +import sys +import time + +import numpy as np + +PROMPT_TOKENS = [ + 2, 18075, 908, 573, 3118, 3963, 578, + 2040, 665, 575, 573, 24655, 108, +] + +TIERS = { + # name: extra kwargs for the FP4 frontend (None => FP8/FP16 base class) + "fp16": None, + "fp8": None, + "nvfp4": {}, + "int4": {"decoder_weight_format": "e0m3", "decoder_act_format": "e0m3"}, + "int4rht": {"decoder_weight_format": "e0m3", "decoder_act_format": "e0m3", + "decoder_rht": True}, +} + + +def build_pipe(mode, checkpoint, num_views): + if mode in ("fp16", "fp8"): + from flash_rt.frontends.torch.pi05_thor import Pi05TorchFrontendThor + return Pi05TorchFrontendThor( + checkpoint, num_views=num_views, autotune=3, use_fa4=True, + use_fp8=(mode == "fp8")) + from flash_rt.frontends.torch.pi05_thor_fp4 import Pi05TorchFrontendThorFP4 + return Pi05TorchFrontendThorFP4( + checkpoint, num_views=num_views, autotune=3, + use_fp4_encoder_ffn=True, fp4_layers=tuple(range(17)), + use_awq=True, awq_alpha=0.8, use_p1_split_gu=True, + use_fp4_decoder=True, use_fa4=True, + use_fp4_encoder_attn=True, use_fp4_siglip_ffn=True, + **TIERS[mode]) + + +def load_observations(fixture, num_views): + data = np.load(fixture) + count = int(data["n"]) + observations = [] + for index in range(count): + obs = {"image": data[f"img_{index}"], "state": data[f"state_{index}"]} + if num_views >= 2: + obs["wrist_image"] = data[f"wrist_{index}"] + if num_views == 3: + obs["wrist_image_right"] = data[f"wrist_right_{index}"] + observations.append(obs) + return observations + + +def run_child(args): + import torch # noqa: F401 (import side effects before pipeline import) + + pipe = build_pipe(args.child_mode, args.checkpoint, args.num_views) + observations = load_observations(args.fixture, args.num_views) + + pipe.set_prompt(PROMPT_TOKENS) + pipe.calibrate(observations, percentile=99.9, verbose=False) + for index in range(args.warmup): + pipe.infer(observations[index % len(observations)]) + + raw_outputs, action_outputs = [], [] + for index, observation in enumerate(observations): + np.random.seed(args.seed + index) + output = pipe.infer(observation) + raw_outputs.append(pipe._g_noise.float().cpu().numpy()) + action_outputs.append(output["actions"]) + + if args.cuda_profile: + torch.cuda.synchronize() + torch.cuda.cudart().cudaProfilerStart() + pipe.infer(observations[0]) + torch.cuda.synchronize() + torch.cuda.cudart().cudaProfilerStop() + + latencies = [] + for index in range(args.iters): + start = time.perf_counter() + pipe.infer(observations[index % len(observations)]) + latencies.append((time.perf_counter() - start) * 1000.0) + torch.cuda.synchronize() + + out_dir = Path(args.output_dir) + out_dir.mkdir(parents=True, exist_ok=True) + out = out_dir / f"{args.child_mode}_actions.npz" + np.savez(out, raw=np.stack(raw_outputs), actions=np.stack(action_outputs)) + result = { + "tier": args.child_mode, + "num_views": args.num_views, + "iters": args.iters, + "p50_ms": statistics.median(latencies), + "p95_ms": float(np.percentile(latencies, 95)), + "min_ms": min(latencies), + "max_ms": max(latencies), + } + print("__CHILD_RESULT__ " + json.dumps(result, sort_keys=True), flush=True) + return 0 + + +def cosines(lhs_all, rhs_all): + per_sample = [] + for index in range(lhs_all.shape[0]): + lhs = lhs_all[index].reshape(-1) + rhs = rhs_all[index].reshape(-1) + per_sample.append(float( + lhs @ rhs / (np.linalg.norm(lhs) * np.linalg.norm(rhs) + 1e-12))) + lhs = lhs_all.reshape(-1) + rhs = rhs_all.reshape(-1) + aggregate = float( + lhs @ rhs / (np.linalg.norm(lhs) * np.linalg.norm(rhs) + 1e-12)) + return aggregate, min(per_sample) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--child-mode", choices=sorted(TIERS)) + parser.add_argument("--checkpoint", required=True) + parser.add_argument("--fixture", required=True) + parser.add_argument("--num-views", type=int, choices=(1, 2, 3), default=3) + parser.add_argument("--output-dir", required=True) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--iters", type=int, default=100) + parser.add_argument( + "--cuda-profile", action="store_true", + help="capture one stable-state infer between " + "cudaProfilerStart/Stop") + parser.add_argument("--seed", type=int, default=20260725) + parser.add_argument("--modes", default="fp16,fp8,nvfp4,int4,int4rht") + args = parser.parse_args() + + # Latency is only meaningful under the same locked-clock discipline the + # strict suite enforces; reuse its verifier rather than duplicating it. + sys.path.insert(0, str(Path(__file__).resolve().parent)) + from bench_pi05_decoder_fp4_e2e import machine_state + machine_state() + + if args.child_mode is not None: + return run_child(args) + + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + modes = args.modes.split(",") + latency = {} + for mode in modes: + command = [ + sys.executable, str(Path(__file__).resolve()), + "--child-mode", mode, + "--checkpoint", args.checkpoint, + "--fixture", args.fixture, + "--num-views", str(args.num_views), + "--output-dir", str(output_dir), + "--warmup", str(args.warmup), + "--iters", str(args.iters), + "--seed", str(args.seed), + ] + child = subprocess.run(command, check=False, capture_output=True, + text=True, env=os.environ.copy()) + if child.returncode != 0: + raise RuntimeError( + f"{mode} child failed rc={child.returncode}\n" + f"{child.stdout[-2000:]}\n{child.stderr[-4000:]}") + lines = [line for line in child.stdout.splitlines() + if line.startswith("__CHILD_RESULT__")] + if len(lines) != 1: + raise RuntimeError( + f"expected one result line from {mode}, got {len(lines)}") + latency[mode] = json.loads(lines[0].split(" ", 1)[1]) + print(f" {mode} done p50={latency[mode]['p50_ms']:.3f} ms", + flush=True) + + ref = np.load(output_dir / "fp16_actions.npz") + raw_ref = ref["raw"].astype(np.float64) + act_ref = ref["actions"].astype(np.float64) + + rows = [] + for mode in modes: + if mode == "fp16": + continue + cur = np.load(output_dir / f"{mode}_actions.npz") + raw_cos, raw_min = cosines(cur["raw"].astype(np.float64), raw_ref) + act_cos, act_min = cosines(cur["actions"].astype(np.float64), act_ref) + rows.append((mode, raw_cos, raw_min, act_cos, act_min)) + + ref_p50 = latency["fp16"]["p50_ms"] if "fp16" in latency else None + print(f"\n=== precision matrix, {args.num_views} view(s) " + "(latency measured, cosine vs FP16) ===") + header = (f"{'tier':10s} {'p50 ms':>8s} {'p95 ms':>8s} {'vs FP16':>8s} " + f"{'raw cos':>9s} {'raw min':>9s} {'act cos':>9s} " + f"{'act min':>9s}") + print(header) + if "fp16" in latency: + f16 = latency["fp16"] + print(f"{'fp16':10s} {f16['p50_ms']:8.3f} {f16['p95_ms']:8.3f} " + f"{1.0:8.3f} {'-':>9s} {'-':>9s} {'-':>9s} {'-':>9s}") + for mode, raw_cos, raw_min, act_cos, act_min in rows: + lat = latency.get(mode, {}) + p50 = lat.get("p50_ms", float("nan")) + p95 = lat.get("p95_ms", float("nan")) + speedup = (ref_p50 / p50) if (ref_p50 and p50 == p50) else float("nan") + print(f"{mode:10s} {p50:8.3f} {p95:8.3f} {speedup:8.3f} " + f"{raw_cos:9.5f} {raw_min:9.5f} {act_cos:9.5f} {act_min:9.5f}") + + payload = { + "num_views": args.num_views, + "latency": latency, + "accuracy_vs_fp16": [ + dict(zip(("tier", "raw_cosine", "raw_min_sample_cosine", + "action_cosine", "action_min_sample_cosine"), r)) + for r in rows], + } + json.dump(payload, open(output_dir / "precision_matrix.json", "w"), + indent=1) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_catalog_pipeline_contracts.py b/tests/test_catalog_pipeline_contracts.py index b914b574..6c30e2b2 100644 --- a/tests/test_catalog_pipeline_contracts.py +++ b/tests/test_catalog_pipeline_contracts.py @@ -89,10 +89,10 @@ def test_pipeline_catalogs_declare_shared_schedule_families(): assert autoregressive.kind == "stage_pipeline" assert autoregressive.family == "autoregressive_decode" - assert autoregressive.version == 2 + assert autoregressive.version == 3 assert _stage_sets(autoregressive) == ( {"prefill", "decode", "token_select"}, - {"input_prepare", "modality_encode"}, + {"input_prepare", "modality_encode", "spec_draft"}, ) assert vla.kind == "stage_pipeline" diff --git a/tests/test_pi05_fp4_fusion_kernels.py b/tests/test_pi05_fp4_fusion_kernels.py index ad7579f2..6ad5083f 100644 --- a/tests/test_pi05_fp4_fusion_kernels.py +++ b/tests/test_pi05_fp4_fusion_kernels.py @@ -190,6 +190,50 @@ def test_fused_geglu_epilogue_matches_split_chain(): f"fused epilogue lost accuracy: {cos_fused:.6f} < {cos_chain:.6f}") +@pytest.mark.parametrize("shape,pair", [ + ((968, 2048, 8192), ("cutlass_fp4_gemm_geglu_il_hw", + "cutlass_fp4_gemm_geglu_il_hw_nod")), + ((10, 1024, 4096), ("cutlass_fp4_gemm_geglu_il_hw_v10", + "cutlass_fp4_gemm_geglu_il_hw_nod_v10")), +]) +def test_geglu_nod_matches_hw_and_skips_dummy(shape, pair): + """The no-D-store GeGLU GEMM must produce bit-identical compact output + to the store-carrying variant, and must never touch the dummy buffer.""" + _require_thor() + torch.manual_seed(20260815) + M, D, H = shape + N_il = 2 * H + + x = torch.randn(M, D, dtype=torch.float16, device='cuda') + w_il = torch.randn(N_il, D, dtype=torch.float16, device='cuda') * 0.02 + q_il = quant_weight_nvfp4(w_il.contiguous()) + act = FP4ActScratch(M, D, device='cuda') + quant_act_nvfp4(x, act, M) + + outputs = {} + for name, fn in (("hw", getattr(fvk_fp4, pair[0])), + ("nod", getattr(fvk_fp4, pair[1]))): + fused = FP4ActScratch(M, H, device='cuda') + fused.packed.fill_(0) + fused.sfa.fill_(0) + dummy = torch.full((M, H), 0xAB, dtype=torch.uint8, device='cuda') + assert fn( + act.packed.data_ptr(), act.sfa.data_ptr(), + q_il['packed'].data_ptr(), q_il['sfb'].data_ptr(), + dummy.data_ptr(), fused.packed.data_ptr(), fused.sfa.data_ptr(), + M, N_il, D, 0) == 0 + torch.cuda.synchronize() + outputs[name] = (fused, dummy) + + hw, nod = outputs["hw"], outputs["nod"] + assert torch.equal(nod[0].packed, hw[0].packed), \ + "nod compact packed differs from hw" + assert torch.equal(nod[0].sfa, hw[0].sfa), "nod compact SFA differs from hw" + assert bool((nod[1] == 0xAB).all()), "nod wrote the dummy buffer" + assert not torch.equal(hw[1], nod[1]), \ + "hw variant did not write the dummy (test would prove nothing)" + + def test_vectorized_siglip_layernorms_match_reference(): """The single-pass LayerNorms must agree with the reference norm + quantize pair (reduction order differs at ulp level)."""