diff --git a/.gitignore b/.gitignore index 3b4758346..f4e696dcc 100644 --- a/.gitignore +++ b/.gitignore @@ -43,4 +43,3 @@ training/_runs/** internal-tests/ internal-docs/ notes/ -third_party/cutlass diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 000000000..281cb2d85 --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "third_party/cutlass"] + path = third_party/cutlass + url = https://github.com/NVIDIA/cutlass.git 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 000000000..74b20a1c2 --- /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 000000000..fd1cab031 --- /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_siglip_ffn_f32out_sm100.cu b/csrc/gemm/fp4/cutlass_fp4_gemm_siglip_ffn_f32out_sm100.cu new file mode 100644 index 000000000..323c56067 --- /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 000000000..98da67e22 --- /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/flash_rt/structures/adapters/ggml/DEVELOPMENT.md b/flash_rt/structures/adapters/ggml/DEVELOPMENT.md new file mode 100644 index 000000000..0c5f47423 --- /dev/null +++ b/flash_rt/structures/adapters/ggml/DEVELOPMENT.md @@ -0,0 +1,113 @@ +# Development guide + +## Architecture + +The adapter has two halves with a hard boundary: + +- **Framework-free half** (`fr_repack.cu`, `fr_quant_act.cu`, + `fr_qkv_post.cu`, `fr_ada.cu`, `fr_decode_attn.cu`, `fr_fa4_vit.cu`): + plain CUDA translation units, entry points declared in `fr_kernels.h` + (raw pointers + `cudaStream_t`, no ggml or CUTLASS types in the + header). GEMMs with fused epilogues live in `csrc/gemm/fp4/` and are + compiled alongside. +- **ggml-facing half** (`fr_ggml.cuh`, `fr_dispatch.cu`): window + predicates (`ggml_cuda_flashrt_should_fuse_*`) and executors that speak + `ggml_tensor`, plus the caches. The host's `ggml-cuda.cu` calls the + predicates from its fuse hook; that call-site code lives in the host + tree, not here. + +Keep the boundary: nothing under `csrc/` or in the framework-free half +may include ggml headers, and `fr_kernels.h` must stay consumable from a +plain C++ translation unit. + +## Caches (all capture-safe) + +- **Weight repack cache** — keyed by weight data pointer, never evicted + (weights are immortal in a loaded model). ggml's split-nibble NVFP4 + blocks are repacked to the CUTLASS wire format (adjacent-pair nibbles, + scale bytes in the Sm1xx atom layout) on first use. +- **Per-evaluation activation cache** — a producer (e.g. the fused adaLN) + can register its already-quantized output; later GEMMs in the same + evaluation reuse it. Keyed by tensor pointer + an evaluation counter so + recycled addresses can never alias. Slots are grow-only so device + addresses stay stable for captured CUDA graphs. +- **One-shot handoffs** (f16 Q from the QKV window to the decode + attention) — single grow-only slot, key cleared on consumption. +- **CUTLASS workspaces** — shape-keyed, grown only outside capture. + +Rule for all of them: no allocation while a CUDA graph is being captured. +Check `cudaStreamIsCapturing` and fall back to the unfused path (or pool +memory) when growth would be needed mid-capture. + +## Adding a fusion window + +1. Express the executor in the framework-free half with a C entry point + in `fr_kernels.h`; consume existing `csrc` GEMMs where possible. +2. Add the predicate/executor pair to `fr_ggml.cuh` / `fr_dispatch.cu`. + The predicate must pin every assumption the kernel makes: dtypes, + shapes, strides (element-exact, not just "contiguous"), op params + (`max_bias`, softcap), and use counts where the window elides + intermediates. +3. Add the call site to the host's fuse hook, and a + `GGML_FLASHRT_NO_` switch in the predicate. +4. Validate per TESTING.md (trigger proof, A/B/A, parity or judge). + +Invariants and pitfalls learned the hard way: + +- **The fuse hook's return contract**: returning 0 means "not fused" and + the anchor node executes normally afterwards — a window that replaces a + single node must also consume the pure-view node that follows it and + return ≥1, or its work is silently overwritten (symptom: identical + results, slower). +- **Overlap checks are allocator-sensitive.** The generic fusion memory + range check vetoes a window when the destination aliases an + outside-window source. The allocator legitimately hands a window's + output the block of an input that dies inside the window; whether that + alias is safe depends on the fused implementation's read-before-write + order, so exemptions are per-window and must be argued in a comment + (see the GeGLU window: the activation is fully consumed by the quantize + kernel before the down GEMM writes). Any change that shifts allocation + (new nodes, another sched) can re-trigger vetoes elsewhere — symptom is + a silent GPU-time regression; diagnose with a kernel census diff. +- **Numeric equivalences must be argued or measured**, e.g. an epilogue + that converts the fp32 accumulator to f16 once is bit-equal to f32 + output plus a separate cast; a fused kernel writing the same values + through the same conversion is bit-identical to the copy chain it + replaces. Anything weaker goes through the real-observation judge. +- **RoPE in `fr_qkv_post.cu` mirrors ggml's `rope_neox`** (yarn + corrections included) and must stay bit-exact with it; the predicate + rejects non-NEOX modes. +- Windows only ever fire on `cc == 1100` (checked at the call site). + +## AOT FlashAttention-4 modules + +`fa4_aot/` holds ahead-of-time exports of the vendored FA4 forward +(vision shape: padded head_dim 80, MHA; prefill shape: head_dim 256, GQA +with one KV head). Regeneration and the export mechanics are documented +in `fa4_aot/README.md`; the short version: + +- CuTe-DSL's `export_to_c` emits a C header (host launch entry, tensor + argument structs, embedded cubin) plus a host object. The tvm-ffi + compile variant only exports a TVM ABI, so the export script strips + `--enable-tvm-ffi` and never executes the resulting object in-process + (its calling convention differs). +- `fr_fa4_shims.c` supplies the small `_cuda*` runtime aliases the object + expects, so neither the build nor the runtime depends on any CuTe-DSL + library. +- Module loading must happen outside CUDA graph capture; the adapter + preloads from `ggml_cuda_flashrt_begin_eval`, which always runs before + a capture can begin. +- The wrapper takes dynamic shapes/strides per tensor, so one export per + (head_dim, GQA config) covers all sequence lengths. The prefill + window's mask handling relies on the pi0.5 prefix-LM property that the + mask is row-uniform pad-only and the real KV length equals the query + count; the padded tail is excluded by the dynamic shape instead of by + a mask. + +## Single-source rule + +Structure changes (GEMM tiles, epilogues, attention decompositions) +belong in `csrc/` or the structures catalog so every host adapter +inherits them; this directory only translates. Nothing here may be +copy-pasted into a host tree, and the host integration must stay behind +its own opt-in build flag so stock builds are unaffected. diff --git a/flash_rt/structures/adapters/ggml/README.md b/flash_rt/structures/adapters/ggml/README.md new file mode 100644 index 000000000..e9dd53a13 --- /dev/null +++ b/flash_rt/structures/adapters/ggml/README.md @@ -0,0 +1,72 @@ +# ggml host adapter (native) + +Native C++/CUDA host adapter that maps FlashRT structures onto ggml's CUDA +backend (llama.cpp family), targeting Jetson AGX Thor (SM110). Unlike the +Python runtime adapters (`vllm_engine.py`, `sglang_engine.py`), this adapter +is consumed at **build time**: the host's CMake compiles these translation +units inside its own build tree and only C symbols cross the boundary. The +host never links against Python, PyTorch, or any FlashRT runtime. + +Documentation: + +- [USAGE.md](USAGE.md) — building a host against this adapter, runtime + switches, deployment notes. +- [TESTING.md](TESTING.md) — operator tests, the qualification gates, and + the benchmarking / parity methodology every change must pass. +- [DEVELOPMENT.md](DEVELOPMENT.md) — layer architecture, how to add a + fusion window, invariants and known pitfalls, AOT FlashAttention-4 + regeneration. + +## What it is + +The adapter is the third host of the `flash_rt/structures` catalog. The +same structures that the torch frontend and the vllm/sglang adapters +consume — block-scaled NVFP4 GEMMs with fused epilogues, fused +norm/modulation producers, the decomposed tiny-M decode attention, the +FlashAttention-4 forward — are mapped here onto ggml's graph executor +through pattern-matched subgraph windows. Heavy math is single-source: + +- **NVFP4 GEMMs** come from `csrc/gemm/fp4/` in this repository + (GeGLU-interleaved, SigLIP-FFN pair, bias/f16-out variants). Nothing is + vendored into the host. +- **FlashAttention-4** is the vendored CuTe-DSL forward under + `csrc/attention/flash_attn_4_src`, consumed as ahead-of-time compiled + modules (see `fa4_aot/`), so the host build needs no CuTe-DSL toolchain. +- The `fr_*.cu` files here are the translation layer only: wire-format + repack (ggml split-nibble NVFP4 → CUTLASS atom layout), activation + quantize, fused RoPE/norm producers, and the dispatch/caching half that + speaks `ggml_tensor`. + +## Layout + +- `fr_kernels.h` — pure C entry points (no ggml, no CUTLASS in the header). +- `fr_gemm_f32out.cu`, `fr_ada.cu`, `fr_qkv_post.cu`, `fr_quant_act.cu`, + `fr_repack.cu`, `fr_decode_attn.cu`, `fr_fa4_vit.cu`, `fr_fa4_shims.c` — + framework-free CUDA translation units. +- `fr_dispatch.cu`, `fr_ggml.cuh` — the ggml-facing half: subgraph window + predicates and executors over `ggml_tensor` chains, weight/activation + caches. Requires ggml-cuda's internal headers on the include path. +- `fa4_aot/` — AOT FlashAttention-4 modules (vision and prefill shapes) + plus their regeneration script and provenance notes. +- `qualification/` — the release gates (see TESTING.md). +- `../../bindings/jetson_pi_edge_pi05.yaml` — the pipeline binding that + maps the host's hot path onto catalog structures under the + complete-hot-path contract. + +## Measured performance (Jetson AGX Thor, pi0.5, 2 camera views) + +| metric | stock llama.cpp (BF16) | with this adapter (NVFP4) | +|---|---|---| +| `llama_encode` + `llama_decode` (host `total_ms`, P50 warm) | 202.7 ms | 35.5 ms (**5.7×**) | +| end-to-end action chunk (ViT + prefill + 10 denoise steps) | — | **42.5 ms** | +| phase split | — | ViT 6.7 + prefill 15.6 + decode 19.8 | + +For context, the FlashRT torch frontend runs the same checkpoint at +36.4 ms end-to-end on the same device; the remaining gap is dominated by +the host graph's fp32 activation dtype (the torch pipeline holds +activations in fp16). + +Numerics: the adapter is bitwise deterministic across processes after +warmup; changes are gated by an exact e2e action golden plus a +real-observation parity protocol against an f16 reference (see +TESTING.md). diff --git a/flash_rt/structures/adapters/ggml/TESTING.md b/flash_rt/structures/adapters/ggml/TESTING.md new file mode 100644 index 000000000..341a20ce7 --- /dev/null +++ b/flash_rt/structures/adapters/ggml/TESTING.md @@ -0,0 +1,76 @@ +# Testing + +Three layers of validation, from operator level to release gate. + +## 1. Operator tests (`test-backend-ops`) + +The host's `test-backend-ops` exercises the NVFP4 mul_mat path against the +CPU reference: + +```bash +GGML_CUDA_FLASHRT_NO_CACHE=1 ./build/bin/test-backend-ops test -o MUL_MAT +``` + +`GGML_CUDA_FLASHRT_NO_CACHE=1` is **required**: the weight repack cache is +keyed by tensor data pointer under the assumption that weights are +immortal, which holds for models but not for the test harness's rapidly +recycled tensors. NVFP4 mismatches up to ~2e-2 are inherent W4A4 +activation-quantization noise (upstream applies the same tolerance to +native FP4 backends), not failures. + +## 2. Qualification gates (`qualification/`) + +`qualification/run_qualification.py` gates a build the way a release +would: + +- **manifest** — the pipeline binding + (`bindings/jetson_pi_edge_pi05.yaml`) must map the host's complete hot + path onto catalog structures. +- **pins** — the structure versions the binding names must match the + catalog (`qualification/pins.yaml`). +- **e2e golden** (`--e2e`, on-device, needs a running server) — drives the + fixed synthetic-input protocol and compares the steady-state action + chunk against `qualification/goldens/pi05_thor_action.json` **exactly**. + The adapter is bitwise deterministic across processes after warmup, so + any bit difference is a real change. + +```bash +python qualification/run_qualification.py # offline gates +python qualification/run_qualification.py --e2e # + on-device golden +python qualification/run_qualification.py --e2e --update-golden +``` + +Take the golden only after at least two warm-up inferences (the first +inference after cold start differs from steady state) and only for +changes whose numerics were judged (below). + +## 3. Benchmark and parity methodology + +Every performance change must pass this protocol on device: + +- **Hot-regime A/B/A sandwich** — run the candidate, the fallback (via its + runtime switch), and the candidate again as three separate server + processes, ≥15 warm-up + ~20 measured inferences each, comparing P50. + Thor drifts ±1–3 ms across long sessions, so only same-session + back-to-back numbers are comparable; single measurements and + cross-session comparisons are not accepted. +- **Bitwise parity** — save the action chunk from each leg. A change that + claims numeric neutrality must be bit-identical to the previous + accepted state. Note the converse trap: bit-identical output *plus* + zero performance delta usually means the window never fired — verify + the window triggers (kernel census, `GGML_FLASHRT_DEBUG`) before + interpreting the A/B. +- **Real-observation judge** — for changes that move numerics, run a set + of real robot observations (gripper-active frames) through the NVFP4 + build and an f16-weights build of the same tree, and compare per-dim + cosine of the action chunks against the f16 reference. The distance to + the reference must not systematically regress. Any bit-level change in + the action path amplifies to ~2e-2 absolute wobble on final actions + through the 10 denoise steps, so raw action diffs are meaningless — + only the distance-to-reference comparison judges accuracy. +- **Kernel-level accounting** — attribute wins with an nsys census + (`GGML_CUDA_DISABLE_GRAPHS=1`, full-lifetime `-t cuda` trace with a + graceful server exit so buffers flush). CUDA-graph replays hide kernels + from the profiler, and profiling on Tegra inflates kernel times, so the + census attributes *where* time went while the non-profiled A/B decides + *whether* the change lands. diff --git a/flash_rt/structures/adapters/ggml/USAGE.md b/flash_rt/structures/adapters/ggml/USAGE.md new file mode 100644 index 000000000..c309ca58f --- /dev/null +++ b/flash_rt/structures/adapters/ggml/USAGE.md @@ -0,0 +1,77 @@ +# Usage + +## Building a host against the adapter + +The reference host is the Jetson-PI-Edge llama.cpp tree, which carries the +integration side (CMake wiring, fuse-hook call sites, pi0 graph changes) +on its FlashRT branch and consumes this repository as a submodule at +`ggml/src/ggml-cuda/flashrt/flashrt-public`: + +```bash +git clone --recursive -b feat/flashrt-thor-kernels +cd Jetson-PI-Edge +cmake -B build -DGGML_CUDA=ON -DGGML_CUDA_FLASHRT=ON -DCMAKE_BUILD_TYPE=Release +cmake --build build --target llama-server -j +``` + +Options: + +- `GGML_CUDA_FLASHRT` (OFF by default) — enables the layer. Without it the + build is stock llama.cpp; every integration point is compiled out. +- `GGML_CUDA_FLASHRT_PUBLIC_DIR` — path to a FlashRT checkout, overriding + the submodule location. +- `GGML_CUDA_FLASHRT_CUTLASS_DIR` — CUTLASS override; defaults to + `third_party/cutlass` inside this repository. + +The adapter is built as a separate CMake OBJECT library with +`-arch=sm_110a` and CUTLASS headers; the rest of ggml-cuda compiles +unchanged. The AOT FlashAttention-4 windows enable automatically when the +`fa4_aot/*.o` artifacts are present (they are checked in; see +`fa4_aot/README.md` to regenerate). + +## Model preparation + +- LLM weights: quantize with the host's `llama-quantize` to the `NVFP4` + target (exposed by the FlashRT branch). Setting `GGML_NVFP4_MSE=1` + during quantization selects per-block scales by reconstruction-MSE + search instead of plain absmax (slower to quantize, more accurate). +- The mmproj (vision tower) is quantized to NVFP4 the same way. + +## Running + +```bash +PI_MODEL=pi05 ./build/bin/llama-server -m --mmproj \ + -ngl 99 --flash-attn on --port +``` + +The server exposes the host's action-chunk HTTP protocol (reset → images → +state → infer). Warm-up matters on Thor: latency reaches its steady state +after roughly 15 inferences. + +## Runtime switches + +All switches are environment variables; unset means enabled/default. + +| variable | effect | +|---|---| +| `GGML_CUDA_FLASHRT_DISABLE=1` | disable the whole layer at runtime (stock kernels) | +| `GGML_FLASHRT_NO_RMS_GEMMA=1` | disable the Gemma norm-chain window | +| `GGML_FLASHRT_NO_QKV_PREFILL=1` | disable the fused prefill QKV window | +| `GGML_FLASHRT_NO_DEC_ATTN=1` | disable the decomposed decode attention | +| `GGML_FLASHRT_NO_VIT_FA4=1` | disable the AOT FA4 vision attention | +| `GGML_FLASHRT_NO_PREFILL_FA4=1` | disable the AOT FA4 prefill attention | +| `GGML_FLASHRT_NO_KV_TAIL=1` | disable the batched persistent-KV tail copies | +| `GGML_FLASHRT_NO_VIS_F16=1` | disable the vision QKV window's direct f16 K/V outputs | +| `GGML_CUDA_FLASHRT_NO_CACHE=1` | disable the pointer-keyed weight repack cache (required for `test-backend-ops`, see TESTING.md) | +| `GGML_FLASHRT_DEBUG=1` | print window-match failure diagnostics | +| `GGML_FLASHRT_DUMP=` / `GGML_FLASHRT_DUMP_MAX=` | dump the first n evaluated graphs' node sequences | + +Host-side switches on the FlashRT branch (outside this repository): +`GGML_PI05_MOD_PRECOMP=0` disables the denoise-schedule modulation +precompute, `LLAMA_GRAPH_REUSE_DISABLE=1` disables graph reuse, +`GGML_CUDA_DISABLE_GRAPHS=1` disables CUDA graphs (useful for profiling: +kernels are invisible to nsys while CUDA graphs replay). + +Every window degrades gracefully: when its predicate does not match (or +its switch is set) the nodes run on stock ggml kernels, so the switches +bisect regressions window by window. diff --git a/flash_rt/structures/adapters/ggml/fa4_aot/README.md b/flash_rt/structures/adapters/ggml/fa4_aot/README.md new file mode 100644 index 000000000..13d5b6abf --- /dev/null +++ b/flash_rt/structures/adapters/ggml/fa4_aot/README.md @@ -0,0 +1,33 @@ +# AOT FlashAttention-4 module (SigLIP vision attention, Thor SM110) + +`fa4_siglip_fwd.h` / `fa4_siglip_fwd.o` are the CuTe-DSL ahead-of-time +export of the vendored FA4 SM100-compatible forward +(`csrc/attention/flash_attn_4_src/flashrt_fa4`) compiled for `sm_110a` +at head_dim 80 — the padded-head layout the ggml adapter's vision path +uses. `fa4_prefill_fwd.h` / `fa4_prefill_fwd.o` are the same export at +the pi0.5 prefill shape (head_dim 256, GQA with one KV head, full +attention): the prefill's row-uniform pad mask is reproduced by passing +the real sequence length as the KV dynamic shape. Sequence length, head count and batch stay dynamic; the softmax +scale is a runtime argument. The `.o` contains the embedded cubin plus +the host launch entry; `fr_fa4_shims.c` provides the small `_cuda*` +runtime aliases the object expects, so no CuTe-DSL runtime library is +needed at build or run time. + +The ggml adapter build enables the FA4 vision-attention window +automatically when these files are present (see the host build's +`GGML_CUDA_FLASHRT` integration); delete them or set +`GGML_FLASHRT_NO_VIT_FA4=1` / `GGML_FLASHRT_NO_PREFILL_FA4=1` to fall +back to the host's own flash attention per site. + +## Regeneration + +Requires the `thor-fa4` runtime deps (`nvidia-cutlass-dsl`, +`quack-kernels`) and PyTorch with CUDA, on the target device: + +```bash +CUTE_DSL_ARCH=sm_110a python export_fa4_siglip.py +``` + +The script compiles the vendored FA4 forward once at the head_dim-80 +shape with `--enable-tvm-ffi` stripped (the plain JIT object carries the +classic C-header exporter) and writes both files into this directory. diff --git a/flash_rt/structures/adapters/ggml/fa4_aot/export_fa4_siglip.py b/flash_rt/structures/adapters/ggml/fa4_aot/export_fa4_siglip.py new file mode 100644 index 000000000..0bb1e3e07 --- /dev/null +++ b/flash_rt/structures/adapters/ggml/fa4_aot/export_fa4_siglip.py @@ -0,0 +1,80 @@ +#!/usr/bin/env python3 +"""AOT-export the vendored FA4 forward for the ggml adapter's vision path. + +Compiles the FA4 SM100-compatible forward (vendored under +csrc/attention/flash_attn_4_src) at the padded SigLIP shape (head_dim 80, +f16, no mask) and writes fa4_siglip_fwd.h / fa4_siglip_fwd.o into this +directory. Run on the target device with the thor-fa4 deps installed: + + CUTE_DSL_ARCH=sm_110a python export_fa4_siglip.py +""" +import os +import sys +from pathlib import Path + +_HERE = Path(__file__).resolve().parent +_REPO = _HERE.parents[4] +sys.path.insert(0, str(_REPO)) + +import torch # noqa: E402 + +from flash_rt.hardware.thor import fa4_backend # noqa: E402 + +assert fa4_backend.is_available(), fa4_backend.status() +fwd = fa4_backend.fa4_fwd() + +import flashrt_fa4.cute.interface_fwd_sm100 as ifw # noqa: E402 + +# Strip --enable-tvm-ffi so the cache holds a plain JitCompiledFunction: +# only that variant carries the classic C-header exporter (embedded cubin +# plus a plain-C host launch entry). The compiled object is exported +# without ever being executed (the tvm-ffi call convention differs). +_holder = [] +_orig_compile = ifw.cute.compile + + +class _NoCall: + def __init__(self, inner): + self._inner = inner + + def __call__(self, *args, **kwargs): + return None + + +def _compile_no_ffi(*args, **kwargs): + kwargs.pop("options", None) + obj = _orig_compile(*args, **kwargs) + _holder.append(obj) + return _NoCall(obj) + + +ifw.cute.compile = _compile_no_ffi +ifw._flash_attn_fwd.compile_cache.clear() + +# vision attention: padded head_dim 80, MHA +NV, SQ, NH, HD = 2, 256, 16, 80 +q = torch.zeros(NV, SQ, NH, HD, dtype=torch.float16, device="cuda") +k = torch.zeros_like(q) +v = torch.zeros_like(q) +out = torch.empty_like(q) +fwd(q, k, v, causal=False, num_splits=1, pack_gqa=False, out=out) +torch.cuda.synchronize() + +assert _holder, "FA4 compile did not run" +_holder[0].export_to_c(str(_HERE), "fa4_siglip_fwd") + +# prefill self-attention: head_dim 256, GQA with one KV head +_holder.clear() +ifw._flash_attn_fwd.compile_cache.clear() +B, SQ2, HQ, HK, HD2 = 1, 559, 8, 1, 256 +q2 = torch.zeros(B, SQ2, HQ, HD2, dtype=torch.float16, device="cuda") +k2 = torch.zeros(B, SQ2, HK, HD2, dtype=torch.float16, device="cuda") +v2 = torch.zeros_like(k2) +out2 = torch.empty_like(q2) +fwd(q2, k2, v2, softmax_scale=HD2 ** -0.5, causal=False, + num_splits=1, pack_gqa=True, out=out2) +torch.cuda.synchronize() + +assert _holder, "FA4 prefill compile did not run" +_holder[0].export_to_c(str(_HERE), "fa4_prefill_fwd") +print("exported:", sorted(p.name for p in _HERE.glob("fa4_*_fwd.*"))) diff --git a/flash_rt/structures/adapters/ggml/fa4_aot/fa4_prefill_fwd.h b/flash_rt/structures/adapters/ggml/fa4_aot/fa4_prefill_fwd.h new file mode 100644 index 000000000..6dd1bc77d --- /dev/null +++ b/flash_rt/structures/adapters/ggml/fa4_aot/fa4_prefill_fwd.h @@ -0,0 +1,101 @@ + +#pragma once + +#include +#include +#include +#include + + +// Macro to check for cuda errors. +#ifndef CUTE_DSL_CUDA_ERROR_CHECK +#define CUTE_DSL_CUDA_ERROR_CHECK(err) { \ + if ((err) != cudaSuccess) { \ + printf("Got Cuda Error %s: %s\n", cudaGetErrorName(err), cudaGetErrorString(err)); \ + } \ +} + +#endif + +typedef struct { + cudaLibrary_t module; +} fa4_prefill_fwd_Kernel_Module_t; + +#ifdef __cplusplus +extern "C" { +#endif +void _mlir_fa4_prefill_fwd_cuda_init(void **); +void _mlir_fa4_prefill_fwd_cuda_load_to_device(void **); +static inline void fa4_prefill_fwd_Kernel_Module_Load(fa4_prefill_fwd_Kernel_Module_t *module) { + cudaLibrary_t *libraryPtr = &(module->module); + cudaError_t ret; + struct { + cudaLibrary_t **libraryPtr; + cudaError_t *ret; + } initArgs = {&libraryPtr, &ret}; + _mlir_fa4_prefill_fwd_cuda_init((void **)(&initArgs)); + CUTE_DSL_CUDA_ERROR_CHECK(ret); + int32_t device_id = 0; + struct { + cudaLibrary_t **library; + int32_t *device_id; + cudaError_t *ret; + } loadArgs = {&libraryPtr, &device_id, &ret}; + int32_t device_count; + CUTE_DSL_CUDA_ERROR_CHECK(cudaGetDeviceCount(&device_count)); + for (int32_t i = 0; i < device_count; i++) { + device_id = i; + _mlir_fa4_prefill_fwd_cuda_load_to_device((void **)(&loadArgs)); + CUTE_DSL_CUDA_ERROR_CHECK(ret); + } +} + +static inline void fa4_prefill_fwd_Kernel_Module_Unload(fa4_prefill_fwd_Kernel_Module_t *module) { + CUTE_DSL_CUDA_ERROR_CHECK(cudaLibraryUnload(module->module)); +} + +#ifdef __cplusplus +} +#endif + +typedef struct { + void *data; + int32_t dynamic_shapes[4]; + int64_t dynamic_strides[3]; +} fa4_prefill_fwd_Tensor_mQ_t; + + +typedef struct { + void *data; + int32_t dynamic_shapes[4]; + int64_t dynamic_strides[3]; +} fa4_prefill_fwd_Tensor_mK_t; + + +typedef struct { + void *data; + int32_t dynamic_shapes[4]; + int64_t dynamic_strides[3]; +} fa4_prefill_fwd_Tensor_mV_t; + + +typedef struct { + void *data; + int32_t dynamic_shapes[4]; + int64_t dynamic_strides[3]; +} fa4_prefill_fwd_Tensor_mO_t; + +#ifdef __cplusplus +extern "C" +#endif +void _mlir_fa4_prefill_fwd__mlir_ciface_cutlass___call___flashrt_fa4cutesm100_hd256_2cta_fmha_forwardBlackwellFusedMultiHeadAttentionForward_object_at__Tensorgmemoi64i64i641_Tensorgmemoi64i64i641_Tensorgmemoi64i64i641_Te(void **args, int32_t num_args); + +static inline int32_t cute_dsl_fa4_prefill_fwd_wrapper(fa4_prefill_fwd_Kernel_Module_t *module, fa4_prefill_fwd_Tensor_mQ_t *mQ, fa4_prefill_fwd_Tensor_mK_t *mK, fa4_prefill_fwd_Tensor_mV_t *mV, fa4_prefill_fwd_Tensor_mO_t *mO, float softmax_scale, cudaStream_t stream) { + int32_t ret; + void *args[7] = { + mQ, mK, mV, mO, &softmax_scale, &stream, + &ret + }; + _mlir_fa4_prefill_fwd__mlir_ciface_cutlass___call___flashrt_fa4cutesm100_hd256_2cta_fmha_forwardBlackwellFusedMultiHeadAttentionForward_object_at__Tensorgmemoi64i64i641_Tensorgmemoi64i64i641_Tensorgmemoi64i64i641_Te(args, 7); + return ret; +} diff --git a/flash_rt/structures/adapters/ggml/fa4_aot/fa4_prefill_fwd.o b/flash_rt/structures/adapters/ggml/fa4_aot/fa4_prefill_fwd.o new file mode 100644 index 000000000..fc9ee97d1 Binary files /dev/null and b/flash_rt/structures/adapters/ggml/fa4_aot/fa4_prefill_fwd.o differ diff --git a/flash_rt/structures/adapters/ggml/fa4_aot/fa4_siglip_fwd.h b/flash_rt/structures/adapters/ggml/fa4_aot/fa4_siglip_fwd.h new file mode 100644 index 000000000..c6fdc58bd --- /dev/null +++ b/flash_rt/structures/adapters/ggml/fa4_aot/fa4_siglip_fwd.h @@ -0,0 +1,101 @@ + +#pragma once + +#include +#include +#include +#include + + +// Macro to check for cuda errors. +#ifndef CUTE_DSL_CUDA_ERROR_CHECK +#define CUTE_DSL_CUDA_ERROR_CHECK(err) { \ + if ((err) != cudaSuccess) { \ + printf("Got Cuda Error %s: %s\n", cudaGetErrorName(err), cudaGetErrorString(err)); \ + } \ +} + +#endif + +typedef struct { + cudaLibrary_t module; +} fa4_siglip_fwd_Kernel_Module_t; + +#ifdef __cplusplus +extern "C" { +#endif +void _mlir_fa4_siglip_fwd_cuda_init(void **); +void _mlir_fa4_siglip_fwd_cuda_load_to_device(void **); +static inline void fa4_siglip_fwd_Kernel_Module_Load(fa4_siglip_fwd_Kernel_Module_t *module) { + cudaLibrary_t *libraryPtr = &(module->module); + cudaError_t ret; + struct { + cudaLibrary_t **libraryPtr; + cudaError_t *ret; + } initArgs = {&libraryPtr, &ret}; + _mlir_fa4_siglip_fwd_cuda_init((void **)(&initArgs)); + CUTE_DSL_CUDA_ERROR_CHECK(ret); + int32_t device_id = 0; + struct { + cudaLibrary_t **library; + int32_t *device_id; + cudaError_t *ret; + } loadArgs = {&libraryPtr, &device_id, &ret}; + int32_t device_count; + CUTE_DSL_CUDA_ERROR_CHECK(cudaGetDeviceCount(&device_count)); + for (int32_t i = 0; i < device_count; i++) { + device_id = i; + _mlir_fa4_siglip_fwd_cuda_load_to_device((void **)(&loadArgs)); + CUTE_DSL_CUDA_ERROR_CHECK(ret); + } +} + +static inline void fa4_siglip_fwd_Kernel_Module_Unload(fa4_siglip_fwd_Kernel_Module_t *module) { + CUTE_DSL_CUDA_ERROR_CHECK(cudaLibraryUnload(module->module)); +} + +#ifdef __cplusplus +} +#endif + +typedef struct { + void *data; + int32_t dynamic_shapes[4]; + int64_t dynamic_strides[3]; +} fa4_siglip_fwd_Tensor_mQ_t; + + +typedef struct { + void *data; + int32_t dynamic_shapes[4]; + int64_t dynamic_strides[3]; +} fa4_siglip_fwd_Tensor_mK_t; + + +typedef struct { + void *data; + int32_t dynamic_shapes[4]; + int64_t dynamic_strides[3]; +} fa4_siglip_fwd_Tensor_mV_t; + + +typedef struct { + void *data; + int32_t dynamic_shapes[4]; + int64_t dynamic_strides[3]; +} fa4_siglip_fwd_Tensor_mO_t; + +#ifdef __cplusplus +extern "C" +#endif +void _mlir_fa4_siglip_fwd__mlir_ciface_cutlass___call___flashrt_fa4cuteflash_fwd_sm100FlashAttentionForwardSm100_object_at__Tensorgmemoi64i64i641_Tensorgmemoi64i64i641_Tensorgmemoi64i64i641_Tensorgmemoi64i64i641_None_01(void **args, int32_t num_args); + +static inline int32_t cute_dsl_fa4_siglip_fwd_wrapper(fa4_siglip_fwd_Kernel_Module_t *module, fa4_siglip_fwd_Tensor_mQ_t *mQ, fa4_siglip_fwd_Tensor_mK_t *mK, fa4_siglip_fwd_Tensor_mV_t *mV, fa4_siglip_fwd_Tensor_mO_t *mO, float softmax_scale, cudaStream_t stream) { + int32_t ret; + void *args[7] = { + mQ, mK, mV, mO, &softmax_scale, &stream, + &ret + }; + _mlir_fa4_siglip_fwd__mlir_ciface_cutlass___call___flashrt_fa4cuteflash_fwd_sm100FlashAttentionForwardSm100_object_at__Tensorgmemoi64i64i641_Tensorgmemoi64i64i641_Tensorgmemoi64i64i641_Tensorgmemoi64i64i641_None_01(args, 7); + return ret; +} diff --git a/flash_rt/structures/adapters/ggml/fa4_aot/fa4_siglip_fwd.o b/flash_rt/structures/adapters/ggml/fa4_aot/fa4_siglip_fwd.o new file mode 100644 index 000000000..d7690a934 Binary files /dev/null and b/flash_rt/structures/adapters/ggml/fa4_aot/fa4_siglip_fwd.o differ diff --git a/flash_rt/structures/adapters/ggml/fr_ada.cu b/flash_rt/structures/adapters/ggml/fr_ada.cu new file mode 100644 index 000000000..60871fdfa --- /dev/null +++ b/flash_rt/structures/adapters/ggml/fr_ada.cu @@ -0,0 +1,371 @@ +// Fused adaLN kernels for the pi0.5 action expert (Thor). +// +// The ggml graph expresses each adaLN application as rms_norm + two +// broadcast repeats + mul + two adds (and the gated residual as repeat + +// mul + add), all on [M, C] tensors with M ~ 10. These kernels collapse +// each chain into one launch. The scale/shift/gate vectors are passed as +// direct pointers (the ggml view tensors' data pointers, which already +// include their byte offsets into the modulation vector). + +#include "fr_kernels.h" + +#include + +#include "cutlass/cutlass.h" +#include "cutlass/detail/sm100_blockscaled_layout.hpp" +#include "cute/tensor.hpp" + +namespace ggml_cuda_flashrt { + +namespace { + +// out[m, c] = norm(x[m])[c] * (1 + scale[c]) + shift[c] +// where norm = rms-normalize when with_rms, identity otherwise. +template +__global__ void kernel_ada_rms(const float * __restrict__ x, + const float * __restrict__ scale, + const float * __restrict__ shift, + float * __restrict__ out, + int C, float eps) { + const int m = blockIdx.x; + const float * xr = x + (int64_t) m * C; + float * orow = out + (int64_t) m * C; + + float inv_rms = 1.0f; + if (with_rms) { + float sumsq = 0.0f; + for (int c = threadIdx.x; c < C; c += blockDim.x) { + const float v = xr[c]; + sumsq += v * v; + } + __shared__ float red[32]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sumsq += __shfl_xor_sync(0xffffffff, sumsq, off); + } + const int warp = threadIdx.x / 32; + if (threadIdx.x % 32 == 0) { + red[warp] = sumsq; + } + __syncthreads(); + if (warp == 0) { + float v = (threadIdx.x < blockDim.x / 32) ? red[threadIdx.x] : 0.0f; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + v += __shfl_xor_sync(0xffffffff, v, off); + } + if (threadIdx.x == 0) { + red[0] = v; + } + } + __syncthreads(); + inv_rms = rsqrtf(red[0] / C + eps); + } + + for (int c = threadIdx.x; c < C; c += blockDim.x) { + const float n = with_rms ? xr[c] * inv_rms : xr[c]; + orow[c] = n * (1.0f + scale[c]) + shift[c]; + } +} + +using AdaCfg = cutlass::detail::Sm1xxBlockScaledConfig<16>; + +__device__ __forceinline__ uint8_t ada_f32_to_e2m1(float x) { + uint8_t sign = (x < 0.f) ? 0x8u : 0x0u; + float ax = fabsf(x); + uint8_t m; + if (ax <= 0.25f) m = 0u; + else if (ax <= 0.75f) m = 1u; + else if (ax <= 1.25f) m = 2u; + else if (ax <= 1.75f) m = 3u; + else if (ax <= 2.5f) m = 4u; + else if (ax <= 3.5f) m = 5u; + else if (ax <= 5.0f) m = 6u; + else m = 7u; + return sign | m; +} + +// Quantize one 16-element block held in registers by 16 consecutive threads? +// Simpler: each thread quantizes one 16-element block it re-reads from the +// just-written f32 output row (L2-hot), writing packed bytes + one SF byte. +template +__device__ __forceinline__ void ada_quant_row(const float * __restrict__ orow, + uint8_t * __restrict__ dst_packed, + uint8_t * __restrict__ dst_sfa, + LayoutSF layout, + int m, int C) { + const int n_blocks = C / 16; + for (int blk = threadIdx.x; blk < n_blocks; blk += blockDim.x) { + const float * v = orow + blk * 16; + float amax = 0.f; + #pragma unroll + for (int i = 0; i < 16; ++i) { + amax = fmaxf(amax, fabsf(v[i])); + } + float desired = amax / 6.f; + if (desired < 1e-12f) desired = 1e-12f; + __nv_fp8_e4m3 q(desired); + dst_sfa[layout(m, blk * 16, 0)] = *reinterpret_cast(&q); + const float inv = 1.f / static_cast(q); + uint2 out; + uint8_t * ob = reinterpret_cast(&out); + #pragma unroll + for (int pIdx = 0; pIdx < 8; ++pIdx) { + const uint8_t lo = ada_f32_to_e2m1(v[2 * pIdx] * inv); + const uint8_t hi = ada_f32_to_e2m1(v[2 * pIdx + 1] * inv); + ob[pIdx] = static_cast(lo | (hi << 4)); + } + reinterpret_cast(dst_packed)[(int64_t) m * n_blocks + blk] = out; + } +} + +// Fused adaLN modulate + NVFP4 quantize of the result. +template +__global__ void kernel_ada_rms_q(const float * __restrict__ x, + const float * __restrict__ scale, + const float * __restrict__ shift, + float * __restrict__ out, + uint8_t * __restrict__ dst_packed, + uint8_t * __restrict__ dst_sfa, + LayoutSF layout, + int C, float eps) { + const int m = blockIdx.x; + const float * xr = x + (int64_t) m * C; + float * orow = out + (int64_t) m * C; + + float inv_rms = 1.0f; + if (with_rms) { + float sumsq = 0.0f; + for (int c = threadIdx.x; c < C; c += blockDim.x) { + const float v = xr[c]; + sumsq += v * v; + } + __shared__ float red[32]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sumsq += __shfl_xor_sync(0xffffffff, sumsq, off); + } + const int warp = threadIdx.x / 32; + if (threadIdx.x % 32 == 0) { + red[warp] = sumsq; + } + __syncthreads(); + if (warp == 0) { + float v = (threadIdx.x < blockDim.x / 32) ? red[threadIdx.x] : 0.0f; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + v += __shfl_xor_sync(0xffffffff, v, off); + } + if (threadIdx.x == 0) { + red[0] = v; + } + } + __syncthreads(); + inv_rms = rsqrtf(red[0] / C + eps); + } + + for (int c = threadIdx.x; c < C; c += blockDim.x) { + const float n = with_rms ? xr[c] * inv_rms : xr[c]; + orow[c] = n * (1.0f + scale[c]) + shift[c]; + } + __syncthreads(); + ada_quant_row(orow, dst_packed, dst_sfa, layout, m, C); +} + +// Fused LayerNorm + affine + NVFP4 quantize of the result. +template +__global__ void kernel_layer_norm_affine_q(const float * __restrict__ x, + const float * __restrict__ w, + const float * __restrict__ b, + float * __restrict__ out, + uint8_t * __restrict__ dst_packed, + uint8_t * __restrict__ dst_sfa, + LayoutSF layout, + int C, float eps) { + const int m = blockIdx.x; + const float * xr = x + (int64_t) m * C; + float * orow = out + (int64_t) m * C; + + float sum = 0.0f, sumsq = 0.0f; + for (int c = threadIdx.x; c < C; c += blockDim.x) { + const float v = xr[c]; + sum += v; + sumsq += v * v; + } + __shared__ float red[2][32]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sumsq += __shfl_xor_sync(0xffffffff, sumsq, off); + } + const int warp = threadIdx.x / 32; + if (threadIdx.x % 32 == 0) { + red[0][warp] = sum; + red[1][warp] = sumsq; + } + __syncthreads(); + if (warp == 0) { + float s = (threadIdx.x < blockDim.x / 32) ? red[0][threadIdx.x] : 0.0f; + float s2 = (threadIdx.x < blockDim.x / 32) ? red[1][threadIdx.x] : 0.0f; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + s += __shfl_xor_sync(0xffffffff, s, off); + s2 += __shfl_xor_sync(0xffffffff, s2, off); + } + if (threadIdx.x == 0) { + red[0][0] = s; + red[1][0] = s2; + } + } + __syncthreads(); + const float mean = red[0][0] / C; + const float var = red[1][0] / C - mean * mean; + const float rstd = rsqrtf(var + eps); + + for (int c = threadIdx.x; c < C; c += blockDim.x) { + orow[c] = (xr[c] - mean) * rstd * w[c] + b[c]; + } + __syncthreads(); + ada_quant_row(orow, dst_packed, dst_sfa, layout, m, C); +} + +// out[m, c] = residual[m, c] + branch[m, c] * gate[c] +__global__ void kernel_gated_residual(const float * __restrict__ residual, + const float * __restrict__ branch, + const float * __restrict__ gate, + float * __restrict__ out, + int C) { + const int m = blockIdx.x; + const int64_t off = (int64_t) m * C; + for (int c = threadIdx.x; c < C; c += blockDim.x) { + out[off + c] = residual[off + c] + branch[off + c] * gate[c]; + } +} + +// out[m, c] = (x[m, c] - mean(x[m])) * rstd(x[m]) * w[c] + b[c] +__global__ void kernel_layer_norm_affine(const float * __restrict__ x, + const float * __restrict__ w, + const float * __restrict__ b, + float * __restrict__ out, + int C, float eps) { + const int m = blockIdx.x; + const float * xr = x + (int64_t) m * C; + float * orow = out + (int64_t) m * C; + + float sum = 0.0f, sumsq = 0.0f; + for (int c = threadIdx.x; c < C; c += blockDim.x) { + const float v = xr[c]; + sum += v; + sumsq += v * v; + } + __shared__ float red[2][32]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + sumsq += __shfl_xor_sync(0xffffffff, sumsq, off); + } + const int warp = threadIdx.x / 32; + if (threadIdx.x % 32 == 0) { + red[0][warp] = sum; + red[1][warp] = sumsq; + } + __syncthreads(); + if (warp == 0) { + float s = (threadIdx.x < blockDim.x / 32) ? red[0][threadIdx.x] : 0.0f; + float s2 = (threadIdx.x < blockDim.x / 32) ? red[1][threadIdx.x] : 0.0f; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + s += __shfl_xor_sync(0xffffffff, s, off); + s2 += __shfl_xor_sync(0xffffffff, s2, off); + } + if (threadIdx.x == 0) { + red[0][0] = s; + red[1][0] = s2; + } + } + __syncthreads(); + const float mean = red[0][0] / C; + const float var = red[1][0] / C - mean * mean; + const float rstd = rsqrtf(var + eps); + + for (int c = threadIdx.x; c < C; c += blockDim.x) { + orow[c] = (xr[c] - mean) * rstd * w[c] + b[c]; + } +} + +// out[c] = a[c] + b[c] +__global__ void kernel_vec_add(const float * __restrict__ a, + const float * __restrict__ b, + float * __restrict__ out, + int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) { + out[i] = a[i] + b[i]; + } +} + +} // namespace + +int ada_rms_mod(const float * x, const float * scale, const float * shift, + float * out, int M, int C, float eps, bool with_rms, + cudaStream_t stream) { + const int threads = 256; + if (with_rms) { + kernel_ada_rms<<>>(x, scale, shift, out, C, eps); + } else { + kernel_ada_rms<<>>(x, scale, shift, out, C, eps); + } + const cudaError_t e = cudaGetLastError(); + return (e == cudaSuccess) ? 0 : -static_cast(e); +} + +int gated_residual(const float * residual, const float * branch, const float * gate, + float * out, int M, int C, cudaStream_t stream) { + kernel_gated_residual<<>>(residual, branch, gate, out, C); + const cudaError_t e = cudaGetLastError(); + return (e == cudaSuccess) ? 0 : -static_cast(e); +} + +int layer_norm_affine(const float * x, const float * w, const float * b, + float * out, int M, int C, float eps, cudaStream_t stream) { + kernel_layer_norm_affine<<>>(x, w, b, out, C, eps); + const cudaError_t e = cudaGetLastError(); + return (e == cudaSuccess) ? 0 : -static_cast(e); +} + +int ada_rms_mod_quant(const float * x, const float * scale, const float * shift, + float * out, void * dst_packed, void * dst_sfa, + int M, int C, float eps, bool with_rms, cudaStream_t stream) { + if (C % 16 != 0) return -1; + auto shape = cute::make_shape(M, 1, C, 1); + auto layout = AdaCfg::tile_atom_to_shape_SFA(shape); + if (with_rms) { + kernel_ada_rms_q<<>>(x, scale, shift, out, + (uint8_t *) dst_packed, (uint8_t *) dst_sfa, layout, C, eps); + } else { + kernel_ada_rms_q<<>>(x, scale, shift, out, + (uint8_t *) dst_packed, (uint8_t *) dst_sfa, layout, C, eps); + } + const cudaError_t e = cudaGetLastError(); + return (e == cudaSuccess) ? 0 : -static_cast(e); +} + +int layer_norm_affine_quant(const float * x, const float * w, const float * b, + float * out, void * dst_packed, void * dst_sfa, + int M, int C, float eps, cudaStream_t stream) { + if (C % 16 != 0) return -1; + auto shape = cute::make_shape(M, 1, C, 1); + auto layout = AdaCfg::tile_atom_to_shape_SFA(shape); + kernel_layer_norm_affine_q<<>>(x, w, b, out, + (uint8_t *) dst_packed, (uint8_t *) dst_sfa, layout, C, eps); + const cudaError_t e = cudaGetLastError(); + return (e == cudaSuccess) ? 0 : -static_cast(e); +} + +int vec_add_f32(const float * a, const float * b, float * out, int n, cudaStream_t stream) { + kernel_vec_add<<<(n + 255) / 256, 256, 0, stream>>>(a, b, out, n); + const cudaError_t e = cudaGetLastError(); + return (e == cudaSuccess) ? 0 : -static_cast(e); +} + +} // namespace ggml_cuda_flashrt diff --git a/flash_rt/structures/adapters/ggml/fr_decode_attn.cu b/flash_rt/structures/adapters/ggml/fr_decode_attn.cu new file mode 100644 index 000000000..0579a5edd --- /dev/null +++ b/flash_rt/structures/adapters/ggml/fr_decode_attn.cu @@ -0,0 +1,167 @@ +// Decomposed tiny-M decode attention for the pi0.5 action expert (Thor). +// +// For q_tokens ≤ 16 over a padded f16 KV of token rows, flash attention's +// stream-k kernel plus its fixup pass is slower than the classic +// decomposition the FlashRT torch pipeline uses: one QK^T GEMM, a masked +// softmax over the KV axis, and one PV GEMM. +// +// All GQA query heads share the single KV head, so instead of a batched +// GEMM per head both contractions run as one wide GEMM over the +// n_head*n_tok query rows. Rows are ordered t-major (row r = t*n_head + h): +// with that ordering the PV output column j lands at byte offset j*hd in +// the flash-attention node's [hd, n_head, n_tok] destination, i.e. the +// GEMM writes the fp32 result contiguously with no strided-C penalty (the +// dispatch layer guarantees the destination is contiguous). The t-major +// order also makes the q gather read the [hd, n_head, n_tok]-contiguous +// Q buffer sequentially. +// +// Numerics follow ggml's fattn contract: scores = scale * q.k + mask (f16 +// mask, slope 1 as max_bias must be 0), softmax in fp32 with running max. + +#include "fr_kernels.h" + +#include +#include + +namespace ggml_cuda_flashrt { + +namespace { + +// gather the permuted f32 Q view into contiguous f16 rows [n_tok*n_head, hd] +// (row r = t*n_head + h), applying nothing else (scale folds into QK alpha) +__global__ void kernel_q_gather_f16(const float * __restrict__ q, + __half * __restrict__ out, + int hd, int n_head, + int64_t s_d, int64_t s_tok, int64_t s_head) { + const int r = blockIdx.x; // t*n_head + h + const int h = r % n_head; + const int t = r / n_head; + const float * src = q + (int64_t) h * s_head + (int64_t) t * s_tok; + __half * dst = out + (int64_t) r * hd; + for (int d = threadIdx.x; d < hd; d += blockDim.x) { + dst[d] = __float2half(src[(int64_t) d * s_d]); + } +} + +// in-place masked softmax over rows of [n_tok*n_head, n_kv] f16 scores. +// mask element for (kv, t) at mask + kv + t*mask_stride (f16, -inf on pads). +__global__ void kernel_mask_softmax_f16(__half * __restrict__ scores, + const __half * __restrict__ mask, + int n_kv, int n_head, int64_t mask_stride) { + const int r = blockIdx.x; // t*n_head + h + const int t = r / n_head; + __half * row = scores + (int64_t) r * n_kv; + const __half * mrow = mask + (int64_t) t * mask_stride; + + float m = -INFINITY; + for (int i = threadIdx.x; i < n_kv; i += blockDim.x) { + const float v = __half2float(row[i]) + __half2float(mrow[i]); + m = fmaxf(m, v); + } + __shared__ float red[32]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + m = fmaxf(m, __shfl_xor_sync(0xffffffff, m, off)); + } + if (threadIdx.x % 32 == 0) red[threadIdx.x / 32] = m; + __syncthreads(); + if (threadIdx.x < 32) { + float v = (threadIdx.x < blockDim.x / 32) ? red[threadIdx.x] : -INFINITY; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + v = fmaxf(v, __shfl_xor_sync(0xffffffff, v, off)); + } + if (threadIdx.x == 0) red[0] = v; + } + __syncthreads(); + m = red[0]; + + float sum = 0.0f; + for (int i = threadIdx.x; i < n_kv; i += blockDim.x) { + const float v = __half2float(row[i]) + __half2float(mrow[i]); + const float e = expf(v - m); + row[i] = __float2half(e); + sum += e; + } + __shared__ float red2[32]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off); + } + if (threadIdx.x % 32 == 0) red2[threadIdx.x / 32] = sum; + __syncthreads(); + if (threadIdx.x < 32) { + float v = (threadIdx.x < blockDim.x / 32) ? red2[threadIdx.x] : 0.0f; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + v += __shfl_xor_sync(0xffffffff, v, off); + } + if (threadIdx.x == 0) red2[0] = v; + } + __syncthreads(); + const float inv = 1.0f / red2[0]; + for (int i = threadIdx.x; i < n_kv; i += blockDim.x) { + row[i] = __float2half(__half2float(row[i]) * inv); + } +} + +} // namespace + +int decode_attn_decomposed(void * cublas_handle, + const float * q, int64_t q_sd, int64_t q_stok, int64_t q_shead, + const void * k_f16_rows, // [n_kv, hd] f16 rows + const void * v_f16_rows, // [n_kv, hd] f16 rows + const void * mask_f16, int64_t mask_stride, + float * dst, int64_t dst_stok, int64_t dst_shead, + void * q16_ws, int q16_ready, void * scores_ws, + int hd, int n_tok, int n_head, int n_kv, + float scale, cudaStream_t stream) { + // the contiguous PV store below requires the [hd, n_head, n_tok] dst + // to be dense; the dispatch layer checks the same before fusing + if (dst_shead != hd || dst_stok != (int64_t) hd * n_head) { + return -1; + } + cublasHandle_t handle = (cublasHandle_t) cublas_handle; + const int R = n_head * n_tok; + + if (!q16_ready) { + kernel_q_gather_f16<<>>( + q, (__half *) q16_ws, hd, n_head, q_sd, q_stok, q_shead); + } + + cublasSetStream(handle, stream); + // scores_col[n_kv, R] = K_col^T [n_kv, hd] x Q16_col [hd, R] + const float beta0 = 0.0f; + cublasStatus_t st = cublasGemmEx( + handle, CUBLAS_OP_T, CUBLAS_OP_N, + n_kv, R, hd, + &scale, + k_f16_rows, CUDA_R_16F, hd, + q16_ws, CUDA_R_16F, hd, + &beta0, + scores_ws, CUDA_R_16F, n_kv, + CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT); + if (st != CUBLAS_STATUS_SUCCESS) return -100 - (int) st; + + kernel_mask_softmax_f16<<>>( + (__half *) scores_ws, (const __half *) mask_f16, n_kv, n_head, mask_stride); + + // dst_col[hd, R] (dense, column r = t*n_head + h at offset r*hd) = + // V_col [hd, n_kv] x P_col [n_kv, R] + const float one = 1.0f; + st = cublasGemmEx( + handle, CUBLAS_OP_N, CUBLAS_OP_N, + hd, R, n_kv, + &one, + v_f16_rows, CUDA_R_16F, hd, + scores_ws, CUDA_R_16F, n_kv, + &beta0, + dst, CUDA_R_32F, hd, + CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT); + if (st != CUBLAS_STATUS_SUCCESS) return -200 - (int) st; + + const cudaError_t e = cudaGetLastError(); + return (e == cudaSuccess) ? 0 : -static_cast(e); +} + +} // namespace ggml_cuda_flashrt diff --git a/flash_rt/structures/adapters/ggml/fr_dispatch.cu b/flash_rt/structures/adapters/ggml/fr_dispatch.cu new file mode 100644 index 000000000..edb039903 --- /dev/null +++ b/flash_rt/structures/adapters/ggml/fr_dispatch.cu @@ -0,0 +1,1993 @@ +// Dispatch glue between ggml-cuda's mul_mat and the FlashRT NVFP4 kernels. +// Host-only logic; the device kernels live in the sibling fr_*.cu files. + +#include "fr_ggml.cuh" +#include "fr_kernels.h" + +// FlashRT kernels consumed directly from the flashrt-public csrc tree +// (GGML_CUDA_FLASHRT_PUBLIC_DIR); no vendored copies. +#include "gemm/fp4/cutlass_fp4_gemm_geglu_il_sm100.cuh" +#include "gemm/fp4/cutlass_fp4_gemm_siglip_ffn_f32out_sm100.cuh" +#include "gemm/fp4/cutlass_fp4_gemm_bias_f32b_f16out_sm100.cuh" + +#include +#include + +namespace { + +// Weights repacked into the CUTLASS wire format, keyed by the ggml tensor's +// device pointer. Weight tensors are immutable and live for the process +// lifetime, so entries are never evicted. +struct repacked_weight { + void * packed = nullptr; + void * sf = nullptr; +}; + +std::unordered_map g_repack_cache; +std::mutex g_repack_mu; + +// The repack allocates with cudaMalloc, which is illegal during CUDA graph +// capture. All weights are repacked during the first (uncaptured) warmup +// evaluation of each graph, so a cache miss while capturing indicates a bug. +const repacked_weight * get_repacked(const ggml_tensor * src0, cudaStream_t stream) { + std::lock_guard lk(g_repack_mu); + + auto it = g_repack_cache.find(src0->data); + if (it != g_repack_cache.end()) { + return &it->second; + } + + const int64_t K = src0->ne[0]; + const int64_t N = src0->ne[1]; + + cudaStreamCaptureStatus cap = cudaStreamCaptureStatusNone; + cudaStreamIsCapturing(stream, &cap); + if (cap != cudaStreamCaptureStatusNone) { + GGML_ABORT("flashrt: weight repack for %s requested during CUDA graph capture", src0->name); + } + + repacked_weight w; + CUDA_CHECK(cudaMalloc(&w.packed, ggml_cuda_flashrt::packed_bytes(N, K))); + CUDA_CHECK(cudaMalloc(&w.sf, ggml_cuda_flashrt::sf_bytes(N, K))); + + const int rc = ggml_cuda_flashrt::repack_weight(src0->data, w.packed, w.sf, (int) N, (int) K, stream); + if (rc != 0) { + GGML_ABORT("flashrt: weight repack failed for %s (N=%lld K=%lld rc=%d)", src0->name, (long long) N, (long long) K, rc); + } + + auto res = g_repack_cache.emplace(src0->data, w); + return &res.first->second; +} + +// Interleaved gate/up weight pairs for the fused GeGLU GEMM, keyed by the +// two tensors' device pointers (same immortality caveat as above). +struct pair_key { + const void * gate; + const void * up; + bool operator==(const pair_key & o) const { return gate == o.gate && up == o.up; } +}; +struct pair_key_hash { + size_t operator()(const pair_key & k) const noexcept { + return std::hash()(k.gate) ^ (std::hash()(k.up) << 1); + } +}; + +std::unordered_map g_pair_cache; + +const repacked_weight * get_repacked_pair(const ggml_tensor * gate_w, const ggml_tensor * up_w, cudaStream_t stream) { + std::lock_guard lk(g_repack_mu); + + pair_key key{gate_w->data, up_w->data}; + auto it = g_pair_cache.find(key); + if (it != g_pair_cache.end()) { + return &it->second; + } + + const int64_t K = gate_w->ne[0]; + const int64_t n_ff = gate_w->ne[1]; + const int64_t N_il = 2 * n_ff; + + cudaStreamCaptureStatus cap = cudaStreamCaptureStatusNone; + cudaStreamIsCapturing(stream, &cap); + if (cap != cudaStreamCaptureStatusNone) { + GGML_ABORT("flashrt: geglu pair repack for %s requested during CUDA graph capture", gate_w->name); + } + + repacked_weight w; + CUDA_CHECK(cudaMalloc(&w.packed, ggml_cuda_flashrt::packed_bytes(N_il, K))); + CUDA_CHECK(cudaMalloc(&w.sf, ggml_cuda_flashrt::sf_bytes(N_il, K))); + + const int rc = ggml_cuda_flashrt::repack_weight_pair_interleaved( + gate_w->data, up_w->data, w.packed, w.sf, (int) n_ff, (int) K, stream); + if (rc != 0) { + GGML_ABORT("flashrt: geglu pair repack failed for %s (n_ff=%lld K=%lld rc=%d)", + gate_w->name, (long long) n_ff, (long long) K, rc); + } + + auto res = g_pair_cache.emplace(key, w); + return &res.first->second; +} + +// Per-evaluation quantized-activation cache. Several ops consume the same +// fp32 activation tensor (q/k/v projections, the adaLN conditioning vector +// across all layers); quantizing it once per graph evaluation removes the +// duplicate quantize launches. Keys use the ggml tensor pointer (unique +// within one evaluation) plus an evaluation counter, so recycled device +// addresses across graphs can never alias. Slot buffers are grow-only and +// never freed, which keeps addresses stable for captured CUDA graphs; a +// replayed graph rewrites any slot before its baked consumers read it. +struct act_slot { + const ggml_tensor * key = nullptr; + uint64_t eval_id = 0; + void * packed = nullptr; + size_t packed_cap = 0; + void * sf = nullptr; + size_t sf_cap = 0; +}; + +act_slot g_act_slots[4]; +int g_act_slot_rr = 0; +uint64_t g_eval_id = 1; + +// Returns cached (packed, sf) for src1 quantized as [M, K], quantizing on a +// miss. Returns false when the cache cannot be used (slot growth needed +// while capturing a CUDA graph); the caller must quantize into pool memory. +bool get_quantized_act(const ggml_tensor * src1, int M, int K, + const void ** out_packed, const void ** out_sf, + cudaStream_t stream) { + for (auto & s : g_act_slots) { + if (s.key == src1 && s.eval_id == g_eval_id) { + *out_packed = s.packed; + *out_sf = s.sf; + return true; + } + } + + act_slot & s = g_act_slots[g_act_slot_rr]; + const size_t need_packed = (size_t) ggml_cuda_flashrt::packed_bytes(M, K); + const size_t need_sf = (size_t) ggml_cuda_flashrt::sf_bytes(M, K); + + if (need_packed > s.packed_cap || need_sf > s.sf_cap) { + cudaStreamCaptureStatus cap = cudaStreamCaptureStatusNone; + cudaStreamIsCapturing(stream, &cap); + if (cap != cudaStreamCaptureStatusNone) { + return false; + } + if (need_packed > s.packed_cap) { + if (s.packed != nullptr) { cudaFree(s.packed); } + CUDA_CHECK(cudaMalloc(&s.packed, need_packed)); + s.packed_cap = need_packed; + } + if (need_sf > s.sf_cap) { + if (s.sf != nullptr) { cudaFree(s.sf); } + CUDA_CHECK(cudaMalloc(&s.sf, need_sf)); + s.sf_cap = need_sf; + } + } + g_act_slot_rr = (g_act_slot_rr + 1) % 4; + + const int rc = ggml_cuda_flashrt::quantize_act_f32( + (const float *) src1->data, s.packed, s.sf, M, K, stream); + if (rc != 0) { + GGML_ABORT("flashrt: activation quantize failed (M=%d K=%d rc=%d)", M, K, rc); + } + s.key = src1; + s.eval_id = g_eval_id; + *out_packed = s.packed; + *out_sf = s.sf; + return true; +} + +// Reserve a cache slot for an activation that a producer kernel will fill +// with already-quantized data (fused quantize). Returns false when slot +// growth would be needed during CUDA graph capture. +bool reserve_quantized_act(const ggml_tensor * out_tensor, int M, int K, + void ** out_packed, void ** out_sf, + cudaStream_t stream) { + act_slot & s = g_act_slots[g_act_slot_rr]; + const size_t need_packed = (size_t) ggml_cuda_flashrt::packed_bytes(M, K); + const size_t need_sf = (size_t) ggml_cuda_flashrt::sf_bytes(M, K); + + if (need_packed > s.packed_cap || need_sf > s.sf_cap) { + cudaStreamCaptureStatus cap = cudaStreamCaptureStatusNone; + cudaStreamIsCapturing(stream, &cap); + if (cap != cudaStreamCaptureStatusNone) { + return false; + } + if (need_packed > s.packed_cap) { + if (s.packed != nullptr) { cudaFree(s.packed); } + CUDA_CHECK(cudaMalloc(&s.packed, need_packed)); + s.packed_cap = need_packed; + } + if (need_sf > s.sf_cap) { + if (s.sf != nullptr) { cudaFree(s.sf); } + CUDA_CHECK(cudaMalloc(&s.sf, need_sf)); + s.sf_cap = need_sf; + } + } + g_act_slot_rr = (g_act_slot_rr + 1) % 4; + + s.key = out_tensor; + s.eval_id = g_eval_id; + *out_packed = s.packed; + *out_sf = s.sf; + return true; +} + +// One-shot f16 Q handoff from the fused decode QKV window to the decomposed +// decode attention: qkv_post writes the rope'd+scaled Q rows as f16 in the +// t-major gather order, and the next attention window consumes them instead +// of running its own gather kernel. A single grow-only slot suffices (the +// producer and consumer alternate strictly within each layer); the key is +// cleared on consumption so a recycled activation address can never alias a +// stale entry. +struct q16_slot { + const void * key = nullptr; // data pointer of the Q tensor written for + uint64_t eval_id = 0; + void * buf = nullptr; + size_t cap = 0; +}; +q16_slot g_q16; + +void * reserve_q16(const void * qdata, size_t bytes, cudaStream_t stream) { + if (bytes > g_q16.cap) { + cudaStreamCaptureStatus cap = cudaStreamCaptureStatusNone; + cudaStreamIsCapturing(stream, &cap); + if (cap != cudaStreamCaptureStatusNone) { + return nullptr; + } + if (g_q16.buf != nullptr) { cudaFree(g_q16.buf); } + CUDA_CHECK(cudaMalloc(&g_q16.buf, bytes)); + g_q16.cap = bytes; + } + g_q16.key = qdata; + g_q16.eval_id = g_eval_id; + return g_q16.buf; +} + +// Grow-only device buffer for the never-written D of the no-D-store GeGLU +// variants (the host-side TMA descriptor still needs a valid allocation). +void * get_dummy_d(size_t bytes) { + static void * buf = nullptr; + static size_t cap = 0; + static std::mutex mu; + std::lock_guard lk(mu); + if (bytes > cap) { + if (buf != nullptr) { + cudaFree(buf); + } + CUDA_CHECK(cudaMalloc(&buf, bytes)); + cap = bytes; + } + return buf; +} + +} // namespace + +bool ggml_cuda_flashrt_should_use(const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst) { + static const bool disabled = getenv("GGML_CUDA_FLASHRT_DISABLE") != nullptr; + if (disabled) { + return false; + } + if (src0->type != GGML_TYPE_NVFP4 || src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) { + return false; + } + if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(src1) || !ggml_is_contiguous(dst)) { + return false; + } + // Batched src1 with unbatched (broadcast) weights folds into a single + // GEMM over all rows because src1/dst are fully contiguous. + if (src0->ne[2] != 1 || src0->ne[3] != 1) { + return false; + } + const int64_t K = src0->ne[0]; + const int64_t N = src0->ne[1]; + if (K % 64 != 0 || N % 16 != 0) { + return false; + } + return true; +} + +void ggml_cuda_flashrt_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { + const int K = (int) src0->ne[0]; + const int N = (int) src0->ne[1]; + const int M = (int) ggml_nrows(src1); // batch dims fold into rows (contiguous, broadcast weights) + + cudaStream_t stream = ctx.stream(); + + // The persistent repack cache is keyed by the weight tensor's device + // pointer, which is only sound when weight tensors are immortal (model + // inference). Tools that create and free tensors at recycled addresses + // (e.g. test-backend-ops) must set GGML_CUDA_FLASHRT_NO_CACHE=1 to + // repack into scratch memory on every call instead. + static const bool no_cache = getenv("GGML_CUDA_FLASHRT_NO_CACHE") != nullptr; + + ggml_cuda_pool_alloc b_packed_scratch(ctx.pool()); + ggml_cuda_pool_alloc b_sf_scratch (ctx.pool()); + + const void * b_packed = nullptr; + const void * b_sf = nullptr; + if (no_cache) { + b_packed_scratch.alloc(ggml_cuda_flashrt::packed_bytes(N, K)); + b_sf_scratch.alloc(ggml_cuda_flashrt::sf_bytes(N, K)); + const int rrc = ggml_cuda_flashrt::repack_weight(src0->data, b_packed_scratch.get(), b_sf_scratch.get(), N, K, stream); + if (rrc != 0) { + GGML_ABORT("flashrt: weight repack failed (N=%d K=%d rc=%d)", N, K, rrc); + } + b_packed = b_packed_scratch.get(); + b_sf = b_sf_scratch.get(); + } else { + const repacked_weight * w = get_repacked(src0, stream); + b_packed = w->packed; + b_sf = w->sf; + } + + ggml_cuda_pool_alloc a_packed(ctx.pool()); + ggml_cuda_pool_alloc a_sf (ctx.pool()); + + const void * q_packed = nullptr; + const void * q_sf = nullptr; + if (!get_quantized_act(src1, M, K, &q_packed, &q_sf, stream)) { + a_packed.alloc(ggml_cuda_flashrt::packed_bytes(M, K)); + a_sf.alloc(ggml_cuda_flashrt::sf_bytes(M, K)); + const int qrc = ggml_cuda_flashrt::quantize_act_f32( + (const float *) src1->data, a_packed.get(), a_sf.get(), M, K, stream); + if (qrc != 0) { + GGML_ABORT("flashrt: activation quantize failed (M=%d K=%d rc=%d)", M, K, qrc); + } + q_packed = a_packed.get(); + q_sf = a_sf.get(); + } + + // ggml's block_nvfp4 scale bytes carry standard e4m3 semantics: its + // dequant table doubles the e2m1 values but its ue4m3 decode halves the + // scale, so the two cancel and no alpha compensation is needed. + const float alpha = 1.0f; + // The 128x128x256 tile beats the wide-N 128x256x128 tile on Thor for + // every production shape measured (including N=16384 prefill FFN). + const bool widen = false; + + const int rc = ggml_cuda_flashrt::gemm_f32out( + q_packed, q_sf, b_packed, b_sf, + (float *) dst->data, M, N, K, alpha, widen, stream); + if (rc != 0) { + GGML_ABORT("flashrt: gemm failed (M=%d N=%d K=%d rc=%d)", M, N, K, rc); + } +} + +bool ggml_cuda_flashrt_should_fuse_ada(const ggml_tensor * rms, const ggml_tensor * mm, const ggml_tensor * bias_add, + const ggml_tensor * view_scale, const ggml_tensor * repeat_scale, + const ggml_tensor * mul, const ggml_tensor * add1, + const ggml_tensor * view_shift, const ggml_tensor * repeat_shift, + const ggml_tensor * add2) { + const ggml_tensor * x = rms != nullptr ? rms->src[0] : mul->src[0]; + const ggml_tensor * normed = rms != nullptr ? rms : x; + const int64_t C = x->ne[0]; + const int64_t M = x->ne[1]; + + if (x->type != GGML_TYPE_F32 || !ggml_is_contiguous(x) || x->ne[2] != 1 || x->ne[3] != 1) { + return false; + } + // modulation projection: [3C, 1] from the shared conditioning vector + if (!ggml_cuda_flashrt_should_use(mm->src[0], mm->src[1], mm) || + mm->ne[0] != 3 * C || mm->ne[1] != 1) { + return false; + } + const ggml_tensor * bias = bias_add->src[1]; + if (bias_add->src[0] != mm || bias->type != GGML_TYPE_F32 || + !ggml_is_contiguous(bias) || bias->ne[0] != 3 * C || !ggml_is_contiguous(bias_add)) { + return false; + } + // scale = mod[0:C], shift = mod[C:2C] + if (view_scale->src[0] != bias_add || view_scale->ne[0] != C || view_scale->ne[1] != 1 || + view_scale->view_offs != 0) { + return false; + } + if (view_shift->src[0] != bias_add || view_shift->ne[0] != C || view_shift->ne[1] != 1 || + view_shift->view_offs != (size_t) C * sizeof(float)) { + return false; + } + if (repeat_scale->src[0] != view_scale || repeat_shift->src[0] != view_shift || + repeat_scale->ne[0] != C || repeat_scale->ne[1] != M) { + return false; + } + // t = normed * scale; t2 = normed + t; out = t2 + shift + if (mul->src[0] != normed || mul->src[1] != repeat_scale || + add1->src[0] != normed || add1->src[1] != mul || + add2->src[0] != add1 || add2->src[1] != repeat_shift || + !ggml_is_contiguous(add2) || add2->ne[0] != C || add2->ne[1] != M) { + return false; + } + return true; +} + +void ggml_cuda_flashrt_ada_norm(ggml_backend_cuda_context & ctx, const ggml_tensor * rms, const ggml_tensor * mm, + ggml_tensor * bias_add, const ggml_tensor * view_scale, const ggml_tensor * mul, + const ggml_tensor * view_shift, ggml_tensor * add2) { + const ggml_tensor * x = rms != nullptr ? rms->src[0] : mul->src[0]; + const ggml_tensor * cond = mm->src[1]; + + const int C = (int) x->ne[0]; + const int M = (int) x->ne[1]; + const int K = (int) mm->src[0]->ne[0]; + const int N = (int) mm->src[0]->ne[1]; // 3C + + cudaStream_t stream = ctx.stream(); + + const repacked_weight * w = get_repacked(mm->src[0], stream); + + ggml_cuda_pool_alloc c_packed(ctx.pool()); + ggml_cuda_pool_alloc c_sf (ctx.pool()); + + const void * q_packed = nullptr; + const void * q_sf = nullptr; + int rc = 0; + if (!get_quantized_act(cond, 1, K, &q_packed, &q_sf, stream)) { + c_packed.alloc(ggml_cuda_flashrt::packed_bytes(1, K)); + c_sf.alloc(ggml_cuda_flashrt::sf_bytes(1, K)); + rc = ggml_cuda_flashrt::quantize_act_f32((const float *) cond->data, c_packed.get(), c_sf.get(), 1, K, stream); + q_packed = c_packed.get(); + q_sf = c_sf.get(); + } + if (rc == 0) { + // bias applied in the GEMM epilogue, writing the biased modulation + // vector (still read later through the gate view) directly + rc = flash_rt::fp4::gemm_bias_f32out(q_packed, q_sf, w->packed, w->sf, + bias_add->src[1]->data, + (float *) bias_add->data, 1, N, K, stream); + } + if (rc == 0) { + const float eps = rms != nullptr ? ggml_get_op_params_f32(rms, 0) : 0.0f; + // emit the quantized form alongside f32 so downstream GEMMs skip + // their activation quantize (registered in the per-eval cache) + void * q_out_packed = nullptr; + void * q_out_sf = nullptr; + if (C % 16 == 0 && reserve_quantized_act(add2, M, C, &q_out_packed, &q_out_sf, stream)) { + rc = ggml_cuda_flashrt::ada_rms_mod_quant((const float *) x->data, + (const float *) view_scale->data, + (const float *) view_shift->data, + (float *) add2->data, q_out_packed, q_out_sf, + M, C, eps, rms != nullptr, stream); + } else { + rc = ggml_cuda_flashrt::ada_rms_mod((const float *) x->data, + (const float *) view_scale->data, + (const float *) view_shift->data, + (float *) add2->data, M, C, eps, rms != nullptr, stream); + } + } + if (rc != 0) { + GGML_ABORT("flashrt: fused adaLN failed (M=%d C=%d rc=%d)", M, C, rc); + } +} + +// Cached-modulation adaLN window: the modulation vector comes from a graph +// input (precomputed per step, see pi0_ae.cpp) instead of a GEMM. Sequence: +// {RMS_NORM, VIEW mod-col, VIEW scale, REPEAT, MUL, ADD, VIEW shift, REPEAT, +// ADD} -> one ada kernel reading scale/shift straight from the input views. +static int g_ada_cached_fail = 0; +#define FR_ADA_FAIL(code) do { g_ada_cached_fail = (code); return false; } while (0) + +bool ggml_cuda_flashrt_should_fuse_ada_cached( + const ggml_tensor * rms, const ggml_tensor * view_col, + const ggml_tensor * view_scale, const ggml_tensor * repeat_scale, + const ggml_tensor * mul, const ggml_tensor * add1, + const ggml_tensor * view_shift, const ggml_tensor * repeat_shift, + const ggml_tensor * add2) { + static const bool disabled = getenv("GGML_CUDA_FLASHRT_DISABLE") != nullptr; + if (disabled) { + FR_ADA_FAIL(1); + } + const ggml_tensor * x = rms->src[0]; + if (x == nullptr || x->type != GGML_TYPE_F32 || !ggml_is_contiguous(x)) { + FR_ADA_FAIL(2); + } + const int64_t C = x->ne[0]; + const int64_t M = x->ne[1]; + if (x->ne[2] != 1 || x->ne[3] != 1 || C % 4 != 0) { + FR_ADA_FAIL(3); + } + if (view_col->type != GGML_TYPE_F32 || view_col->ne[0] != 3 * C || view_col->ne[1] != 1 || + view_col->src[0] == nullptr) { + FR_ADA_FAIL(4); + } + // view_offs accumulates through view-of-view chains down to the ultimate + // source, so compare offsets relative to the enclosing column view + if (view_scale->src[0] != view_col || view_scale->ne[0] != C || view_scale->ne[1] != 1 || + view_scale->view_offs != view_col->view_offs) { + if (getenv("GGML_FLASHRT_DEBUG") != nullptr) { + static int dbg5 = 0; + if (dbg5++ < 4) { + fprintf(stderr, "[fr-ada-c5] src_ok=%d ne0=%lld (C=%lld) ne1=%lld offs=%zu col_offs=%zu\n", + (int) (view_scale->src[0] == view_col), (long long) view_scale->ne[0], + (long long) C, (long long) view_scale->ne[1], + view_scale->view_offs, view_col->view_offs); + } + } + FR_ADA_FAIL(5); + } + if (view_shift->src[0] != view_col || view_shift->ne[0] != C || view_shift->ne[1] != 1 || + view_shift->view_offs != view_col->view_offs + (size_t) C * sizeof(float)) { + FR_ADA_FAIL(6); + } + if (repeat_scale->src[0] != view_scale || repeat_scale->ne[0] != C || repeat_scale->ne[1] != M || + repeat_shift->src[0] != view_shift || repeat_shift->ne[0] != C || repeat_shift->ne[1] != M) { + FR_ADA_FAIL(7); + } + if (mul->src[0] != rms || mul->src[1] != repeat_scale || + add1->src[0] != rms || add1->src[1] != mul || + add2->src[0] != add1 || add2->src[1] != repeat_shift || + !ggml_is_contiguous(add2) || add2->ne[0] != C || add2->ne[1] != M) { + FR_ADA_FAIL(8); + } + g_ada_cached_fail = 0; + return true; +} + +int ggml_cuda_flashrt_ada_cached_fail_code() { return g_ada_cached_fail; } + +void ggml_cuda_flashrt_ada_norm_cached(ggml_backend_cuda_context & ctx, const ggml_tensor * rms, + const ggml_tensor * view_scale, const ggml_tensor * view_shift, + ggml_tensor * add2) { + const ggml_tensor * x = rms->src[0]; + const int C = (int) x->ne[0]; + const int M = (int) x->ne[1]; + const float eps = ggml_get_op_params_f32(rms, 0); + cudaStream_t stream = ctx.stream(); + + int rc; + void * q_out_packed = nullptr; + void * q_out_sf = nullptr; + if (C % 16 == 0 && reserve_quantized_act(add2, M, C, &q_out_packed, &q_out_sf, stream)) { + rc = ggml_cuda_flashrt::ada_rms_mod_quant((const float *) x->data, + (const float *) view_scale->data, + (const float *) view_shift->data, + (float *) add2->data, q_out_packed, q_out_sf, + M, C, eps, true, stream); + } else { + rc = ggml_cuda_flashrt::ada_rms_mod((const float *) x->data, + (const float *) view_scale->data, + (const float *) view_shift->data, + (float *) add2->data, M, C, eps, true, stream); + } + if (rc != 0) { + GGML_ABORT("flashrt: cached adaLN failed (M=%d C=%d rc=%d)", M, C, rc); + } +} + +bool ggml_cuda_flashrt_should_fuse_gated_res(const ggml_tensor * view, const ggml_tensor * repeat, + const ggml_tensor * mul, const ggml_tensor * add) { + const int64_t C = view->ne[0]; + const int64_t M = mul->ne[1]; + if (view->type != GGML_TYPE_F32 || view->ne[1] != 1 || + view->src[0] == nullptr || view->src[0]->type != GGML_TYPE_F32) { + return false; + } + if (repeat->src[0] != view || repeat->ne[0] != C || repeat->ne[1] != M) { + return false; + } + const ggml_tensor * branch = mul->src[0]; + if (mul->src[1] != repeat || branch->type != GGML_TYPE_F32 || !ggml_is_contiguous(branch) || + branch->ne[0] != C || branch->ne[1] != M) { + return false; + } + const ggml_tensor * residual = add->src[0] == mul ? add->src[1] : add->src[0]; + if ((add->src[0] != mul && add->src[1] != mul) || residual == mul || + residual->type != GGML_TYPE_F32 || !ggml_is_contiguous(residual) || + residual->ne[0] != C || residual->ne[1] != M || + !ggml_is_contiguous(add) || add->ne[0] != C || add->ne[1] != M) { + return false; + } + return true; +} + +void ggml_cuda_flashrt_gated_residual(ggml_backend_cuda_context & ctx, const ggml_tensor * view, + const ggml_tensor * mul, ggml_tensor * add) { + const ggml_tensor * branch = mul->src[0]; + const ggml_tensor * residual = add->src[0] == mul ? add->src[1] : add->src[0]; + const int C = (int) view->ne[0]; + const int M = (int) mul->ne[1]; + + const int rc = ggml_cuda_flashrt::gated_residual( + (const float *) residual->data, (const float *) branch->data, + (const float *) view->data, (float *) add->data, M, C, ctx.stream()); + if (rc != 0) { + GGML_ABORT("flashrt: fused gated residual failed (M=%d C=%d rc=%d)", M, C, rc); + } +} + +// SigLIP FFN weights prepared for the fused FP4 pair, keyed by the two +// weight pointers. Up rows are padded to a 32 multiple for the FP4 output; +// the f16 Down weight is quantized with its K padded to the same value. +struct siglip_ffn_weights { + void * up_packed = nullptr; + void * up_sf = nullptr; + void * dn_packed = nullptr; + void * dn_sf = nullptr; + int h_pad = 0; +}; + +std::unordered_map g_sig_ffn_cache; + +const siglip_ffn_weights * get_siglip_ffn_weights(const ggml_tensor * up_w, const ggml_tensor * dn_w, cudaStream_t stream) { + std::lock_guard lk(g_repack_mu); + + pair_key key{up_w->data, dn_w->data}; + auto it = g_sig_ffn_cache.find(key); + if (it != g_sig_ffn_cache.end()) { + return &it->second; + } + + const int64_t K_in = up_w->ne[0]; // model dim (NVFP4 up weight) + const int64_t H = up_w->ne[1]; // hidden dim + const int64_t H_pad = GGML_PAD(H, 64); + const int64_t D_out = dn_w->ne[1]; // model dim (f16 down weight, K = H) + + cudaStreamCaptureStatus cap = cudaStreamCaptureStatusNone; + cudaStreamIsCapturing(stream, &cap); + if (cap != cudaStreamCaptureStatusNone) { + GGML_ABORT("flashrt: siglip ffn weight prep for %s requested during CUDA graph capture", up_w->name); + } + + siglip_ffn_weights w; + w.h_pad = (int) H_pad; + CUDA_CHECK(cudaMalloc(&w.up_packed, ggml_cuda_flashrt::packed_bytes(H_pad, K_in))); + CUDA_CHECK(cudaMalloc(&w.up_sf, ggml_cuda_flashrt::sf_bytes(H_pad, K_in))); + CUDA_CHECK(cudaMalloc(&w.dn_packed, ggml_cuda_flashrt::packed_bytes(D_out, H_pad))); + CUDA_CHECK(cudaMalloc(&w.dn_sf, ggml_cuda_flashrt::sf_bytes(D_out, H_pad))); + + int rc = ggml_cuda_flashrt::repack_weight_rows_padded( + up_w->data, w.up_packed, w.up_sf, (int) H, (int) H_pad, (int) K_in, stream); + if (rc == 0) { + rc = ggml_cuda_flashrt::quantize_weight_f16_padded( + dn_w->data, w.dn_packed, w.dn_sf, (int) D_out, (int) H, (int) H_pad, stream); + } + if (rc != 0) { + GGML_ABORT("flashrt: siglip ffn weight prep failed for %s (H=%lld rc=%d)", up_w->name, (long long) H, rc); + } + + auto res = g_sig_ffn_cache.emplace(key, w); + return &res.first->second; +} + +bool ggml_cuda_flashrt_should_fuse_ln(const ggml_tensor * norm, const ggml_tensor * mul, const ggml_tensor * add) { + const ggml_tensor * x = norm->src[0]; + const int64_t C = norm->ne[0]; + if (x->type != GGML_TYPE_F32 || !ggml_is_contiguous(x) || norm->ne[3] != 1) { + return false; + } + const ggml_tensor * w = mul->src[0] == norm ? mul->src[1] : mul->src[0]; + if ((mul->src[0] != norm && mul->src[1] != norm) || w == norm || + w->type != GGML_TYPE_F32 || !ggml_is_contiguous(w) || + w->ne[0] != C || ggml_nelements(w) != C) { + return false; + } + const ggml_tensor * b = add->src[0] == mul ? add->src[1] : add->src[0]; + if ((add->src[0] != mul && add->src[1] != mul) || b == mul || + b->type != GGML_TYPE_F32 || !ggml_is_contiguous(b) || + b->ne[0] != C || ggml_nelements(b) != C) { + return false; + } + if (!ggml_is_contiguous(add) || add->ne[0] != C || ggml_nrows(add) != ggml_nrows(x)) { + return false; + } + return true; +} + +void ggml_cuda_flashrt_ln_affine(ggml_backend_cuda_context & ctx, const ggml_tensor * norm, const ggml_tensor * mul, ggml_tensor * add) { + const ggml_tensor * x = norm->src[0]; + const ggml_tensor * w = mul->src[0] == norm ? mul->src[1] : mul->src[0]; + const ggml_tensor * b = add->src[0] == mul ? add->src[1] : add->src[0]; + + const int C = (int) norm->ne[0]; + const int M = (int) ggml_nrows(x); + const float eps = ggml_get_op_params_f32(norm, 0); + + int rc; + void * q_out_packed = nullptr; + void * q_out_sf = nullptr; + if (C % 16 == 0 && reserve_quantized_act(add, M, C, &q_out_packed, &q_out_sf, ctx.stream())) { + rc = ggml_cuda_flashrt::layer_norm_affine_quant( + (const float *) x->data, (const float *) w->data, (const float *) b->data, + (float *) add->data, q_out_packed, q_out_sf, M, C, eps, ctx.stream()); + } else { + rc = ggml_cuda_flashrt::layer_norm_affine( + (const float *) x->data, (const float *) w->data, (const float *) b->data, + (float *) add->data, M, C, eps, ctx.stream()); + } + if (rc != 0) { + GGML_ABORT("flashrt: fused layer norm failed (M=%d C=%d rc=%d)", M, C, rc); + } +} + +bool ggml_cuda_flashrt_should_fuse_geglu(const ggml_tensor * gate_mm, const ggml_tensor * up_mm, const ggml_tensor * glu, const ggml_tensor * down_mm) { + if (ggml_get_glu_op(glu) != GGML_GLU_OP_GEGLU) { + return false; + } + // both projections share the same input and shape + if (gate_mm->src[1] != up_mm->src[1]) { + return false; + } + const ggml_tensor * gate_w = gate_mm->src[0]; + const ggml_tensor * up_w = up_mm->src[0]; + const ggml_tensor * down_w = down_mm->src[0]; + if (gate_w->ne[0] != up_w->ne[0] || gate_w->ne[1] != up_w->ne[1]) { + return false; + } + if (!ggml_cuda_flashrt_should_use(gate_w, gate_mm->src[1], gate_mm) || + !ggml_cuda_flashrt_should_use(up_w, up_mm->src[1], up_mm) || + !ggml_cuda_flashrt_should_use(down_w, glu, down_mm)) { + return false; + } + // the down projection must consume the GLU output with K = n_ff + if (down_mm->src[1] != glu || down_w->ne[0] != gate_w->ne[1]) { + return false; + } + return true; +} + +void ggml_cuda_flashrt_geglu_ffn(ggml_backend_cuda_context & ctx, const ggml_tensor * gate_mm, const ggml_tensor * up_mm, const ggml_tensor * glu, ggml_tensor * down_mm) { + GGML_UNUSED(glu); + + const ggml_tensor * src1 = gate_mm->src[1]; + const ggml_tensor * gate_w = gate_mm->src[0]; + const ggml_tensor * down_w = down_mm->src[0]; + + const int K = (int) gate_w->ne[0]; + const int n_ff = (int) gate_w->ne[1]; + const int N_il = 2 * n_ff; + const int M = (int) ggml_nrows(src1); + const int N_out = (int) down_w->ne[1]; + + cudaStream_t stream = ctx.stream(); + + const repacked_weight * w_il = get_repacked_pair(gate_w, up_mm->src[0], stream); + const repacked_weight * w_down = get_repacked(down_w, stream); + + ggml_cuda_pool_alloc a_packed(ctx.pool()); + ggml_cuda_pool_alloc a_sf (ctx.pool()); + + const void * q_packed = nullptr; + const void * q_sf = nullptr; + int rc = 0; + if (!get_quantized_act(src1, M, K, &q_packed, &q_sf, stream)) { + a_packed.alloc(ggml_cuda_flashrt::packed_bytes(M, K)); + a_sf.alloc(ggml_cuda_flashrt::sf_bytes(M, K)); + rc = ggml_cuda_flashrt::quantize_act_f32( + (const float *) src1->data, a_packed.get(), a_sf.get(), M, K, stream); + if (rc != 0) { + GGML_ABORT("flashrt: geglu activation quantize failed (M=%d K=%d rc=%d)", M, K, rc); + } + q_packed = a_packed.get(); + q_sf = a_sf.get(); + } + + ggml_cuda_pool_alloc compact_packed(ctx.pool(), ggml_cuda_flashrt::packed_bytes(M, n_ff)); + ggml_cuda_pool_alloc compact_sfa (ctx.pool(), ggml_cuda_flashrt::sf_bytes(M, n_ff)); + void * dummy_d = get_dummy_d(ggml_cuda_flashrt::packed_bytes(M, N_il)); + + // skinny-M tile for decode-sized batches, default tile otherwise + if (M < 128) { + rc = flash_rt::fp4::cutlass_fp4_gemm_geglu_il_hw_nod_v10( + q_packed, q_sf, w_il->packed, w_il->sf, + dummy_d, compact_packed.get(), compact_sfa.get(), M, N_il, K, stream); + } else { + rc = flash_rt::fp4::cutlass_fp4_gemm_geglu_il_hw_nod( + q_packed, q_sf, w_il->packed, w_il->sf, + dummy_d, compact_packed.get(), compact_sfa.get(), M, N_il, K, stream); + } + if (rc != 0) { + GGML_ABORT("flashrt: geglu gemm failed (M=%d N_il=%d K=%d rc=%d)", M, N_il, K, rc); + } + + rc = ggml_cuda_flashrt::gemm_f32out( + compact_packed.get(), compact_sfa.get(), w_down->packed, w_down->sf, + (float *) down_mm->data, M, N_out, n_ff, /*alpha=*/1.0f, /*widen=*/false, stream); + if (rc != 0) { + GGML_ABORT("flashrt: geglu down gemm failed (M=%d N=%d K=%d rc=%d)", M, N_out, n_ff, rc); + } +} + +// Fused QKV weights (row-concat [k | v | q]) keyed by the three pointers. +struct triple_key { + const void * a; const void * b; const void * c; + bool operator==(const triple_key & o) const { return a == o.a && b == o.b && c == o.c; } +}; +struct triple_key_hash { + size_t operator()(const triple_key & k) const noexcept { + return std::hash()(k.a) ^ (std::hash()(k.b) << 1) ^ (std::hash()(k.c) << 2); + } +}; + +std::unordered_map g_qkv_cache; + +const repacked_weight * get_repacked_qkv(const ggml_tensor * wk, const ggml_tensor * wv, const ggml_tensor * wq, cudaStream_t stream) { + std::lock_guard lk(g_repack_mu); + + triple_key key{wk->data, wv->data, wq->data}; + auto it = g_qkv_cache.find(key); + if (it != g_qkv_cache.end()) { + return &it->second; + } + + const int64_t K = wk->ne[0]; + const int64_t N_tot = wk->ne[1] + wv->ne[1] + wq->ne[1]; + + cudaStreamCaptureStatus cap = cudaStreamCaptureStatusNone; + cudaStreamIsCapturing(stream, &cap); + if (cap != cudaStreamCaptureStatusNone) { + GGML_ABORT("flashrt: qkv repack for %s requested during CUDA graph capture", wq->name); + } + + repacked_weight w; + CUDA_CHECK(cudaMalloc(&w.packed, ggml_cuda_flashrt::packed_bytes(N_tot, K))); + CUDA_CHECK(cudaMalloc(&w.sf, ggml_cuda_flashrt::sf_bytes(N_tot, K))); + + const int rc = ggml_cuda_flashrt::repack_weight_concat3( + wk->data, (int) wk->ne[1], wv->data, (int) wv->ne[1], wq->data, (int) wq->ne[1], + w.packed, w.sf, (int) K, stream); + if (rc != 0) { + GGML_ABORT("flashrt: qkv repack failed for %s (rc=%d)", wq->name, rc); + } + + auto res = g_qkv_cache.emplace(key, w); + return &res.first->second; +} + +bool ggml_cuda_flashrt_should_fuse_siglip_ffn(const ggml_tensor * up_mm, const ggml_tensor * bias1, const ggml_tensor * gelu, + const ggml_tensor * cont1, const ggml_tensor * dn_mm, const ggml_tensor * bias2, + const ggml_tensor * cont2, const ggml_tensor * res_add) { + const ggml_tensor * up_w = up_mm->src[0]; + const ggml_tensor * src1 = up_mm->src[1]; + const ggml_tensor * dn_w = dn_mm->src[0]; + + if (!ggml_cuda_flashrt_should_use(up_w, src1, up_mm)) { + return false; + } + if (dn_w->type != GGML_TYPE_F16 || !ggml_is_contiguous(dn_w) || + dn_w->ne[0] != up_w->ne[1] || dn_w->ne[0] % 16 != 0 || + dn_w->ne[1] % 16 != 0 || dn_w->ne[2] != 1) { + return false; + } + if (ggml_get_unary_op(gelu) != GGML_UNARY_OP_GELU || gelu->src[0] != bias1) { + return false; + } + const ggml_tensor * b1 = bias1->src[0] == up_mm ? bias1->src[1] : bias1->src[0]; + if ((bias1->src[0] != up_mm && bias1->src[1] != up_mm) || b1 == up_mm || + b1->type != GGML_TYPE_F32 || !ggml_is_contiguous(b1) || ggml_nelements(b1) != up_w->ne[1]) { + return false; + } + if (cont1->src[0] != gelu || dn_mm->src[1] != cont1) { + return false; + } + const ggml_tensor * b2 = bias2->src[0] == dn_mm ? bias2->src[1] : bias2->src[0]; + if ((bias2->src[0] != dn_mm && bias2->src[1] != dn_mm) || b2 == dn_mm || + b2->type != GGML_TYPE_F32 || !ggml_is_contiguous(b2) || ggml_nelements(b2) != dn_w->ne[1]) { + return false; + } + if (cont2->src[0] != bias2) { + return false; + } + const ggml_tensor * residual = res_add->src[0] == cont2 ? res_add->src[1] : res_add->src[0]; + if ((res_add->src[0] != cont2 && res_add->src[1] != cont2) || residual == cont2 || + residual->type != GGML_TYPE_F32 || !ggml_is_contiguous(residual) || + !ggml_is_contiguous(res_add) || res_add->ne[0] != dn_w->ne[1] || + ggml_nrows(res_add) != ggml_nrows(src1) || ggml_nrows(residual) != ggml_nrows(src1)) { + return false; + } + return true; +} + +void ggml_cuda_flashrt_siglip_ffn(ggml_backend_cuda_context & ctx, const ggml_tensor * up_mm, const ggml_tensor * bias1, + const ggml_tensor * dn_mm, const ggml_tensor * bias2, + const ggml_tensor * cont2, ggml_tensor * res_add) { + const ggml_tensor * up_w = up_mm->src[0]; + const ggml_tensor * src1 = up_mm->src[1]; + const ggml_tensor * dn_w = dn_mm->src[0]; + const ggml_tensor * b1 = bias1->src[0] == up_mm ? bias1->src[1] : bias1->src[0]; + const ggml_tensor * b2 = bias2->src[0] == dn_mm ? bias2->src[1] : bias2->src[0]; + const ggml_tensor * residual = res_add->src[0] == cont2 ? res_add->src[1] : res_add->src[0]; + + const int K_in = (int) up_w->ne[0]; + const int D_out = (int) dn_w->ne[1]; + const int M = (int) ggml_nrows(src1); + + cudaStream_t stream = ctx.stream(); + + const siglip_ffn_weights * w = get_siglip_ffn_weights(up_w, dn_w, stream); + const int H_pad = w->h_pad; + + const void * q_packed = nullptr; + const void * q_sf = nullptr; + ggml_cuda_pool_alloc a_packed(ctx.pool()); + ggml_cuda_pool_alloc a_sf (ctx.pool()); + int rc = 0; + if (!get_quantized_act(src1, M, K_in, &q_packed, &q_sf, stream)) { + a_packed.alloc(ggml_cuda_flashrt::packed_bytes(M, K_in)); + a_sf.alloc(ggml_cuda_flashrt::sf_bytes(M, K_in)); + rc = ggml_cuda_flashrt::quantize_act_f32((const float *) src1->data, a_packed.get(), a_sf.get(), M, K_in, stream); + q_packed = a_packed.get(); + q_sf = a_sf.get(); + } + + ggml_cuda_pool_alloc hid_packed(ctx.pool(), ggml_cuda_flashrt::packed_bytes(M, H_pad)); + ggml_cuda_pool_alloc hid_sf (ctx.pool(), ggml_cuda_flashrt::sf_bytes(M, H_pad)); + + if (rc == 0) { + rc = flash_rt::fp4::siglip_ffn_up_gelu_fp4out( + q_packed, q_sf, w->up_packed, w->up_sf, b1->data, + hid_packed.get(), hid_sf.get(), M, H_pad, K_in, stream); + } + if (rc == 0) { + rc = flash_rt::fp4::siglip_ffn_down_bias_res_f32( + hid_packed.get(), hid_sf.get(), w->dn_packed, w->dn_sf, b2->data, + residual->data, res_add->data, M, D_out, H_pad, stream, 1.0f); + } + if (rc != 0) { + GGML_ABORT("flashrt: siglip ffn failed (M=%d H_pad=%d rc=%d)", M, H_pad, rc); + } +} + +bool ggml_cuda_flashrt_should_fuse_qkv(const ggml_tensor * k_mm, const ggml_tensor * k_rope, const ggml_tensor * k_cpy, + const ggml_tensor * v_mm, const ggml_tensor * v_cpy, + const ggml_tensor * q_mm, const ggml_tensor * q_rope, const ggml_tensor * q_scale) { + const ggml_tensor * src1 = k_mm->src[1]; + if (v_mm->src[1] != src1 || q_mm->src[1] != src1) { + return false; + } + const ggml_tensor * wk = k_mm->src[0]; + const ggml_tensor * wv = v_mm->src[0]; + const ggml_tensor * wq = q_mm->src[0]; + if (!ggml_cuda_flashrt_should_use(wk, src1, k_mm) || + !ggml_cuda_flashrt_should_use(wv, src1, v_mm) || + !ggml_cuda_flashrt_should_use(wq, src1, q_mm) || + wk->ne[0] != wv->ne[0] || wk->ne[0] != wq->ne[0]) { + return false; + } + // K path: mm -> reshape -> rope -> cpy into an f16 view of the + // persistent KV buffer; single KV head (Nk == head_dim) + const int64_t head_dim = k_rope->src[0]->ne[0]; + if (k_rope->src[0]->op != GGML_OP_RESHAPE || k_rope->src[0]->src[0] != k_mm || + wk->ne[1] != head_dim || k_rope->src[0]->ne[1] != 1) { + return false; + } + if (k_cpy->src[0] != k_rope || k_cpy->src[1] == nullptr || + k_cpy->src[1]->type != GGML_TYPE_F16 || k_cpy->src[1]->op != GGML_OP_VIEW) { + return false; + } + // V path: mm -> reshape -> cpy into f16 view + if (v_cpy->src[0] == nullptr || v_cpy->src[0]->op != GGML_OP_RESHAPE || + v_cpy->src[0]->src[0] != v_mm || wv->ne[1] != head_dim || + v_cpy->src[1] == nullptr || v_cpy->src[1]->type != GGML_TYPE_F16 || v_cpy->src[1]->op != GGML_OP_VIEW) { + return false; + } + // Q path: mm -> reshape -> rope -> scale, head_dim x n_head + if (q_rope->src[0]->op != GGML_OP_RESHAPE || q_rope->src[0]->src[0] != q_mm || + q_rope->src[0]->ne[0] != head_dim || wq->ne[1] % head_dim != 0 || + q_scale->src[0] != q_rope || !ggml_is_contiguous(q_scale)) { + return false; + } + // both ropes must share positions, freq factors and parameters + if (k_rope->src[1] != q_rope->src[1] || k_rope->src[2] != q_rope->src[2] || + memcmp(k_rope->op_params, q_rope->op_params, sizeof(k_rope->op_params)) != 0) { + return false; + } + // NEOX rope only (matches ggml's rope_neox math replicated in qkv_post) + const int mode = ((const int32_t *) k_rope->op_params)[2]; + if (mode != GGML_ROPE_TYPE_NEOX) { + return false; + } + // f16 KV views must be row-contiguous (head_dim elements per token row) + const ggml_tensor * kv = k_cpy->src[1]; + if (kv->nb[0] != sizeof(uint16_t) || kv->nb[1] != head_dim * sizeof(uint16_t)) { + return false; + } + return true; +} + +void ggml_cuda_flashrt_qkv(ggml_backend_cuda_context & ctx, + const ggml_tensor * k_mm, const ggml_tensor * k_rope, const ggml_tensor * k_cpy, + const ggml_tensor * v_mm, const ggml_tensor * v_cpy, + const ggml_tensor * q_mm, const ggml_tensor * q_rope, ggml_tensor * q_scale) { + const ggml_tensor * src1 = k_mm->src[1]; + const ggml_tensor * wk = k_mm->src[0]; + const ggml_tensor * wv = v_mm->src[0]; + const ggml_tensor * wq = q_mm->src[0]; + + const int K = (int) wk->ne[0]; + const int Nk = (int) wk->ne[1]; + const int Nv = (int) wv->ne[1]; + const int Nq = (int) wq->ne[1]; + const int M = (int) ggml_nrows(src1); + const int head_dim = Nk; + + cudaStream_t stream = ctx.stream(); + + const repacked_weight * w = get_repacked_qkv(wk, wv, wq, stream); + + const void * q_packed = nullptr; + const void * q_sf = nullptr; + ggml_cuda_pool_alloc a_packed(ctx.pool()); + ggml_cuda_pool_alloc a_sf (ctx.pool()); + int rc = 0; + if (!get_quantized_act(src1, M, K, &q_packed, &q_sf, stream)) { + a_packed.alloc(ggml_cuda_flashrt::packed_bytes(M, K)); + a_sf.alloc(ggml_cuda_flashrt::sf_bytes(M, K)); + rc = ggml_cuda_flashrt::quantize_act_f32((const float *) src1->data, a_packed.get(), a_sf.get(), M, K, stream); + q_packed = a_packed.get(); + q_sf = a_sf.get(); + } + + const int N_tot = Nk + Nv + Nq; + ggml_cuda_pool_alloc qkv_cat(ctx.pool(), (int64_t) M * N_tot); + + if (rc == 0) { + rc = ggml_cuda_flashrt::gemm_f32out(q_packed, q_sf, w->packed, w->sf, + qkv_cat.get(), M, N_tot, K, 1.0f, false, stream); + } + if (rc == 0) { + // rope parameters, mirrored from ggml-cuda's rope host setup + const int32_t * op = (const int32_t *) k_rope->op_params; + const int n_dims = op[1]; + const int n_ctx_orig = op[4]; + float freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow; + memcpy(&freq_base, op + 5, sizeof(float)); + memcpy(&freq_scale, op + 6, sizeof(float)); + memcpy(&ext_factor, op + 7, sizeof(float)); + memcpy(&attn_factor, op + 8, sizeof(float)); + memcpy(&beta_fast, op + 9, sizeof(float)); + memcpy(&beta_slow, op + 10, sizeof(float)); + float corr_dims[2]; + ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims); + const float theta_scale = powf(freq_base, -2.0f / n_dims); + const float scale_f = ggml_get_op_params_f32(q_scale, 0); + + const ggml_tensor * ff = k_rope->src[2]; + void * q16 = reserve_q16(q_scale->data, + (size_t) M * Nq * sizeof(uint16_t), stream); + rc = ggml_cuda_flashrt::qkv_post( + qkv_cat.get(), (float *) q_scale->data, q16, + k_cpy->src[1]->data, v_cpy->src[1]->data, + (const int32_t *) k_rope->src[1]->data, + ff != nullptr ? (const float *) ff->data : nullptr, + M, Nk, Nv, Nq, head_dim, n_dims, + freq_scale, ext_factor, attn_factor, + corr_dims[0], corr_dims[1], theta_scale, scale_f, stream); + } + if (rc != 0) { + GGML_ABORT("flashrt: fused qkv failed (M=%d N=%d K=%d rc=%d)", M, N_tot, K, rc); + } +} + +// --------------------------------------------------------------------------- +// Vision QKV pad window: {mul_mat, add bias, reshape}x3 + {pad}x3 -> three +// padded-weight GEMMs with the bias in the epilogue, writing the pad buffers +// directly. Widens per-head projections (SigLIP head_dim 72 -> 80) inside the +// repacked weights, so the runtime pad kernels and separate bias adds vanish. + +namespace { + +struct grouppad_weight { + void * packed = nullptr; + void * sf = nullptr; + void * bias = nullptr; // f32 [group_out * n_groups], zero-interleaved +}; + +std::unordered_map g_grouppad_cache; + +const grouppad_weight * get_repacked_grouppad( + const ggml_tensor * w, const ggml_tensor * bias, + int group_in, int group_out, int n_groups, cudaStream_t stream) { + std::lock_guard lk(g_repack_mu); + + auto it = g_grouppad_cache.find(w->data); + if (it != g_grouppad_cache.end()) { + return &it->second; + } + + cudaStreamCaptureStatus cap = cudaStreamCaptureStatusNone; + cudaStreamIsCapturing(stream, &cap); + if (cap != cudaStreamCaptureStatusNone) { + GGML_ABORT("flashrt: grouppad repack for %s requested during CUDA graph capture", w->name); + } + + const int K = (int) w->ne[0]; + const int N_pad = group_out * n_groups; + + grouppad_weight g; + CUDA_CHECK(cudaMalloc(&g.packed, ggml_cuda_flashrt::packed_bytes(N_pad, K))); + CUDA_CHECK(cudaMalloc(&g.sf, ggml_cuda_flashrt::sf_bytes(N_pad, K))); + CUDA_CHECK(cudaMalloc(&g.bias, (size_t) N_pad * sizeof(float))); + + const int rc = ggml_cuda_flashrt::repack_weight_rows_grouppad( + w->data, g.packed, g.sf, group_in, group_out, n_groups, K, stream); + if (rc != 0) { + GGML_ABORT("flashrt: grouppad repack failed for %s (gi=%d go=%d ng=%d K=%d rc=%d)", + w->name, group_in, group_out, n_groups, K, rc); + } + CUDA_CHECK(cudaMemsetAsync(g.bias, 0, (size_t) N_pad * sizeof(float), stream)); + CUDA_CHECK(cudaMemcpy2DAsync(g.bias, (size_t) group_out * sizeof(float), + bias->data, (size_t) group_in * sizeof(float), + (size_t) group_in * sizeof(float), (size_t) n_groups, + cudaMemcpyDeviceToDevice, stream)); + + auto res = g_grouppad_cache.emplace(w->data, g); + return &res.first->second; +} + +// One projection triple of the window: mul_mat -> add(bias) -> reshape -> pad. +bool vis_qkv_pad_leg_ok(const ggml_tensor * mm, const ggml_tensor * add, + const ggml_tensor * resh, const ggml_tensor * pad, + const ggml_tensor * shared_src1) { + if (mm->op != GGML_OP_MUL_MAT || add->op != GGML_OP_ADD || + resh->op != GGML_OP_RESHAPE || pad->op != GGML_OP_PAD) { + return false; + } + const ggml_tensor * w = mm->src[0]; + const ggml_tensor * b = add->src[1]; + if (w == nullptr || w->type != GGML_TYPE_NVFP4 || mm->src[1] != shared_src1) { + return false; + } + if (add->src[0] != mm || resh->src[0] != add || pad->src[0] != resh) { + return false; + } + if (b == nullptr || b->type != GGML_TYPE_F32 || !ggml_is_contiguous(b) || + b->ne[0] != mm->ne[0] || ggml_nrows(b) != 1) { + return false; + } + const int64_t group_in = resh->ne[0]; + const int64_t n_groups = resh->ne[1]; + const int64_t group_out = pad->ne[0]; + if (group_in * n_groups != mm->ne[0] || group_out <= group_in) { + return false; + } + // pad only widens dim0 + if (pad->ne[1] != resh->ne[1] || pad->ne[2] != resh->ne[2] || pad->ne[3] != resh->ne[3]) { + return false; + } + if (pad->type != GGML_TYPE_F32 || !ggml_is_contiguous(pad)) { + return false; + } + const int64_t K = w->ne[0]; + const int64_t N_pad = group_out * n_groups; + if (K % 64 != 0 || N_pad % 16 != 0) { + return false; + } + return true; +} + +} // namespace + +bool ggml_cuda_flashrt_should_fuse_vis_qkv_pad( + const ggml_tensor * mm_q, const ggml_tensor * add_q, const ggml_tensor * resh_q, + const ggml_tensor * mm_k, const ggml_tensor * add_k, const ggml_tensor * resh_k, + const ggml_tensor * mm_v, const ggml_tensor * add_v, const ggml_tensor * resh_v, + const ggml_tensor * pad_q, const ggml_tensor * pad_k, const ggml_tensor * pad_v) { + static const bool disabled = getenv("GGML_CUDA_FLASHRT_DISABLE") != nullptr; + if (disabled) { + return false; + } + const ggml_tensor * src1 = mm_q->src[1]; + if (src1 == nullptr || src1->type != GGML_TYPE_F32 || !ggml_is_contiguous(src1)) { + return false; + } + if (!vis_qkv_pad_leg_ok(mm_q, add_q, resh_q, pad_q, src1) || + !vis_qkv_pad_leg_ok(mm_k, add_k, resh_k, pad_k, src1) || + !vis_qkv_pad_leg_ok(mm_v, add_v, resh_v, pad_v, src1)) { + return false; + } + // identical geometry across the three legs + if (mm_k->ne[0] != mm_q->ne[0] || mm_v->ne[0] != mm_q->ne[0] || + pad_k->ne[0] != pad_q->ne[0] || pad_v->ne[0] != pad_q->ne[0] || + resh_k->ne[0] != resh_q->ne[0] || resh_v->ne[0] != resh_q->ne[0]) { + return false; + } + const int64_t M = ggml_nrows(src1); + if (M <= 0 || M > INT32_MAX) { + return false; + } + return true; +} + +void ggml_cuda_flashrt_vis_qkv_pad(ggml_backend_cuda_context & ctx, + const ggml_tensor * mm_q, const ggml_tensor * add_q, ggml_tensor * pad_q, + const ggml_tensor * mm_k, const ggml_tensor * add_k, ggml_tensor * pad_k, + const ggml_tensor * mm_v, const ggml_tensor * add_v, ggml_tensor * pad_v, + ggml_tensor * k_cast, ggml_tensor * v_cast) { + cudaStream_t stream = ctx.stream(); + + const ggml_tensor * src1 = mm_q->src[1]; + const int K = (int) mm_q->src[0]->ne[0]; + const int M = (int) ggml_nrows(src1); + + const int group_in = (int) (pad_q->src[0]->ne[0]); + const int group_out = (int) pad_q->ne[0]; + const int n_groups = (int) pad_q->ne[1]; + const int N_pad = group_out * n_groups; + + ggml_cuda_pool_alloc a_packed(ctx.pool()); + ggml_cuda_pool_alloc a_sf (ctx.pool()); + const void * q_packed = nullptr; + const void * q_sf = nullptr; + if (!get_quantized_act(src1, M, K, &q_packed, &q_sf, stream)) { + a_packed.alloc(ggml_cuda_flashrt::packed_bytes(M, K)); + a_sf.alloc(ggml_cuda_flashrt::sf_bytes(M, K)); + const int qrc = ggml_cuda_flashrt::quantize_act_f32( + (const float *) src1->data, a_packed.get(), a_sf.get(), M, K, stream); + if (qrc != 0) { + GGML_ABORT("flashrt: vis qkv activation quantize failed (M=%d K=%d rc=%d)", M, K, qrc); + } + q_packed = a_packed.get(); + q_sf = a_sf.get(); + } + + // K/V may go straight to their f16 cast tensors (single rounding from + // the fp32 accumulator, bitwise equal to f32-out + cast); Q stays f32. + const ggml_tensor * legs[3][4] = { + { mm_q, add_q, pad_q, nullptr }, + { mm_k, add_k, pad_k, k_cast }, + { mm_v, add_v, pad_v, v_cast }, + }; + for (auto & leg : legs) { + const grouppad_weight * w = get_repacked_grouppad( + leg[0]->src[0], leg[1]->src[1], group_in, group_out, n_groups, stream); + int rc; + if (leg[3] != nullptr) { + rc = flash_rt::fp4::gemm_bias_f16out( + q_packed, q_sf, w->packed, w->sf, w->bias, + leg[3]->data, M, N_pad, K, stream); + } else { + rc = flash_rt::fp4::gemm_bias_f32out( + q_packed, q_sf, w->packed, w->sf, w->bias, + (float *) leg[2]->data, M, N_pad, K, stream); + } + if (rc != 0) { + GGML_ABORT("flashrt: vis qkv padded gemm failed (M=%d N=%d K=%d rc=%d)", M, N_pad, K, rc); + } + } +} + +// --------------------------------------------------------------------------- +// GEMM + (bias) + residual window: {mul_mat, add bias?, add residual} -> one +// GEMM with the bias and the residual in the epilogue (C may alias D; the +// epilogue reads C before writing D). Covers the encoder/vision o- and +// down-projections whose residual adds were separate bandwidth passes. + +namespace { + +std::unordered_map g_zero_bias_cache; // keyed by N + +const void * get_zero_bias(int N, cudaStream_t stream) { + std::lock_guard lk(g_repack_mu); + auto it = g_zero_bias_cache.find(N); + if (it != g_zero_bias_cache.end()) { + return it->second; + } + cudaStreamCaptureStatus cap = cudaStreamCaptureStatusNone; + cudaStreamIsCapturing(stream, &cap); + if (cap != cudaStreamCaptureStatusNone) { + return nullptr; // caller falls back; allocation is capture-illegal + } + void * p = nullptr; + CUDA_CHECK(cudaMalloc(&p, (size_t) N * sizeof(float))); + CUDA_CHECK(cudaMemsetAsync(p, 0, (size_t) N * sizeof(float), stream)); + g_zero_bias_cache.emplace(N, p); + return p; +} + +const ggml_tensor * mm_res_residual(const ggml_tensor * chain, const ggml_tensor * res_add) { + const ggml_tensor * r = res_add->src[0] == chain ? res_add->src[1] + : res_add->src[1] == chain ? res_add->src[0] : nullptr; + if (r == nullptr || r == chain || r->type != GGML_TYPE_F32 || !ggml_is_contiguous(r) || + !ggml_are_same_shape(r, res_add)) { + return nullptr; + } + return r; +} + +} // namespace + +bool ggml_cuda_flashrt_should_fuse_mm_res(const ggml_tensor * mm, const ggml_tensor * bias_add, + const ggml_tensor * res_add) { + if (!ggml_cuda_flashrt_should_use(mm->src[0], mm->src[1], mm)) { + return false; + } + const int64_t N = mm->ne[0]; + const int64_t K = mm->src[0]->ne[0]; + if (N % 16 != 0 || K % 64 != 0) { + return false; + } + const ggml_tensor * chain = mm; + if (bias_add != nullptr) { + const ggml_tensor * b = bias_add->src[1]; + if (bias_add->src[0] != mm || b == nullptr || b->type != GGML_TYPE_F32 || + !ggml_is_contiguous(b) || b->ne[0] != N || ggml_nrows(b) != 1) { + return false; + } + chain = bias_add; + } + if (mm_res_residual(chain, res_add) == nullptr || + !ggml_is_contiguous(res_add) || !ggml_are_same_shape(res_add, mm)) { + return false; + } + return true; +} + +bool ggml_cuda_flashrt_mm_res(ggml_backend_cuda_context & ctx, const ggml_tensor * mm, + const ggml_tensor * bias_add, ggml_tensor * res_add) { + const ggml_tensor * src1 = mm->src[1]; + const int M = (int) ggml_nrows(src1); + const int N = (int) mm->ne[0]; + const int K = (int) mm->src[0]->ne[0]; + cudaStream_t stream = ctx.stream(); + + const void * bias = bias_add != nullptr ? bias_add->src[1]->data : get_zero_bias(N, stream); + if (bias == nullptr) { + return false; // zero-bias alloc during capture: run unfused + } + const ggml_tensor * residual = mm_res_residual(bias_add != nullptr ? bias_add : mm, res_add); + + const repacked_weight * w = get_repacked(mm->src[0], stream); + + ggml_cuda_pool_alloc a_packed(ctx.pool()); + ggml_cuda_pool_alloc a_sf (ctx.pool()); + const void * q_packed = nullptr; + const void * q_sf = nullptr; + if (!get_quantized_act(src1, M, K, &q_packed, &q_sf, stream)) { + a_packed.alloc(ggml_cuda_flashrt::packed_bytes(M, K)); + a_sf.alloc(ggml_cuda_flashrt::sf_bytes(M, K)); + const int qrc = ggml_cuda_flashrt::quantize_act_f32( + (const float *) src1->data, a_packed.get(), a_sf.get(), M, K, stream); + if (qrc != 0) { + GGML_ABORT("flashrt: mm+res activation quantize failed (M=%d K=%d rc=%d)", M, K, qrc); + } + q_packed = a_packed.get(); + q_sf = a_sf.get(); + } + + const int rc = flash_rt::fp4::siglip_ffn_down_bias_res_f32( + q_packed, q_sf, w->packed, w->sf, bias, + residual->data, (float *) res_add->data, M, N, K, stream, 1.0f); + if (rc != 0) { + GGML_ABORT("flashrt: mm+res fused gemm failed (M=%d N=%d K=%d rc=%d)", M, N, K, rc); + } + return true; +} + +// ── Prefill fused QKV: {q mm→reshape→rope→scale, k mm→reshape→rope→pad, +// v mm→reshape→pad, permute→cpy ×2} ───────────────────────────────────── +// Same fused GEMM + qkv_post as the decode window, but K/V land in per-eval +// padded f16 tensors (token rows of head_dim, KV length padded to the FA KQ +// stride) instead of the persistent KV suffix; the pad rows are zeroed to +// match the graph's PAD semantics (the FA mask multiplies them out, but +// garbage f16 there would poison the softmax with inf/nan). +bool ggml_cuda_flashrt_should_fuse_qkv_prefill( + const ggml_tensor * q_mm, const ggml_tensor * q_rope, const ggml_tensor * q_scale, + const ggml_tensor * k_mm, const ggml_tensor * k_rope, const ggml_tensor * k_pad, + const ggml_tensor * v_mm, const ggml_tensor * v_pad, + const ggml_tensor * k_cpy, const ggml_tensor * v_cpy) { + static const bool disabled = getenv("GGML_FLASHRT_NO_QKV_PREFILL") != nullptr; + if (disabled) { + return false; + } + const ggml_tensor * src1 = q_mm->src[1]; + if (k_mm->src[1] != src1 || v_mm->src[1] != src1) { + return false; + } + const ggml_tensor * wq = q_mm->src[0]; + const ggml_tensor * wk = k_mm->src[0]; + const ggml_tensor * wv = v_mm->src[0]; + if (!ggml_cuda_flashrt_should_use(wq, src1, q_mm) || + !ggml_cuda_flashrt_should_use(wk, src1, k_mm) || + !ggml_cuda_flashrt_should_use(wv, src1, v_mm) || + wk->ne[0] != wv->ne[0] || wk->ne[0] != wq->ne[0]) { + return false; + } + const int64_t head_dim = wk->ne[1]; + if (wv->ne[1] != head_dim || wq->ne[1] % head_dim != 0) { + return false; + } + // Q: mm -> reshape [hd, n_head, M] -> rope -> scale (contiguous f32 out) + if (q_rope->src[0]->op != GGML_OP_RESHAPE || q_rope->src[0]->src[0] != q_mm || + q_rope->src[0]->ne[0] != head_dim || + q_scale->src[0] != q_rope || !ggml_is_contiguous(q_scale) || + q_scale->type != GGML_TYPE_F32) { + return false; + } + // K: mm -> reshape [hd, 1, M] -> rope -> pad along the token dim + if (k_rope->src[0]->op != GGML_OP_RESHAPE || k_rope->src[0]->src[0] != k_mm || + k_rope->src[0]->ne[0] != head_dim || k_rope->src[0]->ne[1] != 1) { + return false; + } + const int64_t M = k_rope->src[0]->ne[2]; + const int64_t kvp = k_pad->ne[2]; + if (k_pad->src[0] != k_rope || k_pad->ne[0] != head_dim || k_pad->ne[1] != 1 || + kvp < M) { + return false; + } + // V: mm -> reshape -> pad, same geometry + if (v_pad->src[0] == nullptr || v_pad->src[0]->op != GGML_OP_RESHAPE || + v_pad->src[0]->src[0] != v_mm || v_pad->src[0]->ne[0] != head_dim || + v_pad->src[0]->ne[1] != 1 || v_pad->src[0]->ne[2] != M || + v_pad->ne[0] != head_dim || v_pad->ne[1] != 1 || v_pad->ne[2] != kvp) { + return false; + } + // both copies materialize [hd, kvp] f16 token rows + for (const ggml_tensor * cpy : { k_cpy, v_cpy }) { + if (cpy->type != GGML_TYPE_F16 || !ggml_is_contiguous(cpy) || + cpy->ne[0] != head_dim || cpy->ne[1] != kvp || cpy->ne[2] != 1 || + cpy->nb[1] != head_dim * sizeof(uint16_t)) { + return false; + } + } + // ropes share positions, freq factors and parameters; NEOX math only + if (k_rope->src[1] != q_rope->src[1] || k_rope->src[2] != q_rope->src[2] || + memcmp(k_rope->op_params, q_rope->op_params, sizeof(k_rope->op_params)) != 0) { + return false; + } + const int mode = ((const int32_t *) k_rope->op_params)[2]; + if (mode != GGML_ROPE_TYPE_NEOX) { + return false; + } + return true; +} + +void ggml_cuda_flashrt_qkv_prefill(ggml_backend_cuda_context & ctx, + const ggml_tensor * q_mm, const ggml_tensor * q_rope, ggml_tensor * q_scale, + const ggml_tensor * k_mm, const ggml_tensor * k_rope, + const ggml_tensor * v_mm, + ggml_tensor * k_cpy, ggml_tensor * v_cpy) { + const ggml_tensor * src1 = q_mm->src[1]; + const ggml_tensor * wq = q_mm->src[0]; + const ggml_tensor * wk = k_mm->src[0]; + const ggml_tensor * wv = v_mm->src[0]; + + const int K = (int) wk->ne[0]; + const int Nk = (int) wk->ne[1]; + const int Nv = (int) wv->ne[1]; + const int Nq = (int) wq->ne[1]; + const int M = (int) ggml_nrows(src1); + const int head_dim = Nk; + const int kvp = (int) k_cpy->ne[1]; + + cudaStream_t stream = ctx.stream(); + + const repacked_weight * w = get_repacked_qkv(wk, wv, wq, stream); + + const void * q_packed = nullptr; + const void * q_sf = nullptr; + ggml_cuda_pool_alloc a_packed(ctx.pool()); + ggml_cuda_pool_alloc a_sf (ctx.pool()); + int rc = 0; + if (!get_quantized_act(src1, M, K, &q_packed, &q_sf, stream)) { + a_packed.alloc(ggml_cuda_flashrt::packed_bytes(M, K)); + a_sf.alloc(ggml_cuda_flashrt::sf_bytes(M, K)); + rc = ggml_cuda_flashrt::quantize_act_f32((const float *) src1->data, a_packed.get(), a_sf.get(), M, K, stream); + q_packed = a_packed.get(); + q_sf = a_sf.get(); + } + + const int N_tot = Nk + Nv + Nq; + ggml_cuda_pool_alloc qkv_cat(ctx.pool(), (int64_t) M * N_tot); + + if (rc == 0) { + rc = ggml_cuda_flashrt::gemm_f32out(q_packed, q_sf, w->packed, w->sf, + qkv_cat.get(), M, N_tot, K, 1.0f, false, stream); + } + if (rc == 0 && kvp > M) { + // zero the pad rows once per eval; qkv_post then fills rows [0, M) + const size_t row_bytes = (size_t) head_dim * sizeof(uint16_t); + cudaMemsetAsync((char *) k_cpy->data + (size_t) M * row_bytes, 0, (size_t) (kvp - M) * row_bytes, stream); + cudaMemsetAsync((char *) v_cpy->data + (size_t) M * row_bytes, 0, (size_t) (kvp - M) * row_bytes, stream); + } + if (rc == 0) { + const int32_t * op = (const int32_t *) k_rope->op_params; + const int n_dims = op[1]; + const int n_ctx_orig = op[4]; + float freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow; + memcpy(&freq_base, op + 5, sizeof(float)); + memcpy(&freq_scale, op + 6, sizeof(float)); + memcpy(&ext_factor, op + 7, sizeof(float)); + memcpy(&attn_factor, op + 8, sizeof(float)); + memcpy(&beta_fast, op + 9, sizeof(float)); + memcpy(&beta_slow, op + 10, sizeof(float)); + float corr_dims[2]; + ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims); + const float theta_scale = powf(freq_base, -2.0f / n_dims); + const float scale_f = ggml_get_op_params_f32(q_scale, 0); + + const ggml_tensor * ff = k_rope->src[2]; + // the rope'd K and plain V rows also land in their graph tensors' + // f32 buffers, feeding the graph-tail persistent-KV store copies + rc = ggml_cuda_flashrt::qkv_post_full( + // no f16 Q handoff on the prefill path (flash attention reads f32 Q) + qkv_cat.get(), (float *) q_scale->data, nullptr, + k_cpy->data, v_cpy->data, + (float *) k_rope->data, (float *) v_mm->data, + (const int32_t *) k_rope->src[1]->data, + ff != nullptr ? (const float *) ff->data : nullptr, + M, Nk, Nv, Nq, head_dim, n_dims, + freq_scale, ext_factor, attn_factor, + corr_dims[0], corr_dims[1], theta_scale, scale_f, stream); + } + if (rc != 0) { + GGML_ABORT("flashrt: fused prefill qkv failed (M=%d N=%d K=%d rc=%d)", M, N_tot, K, rc); + } +} + +// ── Decomposed tiny-M decode attention ────────────────────────────────────── +// FLASH_ATTN_EXT with q_tokens <= 16, single f16 KV head of token rows, an +// f16 mask and no ALiBi/softcap runs as QK-GEMM + masked softmax + PV-GEMM +// (see fr_decode_attn.cu). Faster than the stream-k fattn + fixup pair at +// these shapes. +bool ggml_cuda_flashrt_should_fuse_dec_attn(const ggml_tensor * fa) { + static const bool disabled = getenv("GGML_FLASHRT_NO_DEC_ATTN") != nullptr; + if (disabled) { + return false; + } + const ggml_tensor * q = fa->src[0]; + const ggml_tensor * k = fa->src[1]; + const ggml_tensor * v = fa->src[2]; + const ggml_tensor * mask = fa->src[3]; + if (q == nullptr || k == nullptr || v == nullptr || mask == nullptr || + fa->src[4] != nullptr) { // no attention sinks + return false; + } + const int64_t hd = q->ne[0]; + const int64_t n_tok = q->ne[1]; + const int64_t n_head = q->ne[2]; + const int64_t n_kv = k->ne[1]; + if (q->type != GGML_TYPE_F32 || n_tok > 16 || q->ne[3] != 1 || + hd % 2 != 0 || n_head < 1) { + return false; + } + for (const ggml_tensor * kv : { k, v }) { + if (kv->type != GGML_TYPE_F16 || kv->ne[0] != hd || kv->ne[2] != 1 || + kv->ne[3] != 1 || kv->nb[0] != sizeof(uint16_t) || + (int64_t) kv->nb[1] != hd * (int64_t) sizeof(uint16_t)) { + return false; + } + } + if (v->ne[1] != n_kv) { + return false; + } + if (mask->type != GGML_TYPE_F16 || mask->ne[0] < n_kv || mask->ne[1] < n_tok) { + return false; + } + if (fa->type != GGML_TYPE_F32 || fa->ne[0] != hd || fa->ne[1] != n_head || + fa->ne[2] != n_tok || fa->nb[0] != sizeof(float) || + (int64_t) fa->nb[1] != hd * (int64_t) sizeof(float) || + // the PV GEMM stores its result as one dense column-block + (int64_t) fa->nb[2] != hd * n_head * (int64_t) sizeof(float)) { + return false; + } + float max_bias, softcap; + memcpy(&max_bias, (const float *) fa->op_params + 1, sizeof(float)); + memcpy(&softcap, (const float *) fa->op_params + 2, sizeof(float)); + if (max_bias != 0.0f || softcap != 0.0f) { + return false; + } + return true; +} + +void ggml_cuda_flashrt_dec_attn(ggml_backend_cuda_context & ctx, ggml_tensor * fa) { + const ggml_tensor * q = fa->src[0]; + const ggml_tensor * k = fa->src[1]; + const ggml_tensor * v = fa->src[2]; + const ggml_tensor * mask = fa->src[3]; + const int hd = (int) q->ne[0]; + const int n_tok = (int) q->ne[1]; + const int n_head = (int) q->ne[2]; + const int n_kv = (int) k->ne[1]; + float scale; + memcpy(&scale, (const float *) fa->op_params + 0, sizeof(float)); + cudaStream_t stream = ctx.stream(); + + const int64_t R = (int64_t) n_head * n_tok; + ggml_cuda_pool_alloc scores(ctx.pool(), R * n_kv * sizeof(uint16_t)); + + // consume the f16 Q rows the fused QKV window handed off (one-shot); + // they are valid only when the Q view is the dense t-major layout the + // handoff was written in + const bool q16_hit = g_q16.key != nullptr && g_q16.key == q->data && + g_q16.eval_id == g_eval_id && + q->nb[0] == sizeof(float) && + (int64_t) q->nb[2] == (int64_t) hd * sizeof(float) && + (int64_t) q->nb[1] == (int64_t) hd * n_head * sizeof(float); + ggml_cuda_pool_alloc q16; + void * q16p; + if (q16_hit) { + q16p = g_q16.buf; + g_q16.key = nullptr; + } else { + q16.alloc(ctx.pool(), R * hd * sizeof(uint16_t)); + q16p = q16.get(); + } + + const int rc = ggml_cuda_flashrt::decode_attn_decomposed( + (void *) ctx.cublas_handle(), + (const float *) q->data, + (int64_t) (q->nb[0] / sizeof(float)), + (int64_t) (q->nb[1] / sizeof(float)), + (int64_t) (q->nb[2] / sizeof(float)), + k->data, v->data, + mask->data, (int64_t) (mask->nb[1] / sizeof(uint16_t)), + (float *) fa->data, + (int64_t) (fa->nb[2] / sizeof(float)), + (int64_t) (fa->nb[1] / sizeof(float)), + q16p, q16_hit ? 1 : 0, scores.get(), + hd, n_tok, n_head, n_kv, scale, stream); + if (rc != 0) { + GGML_ABORT("flashrt: decomposed decode attention failed (tok=%d kv=%d rc=%d)", n_tok, n_kv, rc); + } +} + +// ── Batched persistent-KV tail copies ─────────────────────────────────────── +// The prefill graph ends with one f32->f16 row-copy per layer and KV tensor +// into the persistent encoder-KV buffers. Each is a tiny kernel; a run of +// them batches into a single launch with identical rounding. +bool ggml_cuda_flashrt_kv_tail_cpy_ok(const ggml_tensor * cpy, int64_t * hd, int64_t * n_rows) { + static const bool disabled = getenv("GGML_FLASHRT_NO_KV_TAIL") != nullptr; + if (disabled) { + return false; + } + const ggml_tensor * src = cpy->src[0]; + if (cpy->op != GGML_OP_CPY || src == nullptr || cpy->src[1] == nullptr) { + return false; + } + if (src->type != GGML_TYPE_F32 || cpy->type != GGML_TYPE_F16 || + cpy->src[1]->type != GGML_TYPE_F16) { + return false; + } + const int64_t d = src->ne[0]; + const int64_t r = src->ne[2]; + if (src->ne[1] != 1 || src->ne[3] != 1 || d % 2 != 0 || + cpy->ne[0] != d || cpy->ne[1] != 1 || cpy->ne[2] != r || cpy->ne[3] != 1) { + return false; + } + // contiguous rows on both sides (row stride == hd elements) + if (src->nb[0] != sizeof(float) || (int64_t) src->nb[2] != d * (int64_t) sizeof(float) || + cpy->nb[0] != sizeof(uint16_t) || (int64_t) cpy->nb[2] != d * (int64_t) sizeof(uint16_t)) { + return false; + } + if (*hd == 0) { + *hd = d; + *n_rows = r; + } else if (*hd != d || *n_rows != r) { + return false; + } + return true; +} + +bool ggml_cuda_flashrt_kv_tail_cpy(ggml_backend_cuda_context & ctx, ggml_tensor ** cpys, int n) { + if (n < 1 || n > FR_CPY_ROWS_MAX) { + return false; + } + const float * srcs[FR_CPY_ROWS_MAX]; + void * dsts[FR_CPY_ROWS_MAX]; + for (int i = 0; i < n; ++i) { + srcs[i] = (const float *) cpys[i]->src[0]->data; + dsts[i] = cpys[i]->data; + } + const int64_t hd = cpys[0]->ne[0]; + const int64_t n_rows = cpys[0]->ne[2]; + const int rc = ggml_cuda_flashrt::cpy_rows_f32_f16( + srcs, dsts, n, (int) hd, (int) n_rows, ctx.stream()); + if (rc != 0) { + GGML_ABORT("flashrt: batched kv tail copy failed (n=%d hd=%lld rows=%lld rc=%d)", + n, (long long) hd, (long long) n_rows, rc); + } + return true; +} + +// ── SigLIP vision attention via the AOT FlashAttention-4 module ───────────── +// FLASH_ATTN_EXT with head_dim 80, no mask, f32 Q and f16 K/V whose padded +// buffers all share the dense (B, S, H, D) linear layout. Runs as a dense +// f32->f16 Q convert + the FA4 forward + a dense f16->f32 output convert. +bool ggml_cuda_flashrt_should_fuse_vit_fa4(const ggml_tensor * fa, ggml_backend_cuda_context & ctx) { +#ifndef GGML_CUDA_FLASHRT_FA4 + GGML_UNUSED(fa); GGML_UNUSED(ctx); + return false; +#else + static const bool disabled = getenv("GGML_FLASHRT_NO_VIT_FA4") != nullptr; + if (disabled) { + return false; + } + const ggml_tensor * q = fa->src[0]; + const ggml_tensor * k = fa->src[1]; + const ggml_tensor * v = fa->src[2]; + if (q == nullptr || k == nullptr || v == nullptr || + fa->src[3] != nullptr || fa->src[4] != nullptr) { // no mask, no sinks + return false; + } + const int64_t hd = q->ne[0]; + const int64_t S = q->ne[1]; + const int64_t H = q->ne[2]; + const int64_t B = q->ne[3]; + if (hd != 80 || S < 32 || H < 1 || B < 1 || q->type != GGML_TYPE_F32) { + return false; + } + // Q: permuted view over the dense (B, S, H, D) f32 buffer + if (q->nb[0] != sizeof(float) || + (int64_t) q->nb[2] != hd * (int64_t) sizeof(float) || + (int64_t) q->nb[1] != hd * H * (int64_t) sizeof(float) || + (int64_t) q->nb[3] != hd * H * S * (int64_t) sizeof(float)) { + return false; + } + for (const ggml_tensor * kv : { k, v }) { + if (kv->type != GGML_TYPE_F16 || kv->ne[0] != hd || kv->ne[1] != S || + kv->ne[2] != H || kv->ne[3] != B || + kv->nb[0] != sizeof(uint16_t) || + (int64_t) kv->nb[2] != hd * (int64_t) sizeof(uint16_t) || + (int64_t) kv->nb[1] != hd * H * (int64_t) sizeof(uint16_t) || + (int64_t) kv->nb[3] != hd * H * S * (int64_t) sizeof(uint16_t)) { + return false; + } + } + // dst: contiguous [hd, H, S, B] f32 — the same linear layout + if (fa->type != GGML_TYPE_F32 || fa->ne[0] != hd || fa->ne[1] != H || + fa->ne[2] != S || fa->ne[3] != B || fa->nb[0] != sizeof(float) || + (int64_t) fa->nb[1] != hd * (int64_t) sizeof(float) || + (int64_t) fa->nb[2] != hd * H * (int64_t) sizeof(float) || + (int64_t) fa->nb[3] != hd * H * S * (int64_t) sizeof(float)) { + return false; + } + float max_bias, softcap; + memcpy(&max_bias, (const float *) fa->op_params + 1, sizeof(float)); + memcpy(&softcap, (const float *) fa->op_params + 2, sizeof(float)); + if (max_bias != 0.0f || softcap != 0.0f) { + return false; + } + // module load must happen outside CUDA graph capture; fall back until then + return ggml_cuda_flashrt::fa4_vit_ensure_loaded(ctx.stream()) == 0; +#endif // GGML_CUDA_FLASHRT_FA4 +} + +void ggml_cuda_flashrt_vit_fa4(ggml_backend_cuda_context & ctx, ggml_tensor * fa) { +#ifndef GGML_CUDA_FLASHRT_FA4 + GGML_UNUSED(ctx); GGML_UNUSED(fa); + GGML_ABORT("flashrt: FA4 vision attention not built"); +#else + const ggml_tensor * q = fa->src[0]; + const int B = (int) q->ne[3]; + const int S = (int) q->ne[1]; + const int H = (int) q->ne[2]; + const int D = (int) q->ne[0]; + float scale; + memcpy(&scale, (const float *) fa->op_params + 0, sizeof(float)); + + const int64_t n = (int64_t) B * S * H * D; + ggml_cuda_pool_alloc q16(ctx.pool(), n * sizeof(uint16_t)); + ggml_cuda_pool_alloc o16(ctx.pool(), n * sizeof(uint16_t)); + + const int rc = ggml_cuda_flashrt::fa4_vit_attention( + (const float *) q->data, fa->src[1]->data, fa->src[2]->data, + (float *) fa->data, D, q16.get(), o16.get(), + B, S, H, D, scale, ctx.stream()); + if (rc != 0) { + GGML_ABORT("flashrt: FA4 vision attention failed (B=%d S=%d H=%d rc=%d)", B, S, H, rc); + } +#endif // GGML_CUDA_FLASHRT_FA4 +} + +// Variant absorbing the {VIEW (head de-pad), CONT} pair after the FA node: +// the FA4 output converts directly into the CONT's packed destination. +bool ggml_cuda_flashrt_should_fuse_vit_fa4_depad(const ggml_tensor * fa, const ggml_tensor * view, + const ggml_tensor * cont, ggml_backend_cuda_context & ctx) { +#ifndef GGML_CUDA_FLASHRT_FA4 + GGML_UNUSED(fa); GGML_UNUSED(view); GGML_UNUSED(cont); GGML_UNUSED(ctx); + return false; +#else + if (!ggml_cuda_flashrt_should_fuse_vit_fa4(fa, ctx)) { + return false; + } + const int64_t D = fa->ne[0]; + const int64_t H = fa->ne[1]; + const int64_t S = fa->ne[2]; + const int64_t B = fa->ne[3]; + // view: leading d2 <= D slice of the FA output, no offset + if (view->src[0] != fa || view->type != GGML_TYPE_F32 || + view->data != fa->data || + view->ne[0] > D || view->ne[1] != H || view->ne[2] != S || view->ne[3] != B) { + return false; + } + const int64_t D2 = view->ne[0]; + // cont: packed [(H*D2), S, B] contiguous f32 + if (cont->src[0] != view || cont->type != GGML_TYPE_F32 || + cont->ne[0] != H * D2 || cont->ne[1] != S || cont->ne[2] != B || cont->ne[3] != 1 || + !ggml_is_contiguous(cont)) { + return false; + } + return true; +#endif // GGML_CUDA_FLASHRT_FA4 +} + +void ggml_cuda_flashrt_vit_fa4_depad(ggml_backend_cuda_context & ctx, ggml_tensor * fa, + const ggml_tensor * view, ggml_tensor * cont) { +#ifndef GGML_CUDA_FLASHRT_FA4 + GGML_UNUSED(ctx); GGML_UNUSED(fa); GGML_UNUSED(view); GGML_UNUSED(cont); + GGML_ABORT("flashrt: FA4 vision attention not built"); +#else + const ggml_tensor * q = fa->src[0]; + const int B = (int) q->ne[3]; + const int S = (int) q->ne[1]; + const int H = (int) q->ne[2]; + const int D = (int) q->ne[0]; + const int D2 = (int) view->ne[0]; + float scale; + memcpy(&scale, (const float *) fa->op_params + 0, sizeof(float)); + + const int64_t n = (int64_t) B * S * H * D; + ggml_cuda_pool_alloc q16(ctx.pool(), n * sizeof(uint16_t)); + ggml_cuda_pool_alloc o16(ctx.pool(), n * sizeof(uint16_t)); + + const int rc = ggml_cuda_flashrt::fa4_vit_attention( + (const float *) q->data, fa->src[1]->data, fa->src[2]->data, + (float *) cont->data, D2, q16.get(), o16.get(), + B, S, H, D, scale, ctx.stream()); + if (rc != 0) { + GGML_ABORT("flashrt: FA4 vision attention (depad) failed (B=%d S=%d H=%d rc=%d)", B, S, H, rc); + } +#endif // GGML_CUDA_FLASHRT_FA4 +} + +// ── pi0.5 prefill self-attention via the AOT FA4 module ──────────────────── +// FLASH_ATTN_EXT with head_dim 256, one f16 KV head of contiguous token +// rows and an f16 mask. The pi0.5 prefill is a prefix-LM: full attention +// with a row-uniform pad-only mask, and the real KV length equals the +// query count (self-attention), so slicing the padded KV to S reproduces +// the mask exactly. The shape gate (hd 256, S >= 64, padded KV) is +// specific to that graph; GGML_FLASHRT_NO_PREFILL_FA4 disables the window. +bool ggml_cuda_flashrt_should_fuse_prefill_fa4(const ggml_tensor * fa, ggml_backend_cuda_context & ctx) { +#ifndef GGML_CUDA_FLASHRT_FA4 + GGML_UNUSED(fa); GGML_UNUSED(ctx); + return false; +#else + static const bool disabled = getenv("GGML_FLASHRT_NO_PREFILL_FA4") != nullptr; + if (disabled) { + return false; + } + const ggml_tensor * q = fa->src[0]; + const ggml_tensor * k = fa->src[1]; + const ggml_tensor * v = fa->src[2]; + const ggml_tensor * mask = fa->src[3]; + if (q == nullptr || k == nullptr || v == nullptr || mask == nullptr || + fa->src[4] != nullptr) { + return false; + } + const int64_t hd = q->ne[0]; + const int64_t S = q->ne[1]; + const int64_t H = q->ne[2]; + if (hd != 256 || S < 64 || H < 2 || q->ne[3] != 1 || q->type != GGML_TYPE_F32) { + return false; + } + // Q: permuted view over the dense (S, H, D) f32 buffer + if (q->nb[0] != sizeof(float) || + (int64_t) q->nb[2] != hd * (int64_t) sizeof(float) || + (int64_t) q->nb[1] != hd * H * (int64_t) sizeof(float)) { + return false; + } + // K/V: one head of contiguous f16 token rows, padded to a multiple of + // 256 covering exactly S (the pi0.5 prefill padding scheme) + const int64_t SK = k->ne[1]; + if (SK < S || SK % 256 != 0 || SK - S >= 256) { + return false; + } + for (const ggml_tensor * kv : { k, v }) { + if (kv->type != GGML_TYPE_F16 || kv->ne[0] != hd || kv->ne[1] != SK || + kv->ne[2] != 1 || kv->ne[3] != 1 || + kv->nb[0] != sizeof(uint16_t) || + (int64_t) kv->nb[1] != hd * (int64_t) sizeof(uint16_t)) { + return false; + } + } + if (mask->type != GGML_TYPE_F16 || mask->ne[0] < SK || mask->ne[1] < S) { + return false; + } + // dst: contiguous [hd, H, S] f32 — the same dense linear layout + if (fa->type != GGML_TYPE_F32 || fa->ne[0] != hd || fa->ne[1] != H || + fa->ne[2] != S || fa->ne[3] != 1 || fa->nb[0] != sizeof(float) || + (int64_t) fa->nb[1] != hd * (int64_t) sizeof(float) || + (int64_t) fa->nb[2] != hd * H * (int64_t) sizeof(float)) { + return false; + } + float max_bias, softcap; + memcpy(&max_bias, (const float *) fa->op_params + 1, sizeof(float)); + memcpy(&softcap, (const float *) fa->op_params + 2, sizeof(float)); + if (max_bias != 0.0f || softcap != 0.0f) { + return false; + } + return ggml_cuda_flashrt::fa4_vit_ensure_loaded(ctx.stream()) == 0; +#endif // GGML_CUDA_FLASHRT_FA4 +} + +void ggml_cuda_flashrt_prefill_fa4(ggml_backend_cuda_context & ctx, ggml_tensor * fa) { +#ifndef GGML_CUDA_FLASHRT_FA4 + GGML_UNUSED(ctx); GGML_UNUSED(fa); + GGML_ABORT("flashrt: FA4 prefill attention not built"); +#else + const ggml_tensor * q = fa->src[0]; + const int S = (int) q->ne[1]; + const int H = (int) q->ne[2]; + const int D = (int) q->ne[0]; + float scale; + memcpy(&scale, (const float *) fa->op_params + 0, sizeof(float)); + + const int64_t n = (int64_t) S * H * D; + ggml_cuda_pool_alloc q16(ctx.pool(), n * sizeof(uint16_t)); + ggml_cuda_pool_alloc o16(ctx.pool(), n * sizeof(uint16_t)); + + const int rc = ggml_cuda_flashrt::fa4_prefill_attention( + (const float *) q->data, fa->src[1]->data, fa->src[2]->data, + (float *) fa->data, q16.get(), o16.get(), + S, H, D, scale, ctx.stream()); + if (rc != 0) { + GGML_ABORT("flashrt: FA4 prefill attention failed (S=%d H=%d rc=%d)", S, H, rc); + } +#endif // GGML_CUDA_FLASHRT_FA4 +} + +// ── Gemma-style norm fold: {RMS_NORM, MUL(w), ADD(mul, norm)} ──────────────── +// out = rms_norm(x)*w + rms_norm(x) == rms_norm(x)*(1 + w): the adaLN +// modulate kernel with scale = w and shift = 0. ggml's own fused rms_norm +// cannot express this form (the add operand is the norm output itself, which +// no longer exists once the chain is fused), so it never fires on it. +bool ggml_cuda_flashrt_should_fuse_rms_gemma(const ggml_tensor * rms, const ggml_tensor * mul, + const ggml_tensor * add) { + static const bool disabled = getenv("GGML_FLASHRT_NO_RMS_GEMMA") != nullptr; + if (disabled) { + return false; + } + const ggml_tensor * x = rms->src[0]; + if (rms->type != GGML_TYPE_F32 || x == nullptr || x->type != GGML_TYPE_F32 || + !ggml_is_contiguous(rms) || !ggml_is_contiguous(x) || rms->ne[3] != 1) { + return false; + } + if (mul->src[0] != rms) { + return false; + } + const ggml_tensor * w = mul->src[1]; + if (w == nullptr || w->type != GGML_TYPE_F32 || !ggml_is_contiguous(w) || + w->ne[0] != rms->ne[0] || ggml_nrows(w) != 1) { + return false; + } + if (!((add->src[0] == mul && add->src[1] == rms) || + (add->src[0] == rms && add->src[1] == mul))) { + return false; + } + if (add->type != GGML_TYPE_F32 || !ggml_is_contiguous(add) || + !ggml_are_same_shape(add, rms)) { + return false; + } + return true; +} + +bool ggml_cuda_flashrt_rms_gemma(ggml_backend_cuda_context & ctx, const ggml_tensor * rms, + const ggml_tensor * mul, ggml_tensor * add) { + const ggml_tensor * x = rms->src[0]; + const int M = (int) ggml_nrows(x); + const int C = (int) x->ne[0]; + float eps; + memcpy(&eps, rms->op_params, sizeof(float)); + cudaStream_t stream = ctx.stream(); + + const void * zeros = get_zero_bias(C, stream); + if (zeros == nullptr) { + return false; // zero-vector alloc during capture: run unfused + } + const int rc = ggml_cuda_flashrt::ada_rms_mod((const float *) x->data, (const float *) mul->src[1]->data, + (const float *) zeros, (float *) add->data, M, C, eps, + /*with_rms=*/true, stream); + if (rc != 0) { + GGML_ABORT("flashrt: rms_gemma fused kernel failed (M=%d C=%d rc=%d)", M, C, rc); + } + return true; +} + +void ggml_cuda_flashrt_begin_eval() { + g_eval_id++; +#ifdef GGML_CUDA_FLASHRT_FA4 + // begin_eval runs before any CUDA graph capture starts, so the AOT + // module load never has to race a capturing first evaluation + ggml_cuda_flashrt::fa4_vit_ensure_loaded(nullptr); +#endif // GGML_CUDA_FLASHRT_FA4 +} diff --git a/flash_rt/structures/adapters/ggml/fr_fa4_shims.c b/flash_rt/structures/adapters/ggml/fr_fa4_shims.c new file mode 100644 index 000000000..d3186754e --- /dev/null +++ b/flash_rt/structures/adapters/ggml/fr_fa4_shims.c @@ -0,0 +1,42 @@ +// Runtime shims the CuTe-DSL AOT object expects: stable "_cuda*" aliases of +// the CUDA runtime/driver entry points it calls. +#include +#include + +cudaError_t _cudaGetDevice(int * dev) { + return cudaGetDevice(dev); +} + +cudaError_t _cudaDeviceGetAttribute(int * value, enum cudaDeviceAttr attr, int device) { + return cudaDeviceGetAttribute(value, attr, device); +} + +cudaError_t _cudaFuncSetAttribute(const void * func, enum cudaFuncAttribute attr, int value) { + return cudaFuncSetAttribute(func, attr, value); +} + +cudaError_t _cudaKernelSetAttributeForDevice(cudaKernel_t kernel, enum cudaFuncAttribute attr, + int value, int device) { + return cudaKernelSetAttributeForDevice(kernel, attr, value, device); +} + +cudaError_t _cudaLaunchKernelEx(const cudaLaunchConfig_t * config, const void * func, void ** args) { + return cudaLaunchKernelExC(config, func, args); +} + +cudaError_t _cudaLibraryGetKernel(cudaKernel_t * kernel, cudaLibrary_t library, const char * name) { + return cudaLibraryGetKernel(kernel, library, name); +} + +cudaError_t _cudaLibraryLoadData(cudaLibrary_t * library, const void * code, + enum cudaJitOption * jitOptions, void ** jitOptionsValues, + unsigned int numJitOptions, + enum cudaLibraryOption * libraryOptions, + void ** libraryOptionValues, unsigned int numLibraryOptions) { + return cudaLibraryLoadData(library, code, jitOptions, jitOptionsValues, numJitOptions, + libraryOptions, libraryOptionValues, numLibraryOptions); +} + +CUresult _cuKernelGetAttribute(int * pi, CUfunction_attribute attrib, CUkernel kernel, CUdevice dev) { + return cuKernelGetAttribute(pi, attrib, kernel, dev); +} diff --git a/flash_rt/structures/adapters/ggml/fr_fa4_vit.cu b/flash_rt/structures/adapters/ggml/fr_fa4_vit.cu new file mode 100644 index 000000000..9318b61d5 --- /dev/null +++ b/flash_rt/structures/adapters/ggml/fr_fa4_vit.cu @@ -0,0 +1,189 @@ +// FlashAttention-4 (AOT) for the SigLIP vision attention (Thor SM110). +// +// The vendored fa4_aot/ artifacts are the CuTe-DSL AOT export of the FA4 +// SM100-compatible forward compiled for sm_110a at head_dim 80 (the +// padded-head layout this adapter's vision path already uses). The kernel +// takes (batch, seq, heads, head_dim) f16 tensors with arbitrary leading +// strides; softmax scale is a runtime argument. +// +// The ggml flash-attention node's padded Q/K/V/dst all share one linear +// layout (d + h*hd + s*hd*H + b*hd*H*S), which is exactly FA4's +// (B, S, H, D) with strides (S*H*hd, H*hd, hd) — so the f32 Q input and +// f32 output only need dense elementwise converts, and the f16 K/V pass +// straight through. + +#include "fr_kernels.h" + +#include "fa4_aot/fa4_siglip_fwd.h" +#include "fa4_aot/fa4_prefill_fwd.h" + +#include + +namespace ggml_cuda_flashrt { + +namespace { + +__global__ void kernel_f32_to_f16_dense(const float * __restrict__ src, + __half * __restrict__ dst, int64_t n) { + const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) { + dst[i] = __float2half(src[i]); + } +} + +__global__ void kernel_f16_to_f32_dense(const __half * __restrict__ src, + float * __restrict__ dst, int64_t n) { + const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) { + dst[i] = __half2float(src[i]); + } +} + +// convert dropping the head padding: one block per (b, s) token; src rows +// are H heads of D elements, dst rows are H packed slices of D2 (< D) +__global__ void kernel_f16_to_f32_depad(const __half * __restrict__ src, + float * __restrict__ dst, + int H, int D, int D2) { + const int64_t row = blockIdx.x; + const __half * s = src + row * (int64_t) H * D; + float * d = dst + row * (int64_t) H * D2; + for (int t = threadIdx.x; t < H * D2; t += blockDim.x) { + const int h = t / D2; + const int e = t - h * D2; + d[t] = __half2float(s[(int64_t) h * D + e]); + } +} + +fa4_siglip_fwd_Kernel_Module_t g_fa4_module; +bool g_fa4_loaded = false; + +fa4_prefill_fwd_Kernel_Module_t g_fa4p_module; +bool g_fa4p_loaded = false; + +} // namespace + +// Loads the AOT module once; must not run during CUDA graph capture (the +// caller checks and falls back to the unfused path on the first capture). +int fa4_vit_ensure_loaded(cudaStream_t stream) { + if (g_fa4_loaded) { + return 0; + } + cudaStreamCaptureStatus cap = cudaStreamCaptureStatusNone; + cudaStreamIsCapturing(stream, &cap); + if (cap != cudaStreamCaptureStatusNone) { + return -1; + } + fa4_siglip_fwd_Kernel_Module_Load(&g_fa4_module); + cudaError_t e = cudaGetLastError(); + if (e != cudaSuccess) { + return -static_cast(e); + } + g_fa4_loaded = true; + fa4_prefill_fwd_Kernel_Module_Load(&g_fa4p_module); + e = cudaGetLastError(); + if (e != cudaSuccess) { + return -static_cast(e); + } + g_fa4p_loaded = true; + return 0; +} + +// Full (non-causal) self-attention for the pi0.5 prefill: hd-256 GQA FA4. +// q_f32/dst_f32 dense (1, S, H, D); k16/v16 are the first S contiguous +// [D]-rows of the (possibly padded) f16 KV buffers, one KV head. +// The padded tail rows are simply outside the dynamic shape, which is +// exactly the graph's row-uniform pad mask. +int fa4_prefill_attention(const float * q_f32, const void * k16, const void * v16, + float * dst_f32, void * q16_ws, void * o16_ws, + int S, int H, int D, float scale, + cudaStream_t stream) { + if (!g_fa4p_loaded) { + return -1; + } + const int64_t n = (int64_t) S * H * D; + const int threads = 256; + const int64_t blocks = (n + threads - 1) / threads; + + kernel_f32_to_f16_dense<<<(unsigned) blocks, threads, 0, stream>>>( + q_f32, (__half *) q16_ws, n); + + auto fill_q = [&](void * data, auto * t, int heads) { + t->data = data; + t->dynamic_shapes[0] = 1; + t->dynamic_shapes[1] = S; + t->dynamic_shapes[2] = heads; + t->dynamic_shapes[3] = D; + t->dynamic_strides[0] = (int64_t) S * heads * D; + t->dynamic_strides[1] = (int64_t) heads * D; + t->dynamic_strides[2] = D; + }; + fa4_prefill_fwd_Tensor_mQ_t tq; fill_q(q16_ws, &tq, H); + fa4_prefill_fwd_Tensor_mK_t tk; fill_q(const_cast(k16), &tk, 1); + fa4_prefill_fwd_Tensor_mV_t tv; fill_q(const_cast(v16), &tv, 1); + fa4_prefill_fwd_Tensor_mO_t to; fill_q(o16_ws, &to, H); + + const int32_t rc = cute_dsl_fa4_prefill_fwd_wrapper( + &g_fa4p_module, &tq, &tk, &tv, &to, scale, stream); + if (rc != 0) { + return -1000 - rc; + } + + kernel_f16_to_f32_dense<<<(unsigned) blocks, threads, 0, stream>>>( + (const __half *) o16_ws, dst_f32, n); + + const cudaError_t e2 = cudaGetLastError(); + return (e2 == cudaSuccess) ? 0 : -static_cast(e2); +} + +// q_f32: dense (B,S,H,D); k16/v16: dense f16 same layout. When d_out == D +// dst_f32 is the dense padded layout; when d_out < D the head padding is +// dropped and dst_f32 is the packed [(H*d_out), S, B] contiguous tensor. +// q16_ws / o16_ws: workspaces of B*S*H*D halves. +int fa4_vit_attention(const float * q_f32, const void * k16, const void * v16, + float * dst_f32, int d_out, void * q16_ws, void * o16_ws, + int B, int S, int H, int D, float scale, + cudaStream_t stream) { + if (!g_fa4_loaded) { + return -1; + } + const int64_t n = (int64_t) B * S * H * D; + const int threads = 256; + const int64_t blocks = (n + threads - 1) / threads; + + kernel_f32_to_f16_dense<<<(unsigned) blocks, threads, 0, stream>>>( + q_f32, (__half *) q16_ws, n); + + auto fill = [&](void * data, auto * t) { + t->data = data; + t->dynamic_shapes[0] = B; + t->dynamic_shapes[1] = S; + t->dynamic_shapes[2] = H; + t->dynamic_shapes[3] = D; + t->dynamic_strides[0] = (int64_t) S * H * D; + t->dynamic_strides[1] = (int64_t) H * D; + t->dynamic_strides[2] = D; + }; + fa4_siglip_fwd_Tensor_mQ_t tq; fill(q16_ws, &tq); + fa4_siglip_fwd_Tensor_mK_t tk; fill(const_cast(k16), &tk); + fa4_siglip_fwd_Tensor_mV_t tv; fill(const_cast(v16), &tv); + fa4_siglip_fwd_Tensor_mO_t to; fill(o16_ws, &to); + + const int32_t rc = cute_dsl_fa4_siglip_fwd_wrapper( + &g_fa4_module, &tq, &tk, &tv, &to, scale, stream); + if (rc != 0) { + return -1000 - rc; + } + + if (d_out == D) { + kernel_f16_to_f32_dense<<<(unsigned) blocks, threads, 0, stream>>>( + (const __half *) o16_ws, dst_f32, n); + } else { + kernel_f16_to_f32_depad<<<(unsigned) ((int64_t) B * S), threads, 0, stream>>>( + (const __half *) o16_ws, dst_f32, H, D, d_out); + } + + const cudaError_t e = cudaGetLastError(); + return (e == cudaSuccess) ? 0 : -static_cast(e); +} + +} // namespace ggml_cuda_flashrt diff --git a/flash_rt/structures/adapters/ggml/fr_gemm_f32out.cu b/flash_rt/structures/adapters/ggml/fr_gemm_f32out.cu new file mode 100644 index 000000000..0312fbedc --- /dev/null +++ b/flash_rt/structures/adapters/ggml/fr_gemm_f32out.cu @@ -0,0 +1,207 @@ +// Block-scaled NVFP4 GEMM for Thor SM110, fp32 output. +// +// Vendored from FlashRT's cutlass_nvfp4_w4a16_gemm_sm100.cu (Apache-2.0) +// with ElementD changed from bf16 to fp32 so the result lands directly in +// ggml's fp32 dst tensor. The Sm100 CollectiveBuilder dispatch under +// KernelScheduleAuto produces the block-scaled tcgen05 mainloop when built +// for sm_110a. + +#include "fr_kernels.h" + +#include "cute/tensor.hpp" + +#include "cutlass/cutlass.h" +#include "cutlass/numeric_types.h" +#include "cutlass/detail/sm100_blockscaled_layout.hpp" + +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/gemm/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 +#include +#include + +namespace ggml_cuda_flashrt { + +namespace { + +using namespace cute; + +template > +struct FrGemmConfig { + using ElementA = cutlass::float_e2m1_t; + using ElementB = cutlass::float_e2m1_t; + using ElementC = float; + using ElementD = float; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementSF = cutlass::float_ue4m3_t; + + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::RowMajor; + using LayoutD = cutlass::layout::RowMajor; + + using ElementPairA = cutlass::nv_float4_t; + using ElementPairB = cutlass::nv_float4_t; + + static constexpr int AlignmentA = 32; + static constexpr int AlignmentB = 32; + static constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; // 4 + static constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; // 4 + + using ClusterShape = ClusterShapeT; + + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassBlockScaledTensorOp, + TileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, AlignmentC, + ElementD, LayoutD, AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassBlockScaledTensorOp, + ElementPairA, LayoutA, AlignmentA, + ElementPairB, LayoutB, AlignmentB, + ElementAccumulator, + TileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::collective::KernelScheduleAuto + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +}; + +using Sm1xxBlkScaledConfig = cutlass::detail::Sm1xxBlockScaledConfig<16>; + +// Per-shape CUTLASS workspace cache; entries live for the process lifetime +// (weight shapes are fixed per model). +struct ShapeKey { + int M, N, K; + bool operator==(const ShapeKey & o) const { return M == o.M && N == o.N && K == o.K; } +}; +struct ShapeKeyHash { + size_t operator()(const ShapeKey & k) const noexcept { + return (static_cast(k.M) * 1315423911u) + ^ (static_cast(k.N) * 2654435761u) + ^ static_cast(k.K); + } +}; +struct CachedWorkspace { void * ptr = nullptr; size_t size = 0; }; + +std::unordered_map g_ws_cache; +std::mutex g_ws_mu; + +void * get_workspace(int M, int N, int K, size_t needed) { + std::lock_guard lk(g_ws_mu); + ShapeKey key{M, N, K}; + auto it = g_ws_cache.find(key); + if (it != g_ws_cache.end() && it->second.size >= needed) return it->second.ptr; + if (it != g_ws_cache.end()) { cudaFree(it->second.ptr); g_ws_cache.erase(it); } + CachedWorkspace w; w.size = needed; + if (needed > 0) cudaMalloc(&w.ptr, needed); + g_ws_cache[key] = w; + return w.ptr; +} + +template +int run_gemm(const void * A_packed, const void * SFA, + const void * B_packed, const void * SFB, + float * D, int M, int N, int K, + float alpha, cudaStream_t stream) { + using Gemm = typename Config::Gemm; + using ElementSF = typename Config::ElementSF; + using ElementD = typename Config::ElementD; + + using StrideA = typename Gemm::GemmKernel::StrideA; + using StrideB = typename Gemm::GemmKernel::StrideB; + using StrideC = typename Gemm::GemmKernel::StrideC; + using StrideD = typename Gemm::GemmKernel::StrideD; + + StrideA stride_A = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(M, K, 1)); + StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(N, K, 1)); + StrideC stride_C = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(M, N, 1)); + StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(M, N, 1)); + + auto problem_shape_MNKL = cute::make_shape(M, N, K, 1); + auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem_shape_MNKL); + auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem_shape_MNKL); + + using ArrayElementA = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementA; + using ArrayElementB = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementB; + + 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 + }, + { + {alpha, 0.0f}, + nullptr, stride_C, + reinterpret_cast(D), stride_D + } + }; + + Gemm gemm; + size_t ws_size = Gemm::get_workspace_size(args); + void * ws_ptr = get_workspace(M, N, K, ws_size); + + auto status = gemm.can_implement(args); + if (status != cutlass::Status::kSuccess) { + std::fprintf(stderr, "[fr_gemm_f32out] can_implement FAIL M=%d N=%d K=%d (status=%d)\n", + M, N, K, static_cast(status)); + return static_cast(status); + } + status = gemm.initialize(args, ws_ptr, stream); + if (status != cutlass::Status::kSuccess) { + std::fprintf(stderr, "[fr_gemm_f32out] initialize FAIL M=%d N=%d K=%d (status=%d)\n", + M, N, K, static_cast(status)); + return static_cast(status); + } + status = gemm.run(stream); + if (status != cutlass::Status::kSuccess) { + std::fprintf(stderr, "[fr_gemm_f32out] run FAIL M=%d N=%d K=%d (status=%d)\n", + M, N, K, static_cast(status)); + return static_cast(status); + } + return 0; +} + +using ConfigDefault = FrGemmConfig>; +// 2-SM tcgen05 tile: 11-15% faster than the 1-SM tile on Thor for every +// prefill shape measured (M >= ~256); slower at decode-sized M. +using ConfigLargeM = FrGemmConfig, Shape<_2, _1, _1>>; + +} // namespace + +int gemm_f32out(const void * A_packed, const void * SFA, + const void * B_packed, const void * SFB, + float * D, int M, int N, int K, + float alpha, bool widen, cudaStream_t stream) { + (void) widen; // superseded by the M-based tile choice + if (M >= 256) { + return run_gemm(A_packed, SFA, B_packed, SFB, D, M, N, K, alpha, stream); + } + return run_gemm(A_packed, SFA, B_packed, SFB, D, M, N, K, alpha, stream); +} + +} // namespace ggml_cuda_flashrt diff --git a/flash_rt/structures/adapters/ggml/fr_ggml.cuh b/flash_rt/structures/adapters/ggml/fr_ggml.cuh new file mode 100644 index 000000000..38b196306 --- /dev/null +++ b/flash_rt/structures/adapters/ggml/fr_ggml.cuh @@ -0,0 +1,165 @@ +// ggml-facing interface of the FlashRT NVFP4 path (Thor SM110). +// Included from ggml-cuda.cu under #ifdef GGML_CUDA_FLASHRT. +#pragma once + +// Host dependency: ggml-cuda's internal common header. The consuming build +// must put ggml/src/ggml-cuda on the include path (the ggml adapter is +// compiled inside the host's build tree, like the vllm/sglang adapters run +// inside their host's runtime). +#include "common.cuh" + +// Called at the start of every backend graph evaluation; invalidates the +// per-evaluation quantized-activation cache. +void ggml_cuda_flashrt_begin_eval(); + +// True when this mul_mat should be routed to the FlashRT block-scaled NVFP4 +// GEMM: NVFP4 weights, fp32 contiguous activations/dst, no batch dims, +// shapes within kernel alignment. cc must already be checked by the caller. +bool ggml_cuda_flashrt_should_use(const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst); + +// dst = src1 @ src0 via activation quantize + NVFP4 x NVFP4 tcgen05 GEMM. +// Weights are repacked into the CUTLASS wire format on first use and cached +// for the process lifetime. +void ggml_cuda_flashrt_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst); + +// True when the 4-node FFN subgraph {mul_mat gate, mul_mat up, GEGLU, +// mul_mat down} can run as one fused GeGLU GEMM (interleaved gate/up +// weights, FP4 intermediate) followed by the down GEMM. +bool ggml_cuda_flashrt_should_fuse_geglu(const ggml_tensor * gate_mm, const ggml_tensor * up_mm, const ggml_tensor * glu, const ggml_tensor * down_mm); + +// Execute that fused FFN; writes the down mul_mat's dst. +void ggml_cuda_flashrt_geglu_ffn(ggml_backend_cuda_context & ctx, const ggml_tensor * gate_mm, const ggml_tensor * up_mm, const ggml_tensor * glu, ggml_tensor * down_mm); + +// pi0.5 adaLN modulate window: {rms_norm?, mul_mat mod, add bias, view, +// repeat, mul, add, view, repeat, add}. rms is null for the variant whose +// normalized input arrives from a previous graph split. Outputs written by +// the fused execution: the bias add (consumed by the gate view later) and +// the final add. +bool ggml_cuda_flashrt_should_fuse_ada(const ggml_tensor * rms, const ggml_tensor * mm, const ggml_tensor * bias_add, + const ggml_tensor * view_scale, const ggml_tensor * repeat_scale, + const ggml_tensor * mul, const ggml_tensor * add1, + const ggml_tensor * view_shift, const ggml_tensor * repeat_shift, + const ggml_tensor * add2); +void ggml_cuda_flashrt_ada_norm(ggml_backend_cuda_context & ctx, const ggml_tensor * rms, const ggml_tensor * mm, + ggml_tensor * bias_add, const ggml_tensor * view_scale, const ggml_tensor * mul, + const ggml_tensor * view_shift, ggml_tensor * add2); + +// LayerNorm + affine window: {NORM, MUL weight, ADD bias} -> one kernel. +bool ggml_cuda_flashrt_should_fuse_ada_cached( + const ggml_tensor * rms, const ggml_tensor * view_col, + const ggml_tensor * view_scale, const ggml_tensor * repeat_scale, + const ggml_tensor * mul, const ggml_tensor * add1, + const ggml_tensor * view_shift, const ggml_tensor * repeat_shift, + const ggml_tensor * add2); +void ggml_cuda_flashrt_ada_norm_cached(ggml_backend_cuda_context & ctx, const ggml_tensor * rms, + const ggml_tensor * view_scale, const ggml_tensor * view_shift, + ggml_tensor * add2); + +bool ggml_cuda_flashrt_should_fuse_ln(const ggml_tensor * norm, const ggml_tensor * mul, const ggml_tensor * add); +void ggml_cuda_flashrt_ln_affine(ggml_backend_cuda_context & ctx, const ggml_tensor * norm, const ggml_tensor * mul, ggml_tensor * add); + +// SigLIP FFN window: {mul_mat up (NVFP4), add bias, GELU, cont, mul_mat +// down (f16), add bias, cont, add residual} -> fused FP4 Up GEMM (gelu +// epilogue, FP4 hidden) + Down GEMM (bias + residual epilogue). +bool ggml_cuda_flashrt_should_fuse_siglip_ffn(const ggml_tensor * up_mm, const ggml_tensor * bias1, const ggml_tensor * gelu, + const ggml_tensor * cont1, const ggml_tensor * dn_mm, const ggml_tensor * bias2, + const ggml_tensor * cont2, const ggml_tensor * res_add); +void ggml_cuda_flashrt_siglip_ffn(ggml_backend_cuda_context & ctx, const ggml_tensor * up_mm, const ggml_tensor * bias1, + const ggml_tensor * dn_mm, const ggml_tensor * bias2, + const ggml_tensor * cont2, ggml_tensor * res_add); + +// pi0.5 AE fused QKV window: {mm k, reshape, rope, view, cpy, mm v, +// reshape, view, cpy, mm q, reshape, rope, scale} -> one fused GEMM over +// row-concatenated [k|v|q] weights + one post kernel (rope/scale/f16 KV +// suffix stores). +bool ggml_cuda_flashrt_should_fuse_qkv(const ggml_tensor * k_mm, const ggml_tensor * k_rope, const ggml_tensor * k_cpy, + const ggml_tensor * v_mm, const ggml_tensor * v_cpy, + const ggml_tensor * q_mm, const ggml_tensor * q_rope, const ggml_tensor * q_scale); +void ggml_cuda_flashrt_qkv(ggml_backend_cuda_context & ctx, + const ggml_tensor * k_mm, const ggml_tensor * k_rope, const ggml_tensor * k_cpy, + const ggml_tensor * v_mm, const ggml_tensor * v_cpy, + const ggml_tensor * q_mm, const ggml_tensor * q_rope, ggml_tensor * q_scale); + +// pi0.5 gated residual window: {view gate, repeat, mul, add}. +// Vision QKV pad window: {mul_mat, add bias, reshape}x3 + {pad}x3 -> three +// padded-weight GEMMs (bias in epilogue) writing the pad buffers directly. +bool ggml_cuda_flashrt_should_fuse_vis_qkv_pad( + const ggml_tensor * mm_q, const ggml_tensor * add_q, const ggml_tensor * resh_q, + const ggml_tensor * mm_k, const ggml_tensor * add_k, const ggml_tensor * resh_k, + const ggml_tensor * mm_v, const ggml_tensor * add_v, const ggml_tensor * resh_v, + const ggml_tensor * pad_q, const ggml_tensor * pad_k, const ggml_tensor * pad_v); +void ggml_cuda_flashrt_vis_qkv_pad(ggml_backend_cuda_context & ctx, + const ggml_tensor * mm_q, const ggml_tensor * add_q, ggml_tensor * pad_q, + const ggml_tensor * mm_k, const ggml_tensor * add_k, ggml_tensor * pad_k, + const ggml_tensor * mm_v, const ggml_tensor * add_v, ggml_tensor * pad_v, + ggml_tensor * k_cast, ggml_tensor * v_cast); + +// GEMM + optional bias + residual add fused into one epilogue. +bool ggml_cuda_flashrt_should_fuse_mm_res(const ggml_tensor * mm, const ggml_tensor * bias_add, + const ggml_tensor * res_add); +bool ggml_cuda_flashrt_mm_res(ggml_backend_cuda_context & ctx, const ggml_tensor * mm, + const ggml_tensor * bias_add, ggml_tensor * res_add); + +bool ggml_cuda_flashrt_should_fuse_gated_res(const ggml_tensor * view, const ggml_tensor * repeat, + const ggml_tensor * mul, const ggml_tensor * add); +void ggml_cuda_flashrt_gated_residual(ggml_backend_cuda_context & ctx, const ggml_tensor * view, + const ggml_tensor * mul, ggml_tensor * add); + +// Decomposed tiny-M decode attention: replaces a FLASH_ATTN_EXT node with +// QK-GEMM + masked softmax + PV-GEMM for q_tokens <= 16 over a single f16 +// KV head of token rows. +bool ggml_cuda_flashrt_should_fuse_dec_attn(const ggml_tensor * fa); +void ggml_cuda_flashrt_dec_attn(ggml_backend_cuda_context & ctx, ggml_tensor * fa); + +// SigLIP vision attention through the AOT FlashAttention-4 module +// (head_dim 80, no mask). Only available when the adapter was built with +// the fa4_aot artifacts; the first use loads the module (falls back if +// that first use happens during CUDA graph capture). +bool ggml_cuda_flashrt_should_fuse_vit_fa4(const ggml_tensor * fa, ggml_backend_cuda_context & ctx); +void ggml_cuda_flashrt_vit_fa4(ggml_backend_cuda_context & ctx, ggml_tensor * fa); + +// Variant that also absorbs the {VIEW (head de-pad), CONT} pair that +// follows the vision FA node: the FA4 output converts straight into the +// CONT's packed [(H*d), S, B] f32 destination, skipping the padded f32 +// store and the strided copy. +bool ggml_cuda_flashrt_should_fuse_vit_fa4_depad(const ggml_tensor * fa, const ggml_tensor * view, + const ggml_tensor * cont, ggml_backend_cuda_context & ctx); +void ggml_cuda_flashrt_vit_fa4_depad(ggml_backend_cuda_context & ctx, ggml_tensor * fa, + const ggml_tensor * view, ggml_tensor * cont); + +// pi0.5 prefill self-attention (head_dim 256, GQA 1 KV head, full +// attention with a row-uniform pad mask) through the second AOT FA4 +// module. The window assumes the prefix-LM mask semantics of the pi0.5 +// prefill graph (see the dispatch-side checks). +bool ggml_cuda_flashrt_should_fuse_prefill_fa4(const ggml_tensor * fa, ggml_backend_cuda_context & ctx); +void ggml_cuda_flashrt_prefill_fa4(ggml_backend_cuda_context & ctx, ggml_tensor * fa); + +// Prefill fused QKV window: q mm->reshape->rope->scale, k mm->reshape->rope-> +// pad, v mm->reshape->pad, each pad permuted+copied into a padded f16 tensor +// of token rows. One fused GEMM + qkv_post + pad-row zeroing. +bool ggml_cuda_flashrt_should_fuse_qkv_prefill( + const ggml_tensor * q_mm, const ggml_tensor * q_rope, const ggml_tensor * q_scale, + const ggml_tensor * k_mm, const ggml_tensor * k_rope, const ggml_tensor * k_pad, + const ggml_tensor * v_mm, const ggml_tensor * v_pad, + const ggml_tensor * k_cpy, const ggml_tensor * v_cpy); +void ggml_cuda_flashrt_qkv_prefill(ggml_backend_cuda_context & ctx, + const ggml_tensor * q_mm, const ggml_tensor * q_rope, ggml_tensor * q_scale, + const ggml_tensor * k_mm, const ggml_tensor * k_rope, + const ggml_tensor * v_mm, + ggml_tensor * k_cpy, ggml_tensor * v_cpy); + +// Run of terminal f32->f16 row-copy CPY nodes (the persistent encoder-KV +// stores at the prefill graph tail) batched into one kernel launch. A node +// qualifies when it copies [hd, 1, n_rows] contiguous f32 rows into +// contiguous f16 rows and nothing reads the copy back inside the graph; +// all nodes of one batch share hd and n_rows. +bool ggml_cuda_flashrt_kv_tail_cpy_ok(const ggml_tensor * cpy, int64_t * hd, int64_t * n_rows); +bool ggml_cuda_flashrt_kv_tail_cpy(ggml_backend_cuda_context & ctx, ggml_tensor ** cpys, int n); + +// {RMS_NORM, MUL(w), ADD(mul, norm)} -> rms_norm(x)*(1+w) in one kernel. +// The execute returns false (run unfused) only when its zero-vector cache +// cannot allocate during graph capture. +bool ggml_cuda_flashrt_should_fuse_rms_gemma(const ggml_tensor * rms, const ggml_tensor * mul, + const ggml_tensor * add); +bool ggml_cuda_flashrt_rms_gemma(ggml_backend_cuda_context & ctx, const ggml_tensor * rms, + const ggml_tensor * mul, ggml_tensor * add); diff --git a/flash_rt/structures/adapters/ggml/fr_kernels.h b/flash_rt/structures/adapters/ggml/fr_kernels.h new file mode 100644 index 000000000..0ed45a112 --- /dev/null +++ b/flash_rt/structures/adapters/ggml/fr_kernels.h @@ -0,0 +1,189 @@ +// FlashRT NVFP4 kernels for Jetson AGX Thor (SM110). +// +// C-style entry points implemented in the fr_*.cu translation units, which +// are compiled separately for sm_110a with CUTLASS. This header must stay +// free of CUTLASS and ggml includes so it can be consumed from the regular +// ggml-cuda translation units. +// +// Wire format (NVFP4): +// packed: uint8 [rows, K/2], adjacent-pair nibbles (elem 2i low, 2i+1 high) +// scales: e4m3 (positive/ue4m3), one per 16 elements along K, stored in +// the CUTLASS Sm1xx block-scaled atom layout for the GEMM shape +// +// ggml's block_nvfp4 scale bytes carry standard e4m3 semantics (its dequant +// table doubles the e2m1 values and its ue4m3 decode halves the scale, which +// cancel), so they pass through to the GEMM unmodified with alpha = 1.0. +#pragma once + +#include +#include + +namespace ggml_cuda_flashrt { + +static inline int64_t round_up_i64(int64_t x, int64_t m) { return (x + m - 1) / m * m; } + +// Scale-factor buffer size in bytes for a [rows, K] operand: the Sm1xx atom +// layout tiles rows in chunks of 128 and K/16 scale columns in chunks of 4. +static inline int64_t sf_bytes(int64_t rows, int64_t K) { + return round_up_i64(rows, 128) * round_up_i64(K / 16, 4); +} + +static inline int64_t packed_bytes(int64_t rows, int64_t K) { + return rows * (K / 2); +} + +// Block-scaled NVFP4 x NVFP4 GEMM, D = alpha * (A x B), fp32 output. +// A_packed [M, K/2] row-major, B_packed [N, K/2] row-major (used as +// column-major [K, N]), D fp32 [M, N] row-major. +// Requires K % 64 == 0, N % 16 == 0, D 16-byte aligned. +// widen selects a wide-N tile (use for N >= 8192). +// Returns 0 on success. +int gemm_f32out(const void * A_packed, const void * SFA, + const void * B_packed, const void * SFB, + float * D, int M, int N, int K, + float alpha, bool widen, cudaStream_t stream); + +// Quantize fp32 activations [M, K] row-major (contiguous) to NVFP4 packed +// [M, K/2] plus SFA scales in the atom layout for problem (M, x, K). +// Requires K % 16 == 0. Returns 0 on success. +int quantize_act_f32(const float * src, void * dst_packed, void * dst_sfa, + int M, int K, cudaStream_t stream); + +// Repack a ggml GGML_TYPE_NVFP4 weight tensor [N rows, K] into the GEMM's +// B-side wire format: packed [N, K/2] with adjacent-pair nibbles plus SFB +// scales (ggml half-scale bytes, unmodified) in the atom layout. +// Requires K % 64 == 0. Returns 0 on success. +int repack_weight(const void * ggml_blocks, void * dst_packed, void * dst_sf, + int N, int K, cudaStream_t stream); + +// Pairwise-interleave two ggml NVFP4 weight tensors (gate, up; each +// [n_ff rows, K]) into one B operand for the fused GeGLU GEMM: output row 2j +// is gate row j, row 2j+1 is up row j. dst_packed holds 2*n_ff rows; dst_sf +// is sized sf_bytes(2*n_ff, K). +int repack_weight_pair_interleaved(const void * gate_blocks, const void * up_blocks, + void * dst_packed, void * dst_sf, + int n_ff, int K, cudaStream_t stream); + +// Rows-padded repack for the SigLIP FFN Up weight: rows >= N_src are zeros. +int repack_weight_rows_padded(const void * ggml_blocks, void * dst_packed, void * dst_sf, + int N_src, int N_pad, int K, cudaStream_t stream); + +// Quantize an fp16 weight [N, K_src] to NVFP4 wire format, K zero-padded to K_pad. +int quantize_weight_f16_padded(const void * w_f16, void * dst_packed, void * dst_sf, + int N, int K_src, int K_pad, cudaStream_t stream); + +// The SigLIP FFN GEMM pair and gemm_bias_f32out are declared in +// flashrt-public's cutlass_fp4_gemm_siglip_ffn_f32out_sm100.cuh +// (namespace flash_rt::fp4). + +// Group-padded rows repack: n_groups groups widened group_in -> group_out +// rows, pad rows zero (used to widen per-head projections for FA head sizes). +int repack_weight_rows_grouppad(const void * ggml_blocks, void * dst_packed, void * dst_sf, + int group_in, int group_out, int n_groups, int K, cudaStream_t stream); + +// Row-concat repack of three NVFP4 weights (shared K) for the fused QKV GEMM. +int repack_weight_concat3(const void * b0, int N0, const void * b1, int N1, + const void * b2, int N2, + void * dst_packed, void * dst_sf, + int K, cudaStream_t stream); + +// Fused QKV post: RoPE+f16-store K, f16-store V (into the persistent KV +// suffix), RoPE+scale Q (f32 out) from the fused GEMM's [M, Nk+Nv+Nq] rows. +// q16_out (nullable) additionally stores the Q rows as f16 in the same +// [M, Nq] layout, which is the t-major gather order the decomposed decode +// attention consumes. Variant with optional f32 K/V row outputs (for +// graphs whose rope'd K / V feed additional consumers, e.g. persistent-KV +// stores at the graph tail). +int qkv_post_full(const float * qkv_cat, float * q_out, void * q16_out, + void * k_out_f16, void * v_out_f16, + float * k_f32_out, float * v_f32_out, + const int32_t * pos, const float * freq_factors, + int M, int Nk, int Nv, int Nq, int head_dim, int n_dims, + float freq_scale, float ext_factor, float attn_factor, + float corr_low, float corr_high, float theta_scale, float q_scale, + cudaStream_t stream); + +int qkv_post(const float * qkv_cat, float * q_out, void * q16_out, + void * k_out_f16, void * v_out_f16, + const int32_t * pos, const float * freq_factors, + int M, int Nk, int Nv, int Nq, int head_dim, int n_dims, + float freq_scale, float ext_factor, float attn_factor, + float corr_low, float corr_high, float theta_scale, float q_scale, + cudaStream_t stream); + +// Fused adaLN modulate: out[m,c] = norm(x[m])[c] * (1 + scale[c]) + shift[c], +// norm = rms-normalize when with_rms else identity. x/out are [M, C] +// contiguous fp32; scale/shift are [C] vectors. +int ada_rms_mod(const float * x, const float * scale, const float * shift, + float * out, int M, int C, float eps, bool with_rms, + cudaStream_t stream); + +// Fused gated residual: out[m,c] = residual[m,c] + branch[m,c] * gate[c]. +int gated_residual(const float * residual, const float * branch, const float * gate, + float * out, int M, int C, cudaStream_t stream); + +// Quant-emitting variants: additionally write the result quantized to +// NVFP4 packed + SFA (atom layout for an [M, C] activation operand). +int ada_rms_mod_quant(const float * x, const float * scale, const float * shift, + float * out, void * dst_packed, void * dst_sfa, + int M, int C, float eps, bool with_rms, cudaStream_t stream); +int layer_norm_affine_quant(const float * x, const float * w, const float * b, + float * out, void * dst_packed, void * dst_sfa, + int M, int C, float eps, cudaStream_t stream); + +// Fused LayerNorm + affine: out[m,c] = normalize(x[m])[c] * w[c] + b[c]. +int layer_norm_affine(const float * x, const float * w, const float * b, + float * out, int M, int C, float eps, cudaStream_t stream); + +// Decomposed tiny-M attention: one QK^T GEMM over all n_head*n_tok query +// rows (f16, the GQA heads share the single K operand) + masked softmax + +// one PV GEMM whose fp32 output is the dense [hd, n_head, n_tok] dst +// (rows ordered t-major so the store is contiguous; requires +// dst_shead == hd and dst_stok == hd*n_head). q strides are in elements; +// workspaces: q16 [n_tok*n_head, hd] f16, scores [n_tok*n_head, n_kv] +// f16. When q16_ready is nonzero, q16_ws already holds the gathered f16 Q +// rows (produced upstream, e.g. by qkv_post) and the gather kernel is +// skipped. cublas_handle is a cublasHandle_t. +int decode_attn_decomposed(void * cublas_handle, + const float * q, int64_t q_sd, int64_t q_stok, int64_t q_shead, + const void * k_f16_rows, const void * v_f16_rows, + const void * mask_f16, int64_t mask_stride, + float * dst, int64_t dst_stok, int64_t dst_shead, + void * q16_ws, int q16_ready, void * scores_ws, + int hd, int n_tok, int n_head, int n_kv, + float scale, cudaStream_t stream); + +// out[i] = a[i] + b[i] for n fp32 elements. +int vec_add_f32(const float * a, const float * b, float * out, int n, cudaStream_t stream); + +// AOT FlashAttention-4 for the SigLIP vision attention (head_dim 80). +// ensure_loaded loads the vendored module once (fails during CUDA graph +// capture: fall back). attention runs f32->f16 Q convert + FA4 + f16->f32 +// output convert; all tensors dense (B,S,H,D) as one linear buffer. +int fa4_vit_ensure_loaded(cudaStream_t stream); +// d_out == D: dst is the dense padded layout; d_out < D: the head padding +// is dropped and dst is the packed [(H*d_out), S, B] contiguous tensor. +int fa4_vit_attention(const float * q_f32, const void * k16, const void * v16, + float * dst_f32, int d_out, void * q16_ws, void * o16_ws, + int B, int S, int H, int D, float scale, + cudaStream_t stream); + +// Full (non-causal) self-attention for the pi0.5 prefill via a second AOT +// FA4 module (head_dim 256, GQA with one KV head). k16/v16 are the first S +// contiguous rows of the padded f16 KV buffers; the pad rows lie outside +// the dynamic shape, which reproduces the row-uniform pad mask exactly. +int fa4_prefill_attention(const float * q_f32, const void * k16, const void * v16, + float * dst_f32, void * q16_ws, void * o16_ws, + int S, int H, int D, float scale, + cudaStream_t stream); + +// Batched f32->f16 row copies: for each pair p, dst[p][r*hd + i] = +// (half) src[p][r*hd + i] over n_rows rows of hd elements. One launch +// replaces up to FR_CPY_ROWS_MAX individual copy kernels (the persistent +// encoder-KV stores at the prefill graph tail). Rounding matches ggml's +// f32->f16 cpy exactly. +#define FR_CPY_ROWS_MAX 40 +int cpy_rows_f32_f16(const float * const * srcs, void * const * dsts, int n_pairs, + int hd, int n_rows, cudaStream_t stream); + +} // namespace ggml_cuda_flashrt diff --git a/flash_rt/structures/adapters/ggml/fr_qkv_post.cu b/flash_rt/structures/adapters/ggml/fr_qkv_post.cu new file mode 100644 index 000000000..b050dea5f --- /dev/null +++ b/flash_rt/structures/adapters/ggml/fr_qkv_post.cu @@ -0,0 +1,171 @@ +// Fused QKV post-processing for the pi0.5 action expert (Thor). +// +// Consumes the fused QKV GEMM's f32 output [M, Nk + Nv + Nq] (sections +// [k | v | q]) and in one launch: +// - RoPEs K and writes it as f16 into the persistent KV buffer's suffix +// - writes V as f16 into the persistent KV buffer's suffix +// - RoPEs and scales Q, writing the f32 tensor flash attention consumes +// +// The RoPE math mirrors ggml-cuda's rope_neox (yarn corrections included) +// so the fused path is numerically identical to the unfused graph. + +#include "fr_kernels.h" + +#include + +namespace ggml_cuda_flashrt { + +namespace { + +__device__ float qkv_rope_ramp(const float low, const float high, const int i0) { + const float y = (i0 / 2 - low) / max(0.001f, high - low); + return 1.0f - min(1.0f, max(0.0f, y)); +} + +// mirrors ggml-cuda rope_yarn (forward) +__device__ void qkv_rope_yarn(const float theta_extrap, const float freq_scale, + const float corr_low, const float corr_high, + const int i0, const float ext_factor, float mscale, + float & cos_theta, float & sin_theta) { + float theta = freq_scale * theta_extrap; + if (ext_factor != 0.0f) { + const float ramp_mix = qkv_rope_ramp(corr_low, corr_high, i0) * ext_factor; + theta = theta * (1 - ramp_mix) + theta_extrap * ramp_mix; + mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale); + } + cos_theta = cosf(theta) * mscale; + sin_theta = sinf(theta) * mscale; +} + +// one block per token; threads cover the K/Q rope pairs and the V copy +__global__ void kernel_qkv_post(const float * __restrict__ qkv, // [M, Nk+Nv+Nq] + float * __restrict__ q_out, // [M, Nq] f32 (head-major rows) + __half * __restrict__ q16_out, // nullable: same values as f16 + __half * __restrict__ k_out, // suffix rows, head_dim per token + __half * __restrict__ v_out, + float * __restrict__ k_f32_out, // nullable: rope'd K as f32 rows + float * __restrict__ v_f32_out, // nullable: V as f32 rows + const int32_t * __restrict__ pos, + const float * __restrict__ freq_factors, // nullable + int Nk, int Nv, int Nq, + int head_dim, int n_dims, + float freq_scale, float ext_factor, float attn_factor, + float corr_low, float corr_high, + float theta_scale, float q_scale) { + const int t = blockIdx.x; + const float * row = qkv + (int64_t) t * (Nk + Nv + Nq); + const float * krow = row; + const float * vrow = row + Nk; + const float * qrow = row + Nk + Nv; + + const int p = pos[t]; + + // V: plain f16 copy (and optionally the f32 row for downstream readers) + for (int d = threadIdx.x; d < Nv; d += blockDim.x) { + v_out[(int64_t) t * Nv + d] = __float2half(vrow[d]); + if (v_f32_out != nullptr) { + v_f32_out[(int64_t) t * Nv + d] = vrow[d]; + } + } + + // K: rope one head (Nk == head_dim) + for (int i = threadIdx.x; i < Nk / 2; i += blockDim.x) { + const int i0 = 2 * i; // pair index within the head + if (i0 >= n_dims) { + k_out[(int64_t) t * Nk + n_dims + (i0 - n_dims)] = __float2half(krow[n_dims + (i0 - n_dims)]); + k_out[(int64_t) t * Nk + n_dims + (i0 - n_dims) + 1] = __float2half(krow[n_dims + (i0 - n_dims) + 1]); + if (k_f32_out != nullptr) { + k_f32_out[(int64_t) t * Nk + n_dims + (i0 - n_dims)] = krow[n_dims + (i0 - n_dims)]; + k_f32_out[(int64_t) t * Nk + n_dims + (i0 - n_dims) + 1] = krow[n_dims + (i0 - n_dims) + 1]; + } + continue; + } + const float theta_base = p * powf(theta_scale, i0 / 2.0f); + const float freq_factor = freq_factors ? freq_factors[i0 / 2] : 1.0f; + float cos_t, sin_t; + qkv_rope_yarn(theta_base / freq_factor, freq_scale, corr_low, corr_high, + i0, ext_factor, attn_factor, cos_t, sin_t); + const float x0 = krow[i0 / 2]; + const float x1 = krow[i0 / 2 + n_dims / 2]; + const float k0 = x0 * cos_t - x1 * sin_t; + const float k1 = x0 * sin_t + x1 * cos_t; + k_out[(int64_t) t * Nk + i0 / 2] = __float2half(k0); + k_out[(int64_t) t * Nk + i0 / 2 + n_dims / 2] = __float2half(k1); + if (k_f32_out != nullptr) { + k_f32_out[(int64_t) t * Nk + i0 / 2] = k0; + k_f32_out[(int64_t) t * Nk + i0 / 2 + n_dims / 2] = k1; + } + } + + // Q: rope + scale per head + const int n_head = Nq / head_dim; + for (int hp = threadIdx.x; hp < n_head * head_dim / 2; hp += blockDim.x) { + const int h = hp / (head_dim / 2); + const int i0 = 2 * (hp % (head_dim / 2)); + const float * qh = qrow + (int64_t) h * head_dim; + float * oh = q_out + (int64_t) t * Nq + (int64_t) h * head_dim; + __half * oh16 = q16_out != nullptr + ? q16_out + (int64_t) t * Nq + (int64_t) h * head_dim : nullptr; + if (i0 >= n_dims) { + const float v0 = qh[n_dims + (i0 - n_dims)] * q_scale; + const float v1 = qh[n_dims + (i0 - n_dims) + 1] * q_scale; + oh[n_dims + (i0 - n_dims)] = v0; + oh[n_dims + (i0 - n_dims) + 1] = v1; + if (oh16 != nullptr) { + oh16[n_dims + (i0 - n_dims)] = __float2half(v0); + oh16[n_dims + (i0 - n_dims) + 1] = __float2half(v1); + } + continue; + } + const float theta_base = p * powf(theta_scale, i0 / 2.0f); + const float freq_factor = freq_factors ? freq_factors[i0 / 2] : 1.0f; + float cos_t, sin_t; + qkv_rope_yarn(theta_base / freq_factor, freq_scale, corr_low, corr_high, + i0, ext_factor, attn_factor, cos_t, sin_t); + const float x0 = qh[i0 / 2]; + const float x1 = qh[i0 / 2 + n_dims / 2]; + const float v0 = (x0 * cos_t - x1 * sin_t) * q_scale; + const float v1 = (x0 * sin_t + x1 * cos_t) * q_scale; + oh[i0 / 2] = v0; + oh[i0 / 2 + n_dims / 2] = v1; + if (oh16 != nullptr) { + oh16[i0 / 2] = __float2half(v0); + oh16[i0 / 2 + n_dims / 2] = __float2half(v1); + } + } +} + +} // namespace + +int qkv_post_full(const float * qkv_cat, float * q_out, void * q16_out, + void * k_out_f16, void * v_out_f16, + float * k_f32_out, float * v_f32_out, + const int32_t * pos, const float * freq_factors, + int M, int Nk, int Nv, int Nq, int head_dim, int n_dims, + float freq_scale, float ext_factor, float attn_factor, + float corr_low, float corr_high, float theta_scale, float q_scale, + cudaStream_t stream) { + if (n_dims % 2 != 0 || Nk != head_dim || Nq % head_dim != 0) return -1; + kernel_qkv_post<<>>( + qkv_cat, q_out, (__half *) q16_out, (__half *) k_out_f16, (__half *) v_out_f16, + k_f32_out, v_f32_out, + pos, freq_factors, Nk, Nv, Nq, head_dim, n_dims, + freq_scale, ext_factor, attn_factor, corr_low, corr_high, theta_scale, q_scale); + const cudaError_t e = cudaGetLastError(); + return (e == cudaSuccess) ? 0 : -static_cast(e); +} + +int qkv_post(const float * qkv_cat, float * q_out, void * q16_out, + void * k_out_f16, void * v_out_f16, + const int32_t * pos, const float * freq_factors, + int M, int Nk, int Nv, int Nq, int head_dim, int n_dims, + float freq_scale, float ext_factor, float attn_factor, + float corr_low, float corr_high, float theta_scale, float q_scale, + cudaStream_t stream) { + return qkv_post_full(qkv_cat, q_out, q16_out, k_out_f16, v_out_f16, nullptr, nullptr, + pos, freq_factors, M, Nk, Nv, Nq, head_dim, n_dims, + freq_scale, ext_factor, attn_factor, corr_low, corr_high, + theta_scale, q_scale, stream); +} + +} // namespace ggml_cuda_flashrt diff --git a/flash_rt/structures/adapters/ggml/fr_quant_act.cu b/flash_rt/structures/adapters/ggml/fr_quant_act.cu new file mode 100644 index 000000000..5fd82300a --- /dev/null +++ b/flash_rt/structures/adapters/ggml/fr_quant_act.cu @@ -0,0 +1,111 @@ +// fp32 activation quantizer: [M, K] fp32 row-major -> NVFP4 packed + SFA +// scales at the CUTLASS Sm1xx block-scaled atom-layout offsets. +// +// Vendored from FlashRT's quantize_fp4_sfa_bf16.cu (Apache-2.0) with the +// source element type changed from bf16 to fp32. One thread quantizes one +// 16-element block: four 16-byte loads, one 8-byte packed store, one SFA +// byte at the tile-interleaved offset. + +#include "fr_kernels.h" + +#include + +#include "cutlass/cutlass.h" +#include "cutlass/detail/sm100_blockscaled_layout.hpp" +#include "cute/tensor.hpp" + +namespace ggml_cuda_flashrt { + +namespace { + +using CfgVec = cutlass::detail::Sm1xxBlockScaledConfig<16>; + +__device__ __forceinline__ uint8_t fp32_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; +} + +template +__global__ void kernel_quantize_f32( + const float4 * __restrict__ src, // fp32 [M, K] as float4 (4 elements) + uint2 * __restrict__ dst_packed, // [M, K/2] bytes as uint2 (1 block) + uint8_t * __restrict__ dst_sfa, + LayoutSF layout, + int M, int K4) { // K4 = K / 4 float4 chunks per row + const int block_idx = blockIdx.x * blockDim.x + threadIdx.x; + const int row = blockIdx.y; + const int n_blocks = K4 >> 2; // 16 elements per block + if (row >= M || block_idx >= n_blocks) return; + + float vals[16]; + #pragma unroll + for (int c = 0; c < 4; ++c) { + const float4 raw = src[row * K4 + 4 * block_idx + c]; + vals[4 * c + 0] = raw.x; + vals[4 * c + 1] = raw.y; + vals[4 * c + 2] = raw.z; + vals[4 * c + 3] = raw.w; + } + + float amax = 0.f; + #pragma unroll + for (int i = 0; i < 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(desired); + const float bs_dq = static_cast(bs_q); + + dst_sfa[layout(row, block_idx * 16, 0)] = *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 = fp32_to_e2m1(vals[2 * p] * inv_bs); + const uint8_t hi = fp32_to_e2m1(vals[2 * p + 1] * inv_bs); + ob[p] = static_cast(lo | (hi << 4)); + } + dst_packed[row * n_blocks + block_idx] = out; +} + +} // namespace + +int quantize_act_f32(const float * src, void * dst_packed, void * dst_sfa, + int M, int K, cudaStream_t stream) { + if (K % 16 != 0) return -1; + if ((reinterpret_cast(src) & 15) || + (reinterpret_cast(dst_packed) & 7)) return -1; + + const int n_blocks = K / 16; + const int threads = 128; + dim3 grid((n_blocks + threads - 1) / threads, M); + + auto shape = cute::make_shape(M, 1, K, 1); + auto layout = CfgVec::tile_atom_to_shape_SFA(shape); + + kernel_quantize_f32<<>>( + reinterpret_cast(src), + reinterpret_cast(dst_packed), + reinterpret_cast(dst_sfa), + layout, M, K >> 2); + + const cudaError_t e = cudaGetLastError(); + return (e == cudaSuccess) ? 0 : -static_cast(e); +} + +} // namespace ggml_cuda_flashrt diff --git a/flash_rt/structures/adapters/ggml/fr_repack.cu b/flash_rt/structures/adapters/ggml/fr_repack.cu new file mode 100644 index 000000000..b3289863e --- /dev/null +++ b/flash_rt/structures/adapters/ggml/fr_repack.cu @@ -0,0 +1,458 @@ +// One-time repack of a ggml GGML_TYPE_NVFP4 weight tensor into the CUTLASS +// block-scaled GEMM's B-side wire format. +// +// ggml block_nvfp4 (36 bytes / 64 elements): +// uint8 d[4] - one e4m3 scale per 16-element sub-block; standard e4m3 +// semantics (ggml's doubled dequant table and halved ue4m3 +// decode cancel), passed through unmodified +// uint8 qs[32] - e2m1 codes, sub-block s at qs[s*8..s*8+7], byte j holding +// elem[j] in the low nibble and elem[8+j] in the high nibble +// +// Output: packed uint8 [N, K/2] with adjacent-pair nibbles (elem 2i low, +// elem 2i+1 high) and the scale bytes at the Sm1xx atom-layout offsets. +// One thread handles one 16-element sub-block. + +#include "fr_kernels.h" + +#include +#include + +#include "cutlass/cutlass.h" +#include "cutlass/detail/sm100_blockscaled_layout.hpp" +#include "cute/tensor.hpp" + +namespace ggml_cuda_flashrt { + +namespace { + +using CfgVec = cutlass::detail::Sm1xxBlockScaledConfig<16>; + +constexpr int GGML_NVFP4_BLOCK_BYTES = 36; // 4 scale bytes + 32 data bytes + +template +__global__ void kernel_repack( + const uint8_t * __restrict__ src, // ggml block_nvfp4 stream + uint2 * __restrict__ dst_packed, // [N, K/2] bytes as uint2 per sub-block + uint8_t * __restrict__ dst_sf, + LayoutSF layout, + int N, int K16) { // K16 = K / 16 sub-blocks per row + const int t = blockIdx.x * blockDim.x + threadIdx.x; + const int row = blockIdx.y; + if (row >= N || t >= K16) return; + + const int blk = t >> 2; // 64-element ggml block + const int sub = t & 3; // 16-element sub-block within it + + const uint8_t * b = src + (static_cast(row) * (K16 >> 2) + blk) * GGML_NVFP4_BLOCK_BYTES; + const uint8_t scale = b[sub]; + const uint8_t * qs = b + 4 + sub * 8; + + uint2 out; + uint8_t * ob = reinterpret_cast(&out); + #pragma unroll + for (int p = 0; p < 4; ++p) { + // elements 2p, 2p+1 live in the low nibbles of qs[2p], qs[2p+1] + ob[p] = static_cast((qs[2 * p] & 0x0F) | ((qs[2 * p + 1] & 0x0F) << 4)); + // elements 8+2p, 8+2p+1 live in the high nibbles of the same bytes + ob[p + 4] = static_cast((qs[2 * p] >> 4) | ((qs[2 * p + 1] & 0xF0))); + } + + dst_packed[static_cast(row) * K16 + t] = out; + dst_sf[layout(row, t * 16, 0)] = scale; +} + +} // namespace + +namespace { + +// Pairwise-interleaved variant for the fused GeGLU GEMM: output row 2j is +// gate row j, row 2j+1 is up row j (N_il = 2 * n_ff rows total). +template +__global__ void kernel_repack_pair( + const uint8_t * __restrict__ gate, + const uint8_t * __restrict__ up, + uint2 * __restrict__ dst_packed, + uint8_t * __restrict__ dst_sf, + LayoutSF layout, + int n_ff, int K16) { + const int t = blockIdx.x * blockDim.x + threadIdx.x; + const int row = blockIdx.y; // row within gate/up + const int which = blockIdx.z; // 0 = gate, 1 = up + if (row >= n_ff || t >= K16) return; + + const int blk = t >> 2; + const int sub = t & 3; + const uint8_t * src = which ? up : gate; + const uint8_t * b = src + (static_cast(row) * (K16 >> 2) + blk) * GGML_NVFP4_BLOCK_BYTES; + const uint8_t scale = b[sub]; + const uint8_t * qs = b + 4 + sub * 8; + + uint2 out; + uint8_t * ob = reinterpret_cast(&out); + #pragma unroll + for (int p = 0; p < 4; ++p) { + ob[p] = static_cast((qs[2 * p] & 0x0F) | ((qs[2 * p + 1] & 0x0F) << 4)); + ob[p + 4] = static_cast((qs[2 * p] >> 4) | ((qs[2 * p + 1] & 0xF0))); + } + + const int out_row = 2 * row + which; + dst_packed[static_cast(out_row) * K16 + t] = out; + dst_sf[layout(out_row, t * 16, 0)] = scale; +} + +} // namespace + +int repack_weight_pair_interleaved(const void * gate_blocks, const void * up_blocks, + void * dst_packed, void * dst_sf, + int n_ff, int K, cudaStream_t stream) { + if (K % 64 != 0) return -1; + if (reinterpret_cast(dst_packed) & 7) return -1; + + const int K16 = K / 16; + const int N_il = 2 * n_ff; + const int threads = 128; + dim3 grid((K16 + threads - 1) / threads, n_ff, 2); + + auto shape = cute::make_shape(1, N_il, K, 1); + auto layout = CfgVec::tile_atom_to_shape_SFB(shape); + if (static_cast(cute::cosize(layout)) > sf_bytes(N_il, K)) { + return -3; + } + + kernel_repack_pair<<>>( + reinterpret_cast(gate_blocks), + reinterpret_cast(up_blocks), + reinterpret_cast(dst_packed), + reinterpret_cast(dst_sf), + layout, n_ff, K16); + + const cudaError_t e = cudaGetLastError(); + return (e == cudaSuccess) ? 0 : -static_cast(e); +} + +namespace { + +// Rows-padded variant: rows >= N_src emit zero data and zero scales +// (mathematically inert; the extra rows exist only for output alignment). +template +__global__ void kernel_repack_rows_padded( + const uint8_t * __restrict__ src, + uint2 * __restrict__ dst_packed, + uint8_t * __restrict__ dst_sf, + LayoutSF layout, + int N_src, int N_pad, int K16) { + const int t = blockIdx.x * blockDim.x + threadIdx.x; + const int row = blockIdx.y; + if (row >= N_pad || t >= K16) return; + + uint2 out = make_uint2(0, 0); + uint8_t scale = 0; + if (row < N_src) { + const int blk = t >> 2; + const int sub = t & 3; + const uint8_t * b = src + (static_cast(row) * (K16 >> 2) + blk) * GGML_NVFP4_BLOCK_BYTES; + scale = b[sub]; + const uint8_t * qs = b + 4 + sub * 8; + uint8_t * ob = reinterpret_cast(&out); + #pragma unroll + for (int p = 0; p < 4; ++p) { + ob[p] = static_cast((qs[2 * p] & 0x0F) | ((qs[2 * p + 1] & 0x0F) << 4)); + ob[p + 4] = static_cast((qs[2 * p] >> 4) | ((qs[2 * p + 1] & 0xF0))); + } + } + + dst_packed[static_cast(row) * K16 + t] = out; + dst_sf[layout(row, t * 16, 0)] = scale; +} + +// Group-padded variant: output rows are n_groups groups of group_out rows, +// the first group_in of each group copied from the source (group g, lane l -> +// source row g*group_in + l) and the rest zero. Used to widen per-head +// projections (e.g. SigLIP head_dim 72 -> 80) entirely inside the weights so +// no runtime pad kernel is needed. +template +__global__ void kernel_repack_rows_grouppad( + const uint8_t * __restrict__ src, + uint2 * __restrict__ dst_packed, + uint8_t * __restrict__ dst_sf, + LayoutSF layout, + int group_in, int group_out, int n_groups, int K16) { + const int t = blockIdx.x * blockDim.x + threadIdx.x; + const int row = blockIdx.y; + if (row >= group_out * n_groups || t >= K16) return; + + const int lane = row % group_out; + uint2 out = make_uint2(0, 0); + uint8_t scale = 0; + if (lane < group_in) { + const int src_row = (row / group_out) * group_in + lane; + const int blk = t >> 2; + const int sub = t & 3; + const uint8_t * b = src + (static_cast(src_row) * (K16 >> 2) + blk) * GGML_NVFP4_BLOCK_BYTES; + scale = b[sub]; + const uint8_t * qs = b + 4 + sub * 8; + uint8_t * ob = reinterpret_cast(&out); + #pragma unroll + for (int p = 0; p < 4; ++p) { + ob[p] = static_cast((qs[2 * p] & 0x0F) | ((qs[2 * p + 1] & 0x0F) << 4)); + ob[p + 4] = static_cast((qs[2 * p] >> 4) | ((qs[2 * p + 1] & 0xF0))); + } + } + + dst_packed[static_cast(row) * K16 + t] = out; + dst_sf[layout(row, t * 16, 0)] = scale; +} + +__device__ __forceinline__ uint8_t fr_f32_to_e2m1(float x) { + uint8_t sign = (x < 0.f) ? 0x8u : 0x0u; + float ax = fabsf(x); + uint8_t m; + if (ax <= 0.25f) m = 0u; + else if (ax <= 0.75f) m = 1u; + else if (ax <= 1.25f) m = 2u; + else if (ax <= 1.75f) m = 3u; + else if (ax <= 2.5f) m = 4u; + else if (ax <= 3.5f) m = 5u; + else if (ax <= 5.0f) m = 6u; + else m = 7u; + return sign | m; +} + +// Quantize an fp16 weight [N rows, K_src] to the NVFP4 wire format with the +// K dim zero-padded to K_pad (standard e4m3 amax/6 scales, e2m1 nearest). +template +__global__ void kernel_quantize_weight_f16_padded( + const __half * __restrict__ src, + uint2 * __restrict__ dst_packed, + uint8_t * __restrict__ dst_sf, + LayoutSF layout, + int N, int K_src16, int K_pad16) { + const int t = blockIdx.x * blockDim.x + threadIdx.x; + const int row = blockIdx.y; + if (row >= N || t >= K_pad16) return; + + uint2 out = make_uint2(0, 0); + uint8_t sf = 0; + if (t < K_src16) { + const __half * xr = src + (int64_t) row * (K_src16 * 16) + t * 16; + float vals[16]; + float amax = 0.f; + #pragma unroll + for (int i = 0; i < 16; ++i) { + vals[i] = __half2float(xr[i]); + amax = fmaxf(amax, fabsf(vals[i])); + } + float desired = amax / 6.f; + if (desired < 1e-12f) desired = 1e-12f; + __nv_fp8_e4m3 q(desired); + sf = *reinterpret_cast(&q); + const float inv = 1.f / static_cast(q); + uint8_t * ob = reinterpret_cast(&out); + #pragma unroll + for (int p = 0; p < 8; ++p) { + const uint8_t lo = fr_f32_to_e2m1(vals[2 * p] * inv); + const uint8_t hi = fr_f32_to_e2m1(vals[2 * p + 1] * inv); + ob[p] = static_cast(lo | (hi << 4)); + } + } + + dst_packed[static_cast(row) * K_pad16 + t] = out; + dst_sf[layout(row, t * 16, 0)] = sf; +} + +} // namespace + +int repack_weight_rows_padded(const void * ggml_blocks, void * dst_packed, void * dst_sf, + int N_src, int N_pad, int K, cudaStream_t stream) { + if (K % 64 != 0 || N_pad < N_src) return -1; + const int K16 = K / 16; + const int threads = 128; + dim3 grid((K16 + threads - 1) / threads, N_pad); + auto shape = cute::make_shape(1, N_pad, K, 1); + auto layout = CfgVec::tile_atom_to_shape_SFB(shape); + if (static_cast(cute::cosize(layout)) > sf_bytes(N_pad, K)) return -3; + kernel_repack_rows_padded<<>>( + reinterpret_cast(ggml_blocks), + reinterpret_cast(dst_packed), + reinterpret_cast(dst_sf), + layout, N_src, N_pad, K16); + const cudaError_t e = cudaGetLastError(); + return (e == cudaSuccess) ? 0 : -static_cast(e); +} + +int repack_weight_rows_grouppad(const void * ggml_blocks, void * dst_packed, void * dst_sf, + int group_in, int group_out, int n_groups, int K, cudaStream_t stream) { + if (K % 64 != 0 || group_out < group_in || n_groups <= 0) return -1; + const int N_pad = group_out * n_groups; + const int K16 = K / 16; + const int threads = 128; + dim3 grid((K16 + threads - 1) / threads, N_pad); + auto shape = cute::make_shape(1, N_pad, K, 1); + auto layout = CfgVec::tile_atom_to_shape_SFB(shape); + if (static_cast(cute::cosize(layout)) > sf_bytes(N_pad, K)) return -3; + kernel_repack_rows_grouppad<<>>( + reinterpret_cast(ggml_blocks), + reinterpret_cast(dst_packed), + reinterpret_cast(dst_sf), + layout, group_in, group_out, n_groups, K16); + const cudaError_t e = cudaGetLastError(); + return (e == cudaSuccess) ? 0 : -static_cast(e); +} + +int quantize_weight_f16_padded(const void * w_f16, void * dst_packed, void * dst_sf, + int N, int K_src, int K_pad, cudaStream_t stream) { + if (K_src % 16 != 0 || K_pad % 64 != 0 || K_pad < K_src) return -1; + const int K16 = K_pad / 16; + const int threads = 128; + dim3 grid((K16 + threads - 1) / threads, N); + auto shape = cute::make_shape(1, N, K_pad, 1); + auto layout = CfgVec::tile_atom_to_shape_SFB(shape); + if (static_cast(cute::cosize(layout)) > sf_bytes(N, K_pad)) return -3; + kernel_quantize_weight_f16_padded<<>>( + reinterpret_cast(w_f16), + reinterpret_cast(dst_packed), + reinterpret_cast(dst_sf), + layout, N, K_src / 16, K16); + const cudaError_t e = cudaGetLastError(); + return (e == cudaSuccess) ? 0 : -static_cast(e); +} + +namespace { + +// Three-tensor row-concat repack for the fused QKV GEMM: output rows are +// [src0 | src1 | src2] stacked (N0 + N1 + N2 rows, same K). +template +__global__ void kernel_repack_concat3( + const uint8_t * __restrict__ s0, int N0, + const uint8_t * __restrict__ s1, int N1, + const uint8_t * __restrict__ s2, int N2, + uint2 * __restrict__ dst_packed, + uint8_t * __restrict__ dst_sf, + LayoutSF layout, + int K16) { + const int t = blockIdx.x * blockDim.x + threadIdx.x; + const int row = blockIdx.y; + const int N_tot = N0 + N1 + N2; + if (row >= N_tot || t >= K16) return; + + const uint8_t * src; + int src_row; + if (row < N0) { src = s0; src_row = row; } + else if (row < N0 + N1) { src = s1; src_row = row - N0; } + else { src = s2; src_row = row - N0 - N1; } + + const int blk = t >> 2; + const int sub = t & 3; + const uint8_t * b = src + (static_cast(src_row) * (K16 >> 2) + blk) * GGML_NVFP4_BLOCK_BYTES; + const uint8_t scale = b[sub]; + const uint8_t * qs = b + 4 + sub * 8; + + uint2 out; + uint8_t * ob = reinterpret_cast(&out); + #pragma unroll + for (int p = 0; p < 4; ++p) { + ob[p] = static_cast((qs[2 * p] & 0x0F) | ((qs[2 * p + 1] & 0x0F) << 4)); + ob[p + 4] = static_cast((qs[2 * p] >> 4) | ((qs[2 * p + 1] & 0xF0))); + } + + dst_packed[static_cast(row) * K16 + t] = out; + dst_sf[layout(row, t * 16, 0)] = scale; +} + +} // namespace + +int repack_weight_concat3(const void * b0, int N0, const void * b1, int N1, + const void * b2, int N2, + void * dst_packed, void * dst_sf, + int K, cudaStream_t stream) { + if (K % 64 != 0) return -1; + const int K16 = K / 16; + const int N_tot = N0 + N1 + N2; + const int threads = 128; + dim3 grid((K16 + threads - 1) / threads, N_tot); + auto shape = cute::make_shape(1, N_tot, K, 1); + auto layout = CfgVec::tile_atom_to_shape_SFB(shape); + if (static_cast(cute::cosize(layout)) > sf_bytes(N_tot, K)) return -3; + kernel_repack_concat3<<>>( + reinterpret_cast(b0), N0, + reinterpret_cast(b1), N1, + reinterpret_cast(b2), N2, + reinterpret_cast(dst_packed), + reinterpret_cast(dst_sf), + layout, K16); + const cudaError_t e = cudaGetLastError(); + return (e == cudaSuccess) ? 0 : -static_cast(e); +} + +int repack_weight(const void * ggml_blocks, void * dst_packed, void * dst_sf, + int N, int K, cudaStream_t stream) { + if (K % 64 != 0) return -1; + if (reinterpret_cast(dst_packed) & 7) return -1; + + const int K16 = K / 16; + const int threads = 128; + dim3 grid((K16 + threads - 1) / threads, N); + + // SFB layout for a [*, N, K] problem; independent of M. + auto shape = cute::make_shape(1, N, K, 1); + auto layout = CfgVec::tile_atom_to_shape_SFB(shape); + + // The atom layout may address up to the padded (row, k) extents; the + // caller allocates sf_bytes(N, K) which must cover the layout codomain. + if (static_cast(cute::cosize(layout)) > sf_bytes(N, K)) { + return -3; + } + + kernel_repack<<>>( + reinterpret_cast(ggml_blocks), + reinterpret_cast(dst_packed), + reinterpret_cast(dst_sf), + layout, N, K16); + + const cudaError_t e = cudaGetLastError(); + return (e == cudaSuccess) ? 0 : -static_cast(e); +} + +namespace { + +struct cpy_rows_pair { + const float * src; + __half * dst; +}; + +struct cpy_rows_args { + cpy_rows_pair pairs[FR_CPY_ROWS_MAX]; +}; + +// one y-slice of rows per pair; conversion identical to ggml's f32->f16 cpy +__global__ void kernel_cpy_rows_f32_f16(cpy_rows_args args, int hd, int n_rows) { + const cpy_rows_pair p = args.pairs[blockIdx.x]; + for (int r = blockIdx.y; r < n_rows; r += gridDim.y) { + const float * s = p.src + (int64_t) r * hd; + __half * d = p.dst + (int64_t) r * hd; + for (int i = threadIdx.x; i < hd; i += blockDim.x) { + d[i] = __float2half(s[i]); + } + } +} + +} // namespace + +int cpy_rows_f32_f16(const float * const * srcs, void * const * dsts, int n_pairs, + int hd, int n_rows, cudaStream_t stream) { + if (n_pairs < 1 || n_pairs > FR_CPY_ROWS_MAX) { + return -1; + } + cpy_rows_args args; + for (int i = 0; i < n_pairs; ++i) { + args.pairs[i].src = srcs[i]; + args.pairs[i].dst = (__half *) dsts[i]; + } + const int rows_y = n_rows < 64 ? n_rows : 64; + dim3 grid(n_pairs, rows_y); + kernel_cpy_rows_f32_f16<<>>(args, hd, n_rows); + const cudaError_t e = cudaGetLastError(); + return (e == cudaSuccess) ? 0 : -static_cast(e); +} + +} // namespace ggml_cuda_flashrt diff --git a/flash_rt/structures/adapters/ggml/qualification/goldens/base.png b/flash_rt/structures/adapters/ggml/qualification/goldens/base.png new file mode 100644 index 000000000..97a77cdf9 Binary files /dev/null and b/flash_rt/structures/adapters/ggml/qualification/goldens/base.png differ diff --git a/flash_rt/structures/adapters/ggml/qualification/goldens/pi05_thor_action.json b/flash_rt/structures/adapters/ggml/qualification/goldens/pi05_thor_action.json new file mode 100644 index 000000000..87e37c05d --- /dev/null +++ b/flash_rt/structures/adapters/ggml/qualification/goldens/pi05_thor_action.json @@ -0,0 +1 @@ +{"action_final_raw": [[0.09624627977609634, -0.1961977779865265, 0.06703978776931763, -0.5476640462875366, -0.14697180688381195, 0.07153154909610748, -0.002638055244460702, 0.013559071347117424, 0.016760675236582756, 0.007358239963650703, -0.002592590870335698, 0.002914540935307741, 0.006588217802345753, 0.02381635643541813, 0.000132353205117397, 0.004709492437541485, -0.00813357625156641, -0.013270657509565353, 0.009339207783341408, 0.020445771515369415, 0.003282379824668169, -0.002165404614061117, 0.002376250457018614, -0.0039012939669191837, 0.0014441026141867042, 0.0002516416134312749, 0.003512145485728979, -0.007321126293390989, 0.0017400443321093917, 0.004314625635743141, -0.00405048206448555, -0.005579052492976189], [0.12092496454715729, -0.1998528391122818, 0.054504960775375366, -0.5915311574935913, -0.1741236299276352, 0.06810212880373001, -0.003277003997936845, -0.0034880663733929396, 0.003501307452097535, 0.005087920930236578, -0.0022917466703802347, 0.007821531035006046, -0.006648858077824116, 0.027163516730070114, -0.0011810704600065947, -0.00011856241326313466, -0.01714256964623928, 0.004868022631853819, 0.007669864222407341, 0.0008107393514364958, 0.0073386593721807, -0.0016696092206984758, -0.00048284963122569025, -0.009979461319744587, -0.01471013855189085, 0.0012931901728734374, -0.002198341768234968, -0.01144750788807869, 0.010050974786281586, 0.005830400623381138, -0.007881347090005875, 0.0015526266070082784], [0.1221553236246109, -0.20009484887123108, 0.07306937873363495, -0.6259077787399292, -0.17243315279483795, 0.07382355630397797, 0.01399582251906395, 0.008802780881524086, 0.007673643529415131, -0.003587680170312524, -0.00198595249094069, 0.009418884292244911, 0.01879180409014225, 0.02367384359240532, -0.002254959661513567, -0.0034092110581696033, -0.015890508890151978, -0.012552782893180847, 0.010030963458120823, 0.004080003593116999, -0.006269955076277256, -0.004506299272179604, 0.005475796293467283, -0.003492268966510892, -0.005690732505172491, -0.0047112093307077885, -0.007371499668806791, -0.0032012038864195347, 0.00254984968341887, 0.0003990198310930282, 0.006445983424782753, 0.009539195336401463], [0.14189547300338745, -0.20173774659633636, 0.06596016138792038, -0.6811384558677673, -0.1881892830133438, 0.08414757251739502, 0.0160377100110054, 0.02549123391509056, 0.005759424064308405, 0.005022574216127396, 0.014015900902450085, -0.0007282908773049712, 0.010810497216880322, 0.017419779673218727, -0.0012818204704672098, -0.013671827502548695, -0.009128696285188198, -0.002643367275595665, 0.01019663829356432, 0.010841459967195988, -0.002366367494687438, -0.0019233895000070333, 0.0014648950891569257, 0.004127361811697483, 0.0019628419540822506, 0.005564515478909016, -0.004168566316366196, 0.003304382786154747, -0.0038799545727670193, 0.0023837818298488855, -0.008190134540200233, 0.009892459958791733], [0.1522434502840042, -0.20274382829666138, 0.08215358853340149, -0.7039123177528381, -0.19095152616500854, 0.07766478508710861, -0.01269338559359312, 0.00959023181349039, -0.002931674476712942, 0.016207940876483917, -0.009432277642190456, 0.0004770905652549118, 0.004761648364365101, 0.034024689346551895, -0.008643164299428463, 0.0033901091665029526, -0.020978260785341263, -0.010628817602992058, 0.02480119839310646, 0.00730505958199501, 0.0014040278038010001, -0.0018820121185854077, 0.003079682355746627, 0.0030281723011285067, 0.0009533020784147084, 0.0014496227959170938, -0.004266194999217987, -0.007991557009518147, 0.007314159069210291, -0.005737284664064646, -0.0032646663021296263, -0.0032309989910572767], [0.18085116147994995, -0.20060132443904877, 0.06863894313573837, -0.7330130338668823, -0.1915489137172699, 0.09672202914953232, 0.01540116872638464, 0.007338434923440218, 0.017650838941335678, 0.01925886608660221, -0.008773401379585266, 0.013958729803562164, 0.01029625441879034, 0.02788020297884941, -0.0022923294454813004, -0.013784453272819519, -0.016574576497077942, -0.004523152951151133, 0.01866416074335575, 0.013041609898209572, -0.001756405341438949, 0.007274901028722525, -0.0009395781089551747, 0.0024610029067844152, 0.0031260414980351925, 0.003031071275472641, -0.006058946251869202, -0.006632425356656313, -0.0009360113763250411, -0.00719192810356617, -0.001834192662499845, 0.0030427579768002033], [0.18593062460422516, -0.19550056755542755, 0.07160288095474243, -0.7510550022125244, -0.19466954469680786, 0.10721323639154434, -0.013059914112091064, 0.008966893889009953, 0.0030186132062226534, 0.0019175108755007386, 0.0020664960611611605, -0.00020583205332513899, 0.013402078300714493, 0.02108699642121792, -0.002505762968212366, 0.007308184169232845, -0.028859110549092293, -0.005017835181206465, 0.01019445899873972, 0.012598940171301365, -0.004889514297246933, -0.002243961440399289, -0.007915007881820202, 0.002988256746903062, 0.008123243227601051, 0.0008216212736442685, -0.008317344821989536, -0.013078191317617893, -0.01282085757702589, 0.007809976581484079, -0.005311104469001293, -0.0201072059571743], [0.207765594124794, -0.21467824280261993, 0.06326164305210114, -0.7796165943145752, -0.19674517214298248, 0.10123662650585175, 0.00892910547554493, 0.009866893291473389, 0.0011979155242443085, -0.005581721663475037, -0.00986845325678587, -0.00014581959112547338, 0.0097585991024971, 0.0157171580940485, 0.00469200499355793, -0.009907061234116554, -0.01478598453104496, -0.0044782706536352634, 0.015776481479406357, 0.004031747113913298, 0.0027923069428652525, 0.00022482436907012016, 0.0036739581264555454, -0.014940707944333553, -0.009240969084203243, 0.012544935569167137, 0.000710038875695318, -0.0036881084088236094, -0.0074166469275951385, -0.002954066963866353, 0.004561707377433777, 0.00133541040122509], [0.21232953667640686, -0.1888059377670288, 0.06611032038927078, -0.7755351066589355, -0.2044234275817871, 0.10427963733673096, -0.014994261786341667, 0.018358377739787102, -0.0003517218283377588, 0.0059860870242118835, -0.0068831369280815125, 0.003228644607588649, 0.014275294728577137, 0.018018964678049088, 0.0024007963947951794, -0.004222281742841005, -0.025340044870972633, 0.006304721813648939, 0.017566384747624397, 0.0018504681065678596, -0.005991443060338497, -0.0029909336008131504, -0.0013041617348790169, 8.168922795448452e-05, -0.008946222253143787, 0.018412627279758453, -0.003031445201486349, 0.0029337876476347446, 0.0039048572070896626, 0.0065581281669437885, 0.01261038240045309, -0.007260486483573914], [0.21789924800395966, -0.1944425106048584, 0.0784466341137886, -0.8004544973373413, -0.20648646354675293, 0.10525409877300262, -0.012714697048068047, 0.01368754357099533, 0.010657607577741146, 0.009618283249437809, -0.008845558390021324, 0.007256275042891502, -0.011666398495435715, 0.020875461399555206, 0.014044287614524364, 0.0017091069603338838, 0.003284646663814783, -0.004290799144655466, 0.016678009182214737, 0.004360348451882601, 0.006027561146765947, -0.009836606681346893, -0.009942407719790936, 0.001068001496605575, -0.002853973535820842, -0.006787698715925217, -0.012501326389610767, -0.022064579650759697, -0.004612510558217764, 0.007948911748826504, -0.0033561980817466974, -0.008420849218964577]], "action_steps": 10, "action_dim": 32} \ No newline at end of file diff --git a/flash_rt/structures/adapters/ggml/qualification/goldens/wrist.png b/flash_rt/structures/adapters/ggml/qualification/goldens/wrist.png new file mode 100644 index 000000000..6e401b0a8 Binary files /dev/null and b/flash_rt/structures/adapters/ggml/qualification/goldens/wrist.png differ diff --git a/flash_rt/structures/adapters/ggml/qualification/pins.yaml b/flash_rt/structures/adapters/ggml/qualification/pins.yaml new file mode 100644 index 000000000..d38744cf3 --- /dev/null +++ b/flash_rt/structures/adapters/ggml/qualification/pins.yaml @@ -0,0 +1,16 @@ +# Structure versions the ggml adapter's fused windows were written against. +# The qualification runner turns RED when the live catalog moves past a pin: +# that is the signal to re-audit the corresponding window before adopting +# the new structure version, then update the pin. +binding: jetson_pi_edge_pi05 +pins: + vla_tick_pipeline: 2 + decoder_ffn: 1 + vision_ffn: 2 + qkv_pack: 1 + attention_core: 1 + linear_proj: 1 + norm_fused: 1 + adaln_producer: 1 + modnorm_qkv_chain: 3 + cadence_static: 1 diff --git a/flash_rt/structures/adapters/ggml/qualification/run_qualification.py b/flash_rt/structures/adapters/ggml/qualification/run_qualification.py new file mode 100644 index 000000000..c1c61f8e5 --- /dev/null +++ b/flash_rt/structures/adapters/ggml/qualification/run_qualification.py @@ -0,0 +1,177 @@ +#!/usr/bin/env python3 +"""Qualification gates for the ggml host adapter. + +Offline gates (always run): + A. manifest — the pipeline binding validates against the live catalog + (structure renames, removed embedded regions, or a + malformed manifest turn this red). + B. pins — every pinned structure version matches the live catalog + (an upstream version bump turns this red and names the + structure, which is the cue to re-audit the bound window + before adopting the bump). + +On-device gate (opt-in, needs a running llama-server with the pi0.5 model): + C. e2e-golden — drives the fixed synthetic-input protocol against the + server and compares the raw action chunk to a stored + golden. The comparison is exact by default (the adapter + is bitwise deterministic across processes); pass --tol + to allow a max-abs band instead. Any kernel or window + change that shifts numerics turns this red. + +Usage: + python run_qualification.py # gates A+B + python run_qualification.py --e2e --port 8089 # gates A+B+C + python run_qualification.py --e2e --update-golden # refresh the golden +""" + +from __future__ import annotations + +import argparse +import json +import pathlib +import sys + +import yaml + +_HERE = pathlib.Path(__file__).resolve().parent +_REPO = _HERE.parents[4] +if str(_REPO) not in sys.path: + sys.path.insert(0, str(_REPO)) + +from flash_rt.structures.binding import load_binding # noqa: E402 +from flash_rt.structures.registry import load as load_structure # noqa: E402 + +GOLDEN = _HERE / "goldens" / "pi05_thor_action.json" + + +def gate_manifest(binding_name: str) -> tuple[bool, str]: + try: + spec = load_binding(binding_name, require_pipeline_coverage=True) + except Exception as exc: # noqa: BLE001 — any validation failure is red + return False, f"binding failed validation: {exc}" + return True, (f"{spec.name} -> {spec.structure.name}@" + f"{spec.structure.version}, {len(spec.segments)} segments, " + f"contract {spec.coverage_contract}") + + +def gate_pins() -> tuple[bool, str]: + pinned = yaml.safe_load((_HERE / "pins.yaml").read_text())["pins"] + drifted = [] + for name, version in pinned.items(): + try: + live = load_structure(name).version + except KeyError: + drifted.append(f"{name}: pinned @{version}, missing from catalog") + continue + if int(live) != int(version): + drifted.append(f"{name}: pinned @{version}, catalog is @{live}") + if drifted: + return False, "; ".join(drifted) + return True, f"{len(pinned)} structure versions match the catalog" + + +# ---- gate C: end-to-end action golden -------------------------------------- + +def _server_request(base: str, method: str, path: str, body=None): + import urllib.request + opener = urllib.request.build_opener(urllib.request.ProxyHandler({})) + data = json.dumps(body if body is not None else {}).encode() + req = urllib.request.Request(base + path, data=data, method=method, + headers={"Content-Type": "application/json"}) + with opener.open(req, timeout=120) as f: + return json.loads(f.read()) + + +def _synthetic_images(directory: pathlib.Path) -> list[str]: + import numpy as np + from PIL import Image + rng = np.random.default_rng(1234) + paths = [] + for name in ("base.png", "wrist.png"): + p = directory / name + if not p.exists(): + Image.fromarray( + rng.integers(0, 255, (224, 224, 3), dtype=np.uint8)).save(p) + paths.append(str(p)) + return paths + + +def run_pipeline(port: int, image_dir: pathlib.Path, warmup: int = 2) -> dict: + """Fixed synthetic-input protocol; the first evaluations after server + start are cold (cache fill / capture paths) and differ from the steady + state, so the comparison value is taken after ``warmup`` full passes — + steady-state output is bitwise stable across runs and processes.""" + base = f"http://127.0.0.1:{port}" + images = _synthetic_images(image_dir) + state = ",".join(f"{0.01 * i:.4f}" for i in range(32)) + resp = None + for _ in range(warmup + 1): + _server_request(base, "POST", "/foreground/reset") + for p in images: + _server_request(base, "POST", "/foreground/image", {"path": p}) + _server_request(base, "PUT", "/foreground/state", {"state": state}) + resp = _server_request(base, "POST", "/foreground/infer", + {"text": "pick up the object"}) + return {"action_final_raw": resp.get("action_final_raw"), + "action_steps": resp.get("action_steps"), + "action_dim": resp.get("action_dim")} + + +def gate_e2e(port: int, image_dir: pathlib.Path, tol: float, + update_golden: bool) -> tuple[bool, str]: + got = run_pipeline(port, image_dir) + if got["action_final_raw"] is None: + return False, "server returned no action_final_raw" + if update_golden: + GOLDEN.parent.mkdir(parents=True, exist_ok=True) + GOLDEN.write_text(json.dumps(got)) + return True, f"golden updated: {GOLDEN}" + if not GOLDEN.exists(): + return False, f"no golden at {GOLDEN}; run with --update-golden first" + want = json.loads(GOLDEN.read_text()) + if (got["action_steps"] != want["action_steps"] + or got["action_dim"] != want["action_dim"]): + return False, (f"shape drift: {got['action_steps']}x{got['action_dim']}" + f" vs golden {want['action_steps']}x{want['action_dim']}") + flat_got = [x for row in got["action_final_raw"] for x in row] + flat_want = [x for row in want["action_final_raw"] for x in row] + diffs = [abs(a - b) for a, b in zip(flat_got, flat_want)] + worst = max(diffs) + n_diff = sum(1 for d in diffs if d > tol) + if n_diff: + return False, (f"{n_diff}/{len(diffs)} elements beyond tol={tol:g}, " + f"max_abs_diff={worst:.3e}") + return True, (f"{len(diffs)} elements within tol={tol:g} " + f"(max_abs_diff={worst:.3e})") + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--binding", default="jetson_pi_edge_pi05") + ap.add_argument("--e2e", action="store_true", + help="also run the on-device action-golden gate") + ap.add_argument("--port", type=int, default=8089) + ap.add_argument("--image-dir", type=pathlib.Path, + default=_HERE / "goldens") + ap.add_argument("--tol", type=float, default=0.0, + help="max-abs tolerance for the e2e gate (default exact)") + ap.add_argument("--update-golden", action="store_true") + args = ap.parse_args() + + gates = [("manifest", gate_manifest(args.binding)), + ("pins", gate_pins())] + if args.e2e: + args.image_dir.mkdir(parents=True, exist_ok=True) + gates.append(("e2e-golden", + gate_e2e(args.port, args.image_dir, args.tol, + args.update_golden))) + + all_ok = True + for name, (ok, detail) in gates: + print(f"[{'GREEN' if ok else 'RED':5s}] {name}: {detail}") + all_ok = all_ok and ok + return 0 if all_ok else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/flash_rt/structures/bindings/jetson_pi_edge_pi05.yaml b/flash_rt/structures/bindings/jetson_pi_edge_pi05.yaml new file mode 100644 index 000000000..ffa8b90d3 --- /dev/null +++ b/flash_rt/structures/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 +# (flash_rt/structures/adapters/ggml/). 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/third_party/cutlass b/third_party/cutlass new file mode 160000 index 000000000..da5e086da --- /dev/null +++ b/third_party/cutlass @@ -0,0 +1 @@ +Subproject commit da5e086dab31d63815acafdac9a9c5893b1c69e2