diff --git a/CMakeLists.txt b/CMakeLists.txt index 5c5a6fb00..3fdb85f22 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1347,6 +1347,50 @@ else() message(STATUS "SM120 NVFP4 W4A16 fp4_w4a4 kernels: SKIPPED (slim build)") endif() +# Group A addition — interleaved-B warp-split tiers (bind-time repacked B; +# same arch gate and binding guard as the dense warp-split family above). +if(NOT FLASHRT_SLIM_BUILD OR GPU_ARCH STREQUAL "120" OR GPU_ARCH STREQUAL "121") + target_sources(flash_rt_kernels PRIVATE + csrc/kernels/fp4_w4a4_mma_warpsplit_ilv_sm120.cu + csrc/kernels/fp4_w4a4_mma_warpsplit_ilv_mrows_sm120.cu) +endif() + +# Fused SwiGLU + NVFP4 quantize producer (plain CUDA elementwise, no +# CUTLASS dependency; quantize stage transcribed from the production +# bf16 quantize kernel). +target_sources(flash_rt_kernels PRIVATE + csrc/kernels/silu_mul_quantize_fp4_sfa_bf16.cu) + +# Batched 64x64 unit-lower-triangular inverse for the gated-delta WY +# chain (plain CUDA, arch-agnostic). +target_sources(flash_rt_kernels PRIVATE + csrc/kernels/batched_unit_ltri_inv64.cu) + +# Fused (1+w)-form RMSNorm + NVFP4 quantize producer (plain CUDA). +target_sources(flash_rt_kernels PRIVATE + csrc/kernels/rms_norm_quantize_fp4_sfa_bf16.cu) + +# WY-chain norm/pack + chunk-parallel gate cumsum, v2 launch plan. +target_sources(flash_rt_kernels PRIVATE + csrc/kernels/gdn_wy_norm_cumsum_pack_qk_v2.cu) + +# Causal conv1d update with per-thread step batching + GQA split. +target_sources(flash_rt_kernels PRIVATE + csrc/kernels/causal_conv1d_update_steps_gqa_bf16.cu) + +# Gated-delta recurrent decode step, streaming-column form. +target_sources(flash_rt_kernels PRIVATE + csrc/kernels/gdn_recurrent_inout_stream_bf16.cu) + +# Gated norm that also emits its consumer's NVFP4 input. +target_sources(flash_rt_kernels PRIVATE + csrc/kernels/rms_norm_gated_silu_quant_fp4_bf16.cu) + +# Gated-delta chunk core, per-row state stash arm (spec-verify rollback +# by selection). Plain CUDA, no CUTLASS dependency, arch-agnostic. +target_sources(flash_rt_kernels PRIVATE + csrc/kernels/gdn_chunk_from_conv_smem_stash.cu) + # Group B — NVFP4 swizzle/quantize helpers. Their bindings are unguarded today # (exposed on SM89 but runtime-unreachable: every Python loader self-gates on # has_nvfp4()). Gating the sources requires guarding their bindings + includes diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index ade0577e6..e7952a172 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -186,6 +186,16 @@ extern "C" int cutlass_int8_rowwise_bf16out_t64x128( #include "kernels/fp4_w4a4_mma_sm120.cuh" #include "kernels/fp4_w4a4_mma_warpsplit_sm120.cuh" #include "kernels/fp4_w4a4_mma_warpsplit_mrows_sm120.cuh" +#include "kernels/fp4_w4a4_mma_warpsplit_ilv_sm120.cuh" +#include "kernels/fp4_w4a4_mma_warpsplit_ilv_mrows_sm120.cuh" +#include "kernels/gdn_chunk_from_conv_smem_stash.cuh" +#include "kernels/silu_mul_quantize_fp4_sfa_bf16.cuh" +#include "kernels/batched_unit_ltri_inv64.cuh" +#include "kernels/rms_norm_quantize_fp4_sfa_bf16.cuh" +#include "kernels/gdn_wy_norm_cumsum_pack_qk_v2.cuh" +#include "kernels/causal_conv1d_update_steps_gqa_bf16.cuh" +#include "kernels/gdn_recurrent_inout_stream_bf16.cuh" +#include "kernels/rms_norm_gated_silu_quant_fp4_bf16.cuh" #include "quantize/fp8_block128_dequant.cuh" #ifdef FLASHRT_HAVE_NVFP4_SWIZZLE #include "quantize/fp8_block128_to_nvfp4_swizzled.cuh" @@ -6630,10 +6640,225 @@ a full 16-row tile, so M<=16 rows cost the same weight HBM as M=1; combined with warp-split-K (K split across `warps`, partials summed in shared memory -> graph-replay safe) to fill the SMs on long K. M in 1..16; N%8==0; K%64==0; (K/64)%warps==0; warps in {2,4,8}; stages in {3,4,6}. +)pbdoc"); + + m.def("fp4_w4a4_mma_sm120_warpsplit_ilv_bf16out", + [](uintptr_t A, uintptr_t B, uintptr_t D, int N, int K, uintptr_t SFA, + uintptr_t SFB, float alpha, int warps, int stages, + uintptr_t stream) -> int { + return flash_rt::gemm::fp4_w4a4_mma_sm120_warpsplit_ilv_bf16out( + to_ptr(A), to_ptr(B), to_ptr(D), N, K, to_ptr(SFA), + to_ptr(SFB), alpha, warps, stages, to_stream(stream)); + }, + py::arg("A"), py::arg("B"), py::arg("D"), py::arg("N"), py::arg("K"), + py::arg("SFA"), py::arg("SFB"), py::arg("alpha") = 1.0f, + py::arg("warps") = 4, py::arg("stages") = 4, py::arg("stream") = 0, + R"pbdoc( +NVFP4 W4A4 warp-split-K M=1 GEMV over the interleaved B layout (SM120). +B must come from fp4_w4a4_repack_b_ilv_sm120 (bind-time repack): each +8-column group's K-tiles interleave so the block's global reads are fully +contiguous. Bit-exact vs the dense warp-split GEMV (identical reduction +order). N%8==0; K%64==0; (K/64)%warps==0; warps in {2,4,8}; stages {3,4,6}. +)pbdoc"); + + m.def("fp4_w4a4_mma_sm120_warpsplit_ilv_mrows_bf16out", + [](uintptr_t A, uintptr_t B, uintptr_t D, int M, int N, int K, + uintptr_t SFA, uintptr_t SFB, float alpha, int warps, int stages, + uintptr_t stream) -> int { + return flash_rt::gemm:: + fp4_w4a4_mma_sm120_warpsplit_ilv_mrows_bf16out( + to_ptr(A), to_ptr(B), to_ptr(D), M, N, K, to_ptr(SFA), + to_ptr(SFB), alpha, warps, stages, to_stream(stream)); + }, + py::arg("A"), py::arg("B"), py::arg("D"), py::arg("M"), py::arg("N"), + py::arg("K"), py::arg("SFA"), py::arg("SFB"), py::arg("alpha") = 1.0f, + py::arg("warps") = 4, py::arg("stages") = 4, py::arg("stream") = 0, + R"pbdoc( +NVFP4 W4A4 small-M (M<=16) warp-split-K GEMM over the interleaved B layout +(SM120). Shares the repacked B with the M=1 interleaved entry; bit-exact per +row vs it (identical reduction order). M in 1..16; N%8==0; K%64==0; +(K/64)%warps==0; warps in {2,4,8}; stages in {3,4,6}. +)pbdoc"); + + m.def("fp4_w4a4_repack_b_ilv_sm120", + [](uintptr_t B_packed, uintptr_t B_ilv, int N, int K, + uintptr_t stream) -> int { + return flash_rt::gemm::fp4_w4a4_repack_b_ilv_sm120( + to_ptr(B_packed), to_ptr(B_ilv), N, K, to_stream(stream)); + }, + py::arg("B_packed"), py::arg("B_ilv"), py::arg("N"), py::arg("K"), + py::arg("stream") = 0, + R"pbdoc( +Bind-time repack: dense packed B (N, K/2 bytes row-major) -> interleaved +[N/8 groups][K/64 tiles][8 cols x 32B] layout for the warpsplit_ilv entries. +dst byte size equals src (N*K/2). Pure byte permutation on stream. )pbdoc"); #endif + m.def("silu_mul_quantize_fp4_sfa_bf16", + [](uintptr_t merged, uintptr_t packed, uintptr_t sfa, int N, int H, + uintptr_t stream) -> int { + return flash_rt::fp4::silu_mul_quantize_fp4_sfa_bf16( + to_ptr(merged), to_ptr(packed), to_ptr(sfa), N, H, + to_stream(stream)); + }, + py::arg("merged"), py::arg("packed"), py::arg("sfa"), py::arg("N"), + py::arg("H"), py::arg("stream") = 0, + R"pbdoc( +Fused SwiGLU activation + NVFP4 quantize producer. merged: bf16 (N, 2H) +row-major, halves ordered [gate|up], H%16==0. Emits packed (N, H/2 bytes) ++ SFA in the 128x64-atom layout ((N+127)/128 * (H+63)/64 * 512 bytes). +silu·mul in fp32 rounded through bf16, then the production quantize path +verbatim - bit-exact vs the split (mul kernel -> quantize kernel) chain. +)pbdoc"); + + m.def("batched_unit_ltri_inv64_f32", + [](uintptr_t A, uintptr_t X, int B, uintptr_t stream) -> int { + return flash_rt::kernels::batched_unit_ltri_inv64_f32( + to_ptr(A), to_ptr(X), B, to_stream(stream)); + }, + py::arg("A"), py::arg("X"), py::arg("B"), py::arg("stream") = 0, + R"pbdoc( +Batched 64x64 unit-lower-triangular inverse, fp32: X = inv(I + +strict_tril(A)). A, X: (B, 64, 64) row-major; only A's strict lower +triangle is read. Column-independent forward substitution - the same +recurrence class as the batched cuBLAS solve it replaces. +)pbdoc"); + + m.def("rms_norm_quantize_fp4_sfa_bf16", + [](uintptr_t x, uintptr_t w, float eps, uintptr_t normed, + uintptr_t packed, uintptr_t sfa, int N, int D, + uintptr_t stream) -> int { + return flash_rt::fp4::rms_norm_quantize_fp4_sfa_bf16( + to_ptr(x), to_ptr(w), eps, to_ptr(normed), to_ptr(packed), + to_ptr(sfa), N, D, to_stream(stream)); + }, + py::arg("x"), py::arg("w"), py::arg("eps"), py::arg("normed"), + py::arg("packed"), py::arg("sfa"), py::arg("N"), py::arg("D"), + py::arg("stream") = 0, + R"pbdoc( +Fused (1+w)-form RMSNorm + NVFP4 quantize producer. x: bf16 (N, D); +w: bf16 (D) residual-form weight (kernel applies 1+w). Emits the normed +bf16 rows AND their packed FP4 (N, D/2 bytes) + SFA (128x64-atom layout) +off one read of x - the norm the host computes, plus the quantization +its bound consumers would each redo. D%16==0, D<=8192. +)pbdoc"); + + m.def("gdn_wy_norm_cumsum_pack_qk_v2_bf16", + [](uintptr_t q16, uintptr_t k16, uintptr_t g, uintptr_t q16_l2, + uintptr_t k16_l2, uintptr_t q_pack_hv, uintptr_t k_pack_hk, + uintptr_t g_cumsum, int S, uintptr_t stream) -> int { + return flash_rt::kernels::gdn_wy_norm_cumsum_pack_qk_v2_bf16( + to_ptr(q16), to_ptr(k16), to_ptr(g), to_ptr(q16_l2), + to_ptr(k16_l2), to_ptr(q_pack_hv), to_ptr(k_pack_hk), + to_ptr(g_cumsum), S, to_stream(stream)); + }, + py::arg("q16"), py::arg("k16"), py::arg("g"), py::arg("q16_l2"), + py::arg("k16_l2"), py::arg("q_pack_hv"), py::arg("k_pack_hk"), + py::arg("g_cumsum"), py::arg("S"), py::arg("stream") = 0, + R"pbdoc( +WY-chain q/k L2-norm + pack + per-chunk gate cumsum for the fixed +16/48/128/64 family, v2 launch plan: the norm/pack math transcribed +verbatim from the packaged fast arm, the gate cumsum parallelized over +the independent (chunk, head) pairs with the packaged serial order kept +inside each chunk - bit-exact against the packaged pair. +)pbdoc"); + + m.def("causal_conv1d_update_steps_gqa_bf16", + [](uintptr_t x, uintptr_t w, uintptr_t bias, uintptr_t state, + uintptr_t q16, uintptr_t k16, uintptr_t v48, int S, + bool apply_silu, uintptr_t stream) -> int { + return flash_rt::kernels::causal_conv1d_update_steps_gqa_bf16( + to_ptr(x), to_ptr(w), to_ptr(bias), to_ptr(state), + to_ptr(q16), to_ptr(k16), to_ptr(v48), S, apply_silu, + to_stream(stream)); + }, + py::arg("x"), py::arg("w"), py::arg("bias"), py::arg("state"), + py::arg("q16"), py::arg("k16"), py::arg("v48"), py::arg("S"), + py::arg("apply_silu") = true, py::arg("stream") = 0, + R"pbdoc( +Chunk-parallel causal conv1d update (K=4, the 2048/2048/6144 q/k/v +channel family) with per-thread step batching and GQA split outputs. +Taps roll through registers - (STEPS+3)/STEPS read amplification +instead of 4x - with the packaged kernel's tap order and fma chain, +bit-exact. x: (S, 10240); w: (10240, 4); state: (10240, 3) last raw +inputs; q16/k16: (S, 2048), v48: (S, 6144), silu applied. +)pbdoc"); + + m.def("rms_norm_gated_silu_quant_fp4_bf16", + [](uintptr_t x, uintptr_t gate, uintptr_t weight, uintptr_t out, + uintptr_t packed, uintptr_t sfa, int M, int dim, float eps, + uintptr_t stream) -> int { + return flash_rt::kernels::rms_norm_gated_silu_quant_fp4_bf16( + to_ptr(x), to_ptr(gate), to_ptr(weight), to_ptr(out), + to_ptr(packed), to_ptr(sfa), M, dim, eps, + to_stream(stream)); + }, + py::arg("x"), py::arg("gate"), py::arg("weight"), + py::arg("out"), py::arg("packed"), py::arg("sfa"), py::arg("M"), + py::arg("dim"), py::arg("eps") = 1e-6f, py::arg("stream") = 0, + R"pbdoc( +Fused RMSNorm + weight + silu(gate) that also emits the row's NVFP4 +packed bytes and SFA, read by the output projection as one (1, M*128) +activation. Norm arithmetic transcribed from the packaged gated-norm +kernel; quantize stage is the production path verbatim. dim must be 128. +)pbdoc"); + + m.def("gdn_recurrent_inout_stream_bf16", + [](uintptr_t q, uintptr_t k, uintptr_t v, uintptr_t g, + uintptr_t beta, uintptr_t state_in, uintptr_t state_out, + uintptr_t out, int B, int num_v_heads, int head_dim, + bool use_qk_l2norm, uintptr_t stream) -> int { + return flash_rt::kernels::gdn_recurrent_inout_stream_bf16( + to_ptr(q), to_ptr(k), to_ptr(v), to_ptr(g), + to_ptr(beta), to_ptr(state_in), to_ptr(state_out), + to_ptr(out), B, num_v_heads, head_dim, use_qk_l2norm, + to_stream(stream)); + }, + py::arg("q"), py::arg("k"), py::arg("v"), py::arg("g"), + py::arg("beta"), py::arg("state_in"), py::arg("state_out"), + py::arg("out"), py::arg("B"), py::arg("num_v_heads"), + py::arg("head_dim"), py::arg("use_qk_l2norm") = true, + py::arg("stream") = 0, + R"pbdoc( +Gated-delta recurrent decode step over a streaming-column form: one warp +per 32 value columns instead of one block per head, so the same total +thread count spreads over 4x the blocks (the packaged kernel leaves +three quarters of a 170-SM part idle). Per-column arithmetic is the +packaged kernel's, unchanged; the q/k L2 norm reduces over a warp +rather than a 128-thread block, so its fp32 rounding may differ. +head_dim must be 128. +)pbdoc"); + + m.def("gdn_chunk_from_conv_smem_h_stash_bf16", + [](uintptr_t conv_out, uintptr_t a, uintptr_t b_, uintptr_t neg_exp_a, + uintptr_t dt_bias, uintptr_t state, uintptr_t out, uintptr_t stash, + int S, int num_v_heads, int num_k_heads, int head_dim, + int a_stride, int b_stride, bool use_qk_l2norm, + uintptr_t stream) { + flash_rt::gdn::gdn_chunk_from_conv_smem_h_stash_bf16( + to_ptr(conv_out), to_ptr(a), to_ptr(b_), + reinterpret_cast(neg_exp_a), + reinterpret_cast(dt_bias), + to_ptr(state), to_ptr(out), to_ptr(stash), S, + num_v_heads, num_k_heads, head_dim, a_stride, b_stride, + use_qk_l2norm, to_stream(stream)); + }, + py::arg("conv_out"), py::arg("a"), py::arg("b"), + py::arg("neg_exp_a"), py::arg("dt_bias"), py::arg("state"), + py::arg("out"), py::arg("stash"), py::arg("S"), + py::arg("num_v_heads"), py::arg("num_k_heads"), + py::arg("head_dim"), py::arg("a_stride"), py::arg("b_stride"), + py::arg("use_qk_l2norm") = true, py::arg("stream") = 0, + R"pbdoc( +Gated-delta from-conv chunk core with a per-row state stash (spec verify). +Identical recurrence and per-row bf16 state requantisation to the plain +chunk kernel; stash row s additionally records the carried state after row +s, bit-equal to the final state of a re-advance over rows 0..s. A rejected +speculative round rolls back by selecting a stash row. head_dim must be 128. +)pbdoc"); + // ───────────────────────────────────────────────────────────────────── // SM100 NVFP4 W4A16 GEMM bindings (Thor SM110). // diff --git a/csrc/gemm/fp4/cutlass_nvfp4_gemm_m256_sm120.cu b/csrc/gemm/fp4/cutlass_nvfp4_gemm_m256_sm120.cu new file mode 100644 index 000000000..b3f76830a --- /dev/null +++ b/csrc/gemm/fp4/cutlass_nvfp4_gemm_m256_sm120.cu @@ -0,0 +1,111 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Large-M NVFP4 W4A4 GEMM tier for sm_120: MmaTileShape 256x128x128, +// cooperative schedule, cluster 1x1. On the M~2048 prefill shapes this +// tile wins every family over the 128x128x128 baseline (measured +// 1428/1588/1430/1360 TFLOPS on 17408x5120 / 5120x17408 / 12288x5120 / +// 16384x5120 vs 1303/1361/1311/1289; FP4 MMA roof ~2020). K traversal +// order is unchanged, so per-element accumulation matches the baseline. +// Dispatch intent: route M >= 512 here, keep the small-M tier as is. +// +// Requires CUTLASS >= 4.5 (Sm120 blockscaled collective; the Sm100-tagged +// builder path refuses initialize on sm_120 from 4.5.x). Build with +// -gencode arch=compute_120a,code=sm_120a --expt-relaxed-constexpr. +// +// Header: cutlass_nvfp4_gemm_m256_sm120.cuh. +#include "cutlass_nvfp4_gemm_m256_sm120.cuh" + +#include + +#include "cutlass/cutlass.h" +#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" + +namespace flash_rt { +namespace gemm { +namespace { + +using namespace cute; + +using ElementA = cutlass::nv_float4_t; +using ElementB = cutlass::nv_float4_t; +using ElementD = cutlass::bfloat16_t; +using ElementC = cutlass::bfloat16_t; +using EAcc = float; +using Arch = cutlass::arch::Sm120; +using Op = cutlass::arch::OpClassBlockScaledTensorOp; +using MmaTile = Shape<_256, _128, _128>; +using Cluster = Shape<_1, _1, _1>; + +using CE = typename cutlass::epilogue::collective::CollectiveBuilder< + Arch, Op, MmaTile, Cluster, + cutlass::epilogue::collective::EpilogueTileAuto, EAcc, EAcc, + ElementC, cutlass::layout::RowMajor, 8, + ElementD, cutlass::layout::RowMajor, 8, + cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp; +using CM = typename cutlass::gemm::collective::CollectiveBuilder< + Arch, Op, + ElementA, cutlass::layout::RowMajor, 32, + ElementB, cutlass::layout::ColumnMajor, 32, + EAcc, MmaTile, Cluster, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CE::SharedStorage))>, + cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp; +using GK = cutlass::gemm::kernel::GemmUniversal, + CM, CE, void>; +using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + +} // namespace + +size_t nvfp4_gemm_m256_sm120_workspace_size(int M, int N, int K) { + using SA = typename Gemm::GemmKernel::StrideA; + using SB = typename Gemm::GemmKernel::StrideB; + using SD = typename Gemm::GemmKernel::StrideD; + using Cfg = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + auto sa = cutlass::make_cute_packed_stride(SA{}, {M, K, 1}); + auto sb = cutlass::make_cute_packed_stride(SB{}, {N, K, 1}); + auto sd = cutlass::make_cute_packed_stride(SD{}, {M, N, 1}); + auto lsfa = Cfg::tile_atom_to_shape_SFA(make_shape(M, N, K, 1)); + auto lsfb = Cfg::tile_atom_to_shape_SFB(make_shape(M, N, K, 1)); + typename Gemm::Arguments args{ + cutlass::gemm::GemmUniversalMode::kGemm, {M, N, K, 1}, + {nullptr, sa, nullptr, sb, nullptr, lsfa, nullptr, lsfb}, + {{1.f, 0.f}, nullptr, sd, nullptr, sd}}; + return Gemm::get_workspace_size(args); +} + +int nvfp4_gemm_m256_sm120_bf16(const void* A_packed, const void* SFA, + const void* B_packed, const void* SFB, + void* D_bf16, int M, int N, int K, + float alpha, void* workspace, + cudaStream_t stream) { + using SA = typename Gemm::GemmKernel::StrideA; + using SB = typename Gemm::GemmKernel::StrideB; + using SD = typename Gemm::GemmKernel::StrideD; + using Cfg = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + auto sa = cutlass::make_cute_packed_stride(SA{}, {M, K, 1}); + auto sb = cutlass::make_cute_packed_stride(SB{}, {N, K, 1}); + auto sd = cutlass::make_cute_packed_stride(SD{}, {M, N, 1}); + auto lsfa = Cfg::tile_atom_to_shape_SFA(make_shape(M, N, K, 1)); + auto lsfb = Cfg::tile_atom_to_shape_SFB(make_shape(M, N, K, 1)); + typename Gemm::Arguments args{ + cutlass::gemm::GemmUniversalMode::kGemm, {M, N, K, 1}, + {reinterpret_cast(A_packed), sa, + reinterpret_cast(B_packed), sb, + reinterpret_cast(SFA), lsfa, + reinterpret_cast(SFB), lsfb}, + {{alpha, 0.f}, nullptr, sd, + reinterpret_cast(D_bf16), sd}}; + Gemm gemm; + if (gemm.can_implement(args) != cutlass::Status::kSuccess) return 1; + if (gemm.initialize(args, workspace, stream) != cutlass::Status::kSuccess) + return 2; + if (gemm.run(stream) != cutlass::Status::kSuccess) return 3; + return 0; +} + +} // namespace gemm +} // namespace flash_rt diff --git a/csrc/gemm/fp4/cutlass_nvfp4_gemm_m256_sm120.cuh b/csrc/gemm/fp4/cutlass_nvfp4_gemm_m256_sm120.cuh new file mode 100644 index 000000000..1ade2ae64 --- /dev/null +++ b/csrc/gemm/fp4/cutlass_nvfp4_gemm_m256_sm120.cuh @@ -0,0 +1,22 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Large-M (M >= ~512) NVFP4 W4A4 GEMM tier for sm_120: 256x128x128 +// cooperative tile, bf16 output. See the .cu for measurements and the +// CUTLASS >= 4.5 requirement. Additive. +#pragma once +#include +#include +namespace flash_rt { +namespace gemm { +size_t nvfp4_gemm_m256_sm120_workspace_size(int M, int N, int K); +// A (M, K/2) row-major packed, B (N, K/2) column-major-K packed, SFA/SFB +// in the CUTLASS Sm1xx block-scaled layouts, D (M, N) bf16 row-major. +// workspace from nvfp4_gemm_m256_sm120_workspace_size. Returns 0 on +// success, 1 refused (can_implement), 2 initialize, 3 run. +int nvfp4_gemm_m256_sm120_bf16(const void* A_packed, const void* SFA, + const void* B_packed, const void* SFB, + void* D_bf16, int M, int N, int K, + float alpha, void* workspace, + cudaStream_t stream); +} // namespace gemm +} // namespace flash_rt diff --git a/csrc/kernels/batched_unit_ltri_inv64.cu b/csrc/kernels/batched_unit_ltri_inv64.cu new file mode 100644 index 000000000..e876b2bc1 --- /dev/null +++ b/csrc/kernels/batched_unit_ltri_inv64.cu @@ -0,0 +1,58 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Batched 64x64 unit-lower-triangular inverse. See header for the +// contract; same forward-substitution recurrence the batched cuBLAS +// solve runs, so the fp32 error class matches the path it replaces. +#include "kernels/batched_unit_ltri_inv64.cuh" + +namespace flash_rt { +namespace kernels { +namespace { + +template // matrices per block; blockDim.x = 64*MPB +__global__ void unit_ltri_inv64_kernel(const float* __restrict__ A, + float* __restrict__ X, int B) { + __shared__ float Ls[MPB][64][65]; + const int s = threadIdx.x >> 6; + const int c = threadIdx.x & 63; + const int mi = blockIdx.x * MPB + s; + if (mi >= B) return; + const float* Ab = A + (size_t)mi * 4096; + for (int idx = c; idx < 4096; idx += 64) { + const int i = idx >> 6, j = idx & 63; + Ls[s][i][j] = (j < i) ? Ab[idx] : 0.f; + } + __syncwarp(); + + float x[64]; + #pragma unroll + for (int i = 0; i < 64; ++i) { + float acc = (i == c) ? 1.f : 0.f; + #pragma unroll + for (int j = 0; j < i; ++j) + acc -= Ls[s][i][j] * x[j]; + x[i] = acc; + } + + float* Xb = X + (size_t)mi * 4096; + #pragma unroll + for (int i = 0; i < 64; ++i) + Xb[i * 64 + c] = x[i]; +} + +} // namespace + +int batched_unit_ltri_inv64_f32(const void* A, void* X, int B, + cudaStream_t stream) { + if (!A || !X || B <= 0) return 1; + constexpr int MPB = 2; + const int blocks = (B + MPB - 1) / MPB; + unit_ltri_inv64_kernel<<>>( + reinterpret_cast(A), + reinterpret_cast(X), B); + const cudaError_t e = cudaGetLastError(); + return (e == cudaSuccess) ? 0 : -static_cast(e); +} + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/batched_unit_ltri_inv64.cuh b/csrc/kernels/batched_unit_ltri_inv64.cuh new file mode 100644 index 000000000..d18ff4fd9 --- /dev/null +++ b/csrc/kernels/batched_unit_ltri_inv64.cuh @@ -0,0 +1,22 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Batched 64x64 unit-lower-triangular inverse: X = inv(I + strict_tril(A)) +// in fp32. Column-independent forward substitution — one lane per column, +// per-column history in registers (fully unrolled), the L tile broadcast +// from shared memory, no cross-lane synchronization. The identity/tril +// preparation folds into the kernel (the strict lower triangle is read +// straight from A, the unit diagonal is implicit), removing the eye/ +// expand/tril materializations the host-side solve needed. Additive. +#pragma once +#include + +namespace flash_rt { +namespace kernels { + +// A: (B, 64, 64) fp32 row-major (only the strict lower triangle is +// read). X: (B, 64, 64) fp32. Returns 0 on success. +int batched_unit_ltri_inv64_f32(const void* A, void* X, int B, + cudaStream_t stream); + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/causal_conv1d_update_steps_gqa_bf16.cu b/csrc/kernels/causal_conv1d_update_steps_gqa_bf16.cu new file mode 100644 index 000000000..7e00885ae --- /dev/null +++ b/csrc/kernels/causal_conv1d_update_steps_gqa_bf16.cu @@ -0,0 +1,105 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Chunk-parallel causal conv1d update with per-thread step batching +// and GQA split outputs. See header for the contract. +#include "kernels/causal_conv1d_update_steps_gqa_bf16.cuh" + +#include + +namespace flash_rt { +namespace kernels { +namespace { + +constexpr int kConvDim = 10240; +constexpr int kK = 4; +constexpr int kSteps = 8; +constexpr int kThreads = 128; + +__device__ __forceinline__ float conv_silu(float x) { + // matches the packaged kernel: x * sigmoid(x) via the fast exp + return x / (1.0f + __expf(-x)); +} + +__global__ void conv1d_steps_gqa_kernel( + const __nv_bfloat16* __restrict__ x, + const __nv_bfloat16* __restrict__ w, + const __nv_bfloat16* __restrict__ bias, + const __nv_bfloat16* __restrict__ state, + __nv_bfloat16* __restrict__ q16, + __nv_bfloat16* __restrict__ k16, + __nv_bfloat16* __restrict__ v48, + int S, bool apply_silu) { + const int c = blockIdx.x * kThreads + threadIdx.x; + if (c >= kConvDim) return; + const int s0 = blockIdx.y * kSteps; + + // tap weights and bias for this channel + float wv[kK]; + #pragma unroll + for (int i = 0; i < kK; ++i) + wv[i] = static_cast(w[c * kK + i]); + const float bv = (bias != nullptr) ? static_cast(bias[c]) : 0.f; + + // rolling window xv[0..2] = taps t = s-3..s-1, xv[3] = t = s + float xv[kK]; + #pragma unroll + for (int i = 0; i < kK - 1; ++i) { + const int t = s0 - (kK - 1) + i; + float v = 0.f; + if (t >= 0) { + v = static_cast(x[(size_t)t * kConvDim + c]); + } else if (t >= -(kK - 1)) { + v = static_cast(state[c * (kK - 1) + (t + kK - 1)]); + } + xv[i] = v; + } + + #pragma unroll + for (int j = 0; j < kSteps; ++j) { + const int s = s0 + j; + if (s >= S) return; + xv[kK - 1] = static_cast(x[(size_t)s * kConvDim + c]); + // same accumulation order as the packaged kernel: bias first, + // then taps in ascending-t order + float acc = bv; + #pragma unroll + for (int i = 0; i < kK; ++i) + acc = fmaf(xv[i], wv[i], acc); + if (apply_silu) acc = conv_silu(acc); + const __nv_bfloat16 y = __float2bfloat16(acc); + if (c < 2048) { + q16[(size_t)s * 2048 + c] = y; + } else if (c < 4096) { + k16[(size_t)s * 2048 + (c - 2048)] = y; + } else { + v48[(size_t)s * 6144 + (c - 4096)] = y; + } + #pragma unroll + for (int i = 0; i < kK - 1; ++i) + xv[i] = xv[i + 1]; + } +} + +} // namespace + +int causal_conv1d_update_steps_gqa_bf16( + const void* x, const void* w, const void* bias, const void* state, + void* q16, void* k16, void* v48, int S, bool apply_silu, + cudaStream_t stream) { + if (!x || !w || !state || !q16 || !k16 || !v48) return 1; + if (S <= 0) return 2; + dim3 grid(kConvDim / kThreads, (S + kSteps - 1) / kSteps); + conv1d_steps_gqa_kernel<<>>( + reinterpret_cast(x), + reinterpret_cast(w), + reinterpret_cast(bias), + reinterpret_cast(state), + reinterpret_cast<__nv_bfloat16*>(q16), + reinterpret_cast<__nv_bfloat16*>(k16), + reinterpret_cast<__nv_bfloat16*>(v48), S, apply_silu); + const cudaError_t e = cudaGetLastError(); + return (e == cudaSuccess) ? 0 : -static_cast(e); +} + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/causal_conv1d_update_steps_gqa_bf16.cuh b/csrc/kernels/causal_conv1d_update_steps_gqa_bf16.cuh new file mode 100644 index 000000000..ed12a596a --- /dev/null +++ b/csrc/kernels/causal_conv1d_update_steps_gqa_bf16.cuh @@ -0,0 +1,26 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Chunk-parallel causal conv1d update (K=4) with per-thread step +// batching and GQA split outputs. The packaged chunk-parallel kernel +// computes one (token, channel) per thread, so every input element is +// re-read K times from DRAM; batching STEPS consecutive tokens per +// thread rolls the taps through registers ((STEPS+K-1)/STEPS read +// amplification instead of K). Tap order and fp32 fma chain match the +// packaged kernel exactly — bit-exact outputs. Additive. +#pragma once +#include + +namespace flash_rt { +namespace kernels { + +// x: (S, 10240) bf16 (the fixed 2048/2048/6144 q/k/v channel family, +// conv K=4). w: (10240, 4) bf16. state: (1, 10240, 3) bf16 (the last +// K-1 raw inputs). bias: (10240) bf16 or null. q16/k16: (S, 2048), +// v48: (S, 6144) bf16, silu applied. Returns 0 on success. +int causal_conv1d_update_steps_gqa_bf16( + const void* x, const void* w, const void* bias, const void* state, + void* q16, void* k16, void* v48, int S, bool apply_silu, + cudaStream_t stream); + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/checks/bench_gemm_m256.cu b/csrc/kernels/checks/bench_gemm_m256.cu new file mode 100644 index 000000000..bee77ca7d --- /dev/null +++ b/csrc/kernels/checks/bench_gemm_m256.cu @@ -0,0 +1,37 @@ +// Large-M NVFP4 GEMM tier: TFLOPS receipt on the prefill shape family. +// Build: nvcc -gencode arch=compute_120a,code=sm_120a -O3 -std=c++17 \ +// --expt-relaxed-constexpr -I /include \ +// -I /tools/util/include -I ../../csrc/gemm/fp4 \ +// bench_gemm_m256.cu -o bench_gemm_m256 (CUTLASS >= 4.5) +#include "../../gemm/fp4/cutlass_nvfp4_gemm_m256_sm120.cu" +#include +#include +#include +#define CK(x) do { auto e=(x); if(e!=cudaSuccess){printf("err %s @%d\n",cudaGetErrorString(e),__LINE__);exit(1);} } while(0) +int main() { + const int M = 2044; + int shapes[][2] = {{17408,5120},{5120,17408},{12288,5120},{16384,5120}}; + uint8_t *A,*B,*SFA,*SFB,*D; + CK(cudaMalloc(&A,(size_t)M*17408/2)); CK(cudaMalloc(&B,(size_t)17408ull*17408/2)); + CK(cudaMalloc(&SFA,(size_t)M*17408/8)); CK(cudaMalloc(&SFB,(size_t)17408ull*17408/8)); + CK(cudaMalloc(&D,(size_t)M*17408*2)); + std::vector h(1<<24); for (auto& x : h) x = rand() & 0xff; + cudaMemcpy(A,h.data(),std::min((size_t)M*17408/2,h.size()),cudaMemcpyHostToDevice); + cudaMemset(SFA,0x3f,(size_t)M*17408/8); cudaMemset(SFB,0x3f,(size_t)17408ull*17408/8); + for (auto& sh : shapes) { + int N = sh[0], K = sh[1]; + size_t ws = flash_rt::gemm::nvfp4_gemm_m256_sm120_workspace_size(M,N,K); + void* wk; CK(cudaMalloc(&wk, ws ? ws : 4)); + int rc = flash_rt::gemm::nvfp4_gemm_m256_sm120_bf16(A,SFA,B,SFB,D,M,N,K,1.f,wk,0); + CK(cudaDeviceSynchronize()); + if (rc) { printf("N=%d K=%d rc=%d\n",N,K,rc); cudaFree(wk); continue; } + cudaEvent_t e0,e1; cudaEventCreate(&e0); cudaEventCreate(&e1); + cudaEventRecord(e0); + for (int i=0;i<20;++i) flash_rt::gemm::nvfp4_gemm_m256_sm120_bf16(A,SFA,B,SFB,D,M,N,K,1.f,wk,0); + cudaEventRecord(e1); CK(cudaEventSynchronize(e1)); + float ms; cudaEventElapsedTime(&ms,e0,e1); ms/=20; + printf("M=%d N=%5d K=%5d: %7.0f TFLOPS (roof ~2020)\n", M,N,K, 2.0*M*N*K/ms/1e9); + cudaFree(wk); + } + return 0; +} diff --git a/csrc/kernels/checks/check_warpsplit_ilv.cu b/csrc/kernels/checks/check_warpsplit_ilv.cu new file mode 100644 index 000000000..11af9d07d --- /dev/null +++ b/csrc/kernels/checks/check_warpsplit_ilv.cu @@ -0,0 +1,68 @@ +#include +#include +#include +// Interleaved-B GEMV bit-exactness: random packed data, host-side repack +// reference, base vs ilv outputs compared bit-for-bit across shapes and +// warp configs; the device repack entry is checked against the host +// reference bytes as well. +// Build: nvcc -gencode arch=compute_120a,code=sm_120a -O3 -std=c++17 \ +// -I /include check_warpsplit_ilv.cu -o check_warpsplit_ilv +#include "../fp4_w4a4_mma_warpsplit_sm120.cuh" +#include "../fp4_w4a4_mma_warpsplit_ilv_sm120.cuh" +#include +#include +#include +#include +#define CK(x) do { auto e=(x); if(e!=cudaSuccess){printf("err %s @%d\n",cudaGetErrorString(e),__LINE__);exit(1);} } while(0) +int main() { + int fails = 0; + int shapes[][2] = {{17408,5120},{5120,17408},{12288,5120},{1024,5120},{5120,6144},{16384,5120}}; + for (auto& sh : shapes) { + int N = sh[0], K = sh[1], KH = K/2, KI = K/64; + size_t bBytes=(size_t)N*KH, sfbBytes=((size_t)((N+127)/128))*((K/16+3)/4)*512+4096; + std::vector hB(bBytes), hSFB(sfbBytes), hA(KH), hSFA((size_t)KI*512+4096), hBi(bBytes); + srand(N ^ K); + for (auto& x : hB) x = rand() & 0xff; + for (auto& x : hSFB) x = 0x30 + (rand() & 0xf); + for (auto& x : hA) x = rand() & 0xff; + for (auto& x : hSFA) x = 0x30 + (rand() & 0xf); + // host repack: Bi[g*KH*8 + kt*256 + col*32 + off*4 .. +4] = B[(g*8+col)*KH + kt*32 + off*4] + for (int g = 0; g < N/8; ++g) + for (int kt = 0; kt < KI; ++kt) + for (int col = 0; col < 8; ++col) + memcpy(&hBi[(size_t)g*KH*8 + (size_t)kt*256 + col*32], + &hB[(size_t)(g*8+col)*KH + (size_t)kt*32], 32); + uint8_t *A,*SFA,*B,*Bi,*Bi_dev,*SFB; __nv_bfloat16 *D1,*D2; + CK(cudaMalloc(&A,KH)); CK(cudaMalloc(&SFA,hSFA.size())); CK(cudaMalloc(&B,bBytes)); + CK(cudaMalloc(&Bi,bBytes)); CK(cudaMalloc(&Bi_dev,bBytes)); CK(cudaMalloc(&SFB,sfbBytes)); + CK(cudaMalloc(&D1,N*2)); CK(cudaMalloc(&D2,N*2)); + cudaMemcpy(A,hA.data(),KH,cudaMemcpyHostToDevice); + cudaMemcpy(SFA,hSFA.data(),hSFA.size(),cudaMemcpyHostToDevice); + cudaMemcpy(B,hB.data(),bBytes,cudaMemcpyHostToDevice); + cudaMemcpy(Bi,hBi.data(),bBytes,cudaMemcpyHostToDevice); + cudaMemcpy(SFB,hSFB.data(),sfbBytes,cudaMemcpyHostToDevice); + flash_rt::gemm::fp4_w4a4_repack_b_ilv_sm120(B, Bi_dev, N, K, 0); + CK(cudaDeviceSynchronize()); + { std::vector dv(bBytes); + cudaMemcpy(dv.data(), Bi_dev, bBytes, cudaMemcpyDeviceToHost); + size_t rd = 0; for (size_t i2 = 0; i2 < bBytes; ++i2) rd += (dv[i2] != hBi[i2]); + printf("N=%d K=%d device-repack vs host: byte-diff=%zu %s\n", N, K, rd, rd? "FAIL":"OK"); + fails += (rd != 0); } + for (int w : {2, 4, 8}) { + if (KI % w) continue; + cudaMemset(D1,0,N*2); cudaMemset(D2,0,N*2); + int r1 = flash_rt::gemm::fp4_w4a4_mma_sm120_warpsplit_bf16out(A,B,D1,N,K,SFA,SFB,1.f,w,3,0); + int r2 = flash_rt::gemm::fp4_w4a4_mma_sm120_warpsplit_ilv_bf16out(A,Bi,D2,N,K,SFA,SFB,1.f,w,3,0); + CK(cudaDeviceSynchronize()); + std::vector o1(N), o2(N); + cudaMemcpy(o1.data(),D1,N*2,cudaMemcpyDeviceToHost); + cudaMemcpy(o2.data(),D2,N*2,cudaMemcpyDeviceToHost); + int diff = 0; for (int i = 0; i < N; ++i) diff += (o1[i] != o2[i]); + printf("N=%d K=%d w%d: rc=%d/%d bit-diff=%d %s\n", N,K,w,r1,r2,diff, diff? "FAIL":"OK"); + fails += (diff != 0) + r1 + r2; + } + cudaFree(A);cudaFree(SFA);cudaFree(B);cudaFree(Bi);cudaFree(SFB);cudaFree(D1);cudaFree(D2); + } + printf(fails ? "CHECK FAILED\n" : "ALL BIT-EXACT\n"); + return fails != 0; +} diff --git a/csrc/kernels/checks/check_warpsplit_ilv_mrows.cu b/csrc/kernels/checks/check_warpsplit_ilv_mrows.cu new file mode 100644 index 000000000..d156ba448 --- /dev/null +++ b/csrc/kernels/checks/check_warpsplit_ilv_mrows.cu @@ -0,0 +1,83 @@ +#include +#include +#include +// Multi-row interleaved GEMM bit-exactness: random packed data, the +// multi-row kernel's M rows compared bit-for-bit against M independent +// M=1 interleaved GEMV calls (row r's scales staged at the production +// 512B-block row offset for the multi-row read, at block offset 0 for +// the per-row reference). Rows are independent under the shared MMA atom +// and the reduction order is identical, so any difference is a defect. +// Build: nvcc -gencode arch=compute_120a,code=sm_120a -O3 -std=c++17 \ +// --expt-relaxed-constexpr -I /include \ +// check_warpsplit_ilv_mrows.cu ../fp4_w4a4_mma_warpsplit_ilv_sm120.cu \ +// ../fp4_w4a4_mma_warpsplit_ilv_mrows_sm120.cu -o check_warpsplit_ilv_mrows +#include "../fp4_w4a4_mma_warpsplit_ilv_sm120.cuh" +#include "../fp4_w4a4_mma_warpsplit_ilv_mrows_sm120.cuh" +#include +#include +#include +#define CK(x) do { auto e=(x); if(e!=cudaSuccess){printf("err %s @%d\n",cudaGetErrorString(e),__LINE__);exit(1);} } while(0) +int main() { + int fails = 0; + int shapes[][2] = {{17408,5120},{5120,17408},{12288,5120},{1024,5120},{5120,6144},{16384,5120}}; + for (auto& sh : shapes) { + int N = sh[0], K = sh[1], KH = K/2, KI = K/64; + size_t bBytes=(size_t)N*KH, sfbBytes=((size_t)((N+127)/128))*((K/16+3)/4)*512+4096; + std::vector hB(bBytes), hSFB(sfbBytes), hBi(bBytes); + srand(N ^ K); + for (auto& x : hB) x = rand() & 0xff; + for (auto& x : hSFB) x = 0x30 + (rand() & 0xf); + for (int g = 0; g < N/8; ++g) + for (int kt = 0; kt < KI; ++kt) + for (int col = 0; col < 8; ++col) + memcpy(&hBi[(size_t)g*KH*8 + (size_t)kt*256 + col*32], + &hB[(size_t)(g*8+col)*KH + (size_t)kt*32], 32); + uint8_t *Bi,*SFB; CK(cudaMalloc(&Bi,bBytes)); CK(cudaMalloc(&SFB,sfbBytes)); + cudaMemcpy(Bi,hBi.data(),bBytes,cudaMemcpyHostToDevice); + cudaMemcpy(SFB,hSFB.data(),sfbBytes,cudaMemcpyHostToDevice); + for (int M : {1, 2, 4, 7, 8, 16}) { + std::vector hA((size_t)M*KH), hS((size_t)M*KI*4); + for (auto& x : hA) x = rand() & 0xff; + for (auto& x : hS) x = 0x30 + (rand() & 0xf); + std::vector hSFAm((size_t)KI*512+4096, 0); + for (int r = 0; r < M; ++r) + for (int kt = 0; kt < KI; ++kt) + memcpy(&hSFAm[(size_t)kt*512 + r*16], &hS[((size_t)r*KI+kt)*4], 4); + uint8_t *A,*SFAm,*SFA1; __nv_bfloat16 *Dm,*D1; + CK(cudaMalloc(&A,hA.size())); CK(cudaMalloc(&SFAm,hSFAm.size())); + CK(cudaMalloc(&SFA1,(size_t)KI*512+4096)); + CK(cudaMalloc(&Dm,(size_t)M*N*2)); CK(cudaMalloc(&D1,N*2)); + cudaMemcpy(A,hA.data(),hA.size(),cudaMemcpyHostToDevice); + cudaMemcpy(SFAm,hSFAm.data(),hSFAm.size(),cudaMemcpyHostToDevice); + for (int w : {2, 4}) { + if (KI % w) continue; + cudaMemset(Dm,0,(size_t)M*N*2); + int rcm = flash_rt::gemm::fp4_w4a4_mma_sm120_warpsplit_ilv_mrows_bf16out( + A,Bi,Dm,M,N,K,SFAm,SFB,1.f,w,3,0); + CK(cudaDeviceSynchronize()); + std::vector om((size_t)M*N), o1(N); + cudaMemcpy(om.data(),Dm,(size_t)M*N*2,cudaMemcpyDeviceToHost); + int diff = 0, rc1sum = 0; + for (int r = 0; r < M; ++r) { + std::vector hSFA1((size_t)KI*512+4096, 0); + for (int kt = 0; kt < KI; ++kt) + memcpy(&hSFA1[(size_t)kt*512], &hS[((size_t)r*KI+kt)*4], 4); + cudaMemcpy(SFA1,hSFA1.data(),hSFA1.size(),cudaMemcpyHostToDevice); + cudaMemset(D1,0,N*2); + rc1sum += flash_rt::gemm::fp4_w4a4_mma_sm120_warpsplit_ilv_bf16out( + A + (size_t)r*KH,Bi,D1,N,K,SFA1,SFB,1.f,w,3,0); + CK(cudaDeviceSynchronize()); + cudaMemcpy(o1.data(),D1,N*2,cudaMemcpyDeviceToHost); + for (int i = 0; i < N; ++i) diff += (om[(size_t)r*N+i] != o1[i]); + } + printf("N=%d K=%d M=%d w%d: rc=%d/%d bit-diff=%d %s\n", + N,K,M,w,rcm,rc1sum,diff, diff? "FAIL":"OK"); + fails += (diff != 0) + rcm + rc1sum; + } + cudaFree(A);cudaFree(SFAm);cudaFree(SFA1);cudaFree(Dm);cudaFree(D1); + } + cudaFree(Bi);cudaFree(SFB); + } + printf(fails ? "CHECK FAILED\n" : "ALL BIT-EXACT\n"); + return fails != 0; +} diff --git a/csrc/kernels/checks/check_wy_kkt_mma.cu b/csrc/kernels/checks/check_wy_kkt_mma.cu new file mode 100644 index 000000000..a2f7e7ef3 --- /dev/null +++ b/csrc/kernels/checks/check_wy_kkt_mma.cu @@ -0,0 +1,82 @@ +// WY K*KT: scalar reference (transcribed from the production kernel) +// vs the MMA entry - numeric band + speed. +// Build: nvcc -gencode arch=compute_120a,code=sm_120a -O3 -std=c++17 \ +// check_wy_kkt_mma.cu -o check_wy_kkt_mma +#include "../gated_delta_wy_kkt_mma.cu" +#include +#include +#include +#include +#include +#define CK(x) do { auto e=(x); if(e!=cudaSuccess){printf("err %s @%d\n",cudaGetErrorString(e),__LINE__);exit(1);} } while(0) +constexpr int kWyChunk = 64, kHD = 128; + +__global__ void kkt_v1_kernel( + const __nv_bfloat16* __restrict__ k16_l2, const __nv_bfloat16* __restrict__ beta, + const __nv_bfloat16* __restrict__ g_cumsum, float* __restrict__ A, + int S, int num_k_heads, int num_v_heads, int head_group_size) { + const int pair = blockIdx.x * blockDim.x + threadIdx.x; + if (pair >= kWyChunk * kWyChunk) return; + const int i = pair / kWyChunk, j = pair - i * kWyChunk; + const int vh = blockIdx.y, chunk = blockIdx.z; + const int si = chunk * kWyChunk + i, sj = chunk * kWyChunk + j; + const size_t a_off = (((static_cast(chunk) * num_v_heads + vh) * kWyChunk + i) * kWyChunk + j); + if (i <= j || si >= S || sj >= S) { A[a_off] = 0.0f; return; } + const int kh = vh / head_group_size; + const size_t ki = (static_cast(si) * num_k_heads + kh) * kHD; + const size_t kj = (static_cast(sj) * num_k_heads + kh) * kHD; + float dot = 0.0f; + #pragma unroll 16 + for (int d = 0; d < kHD; ++d) + dot = fmaf((float)k16_l2[ki + d], (float)k16_l2[kj + d], dot); + const float bi = (float)beta[(size_t)si * num_v_heads + vh]; + const float gi = (float)g_cumsum[(size_t)si * num_v_heads + vh]; + const float gj = (float)g_cumsum[(size_t)sj * num_v_heads + vh]; + A[a_off] = bi * dot * __expf(gi - gj); +} + +int main() { + int S = 2048, KH = 16, VH = 48, GRP = 3; + int chunks = S / 64; + size_t nK = (size_t)S * KH * kHD, nBG = (size_t)S * VH, nA = (size_t)chunks * VH * 64 * 64; + std::vector hK(nK), hB(nBG), hG(nBG); + srand(11); + auto rb = [&](float scale){ float f = ((rand()%2000)-1000)/1000.0f*scale; __nv_bfloat16 b=__float2bfloat16(f); return *reinterpret_cast(&b); }; + for (auto& x : hK) x = rb(1.0f); + for (auto& x : hB) x = rb(0.9f); + for (auto& x : hG) x = rb(2.0f); + __nv_bfloat16 *K, *B, *G; float *A1, *A2; + CK(cudaMalloc(&K, nK*2)); CK(cudaMalloc(&B, nBG*2)); CK(cudaMalloc(&G, nBG*2)); + CK(cudaMalloc(&A1, nA*4)); CK(cudaMalloc(&A2, nA*4)); + cudaMemcpy(K, hK.data(), nK*2, cudaMemcpyHostToDevice); + cudaMemcpy(B, hB.data(), nBG*2, cudaMemcpyHostToDevice); + cudaMemcpy(G, hG.data(), nBG*2, cudaMemcpyHostToDevice); + dim3 g1((64*64+255)/256, VH, chunks); + kkt_v1_kernel<<>>(K, B, G, A1, S, KH, VH, GRP); + qwen36_gdn_wy_kkt_b64_mma_bf16(K, B, G, A2, S, 0); + CK(cudaDeviceSynchronize()); + std::vector o1(nA), o2(nA); + cudaMemcpy(o1.data(), A1, nA*4, cudaMemcpyDeviceToHost); + cudaMemcpy(o2.data(), A2, nA*4, cudaMemcpyDeviceToHost); + double maxrel = 0, maxabs = 0; size_t bad = 0; + for (size_t i = 0; i < nA; ++i) { + double d = fabs((double)o1[i] - o2[i]); + maxabs = fmax(maxabs, d); + double den = fmax(fabs((double)o1[i]), 1e-3); + maxrel = fmax(maxrel, d / den); + if (d / den > 1e-2 && d > 1e-3) ++bad; + } + printf("numeric: maxabs=%.3e maxrel=%.3e bad=%zu/%zu\n", maxabs, maxrel, bad, nA); + cudaEvent_t e0, e1; cudaEventCreate(&e0); cudaEventCreate(&e1); + for (auto which : {1, 2}) { + cudaEventRecord(e0); + for (int it = 0; it < 30; ++it) { + if (which == 1) kkt_v1_kernel<<>>(K, B, G, A1, S, KH, VH, GRP); + else qwen36_gdn_wy_kkt_b64_mma_bf16(K, B, G, A2, S, 0); + } + cudaEventRecord(e1); CK(cudaEventSynchronize(e1)); + float ms; cudaEventElapsedTime(&ms, e0, e1); ms /= 30; + printf("v%d: %8.1f us/layer (x48 layers = %.1f ms)\n", which, ms*1e3, ms*48); + } + return 0; +} diff --git a/csrc/kernels/fp4_w4a4_mma_warpsplit_ilv_mrows_sm120.cu b/csrc/kernels/fp4_w4a4_mma_warpsplit_ilv_mrows_sm120.cu new file mode 100644 index 000000000..9d483aa25 --- /dev/null +++ b/csrc/kernels/fp4_w4a4_mma_warpsplit_ilv_mrows_sm120.cu @@ -0,0 +1,159 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Multi-row (M<=16) interleaved-B warp-split-K NVFP4 W4A4 GEMM for sm_120. +// See the header for the contract. The mainloop is the M=1 interleaved +// GEMV's, unchanged — same MMA atom, same pipeline, same reduction order — +// with A/SFA rows 0..M-1 loaded live (rows M..15 zero ballast) and the +// epilogue reducing and writing all M rows. One weight stream serves every +// row: measured 74-80% of the DRAM read roof at M=8 on the wide decode +// shapes, 5.3-6.0x over M sequential GEMV calls, 1.16-1.37x over the +// plain-B multi-row variant at M=8 (the interleave advantage grows with M). +// Additive: new file + new entry point. + +#include "fp4_w4a4_mma_warpsplit_ilv_mrows_sm120.cuh" + +#include +#include +#include + +#include "cute/arch/mma_sm120.hpp" +#include "cutlass/numeric_types.h" + +namespace flash_rt { +namespace gemm { +namespace { + +using AtomType = cute::SM120::BLOCKSCALED::SM120_16x8x64_TN_VS< + cutlass::float_e2m1_t, cutlass::float_e2m1_t, float, + cutlass::float_ue4m3_t, 16>; + +__device__ __forceinline__ uint32_t fa(const uint8_t* s, int t0, int t1, int r) { + int ro = ((r & 1) ? (t1 + 8) : t1) * 32; + return *reinterpret_cast(s + ro + t0 * 4 + ((r >> 1) & 1) * 16); +} +__device__ __forceinline__ uint32_t fb(const uint8_t* s, int t0, int t1, int r) { + return *reinterpret_cast(s + t1 * 32 + t0 * 4 + r * 16); +} +__device__ __forceinline__ uint32_t fsa(const uint8_t* p, int u) { + return *reinterpret_cast(p + u * 4); +} +__device__ __forceinline__ void cpa(uint8_t* d, const uint8_t* s) { + uint32_t i = __cvta_generic_to_shared(d); + asm volatile("cp.async.ca.shared.global.L2::128B [%0], [%1], 4;\n" :: "r"(i), "l"(s)); +} +__device__ __forceinline__ void commit() { asm volatile("cp.async.commit_group;\n" ::); } +template __device__ __forceinline__ void waitg() { + asm volatile("cp.async.wait_group %0;\n" :: "n"(N)); +} + +template +__global__ void warpsplit_ilv_mrows_kernel( + const uint8_t* __restrict__ A, const uint8_t* __restrict__ B, + const uint8_t* __restrict__ SFA, const uint8_t* __restrict__ SFB, + __nv_bfloat16* __restrict__ D, float alpha, int M, int N, int K) { + __shared__ uint8_t sA[WARPS][STAGES][16 * 32]; + __shared__ uint8_t sSFA[WARPS][STAGES][16 * 4]; + __shared__ uint8_t sB[WARPS][STAGES][8 * 32]; + __shared__ uint8_t sSFB[WARPS][STAGES][8 * 4]; + __shared__ float s_red[WARPS][16][8]; // per-warp (row, col) partials + + int tid = threadIdx.x, warp = tid >> 5, lane = tid & 31; + int my_n = blockIdx.x * 8; + const int KI = K / 64, KIw = KI / WARPS; + const int kt0 = warp * KIw; + const int KH = K / 2, ncs = (K / 16 + 3) / 4; + int t0 = lane & 3, t1 = lane >> 2, sau = (lane & 1) * 8 + (lane >> 2), sbu = lane >> 2; + float c0 = 0, c1 = 0, c2 = 0, c3 = 0; + + uint8_t (*mA)[16 * 32] = sA[warp]; + uint8_t (*mSFA)[16 * 4] = sSFA[warp]; + uint8_t (*mB)[8 * 32] = sB[warp]; + uint8_t (*mSFB)[8 * 4] = sSFB[warp]; + + // rows >= M stay zero ballast, exactly the v9 discipline for rows >= 1 + if (lane >= M && lane < 16) { + #pragma unroll + for (int st = 0; st < STAGES; ++st) { + int4* av = reinterpret_cast(mA[st]); int4 z{0, 0, 0, 0}; + av[lane * 2] = z; av[lane * 2 + 1] = z; + } + } + if (lane < 4) + for (int st = 0; st < STAGES; ++st) + for (int i = M * 4 + lane; i < 64; i += 4) mSFA[st][i] = 0; + + auto ld = [&](int bf, int kt) { + int bo = kt * 32; + // A: M live rows, 32B each — row r of the smem tile is row r of A + for (int idx = lane; idx < M * 8; idx += 32) { + int rr = idx >> 3, off = idx & 7; + cpa(mA[bf] + rr * 32 + off * 4, A + (size_t)rr * KH + bo + off * 4); + } + if (lane < M) cpa(mSFA[bf] + lane * 4, SFA + (size_t)kt * 512 + lane * 16); + { + const uint8_t* Bg = B + (size_t)blockIdx.x * KH * 8 + (size_t)kt * 256; + for (int c = 0; c < 2; ++c) { int ch = lane + c * 32; + cpa(mB[bf] + ch * 4, Bg + ch * 4); } + } + if (lane < 8) { int col = my_n + lane, rb = col >> 7, ri = col & 127; + int si = rb * ncs + kt, ib = (ri & 31) * 16 + ((ri >> 5) & 3) * 4; + cpa(mSFB[bf] + lane * 4, SFB + si * 512 + ib); } + }; + #pragma unroll + for (int st = 0; st < STAGES - 1; ++st) { if (st < KIw) ld(st, kt0 + st); commit(); } + for (int j = 0; j < KIw; ++j) { + int cb = j % STAGES, jp = j + STAGES - 1; + if (jp < KIw) ld(jp % STAGES, kt0 + jp); + commit(); waitg(); __syncwarp(); + uint32_t a0 = fa(mA[cb], t0, t1, 0), a1 = fa(mA[cb], t0, t1, 1); + uint32_t a2 = fa(mA[cb], t0, t1, 2), a3 = fa(mA[cb], t0, t1, 3); + uint32_t b0 = fb(mB[cb], t0, t1, 0), b1 = fb(mB[cb], t0, t1, 1); + uint32_t sfa = fsa(mSFA[cb], sau), sfb = fsa(mSFB[cb], sbu); + float d0, d1, d2, d3; + AtomType::fma(d0, d1, d2, d3, a0, a1, a2, a3, b0, b1, c0, c1, c2, c3, sfa, sfb); + c0 = d0; c1 = d1; c2 = d2; c3 = d3; + } + // fragment: lane q=lane>>2 holds row q (c0,c1) and row q+8 (c2,c3), + // cols 2r / 2r+1 with r=lane&3 + int q = lane >> 2, r = lane & 3; + if (q < M) { s_red[warp][q][r * 2] = c0; s_red[warp][q][r * 2 + 1] = c1; } + if (q + 8 < M) { s_red[warp][q + 8][r * 2] = c2; s_red[warp][q + 8][r * 2 + 1] = c3; } + __syncthreads(); + if (warp == 0) { + for (int out = lane; out < M * 8; out += 32) { + int row = out >> 3, col8 = out & 7; + float acc = 0.f; + #pragma unroll + for (int w = 0; w < WARPS; ++w) acc += s_red[w][row][col8]; + int col = my_n + col8; + if (col < N) D[(size_t)row * N + col] = __float2bfloat16(acc * alpha); + } + } +} + +} // namespace + +int fp4_w4a4_mma_sm120_warpsplit_ilv_mrows_bf16out( + const void* A_packed, const void* B_packed, void* D_bf16, int M, int N, + int K, const void* SFA, const void* SFB, float alpha, int warps, + int stages, cudaStream_t stream) { + if (!A_packed || !B_packed || !D_bf16 || !SFA || !SFB) return 1; + if (K <= 0 || (K % 64) != 0 || ((K / 64) % warps) != 0) return 2; + if (N <= 0 || (N % 8) != 0) return 3; + if (M < 1 || M > 16) return 4; + dim3 grid(N / 8); + auto a = reinterpret_cast(A_packed); + auto b = reinterpret_cast(B_packed); + auto sa = reinterpret_cast(SFA); + auto sb = reinterpret_cast(SFB); + auto d = reinterpret_cast<__nv_bfloat16*>(D_bf16); + #define WSIM_L(ST, WP) warpsplit_ilv_mrows_kernel<<>>(a, b, sa, sb, d, alpha, M, N, K) + if (warps == 2) { if (stages == 3) WSIM_L(3, 2); else if (stages == 4) WSIM_L(4, 2); else if (stages == 6) WSIM_L(6, 2); else return 5; } + else if (warps == 4) { if (stages == 3) WSIM_L(3, 4); else if (stages == 4) WSIM_L(4, 4); else if (stages == 6) WSIM_L(6, 4); else return 5; } + else if (warps == 8) { if (stages == 3) WSIM_L(3, 8); else if (stages == 4) WSIM_L(4, 8); else return 5; } + else return 6; + return 0; +} + +} // namespace gemm +} // namespace flash_rt diff --git a/csrc/kernels/fp4_w4a4_mma_warpsplit_ilv_mrows_sm120.cuh b/csrc/kernels/fp4_w4a4_mma_warpsplit_ilv_mrows_sm120.cuh new file mode 100644 index 000000000..572559648 --- /dev/null +++ b/csrc/kernels/fp4_w4a4_mma_warpsplit_ilv_mrows_sm120.cuh @@ -0,0 +1,27 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Multi-row (M<=16) interleaved-B warp-split-K NVFP4 W4A4 GEMM for sm_120 +// — the spec-verify shapes (M = draft block + 1, and the shorter re-advance +// prefixes). The 16x8x64 block-scaled MMA atom already produces a 16-row +// output tile; the M=1 interleaved GEMV feeds it one live row and fifteen +// rows of zero ballast. This entry feeds up to 16 live rows through the +// same atom, same shared-memory plan, same accumulator budget: B stays the +// only DRAM-bound stream (A is M*K/2 bytes, L2-resident across the grid), +// so a multi-row call lands near single-GEMV cost. Bit-exact per row vs +// the M=1 interleaved kernel (identical reduction order). Additive. +#pragma once +#include +namespace flash_rt { +namespace gemm { +// A_packed (M, K/2) row-major; B_ilv = interleaved B from +// fp4_w4a4_repack_b_ilv_sm120 (shared with the M=1 entry). D_bf16 (M, N) +// row-major. SFA: the production 512B-per-K64-tile block, row r's four +// scales at (tile*512 + r*16). SFB keeps the base swizzled layout. +// warps in {2,4,8}, stages in {3,4,6}. 1<=M<=16, N%8==0, K%64==0, +// (K/64)%warps==0. Returns 0 on success. +int fp4_w4a4_mma_sm120_warpsplit_ilv_mrows_bf16out( + const void* A_packed, const void* B_ilv, void* D_bf16, int M, int N, + int K, const void* SFA, const void* SFB, float alpha, int warps, + int stages, cudaStream_t stream); +} // namespace gemm +} // namespace flash_rt diff --git a/csrc/kernels/fp4_w4a4_mma_warpsplit_ilv_sm120.cu b/csrc/kernels/fp4_w4a4_mma_warpsplit_ilv_sm120.cu new file mode 100644 index 000000000..2afb4d40a --- /dev/null +++ b/csrc/kernels/fp4_w4a4_mma_warpsplit_ilv_sm120.cu @@ -0,0 +1,187 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Interleaved-B variant of the warp-split-K NVFP4 W4A4 M=1 GEMV (sm_120). +// The dense B walk of the base kernel issues, per K-tile, eight 32B reads +// strided by K/2 bytes (one per output column) - a many-stream pattern that +// plateaus ~86% of achievable DRAM read bandwidth on the wide decode +// shapes. This variant reads B in a bind-time interleaved layout, +// addr = group*(K/2*8) + kt*256 + col*32 + off*4, +// so each block's global reads are fully sequential; measured 89-91% of +// the achievable read roof on the 17408x5120 / 16384x5120 / 5120x17408 +// families, bit-exact against the base kernel (identical reduction order). +// The layout is produced once at bind time by +// fp4_w4a4_repack_b_ilv_sm120 below. Additive: new file + new entries. +// +// Header: fp4_w4a4_mma_warpsplit_ilv_sm120.cuh. +#include +#include +#include + +#include "cute/arch/mma_sm120.hpp" +#include "cutlass/numeric_types.h" + +namespace flash_rt { +namespace gemm { +namespace { + +using AtomType = cute::SM120::BLOCKSCALED::SM120_16x8x64_TN_VS< + cutlass::float_e2m1_t, cutlass::float_e2m1_t, float, + cutlass::float_ue4m3_t, 16>; + +__device__ __forceinline__ uint32_t fa(const uint8_t* s, int t0, int t1, int r) { + int ro = ((r & 1) ? (t1 + 8) : t1) * 32; + return *reinterpret_cast(s + ro + t0 * 4 + ((r >> 1) & 1) * 16); +} +__device__ __forceinline__ uint32_t fb(const uint8_t* s, int t0, int t1, int r) { + return *reinterpret_cast(s + t1 * 32 + t0 * 4 + r * 16); +} +__device__ __forceinline__ uint32_t fsa(const uint8_t* p, int u) { + return *reinterpret_cast(p + u * 4); +} +__device__ __forceinline__ void cpa(uint8_t* d, const uint8_t* s) { + uint32_t i = __cvta_generic_to_shared(d); + asm volatile("cp.async.ca.shared.global.L2::128B [%0], [%1], 4;\n" :: "r"(i), "l"(s)); +} +__device__ __forceinline__ void commit() { asm volatile("cp.async.commit_group;\n" ::); } +template __device__ __forceinline__ void waitg() { + asm volatile("cp.async.wait_group %0;\n" :: "n"(N)); +} + +template +__global__ void warpsplit_kernel( + const uint8_t* __restrict__ A, const uint8_t* __restrict__ B, + const uint8_t* __restrict__ SFA, const uint8_t* __restrict__ SFB, + __nv_bfloat16* __restrict__ D, float alpha, int N, int K) { + // per-warp pipeline buffers + __shared__ uint8_t sA[WARPS][STAGES][16 * 32]; + __shared__ uint8_t sSFA[WARPS][STAGES][16 * 4]; + __shared__ uint8_t sB[WARPS][STAGES][8 * 32]; + __shared__ uint8_t sSFB[WARPS][STAGES][8 * 4]; + __shared__ float s_red[WARPS][8]; // each warp's 8 col partials + + int tid = threadIdx.x, warp = tid >> 5, lane = tid & 31; + int my_n = blockIdx.x * 8; + const int KI = K / 64, KIw = KI / WARPS; // K-tiles per warp + const int kt0 = warp * KIw; + const int KH = K / 2, ncs = (K / 16 + 3) / 4; + int t0 = lane & 3, t1 = lane >> 2, sau = (lane & 1) * 8 + (lane >> 2), sbu = lane >> 2; + float c0 = 0, c1 = 0, c2 = 0, c3 = 0; + + uint8_t (*mA)[16 * 32] = sA[warp]; + uint8_t (*mSFA)[16 * 4] = sSFA[warp]; + uint8_t (*mB)[8 * 32] = sB[warp]; + uint8_t (*mSFB)[8 * 4] = sSFB[warp]; + + if (lane >= 1 && lane < 16) { + #pragma unroll + for (int st = 0; st < STAGES; ++st) { + int4* av = reinterpret_cast(mA[st]); int4 z{0, 0, 0, 0}; + av[lane * 2] = z; av[lane * 2 + 1] = z; + } + if (lane < 4) for (int st = 0; st < STAGES; ++st) + for (int i = 4 + lane; i < 64; i += 4) mSFA[st][i] = 0; + } + auto ld = [&](int bf, int kt) { + int bo = kt * 32; + if (lane < 8) cpa(mA[bf] + lane * 4, A + bo + lane * 4); + if (lane == 0) cpa(mSFA[bf], SFA + kt * 512); + { + const uint8_t* Bg = B + (size_t)blockIdx.x * KH * 8 + (size_t)kt * 256; + for (int c = 0; c < 2; ++c) { int ch = lane + c * 32; + cpa(mB[bf] + ch * 4, Bg + ch * 4); } + } + if (lane < 8) { int col = my_n + lane, rb = col >> 7, ri = col & 127; + int si = rb * ncs + kt, ib = (ri & 31) * 16 + ((ri >> 5) & 3) * 4; + cpa(mSFB[bf] + lane * 4, SFB + si * 512 + ib); } + }; + #pragma unroll + for (int st = 0; st < STAGES - 1; ++st) { if (st < KIw) ld(st, kt0 + st); commit(); } + for (int j = 0; j < KIw; ++j) { + int cb = j % STAGES, jp = j + STAGES - 1; + if (jp < KIw) ld(jp % STAGES, kt0 + jp); + commit(); waitg(); __syncwarp(); + uint32_t a0 = fa(mA[cb], t0, t1, 0), a1 = fa(mA[cb], t0, t1, 1); + uint32_t a2 = fa(mA[cb], t0, t1, 2), a3 = fa(mA[cb], t0, t1, 3); + uint32_t b0 = fb(mB[cb], t0, t1, 0), b1 = fb(mB[cb], t0, t1, 1); + uint32_t sfa = fsa(mSFA[cb], sau), sfb = fsa(mSFB[cb], sbu); + float d0, d1, d2, d3; + AtomType::fma(d0, d1, d2, d3, a0, a1, a2, a3, b0, b1, c0, c1, c2, c3, sfa, sfb); + c0 = d0; c1 = d1; c2 = d2; c3 = d3; + } + // each warp: lanes 0..3 hold row-0 partials c0 (col 2r) / c1 (col 2r+1) + int q = lane >> 2, r = lane & 3; + if (q == 0) { s_red[warp][r * 2] = c0; s_red[warp][r * 2 + 1] = c1; } + __syncthreads(); + // warp 0 sums the WARPS partials per col and writes the bf16 output + if (warp == 0 && lane < 8) { + float acc = 0.f; + #pragma unroll + for (int w = 0; w < WARPS; ++w) acc += s_red[w][lane]; + int col = my_n + lane; + if (col < N) D[col] = __float2bfloat16(acc * alpha); + } +} + +} // namespace + +int fp4_w4a4_mma_sm120_warpsplit_ilv_bf16out( + const void* A_packed, const void* B_packed, void* D_bf16, int N, int K, + const void* SFA, const void* SFB, float alpha, int warps, int stages, + cudaStream_t stream) { + if (!A_packed || !B_packed || !D_bf16 || !SFA || !SFB) return 1; + if (K <= 0 || (K % 64) != 0 || ((K / 64) % warps) != 0) return 2; + if (N <= 0 || (N % 8) != 0) return 3; + dim3 grid(N / 8); + auto a = reinterpret_cast(A_packed); + auto b = reinterpret_cast(B_packed); + auto sa = reinterpret_cast(SFA); + auto sb = reinterpret_cast(SFB); + auto d = reinterpret_cast<__nv_bfloat16*>(D_bf16); + #define WS_L(ST, WP) warpsplit_kernel<<>>(a, b, sa, sb, d, alpha, N, K) + if (warps == 2) { if (stages == 3) WS_L(3, 2); else if (stages == 4) WS_L(4, 2); else if (stages == 6) WS_L(6, 2); else return 5; } + else if (warps == 4) { if (stages == 3) WS_L(3, 4); else if (stages == 4) WS_L(4, 4); else if (stages == 6) WS_L(6, 4); else return 5; } + else if (warps == 8) { if (stages == 3) WS_L(3, 8); else if (stages == 4) WS_L(4, 8); else return 5; } + else return 6; + return 0; +} + + +namespace { + +// bind-time repack: dense packed B (N, K/2 bytes row-major) -> interleaved +// [N/8 groups][K/64 tiles][8 cols x 32B]. Pure byte permutation; one thread +// moves one 32B tile-column. +__global__ void repack_b_ilv_kernel(const uint8_t* __restrict__ src, + uint8_t* __restrict__ dst, + int N, int K) { + const int KH = K / 2, KI = K / 64; + const size_t idx = (size_t)blockIdx.x * blockDim.x + threadIdx.x; + const size_t total = (size_t)(N / 8) * KI * 8; + if (idx >= total) return; + const int col = idx % 8; + const int kt = (idx / 8) % KI; + const size_t g = idx / 8 / KI; + const uint8_t* s = src + (g * 8 + col) * (size_t)KH + (size_t)kt * 32; + uint8_t* d = dst + g * (size_t)KH * 8 + (size_t)kt * 256 + col * 32; + #pragma unroll + for (int i = 0; i < 8; ++i) + reinterpret_cast(d)[i] = + reinterpret_cast(s)[i]; +} + +} // namespace + +int fp4_w4a4_repack_b_ilv_sm120(const void* B_packed, void* B_ilv, int N, + int K, cudaStream_t stream) { + if (N <= 0 || (N % 8) != 0 || K <= 0 || (K % 64) != 0) return 2; + const size_t total = (size_t)(N / 8) * (K / 64) * 8; + const int threads = 256; + repack_b_ilv_kernel<<<(unsigned)((total + threads - 1) / threads), + threads, 0, stream>>>( + reinterpret_cast(B_packed), + reinterpret_cast(B_ilv), N, K); + return 0; +} + +} // namespace gemm +} // namespace flash_rt diff --git a/csrc/kernels/fp4_w4a4_mma_warpsplit_ilv_sm120.cuh b/csrc/kernels/fp4_w4a4_mma_warpsplit_ilv_sm120.cuh new file mode 100644 index 000000000..a33649f9d --- /dev/null +++ b/csrc/kernels/fp4_w4a4_mma_warpsplit_ilv_sm120.cuh @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Interleaved-B warp-split-K NVFP4 W4A4 M=1 GEMV for sm_120, plus the +// bind-time B repack that produces its layout. Bit-exact vs the base +// warpsplit kernel (identical reduction order); serves the wide decode +// shapes at 89-91% of the achievable DRAM read roof. Additive. +#pragma once +#include +namespace flash_rt { +namespace gemm { +// A_packed (K/2,), B_ilv = interleaved B from fp4_w4a4_repack_b_ilv_sm120, +// D_bf16 (N,). SFA (K/16,) and SFB (N, K/16) keep the base kernel's +// swizzled layouts. warps in {2,4,8}, stages in {3,4,6}. N%8==0, K%64==0, +// (K/64)%warps==0. Returns 0 on success. +int fp4_w4a4_mma_sm120_warpsplit_ilv_bf16out( + const void* A_packed, const void* B_ilv, void* D_bf16, int N, int K, + const void* SFA, const void* SFB, float alpha, int warps, int stages, + cudaStream_t stream); +// dense packed B (N, K/2 row-major) -> interleaved layout +// [N/8][K/64][8 cols x 32B]. dst size equals src size (N*K/2 bytes). +int fp4_w4a4_repack_b_ilv_sm120(const void* B_packed, void* B_ilv, int N, + int K, cudaStream_t stream); +} // namespace gemm +} // namespace flash_rt diff --git a/csrc/kernels/gated_delta_wy_kkt_mma.cu b/csrc/kernels/gated_delta_wy_kkt_mma.cu new file mode 100644 index 000000000..db0ad0f13 --- /dev/null +++ b/csrc/kernels/gated_delta_wy_kkt_mma.cu @@ -0,0 +1,146 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// MMA rewrite of the WY K*K^T chunk kernel. The scalar kernel assigns one +// (i,j) pair per thread - a 128-wide fmaf chain with no smem reuse, half +// the threads writing only zeros, and the same K^T K dot recomputed for +// every v-head in a K-head group. Here one block owns one (chunk, k-head): +// the 64x128 K slab loads to shared memory once, eight warps produce the +// 64x64 Gram tile with wmma (bf16 in, f32 accumulate), and the epilogue +// applies each group member's beta_i * exp(gi - gj) scaling. Measured +// 32.8x on the production S=2048 shape (876.8 -> 26.7 us/layer); numerics +// sit in the bf16 reduction-order band (max rel ~2e-3) - the consumer's +// teacher-forced gate is the judge, as with the sibling mma_fla kernels. +// Any S (in-bounds guard + zero padding), covers short continuation and +// deep buckets. Additive: new file + new entry; same argument surface and +// A layout as qwen36_gdn_wy_kkt_b64_bf16. +#include "gated_delta_wy_kkt_mma.cuh" + +#include +#include +#include + +namespace flash_rt { +namespace kernels { +namespace wy_kkt_mma { +namespace { + +constexpr int kHD = 128; +constexpr int kQHeads = 16; +constexpr int kVHeads = 48; +constexpr int kWyChunk = 64; + +using namespace nvcuda; + +constexpr int CH = kWyChunk; +constexpr int HD = kHD; + +// grid: (num_kh, chunks); block: 256 (8 warps) +__global__ void qwen36_gdn_wy_kkt_b64_mma_kernel( + const __nv_bfloat16* __restrict__ k16_l2, // [S, KH, HD] + const __nv_bfloat16* __restrict__ beta, // [S, VH] + const __nv_bfloat16* __restrict__ g_cumsum, // [S, VH] + float* __restrict__ A, // [chunks, VH, CH, CH] + int S, int num_kh, int num_vh, int group) { + const int kh = blockIdx.x; + const int chunk = blockIdx.y; + const int s0 = chunk * CH; + const int tid = threadIdx.x, warp = tid >> 5; + + __shared__ __align__(16) __nv_bfloat16 sK[CH][HD + 8]; // +8 pad 防冲突 + __shared__ __align__(16) float sD[CH][CH]; + + // K 块装载: 行 si 超界补零 (S 非 64 倍时) + for (int idx = tid * 8; idx < CH * HD; idx += blockDim.x * 8) { + int r = idx / HD, c = idx % HD; + int si = s0 + r; + if (si < S) { + *reinterpret_cast(&sK[r][c]) = *reinterpret_cast( + &k16_l2[(static_cast(si) * num_kh + kh) * HD + c]); + } else { + int4 z{0, 0, 0, 0}; + *reinterpret_cast(&sK[r][c]) = z; + } + } + __syncthreads(); + + // 8 warp: warp w 负责行块 (w>>1)*16, 列块 (w&1)*32 (2 个 16x16 tile) + { + const int rt = warp >> 1; // 0..3 → 行 16 块 + const int ct = warp & 1; // 0..1 → 列 32 半区 + wmma::fragment acc[2]; + wmma::fill_fragment(acc[0], 0.f); + wmma::fill_fragment(acc[1], 0.f); + wmma::fragment fa; + wmma::fragment fb; + for (int kk = 0; kk < HD; kk += 16) { + wmma::load_matrix_sync(fa, &sK[rt * 16][kk], HD + 8); + #pragma unroll + for (int t = 0; t < 2; ++t) { + // B = K^T: 列 j 块 = sK 行 (ct*32 + t*16) 作 col_major + wmma::load_matrix_sync(fb, &sK[ct * 32 + t * 16][kk], HD + 8); + wmma::mma_sync(acc[t], fa, fb, acc[t]); + } + } + #pragma unroll + for (int t = 0; t < 2; ++t) + wmma::store_matrix_sync(&sD[rt * 16][ct * 32 + t * 16], acc[t], + CH, wmma::mem_row_major); + } + __syncthreads(); + + // epilogue: 组内每 vh 独立缩放写出 (i>j 且界内, 否则 0) + const int vh0 = kh * group; + for (int g = 0; g < group; ++g) { + const int vh = vh0 + g; + const size_t base = + ((static_cast(chunk) * num_vh + vh) * CH) * CH; + for (int idx = tid; idx < CH * CH; idx += blockDim.x) { + const int i = idx >> 6, j = idx & 63; + const int si = s0 + i, sj = s0 + j; + float out = 0.f; + if (i > j && si < S && sj < S) { + const float bi = __bfloat162float( + beta[static_cast(si) * num_vh + vh]); + const float gi = __bfloat162float( + g_cumsum[static_cast(si) * num_vh + vh]); + const float gj = __bfloat162float( + g_cumsum[static_cast(sj) * num_vh + vh]); + out = bi * sD[i][j] * __expf(gi - gj); + } + A[base + idx] = out; + } + } +} + + +} // namespace + +void run(const void* k16_l2, const void* beta, const void* g_cumsum, + void* A, int S, cudaStream_t stream) { + if (S <= 0) return; + const int chunks = (S + kWyChunk - 1) / kWyChunk; + qwen36_gdn_wy_kkt_b64_mma_kernel<<>>( + reinterpret_cast(k16_l2), + reinterpret_cast(beta), + reinterpret_cast(g_cumsum), + reinterpret_cast(A), S, kQHeads, kVHeads, + kVHeads / kQHeads); +} + +} // namespace wy_kkt_mma +} // namespace kernels +} // namespace flash_rt + +void qwen36_gdn_wy_kkt_b64_mma_bf16( + const void* k16_l2, + const void* beta, + const void* g_cumsum, + void* A, + int S, + cudaStream_t stream) +{ + flash_rt::kernels::wy_kkt_mma::run(k16_l2, beta, g_cumsum, A, S, stream); +} diff --git a/csrc/kernels/gated_delta_wy_kkt_mma.cuh b/csrc/kernels/gated_delta_wy_kkt_mma.cuh new file mode 100644 index 000000000..3f2194a06 --- /dev/null +++ b/csrc/kernels/gated_delta_wy_kkt_mma.cuh @@ -0,0 +1,15 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// MMA rewrite of the WY K*K^T chunk kernel (see gated_delta_wy_kkt_mma.cu). +// Same argument surface and A layout as qwen36_gdn_wy_kkt_b64_bf16, so a +// consumer switches by entry name alone. Additive. +#pragma once +#include + +void qwen36_gdn_wy_kkt_b64_mma_bf16( + const void* k16_l2, + const void* beta, + const void* g_cumsum, + void* A, + int S, + cudaStream_t stream); diff --git a/csrc/kernels/gdn_chunk_from_conv_smem_stash.cu b/csrc/kernels/gdn_chunk_from_conv_smem_stash.cu new file mode 100644 index 000000000..0b9203a7d --- /dev/null +++ b/csrc/kernels/gdn_chunk_from_conv_smem_stash.cu @@ -0,0 +1,208 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Per-row-stash arm of the gated-delta from-conv chunk core. The +// recurrence, gating math, reduction order, and per-row bf16 state +// requantisation are the plain kernel's, verbatim — the only addition +// is the stash write after each row's update, quantised exactly like +// the carried state, so stash row s is bit-equal to the final state a +// re-advance over rows 0..s would store. One block per value head, +// state resident in shared memory across the row walk. +#include "gdn_chunk_from_conv_smem_stash.cuh" + +#include +#include + +namespace flash_rt { +namespace gdn { +namespace { + +constexpr float kEps = 1e-6f; +constexpr int kHD = 128; + +template +__device__ __forceinline__ float block_reduce_sum(float val, float* smem) { + const int t = threadIdx.x; + const int lane = t & 31; + const int warp = t >> 5; + for (int off = 16; off > 0; off >>= 1) + val += __shfl_down_sync(0xffffffffu, val, off); + if (lane == 0) smem[warp] = val; + __syncthreads(); + if (warp == 0) { + val = (t < HD / 32) ? smem[lane] : 0.0f; + for (int off = 16; off > 0; off >>= 1) + val += __shfl_down_sync(0xffffffffu, val, off); + if (lane == 0) smem[0] = val; + } + __syncthreads(); + return smem[0]; +} + +template +__global__ void gdn_chunk_from_conv_smem_stash_kernel( + const __nv_bfloat16* __restrict__ conv_out, + const __nv_bfloat16* __restrict__ a_in, + const __nv_bfloat16* __restrict__ b_in, + const float* __restrict__ neg_exp_A_log, + const float* __restrict__ dt_bias, + __nv_bfloat16* __restrict__ state, + __nv_bfloat16* __restrict__ out_, + __nv_bfloat16* __restrict__ stash, + int S, + int num_v_heads, + int num_k_heads, + int a_stride, + int b_stride, + bool use_qk_l2norm) +{ + static_assert(HD == 128, "HD must be 128 for this host family"); + const int h = blockIdx.x; + const int t = threadIdx.x; + if (t >= HD) return; + + extern __shared__ float smem[]; + float* state_s = smem; + float* qs = state_s + HD * HD; + float* ks = qs + HD; + float* scratch = ks + HD; + float* gate_values = scratch + 32; + + const size_t state_h_off = (size_t)h * HD * HD; + #pragma unroll 16 + for (int i = 0; i < HD; ++i) { + state_s[i * HD + t] = static_cast( + state[state_h_off + (size_t)i * HD + t]); + } + __syncthreads(); + + const int broadcast = num_v_heads / num_k_heads; + const int src_h = h / broadcast; + const int qk_width = num_k_heads * HD; + const int row_width = (2 * num_k_heads + num_v_heads) * HD; + for (int s = 0; s < S; ++s) { + const size_t row = static_cast(s) * row_width; + const size_t out_off = ((size_t)s * num_v_heads + h) * HD + t; + qs[t] = static_cast(conv_out[row + src_h * HD + t]); + ks[t] = static_cast(conv_out[row + qk_width + src_h * HD + t]); + __syncthreads(); + + if (use_qk_l2norm) { + float q_sq = qs[t] * qs[t]; + float k_sq = ks[t] * ks[t]; + q_sq = block_reduce_sum(q_sq, scratch); + // barrier between reduce calls sharing scratch: without it warp 0 + // begins the second reduce's scratch writes while a slower warp + // still reads the first result (the plain kernel's receipt) + __syncthreads(); + k_sq = block_reduce_sum(k_sq, scratch); + const float q_inv = rsqrtf(q_sq + kEps); + const float k_inv = rsqrtf(k_sq + kEps); + qs[t] *= q_inv; + ks[t] *= k_inv; + __syncthreads(); + } + + qs[t] *= rsqrtf(static_cast(HD)); + __syncthreads(); + + if (t == 0) { + const float av = + static_cast(a_in[s * a_stride + h]) + dt_bias[h]; + const float sp = log1pf(__expf(av)); + const float g_log = static_cast( + __float2bfloat16(neg_exp_A_log[h] * sp)); + gate_values[0] = __expf(g_log); + const float bv = static_cast(b_in[s * b_stride + h]); + gate_values[1] = static_cast( + __float2bfloat16(1.0f / (1.0f + __expf(-bv)))); + } + __syncthreads(); + const float g_t = gate_values[0]; + const float beta_t = gate_values[1]; + + #pragma unroll 16 + for (int i = 0; i < HD; ++i) { + state_s[i * HD + t] *= g_t; + } + + float kv_mem = 0.0f; + #pragma unroll 16 + for (int i = 0; i < HD; ++i) { + kv_mem = fmaf(state_s[i * HD + t], ks[i], kv_mem); + } + + const float v_t = + static_cast(conv_out[row + 2 * qk_width + h * HD + t]); + const float delta = (v_t - kv_mem) * beta_t; + + #pragma unroll 16 + for (int i = 0; i < HD; ++i) { + state_s[i * HD + t] = + fmaf(ks[i], delta, state_s[i * HD + t]); + } + + float out_t = 0.0f; + #pragma unroll 16 + for (int i = 0; i < HD; ++i) { + out_t = fmaf(state_s[i * HD + t], qs[i], out_t); + } + out_[out_off] = __float2bfloat16(out_t); + + // the stash write IS the carried requantisation: row s records + // bf16(state after row s), exactly the value a re-advance ending + // here would store, and exactly the value row s+1 resumes from + const size_t stash_off = + (((size_t)s * num_v_heads + h)) * HD * HD; + #pragma unroll 16 + for (int i = 0; i < HD; ++i) { + const __nv_bfloat16 q16 = __float2bfloat16(state_s[i * HD + t]); + stash[stash_off + (size_t)i * HD + t] = q16; + state_s[i * HD + t] = static_cast(q16); + } + __syncthreads(); + } + + #pragma unroll 16 + for (int i = 0; i < HD; ++i) { + state[state_h_off + (size_t)i * HD + t] = + __float2bfloat16(state_s[i * HD + t]); + } +} + +} // namespace + +void gdn_chunk_from_conv_smem_h_stash_bf16( + const void* conv_out, const void* a, const void* b, + const float* neg_exp_A_log, const float* dt_bias, void* state, + void* out, void* stash, int S, int num_v_heads, int num_k_heads, + int head_dim, int a_stride, int b_stride, bool use_qk_l2norm, + cudaStream_t stream) +{ + if (S <= 0 || num_v_heads <= 0) return; + if (head_dim != kHD) return; + dim3 grid(num_v_heads, 1); + dim3 block(kHD); + constexpr size_t kSmemBytes = + (kHD * kHD + 2 * kHD + 34) * sizeof(float); + static bool attr_set = false; + if (!attr_set) { + cudaFuncSetAttribute( + gdn_chunk_from_conv_smem_stash_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(kSmemBytes)); + attr_set = true; + } + gdn_chunk_from_conv_smem_stash_kernel<<< + grid, block, kSmemBytes, stream>>>( + reinterpret_cast(conv_out), + reinterpret_cast(a), + reinterpret_cast(b), + neg_exp_A_log, dt_bias, + reinterpret_cast<__nv_bfloat16*>(state), + reinterpret_cast<__nv_bfloat16*>(out), + reinterpret_cast<__nv_bfloat16*>(stash), + S, num_v_heads, num_k_heads, a_stride, b_stride, use_qk_l2norm); +} + +} // namespace gdn +} // namespace flash_rt diff --git a/csrc/kernels/gdn_chunk_from_conv_smem_stash.cuh b/csrc/kernels/gdn_chunk_from_conv_smem_stash.cuh new file mode 100644 index 000000000..cb634e421 --- /dev/null +++ b/csrc/kernels/gdn_chunk_from_conv_smem_stash.cuh @@ -0,0 +1,26 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Gated-delta chunk core with per-row state stash — the spec-verify +// arm. Identical recurrence to the plain from-conv chunk kernel (same +// per-row bf16 state requantisation, same reduction order), with the +// post-row state additionally written to a stash slab. A rejected +// speculative round then rolls back by *selecting* the stash row at +// the accepted length instead of re-driving the state sublayers: the +// stash row is bit-equal to what a re-advance over that prefix would +// have stored. Additive: new file + new entry point. +#pragma once +#include +namespace flash_rt { +namespace gdn { +// conv_out (S, (2*Hk+Hv)*128) packed q|k|v rows from the conv update; +// a/b (S, *_stride) raw gating projections; state (Hv, 128, 128) bf16, +// carried in place; out (S, Hv, 128); stash (S, Hv, 128, 128) bf16 — +// row s holds the carried state *after* row s. head_dim must be 128. +void gdn_chunk_from_conv_smem_h_stash_bf16( + const void* conv_out, const void* a, const void* b, + const float* neg_exp_A_log, const float* dt_bias, void* state, + void* out, void* stash, int S, int num_v_heads, int num_k_heads, + int head_dim, int a_stride, int b_stride, bool use_qk_l2norm, + cudaStream_t stream); +} // namespace gdn +} // namespace flash_rt diff --git a/csrc/kernels/gdn_recurrent_inout_stream_bf16.cu b/csrc/kernels/gdn_recurrent_inout_stream_bf16.cu new file mode 100644 index 000000000..be5104161 --- /dev/null +++ b/csrc/kernels/gdn_recurrent_inout_stream_bf16.cu @@ -0,0 +1,142 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Gated-delta recurrent decode step, streaming-column form. See header. +#include "kernels/gdn_recurrent_inout_stream_bf16.cuh" + +#include + +namespace flash_rt { +namespace kernels { +namespace { + +constexpr int kHD = 128; +constexpr float kEps = 1e-6f; + +// the packaged kernel's block reduction, transcribed: splitting the +// value columns across smaller blocks would reduce the q/k norms over +// a warp instead, and that changes their fp32 summation order. The +// measured win is not in the split — it is in not spilling the state +// column — so the block width stays 128 and the result stays bitwise +__device__ __forceinline__ float block_reduce_sum(float val, + float* smem) { + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + val += __shfl_xor_sync(0xffffffffu, val, off); + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + if (lane == 0) smem[warp] = val; + __syncthreads(); + if (warp == 0) { + val = (lane < (kHD / 32)) ? smem[lane] : 0.0f; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + val += __shfl_xor_sync(0xffffffffu, val, off); + if (lane == 0) smem[0] = val; + } + __syncthreads(); + return smem[0]; +} + +__global__ void recurrent_stream_kernel( + const __nv_bfloat16* __restrict__ q_in, + const __nv_bfloat16* __restrict__ k_in, + const __nv_bfloat16* __restrict__ v_in, + const __nv_bfloat16* __restrict__ g_in, + const __nv_bfloat16* __restrict__ beta_in, + const __nv_bfloat16* __restrict__ state_in, + __nv_bfloat16* __restrict__ state_out, + __nv_bfloat16* __restrict__ out_, + int num_v_heads, bool use_qk_l2norm) { + const int h = blockIdx.x; + const int b = blockIdx.y; + const int t = threadIdx.x; // this thread's value column + if (t >= kHD) return; + + const size_t hv_off = ((size_t)b * num_v_heads + h) * kHD; + + __shared__ float smem[2 * kHD + 32]; + float* sq = smem; + float* sk = smem + kHD; + float* scratch = smem + 2 * kHD; + sq[t] = static_cast(q_in[hv_off + t]); + sk[t] = static_cast(k_in[hv_off + t]); + __syncthreads(); + + if (use_qk_l2norm) { + float q_sq = block_reduce_sum(sq[t] * sq[t], scratch); + __syncthreads(); + float k_sq = block_reduce_sum(sk[t] * sk[t], scratch); + sq[t] *= rsqrtf(q_sq + kEps); + sk[t] *= rsqrtf(k_sq + kEps); + __syncthreads(); + } + sq[t] *= rsqrtf(static_cast(kHD)); + __syncthreads(); + + const float g_t = + __expf(static_cast(g_in[b * num_v_heads + h])); + const float beta_t = + static_cast(beta_in[b * num_v_heads + h]); + + // Two streaming passes rather than a 128-entry per-thread array. + // That array is 128 registers on top of everything else, past what + // a thread can hold, so the column it is meant to keep resident + // spills to local memory and the step reads its own state twice + // through DRAM anyway. Reading state twice explicitly costs the + // same traffic on the first pass and hits cache on the second (a + // block's slice is 32KB), while the register budget drops far + // enough for the scheduler to hide the latency. Per-column + // arithmetic and its order are unchanged. + const size_t state_h_off = hv_off * kHD; + float kv_mem = 0.0f; + #pragma unroll 16 + for (int i = 0; i < kHD; ++i) { + const float c = static_cast( + state_in[state_h_off + (size_t)i * kHD + t]) * g_t; + kv_mem = fmaf(c, sk[i], kv_mem); + } + + const float v_t = static_cast(v_in[hv_off + t]); + const float delta = (v_t - kv_mem) * beta_t; + + float out_t = 0.0f; + #pragma unroll 16 + for (int i = 0; i < kHD; ++i) { + const float c = fmaf( + sk[i], delta, + static_cast( + state_in[state_h_off + (size_t)i * kHD + t]) * g_t); + state_out[state_h_off + (size_t)i * kHD + t] = + __float2bfloat16(c); + out_t = fmaf(c, sq[i], out_t); + } + out_[hv_off + t] = __float2bfloat16(out_t); +} + +} // namespace + +int gdn_recurrent_inout_stream_bf16( + const void* q, const void* k, const void* v, const void* g, + const void* beta, const void* state_in, void* state_out, void* out, + int B, int num_v_heads, int head_dim, bool use_qk_l2norm, + cudaStream_t stream) { + if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out) + return 1; + if (head_dim != kHD || B <= 0 || num_v_heads <= 0) return 2; + dim3 grid(num_v_heads, B); + recurrent_stream_kernel<<>>( + reinterpret_cast(q), + reinterpret_cast(k), + reinterpret_cast(v), + reinterpret_cast(g), + reinterpret_cast(beta), + reinterpret_cast(state_in), + reinterpret_cast<__nv_bfloat16*>(state_out), + reinterpret_cast<__nv_bfloat16*>(out), num_v_heads, + use_qk_l2norm); + const cudaError_t e = cudaGetLastError(); + return (e == cudaSuccess) ? 0 : -static_cast(e); +} + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/gdn_recurrent_inout_stream_bf16.cuh b/csrc/kernels/gdn_recurrent_inout_stream_bf16.cuh new file mode 100644 index 000000000..713db0f7a --- /dev/null +++ b/csrc/kernels/gdn_recurrent_inout_stream_bf16.cuh @@ -0,0 +1,37 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Gated-delta recurrent decode step, streaming-column form. +// +// The recurrence gives every value column its own thread: that thread +// holds state[:, v] and reduces over K entirely on its own. The +// packaged kernel keeps that column in a 128-entry per-thread array, +// which no thread can hold — it spills, and the step pays DRAM for +// state it believes is resident. This entry streams the column +// instead. +// +// The whole kernel is bit-identical to the packaged one — same block +// reduction for the q/k norms, same per-column arithmetic in the same +// order. The only change is that the column no longer lives in a +// per-thread array: 128 registers on top of everything else is past +// what a thread holds, so that array spills to local memory and the +// step reads its own state through DRAM anyway. Streaming the column +// in two passes costs the same traffic on the first and hits cache on +// the second (a head's slice is 32KB), while the register budget +// drops far enough for the scheduler to hide the latency. Additive. +#pragma once +#include + +namespace flash_rt { +namespace kernels { + +// q/k/v: (B, H, 128) bf16. g/beta: (B, H) bf16. state_in/out: +// (B, H, 128, 128) bf16 (may alias). out: (B, H, 128) bf16. +// head_dim must be 128. Returns 0 on success. +int gdn_recurrent_inout_stream_bf16( + const void* q, const void* k, const void* v, const void* g, + const void* beta, const void* state_in, void* state_out, void* out, + int B, int num_v_heads, int head_dim, bool use_qk_l2norm, + cudaStream_t stream); + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/gdn_wy_norm_cumsum_pack_qk_v2.cu b/csrc/kernels/gdn_wy_norm_cumsum_pack_qk_v2.cu new file mode 100644 index 000000000..8da3ea375 --- /dev/null +++ b/csrc/kernels/gdn_wy_norm_cumsum_pack_qk_v2.cu @@ -0,0 +1,138 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// WY-chain q/k L2-norm + pack + per-chunk gate cumsum, v2 launch plan. +// See header. Norm kernel body and reduction transcribed verbatim from +// the packaged fast arm; the cumsum keeps the packaged serial order +// inside each 64-token chunk (bit-exact) and parallelizes across the +// independent (chunk, head) pairs. +#include "kernels/gdn_wy_norm_cumsum_pack_qk_v2.cuh" + +#include + +namespace flash_rt { +namespace kernels { +namespace { + +constexpr int kHD = 128; +constexpr int kQHeads = 16; +constexpr int kVHeads = 48; +constexpr int kWyChunk = 64; +constexpr float kEps = 1e-6f; + +template +__device__ __forceinline__ float block_reduce_sum(float val, float* smem) { + for (int off = 16; off > 0; off >>= 1) { + val += __shfl_xor_sync(0xffffffff, val, off); + } + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + if (lane == 0) smem[warp] = val; + __syncthreads(); + if (warp == 0) { + val = (lane < (HD / 32)) ? smem[lane] : 0.0f; + for (int off = 16; off > 0; off >>= 1) { + val += __shfl_xor_sync(0xffffffff, val, off); + } + if (lane == 0) smem[0] = val; + } + __syncthreads(); + return smem[0]; +} + +__global__ void norm_qk_v2_kernel( + const __nv_bfloat16* __restrict__ q16, + const __nv_bfloat16* __restrict__ k16, + __nv_bfloat16* __restrict__ q16_l2, + __nv_bfloat16* __restrict__ k16_l2, + __nv_bfloat16* __restrict__ q_pack_hv, + __nv_bfloat16* __restrict__ k_pack_hk, + int S, int num_k_heads, int num_v_heads, int head_group_size) { + const int t = threadIdx.x; + const int h = blockIdx.x; + const int s = blockIdx.y; + if (t >= kHD || h >= num_k_heads || s >= S) return; + + __shared__ float scratch[32]; + const size_t off = (static_cast(s) * num_k_heads + h) * kHD + t; + const float qv = static_cast(q16[off]); + const float kv = static_cast(k16[off]); + float q_sq = qv * qv; + float k_sq = kv * kv; + q_sq = block_reduce_sum(q_sq, scratch); + __syncthreads(); + k_sq = block_reduce_sum(k_sq, scratch); + __syncthreads(); + const float q_inv = rsqrtf(q_sq + kEps); + const float k_inv = rsqrtf(k_sq + kEps); + const __nv_bfloat16 q_norm = __float2bfloat16(qv * q_inv); + const __nv_bfloat16 k_norm = __float2bfloat16(kv * k_inv); + q16_l2[off] = q_norm; + k16_l2[off] = k_norm; + if (k_pack_hk != nullptr) { + const int chunk = s / kWyChunk; + const int tt = s - chunk * kWyChunk; + k_pack_hk[ + ((static_cast(chunk) * num_k_heads + h) * kWyChunk + tt) + * kHD + t] = k_norm; + } + if (q_pack_hv != nullptr) { + const int chunk = s / kWyChunk; + const int tt = s - chunk * kWyChunk; + #pragma unroll + for (int r = 0; r < head_group_size; ++r) { + const int vh = h * head_group_size + r; + q_pack_hv[ + ((static_cast(chunk) * num_v_heads + vh) * kWyChunk + tt) + * kHD + t] = q_norm; + } + } +} + +// one thread per (chunk, head): the packaged serial order inside the +// chunk, chunks in parallel (the accumulator resets at every boundary) +__global__ void cumsum_g_v2_kernel( + const __nv_bfloat16* __restrict__ g, + __nv_bfloat16* __restrict__ g_cumsum, + int S, int num_v_heads, int num_chunks) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + const int h = idx % num_v_heads; + const int chunk = idx / num_v_heads; + if (chunk >= num_chunks) return; + const int s0 = chunk * kWyChunk; + const int s1 = min(s0 + kWyChunk, S); + float acc = 0.0f; + for (int s = s0; s < s1; ++s) { + const size_t off = static_cast(s) * num_v_heads + h; + acc += static_cast(g[off]); + g_cumsum[off] = __float2bfloat16(acc); + } +} + +} // namespace + +int gdn_wy_norm_cumsum_pack_qk_v2_bf16( + const void* q16, const void* k16, const void* g, void* q16_l2, + void* k16_l2, void* q_pack_hv, void* k_pack_hk, void* g_cumsum, + int S, cudaStream_t stream) { + if (!q16 || !k16 || !g || !q16_l2 || !k16_l2 || !g_cumsum) return 1; + if (S <= 0) return 2; + norm_qk_v2_kernel<<>>( + reinterpret_cast(q16), + reinterpret_cast(k16), + reinterpret_cast<__nv_bfloat16*>(q16_l2), + reinterpret_cast<__nv_bfloat16*>(k16_l2), + reinterpret_cast<__nv_bfloat16*>(q_pack_hv), + reinterpret_cast<__nv_bfloat16*>(k_pack_hk), + S, kQHeads, kVHeads, kVHeads / kQHeads); + const int chunks = (S + kWyChunk - 1) / kWyChunk; + const int total = chunks * kVHeads; + cumsum_g_v2_kernel<<<(total + 127) / 128, 128, 0, stream>>>( + reinterpret_cast(g), + reinterpret_cast<__nv_bfloat16*>(g_cumsum), + S, kVHeads, chunks); + const cudaError_t e = cudaGetLastError(); + return (e == cudaSuccess) ? 0 : -static_cast(e); +} + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/gdn_wy_norm_cumsum_pack_qk_v2.cuh b/csrc/kernels/gdn_wy_norm_cumsum_pack_qk_v2.cuh new file mode 100644 index 000000000..d79da5a84 --- /dev/null +++ b/csrc/kernels/gdn_wy_norm_cumsum_pack_qk_v2.cuh @@ -0,0 +1,28 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// WY-chain q/k L2-norm + pack + per-chunk gate cumsum, v2 launch plan. +// The math is a verbatim transcription of the packaged +// norm_cumsum_pack_qk pair (same block reduction, same rounding); the +// v2 is the cumsum's launch: the packaged kernel walks the whole +// prompt from a single 64-thread block, while the per-64-token chunks +// are independent by construction (the accumulator resets at every +// chunk boundary), so this plan gives each (chunk, head) its own lane +// — bit-exact per chunk, parallel across them. Additive. +#pragma once +#include + +namespace flash_rt { +namespace kernels { + +// Fixed 16 k-head / 48 v-head / 128 head-dim / 64-chunk family (the +// same constants the packaged fast arm serves). q16/k16: (S, 16, 128) +// bf16. q16_l2/k16_l2: same shape. q_pack_hv: (NT, 48, 64, 128). +// k_pack_hk: (NT, 16, 64, 128). g: (S, 48) bf16 -> g_cumsum (S, 48). +// Returns 0 on success. +int gdn_wy_norm_cumsum_pack_qk_v2_bf16( + const void* q16, const void* k16, const void* g, void* q16_l2, + void* k16_l2, void* q_pack_hv, void* k_pack_hk, void* g_cumsum, + int S, cudaStream_t stream); + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/rms_norm_gated_silu_quant_fp4_bf16.cu b/csrc/kernels/rms_norm_gated_silu_quant_fp4_bf16.cu new file mode 100644 index 000000000..cb5bce9b5 --- /dev/null +++ b/csrc/kernels/rms_norm_gated_silu_quant_fp4_bf16.cu @@ -0,0 +1,142 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Fused RMSNorm + weight + silu(gate) + NVFP4 quantize. See header. +#include "kernels/rms_norm_gated_silu_quant_fp4_bf16.cuh" + +#include +#include +#include + +namespace flash_rt { +namespace kernels { +namespace { + +constexpr int kDim = 128; + +__device__ __forceinline__ int rgs_sfa_offset_128x64( + int row, int k, int dim) { + const int row_block = row >> 7; + const int row_in_block = row & 127; + const int k_block = k >> 6; + const int k_in_block = k & 63; + const int k_blocks = (dim + 63) >> 6; + return row_block * k_blocks * 512 + k_block * 512 + + (row_in_block & 31) * 16 + (row_in_block >> 5) * 4 + + (k_in_block >> 4); +} + +__device__ __forceinline__ uint8_t rgs_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; +} + +__global__ void rms_norm_gated_silu_quant_kernel( + const __nv_bfloat16* __restrict__ x, + const __nv_bfloat16* __restrict__ gate, + const __nv_bfloat16* __restrict__ weight, + __nv_bfloat16* __restrict__ out, + uint2* __restrict__ packed, + uint8_t* __restrict__ sfa, + int M, int D, float eps) { + const int m = blockIdx.x; + const int t = threadIdx.x; + if (m >= M || t >= kDim) return; + + const size_t row_off = (size_t)m * kDim + t; + const float xv = __bfloat162float(x[row_off]); + const float gv = __bfloat162float(gate[row_off]); + + // block-reduce sum-of-squares, transcribed from the packaged kernel + float sq = xv * xv; + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffff, sq, off); + __shared__ float warp_sq[4]; + __shared__ float reduced; + const int lane = t & 31; + const int warp = t >> 5; + if (lane == 0) warp_sq[warp] = sq; + __syncthreads(); + if (warp == 0) { + float v = (lane < 4) ? warp_sq[lane] : 0.0f; + v += __shfl_xor_sync(0xffffffff, v, 1); + v += __shfl_xor_sync(0xffffffff, v, 2); + if (lane == 0) reduced = v; + } + __syncthreads(); + + const float rms_inv = rsqrtf(reduced / static_cast(kDim) + eps); + const float wv = __bfloat162float(weight[t]); + const __nv_bfloat16 norm_bf = __float2bfloat16(xv * rms_inv); + const __nv_bfloat16 weighted_bf = + __float2bfloat16(wv * __bfloat162float(norm_bf)); + const float silu_g = gv / (1.0f + __expf(-gv)); + const __nv_bfloat16 out_bf = + __float2bfloat16(__bfloat162float(weighted_bf) * silu_g); + out[row_off] = out_bf; + + // the consumer reads this row as one (1, M*128) activation: its + // 16-element quantize blocks tile a head's lanes exactly, so the + // eight blocks of this row quantize here, in the production path's + // own arithmetic + __shared__ float vals[kDim]; + vals[t] = __bfloat162float(out_bf); + __syncthreads(); + if (t >= kDim / 16) return; + const int blk = t; // 0..7 + const int base = blk * 16; + float amax = 0.f; + #pragma unroll + for (int i = 0; i < 16; ++i) { + const float a = fabsf(vals[base + i]); + if (a > amax) amax = a; + } + float desired = amax / 6.f; + if (desired < 1e-12f) desired = 1e-12f; + __nv_fp8_e4m3 bs_q = __nv_fp8_e4m3(fmaxf(desired, 0.f)); + const float bs_dq = static_cast(bs_q); + const int col = m * kDim + base; // column in the (1, D) row + sfa[rgs_sfa_offset_128x64(0, col, D)] = + *reinterpret_cast(&bs_q); + const float inv_bs = 1.f / bs_dq; + uint2 o; + uint8_t* ob = reinterpret_cast(&o); + #pragma unroll + for (int p = 0; p < 8; ++p) { + const uint8_t lo = rgs_fp32_to_e2m1(vals[base + 2 * p] * inv_bs); + const uint8_t hi = rgs_fp32_to_e2m1(vals[base + 2 * p + 1] * inv_bs); + ob[p] = static_cast(lo | (hi << 4)); + } + packed[(size_t)m * (kDim / 16) + blk] = o; +} + +} // namespace + +int rms_norm_gated_silu_quant_fp4_bf16( + const void* x, const void* gate, const void* weight, void* out, + void* packed, void* sfa, int M, int dim, float eps, + cudaStream_t stream) { + if (!x || !gate || !weight || !out || !packed || !sfa) return 1; + if (dim != kDim || M <= 0) return 2; + rms_norm_gated_silu_quant_kernel<<>>( + reinterpret_cast(x), + reinterpret_cast(gate), + reinterpret_cast(weight), + reinterpret_cast<__nv_bfloat16*>(out), + reinterpret_cast(packed), + reinterpret_cast(sfa), M, M * kDim, eps); + const cudaError_t e = cudaGetLastError(); + return (e == cudaSuccess) ? 0 : -static_cast(e); +} + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/rms_norm_gated_silu_quant_fp4_bf16.cuh b/csrc/kernels/rms_norm_gated_silu_quant_fp4_bf16.cuh new file mode 100644 index 000000000..8740cfacc --- /dev/null +++ b/csrc/kernels/rms_norm_gated_silu_quant_fp4_bf16.cuh @@ -0,0 +1,27 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Fused RMSNorm + weight + silu(gate) that also emits its consumer's +// NVFP4 input. The gated norm's output has exactly one consumer — the +// output projection — whose first act is to quantize it, and the +// quantizer's 16-element blocks tile a head's 128 lanes exactly, so +// the whole quantization fits inside the block that just produced the +// row. The norm arithmetic is transcribed unchanged from the packaged +// kernel; the quantize stage is the production path verbatim. +// Additive. +#pragma once +#include + +namespace flash_rt { +namespace kernels { + +// x, gate: (M, 128) bf16. weight: (128) bf16. out: (M, 128) bf16 — +// the normed rows, still written for hosts that read them. packed: +// (1, M*128/2) bytes; sfa: the 128x64-atom block for (1, M*128). +// dim must be 128. Returns 0 on success. +int rms_norm_gated_silu_quant_fp4_bf16( + const void* x, const void* gate, const void* weight, void* out, + void* packed, void* sfa, int M, int dim, float eps, + cudaStream_t stream); + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/rms_norm_quantize_fp4_sfa_bf16.cu b/csrc/kernels/rms_norm_quantize_fp4_sfa_bf16.cu new file mode 100644 index 000000000..8bef58b0c --- /dev/null +++ b/csrc/kernels/rms_norm_quantize_fp4_sfa_bf16.cu @@ -0,0 +1,166 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Fused (1+w)-form RMSNorm + NVFP4 quantize producer. See header. One +// block per row: the fp32 row is staged in shared memory (one global +// read), reduced for mean(x^2), then each thread quantizes 16-element +// blocks with the production quantize path verbatim while also writing +// the normed bf16 row. +#include "kernels/rms_norm_quantize_fp4_sfa_bf16.cuh" + +#include +#include +#include + +namespace flash_rt { +namespace fp4 { +namespace { + +__device__ __forceinline__ int rnq_sfa_offset_128x64( + int row, int k, int dim) { + const int row_block = row >> 7; + const int row_in_block = row & 127; + const int k_block = k >> 6; + const int k_in_block = k & 63; + const int k_blocks = (dim + 63) >> 6; + return row_block * k_blocks * 512 + k_block * 512 + + (row_in_block & 31) * 16 + (row_in_block >> 5) * 4 + + (k_in_block >> 4); +} + +__device__ __forceinline__ uint8_t rnq_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_rms_norm_quantize_fp4_sfa_bf16( + const int4* __restrict__ x, // bf16 (N, D) as int4 (8 elems) + const __nv_bfloat16* __restrict__ w, + float eps, + int4* __restrict__ normed, // bf16 (N, D) as int4 + uint2* __restrict__ dst_packed, + uint8_t* __restrict__ dst_sfa, + int N, int D) { + __shared__ float red[THREADS / 32]; + const int row = blockIdx.x; + if (row >= N) return; + const int D8 = D >> 3; + const int4* xr = x + (size_t)row * D8; + + // pass 1: vectorized sum of squares (the row stays L2-resident for + // pass 2 — no shared-memory staging, so occupancy is thread-bound) + float ssq = 0.f; + for (int i = threadIdx.x; i < D8; i += THREADS) { + const int4 raw = xr[i]; + const __nv_bfloat16* h = reinterpret_cast(&raw); + #pragma unroll + for (int j = 0; j < 8; ++j) { + const float v = __bfloat162float(h[j]); + ssq += v * v; + } + } + #pragma unroll + for (int o = 16; o; o >>= 1) + ssq += __shfl_down_sync(0xffffffffu, ssq, o); + if ((threadIdx.x & 31) == 0) red[threadIdx.x >> 5] = ssq; + __syncthreads(); + if (threadIdx.x < 32) { + float v = (threadIdx.x < THREADS / 32) ? red[threadIdx.x] : 0.f; + #pragma unroll + for (int o = 16; o; o >>= 1) + v += __shfl_down_sync(0xffffffffu, v, o); + if (threadIdx.x == 0) red[0] = v; + } + __syncthreads(); + const float rstd = rsqrtf(red[0] / (float)D + eps); + + // pass 2: one 16-element block per iteration — vectorized re-read + // (L2), norm, vectorized bf16 store, production quantize path + const int n_blocks = D / 16; + const int4* wr = reinterpret_cast(w); + for (int b = threadIdx.x; b < n_blocks; b += THREADS) { + const int4 r0 = xr[2 * b], r1 = xr[2 * b + 1]; + const int4 w0 = wr[2 * b], w1 = wr[2 * b + 1]; + const __nv_bfloat16* h0 = reinterpret_cast(&r0); + const __nv_bfloat16* h1 = reinterpret_cast(&r1); + const __nv_bfloat16* g0 = reinterpret_cast(&w0); + const __nv_bfloat16* g1 = reinterpret_cast(&w1); + float vals[16]; + int4 nb0, nb1; + __nv_bfloat16* nh0 = reinterpret_cast<__nv_bfloat16*>(&nb0); + __nv_bfloat16* nh1 = reinterpret_cast<__nv_bfloat16*>(&nb1); + float amax = 0.f; + #pragma unroll + for (int i = 0; i < 8; ++i) { + const float y = __bfloat162float(h0[i]) * rstd * + (1.f + __bfloat162float(g0[i])); + nh0[i] = __float2bfloat16(y); + vals[i] = __bfloat162float(nh0[i]); + const float a = fabsf(vals[i]); + if (a > amax) amax = a; + } + #pragma unroll + for (int i = 0; i < 8; ++i) { + const float y = __bfloat162float(h1[i]) * rstd * + (1.f + __bfloat162float(g1[i])); + nh1[i] = __float2bfloat16(y); + vals[8 + i] = __bfloat162float(nh1[i]); + const float a = fabsf(vals[8 + i]); + if (a > amax) amax = a; + } + normed[(size_t)row * D8 + 2 * b] = nb0; + normed[(size_t)row * D8 + 2 * b + 1] = nb1; + float desired = amax / 6.f; + if (desired < 1e-12f) desired = 1e-12f; + __nv_fp8_e4m3 bs_q = __nv_fp8_e4m3(fmaxf(desired, 0.f)); + const float bs_dq = static_cast(bs_q); + dst_sfa[rnq_sfa_offset_128x64(row, b * 16, D)] = + *reinterpret_cast(&bs_q); + const float inv_bs = 1.f / bs_dq; + uint2 out; + uint8_t* ob = reinterpret_cast(&out); + #pragma unroll + for (int p = 0; p < 8; ++p) { + const uint8_t lo = rnq_fp32_to_e2m1(vals[2 * p] * inv_bs); + const uint8_t hi = rnq_fp32_to_e2m1(vals[2 * p + 1] * inv_bs); + ob[p] = static_cast(lo | (hi << 4)); + } + dst_packed[(size_t)row * n_blocks + b] = out; + } +} + +} // namespace + +int rms_norm_quantize_fp4_sfa_bf16( + const void* x_bf16, const void* w_bf16, float eps, void* normed_bf16, + void* dst_packed, void* dst_sfa, int N, int D, cudaStream_t stream) { + if (!x_bf16 || !w_bf16 || !normed_bf16 || !dst_packed || !dst_sfa) + return 1; + if (N <= 0 || D <= 0 || (D % 16) != 0 || D > 8192) return 2; + if ((reinterpret_cast(x_bf16) & 15) || + (reinterpret_cast(w_bf16) & 15) || + (reinterpret_cast(normed_bf16) & 15)) return 3; + constexpr int THREADS = 256; + kernel_rms_norm_quantize_fp4_sfa_bf16 + <<>>( + reinterpret_cast(x_bf16), + reinterpret_cast(w_bf16), eps, + reinterpret_cast(normed_bf16), + reinterpret_cast(dst_packed), + reinterpret_cast(dst_sfa), N, D); + const cudaError_t e = cudaGetLastError(); + return (e == cudaSuccess) ? 0 : -static_cast(e); +} + +} // namespace fp4 +} // namespace flash_rt diff --git a/csrc/kernels/rms_norm_quantize_fp4_sfa_bf16.cuh b/csrc/kernels/rms_norm_quantize_fp4_sfa_bf16.cuh new file mode 100644 index 000000000..e74c70e8c --- /dev/null +++ b/csrc/kernels/rms_norm_quantize_fp4_sfa_bf16.cuh @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Fused (1+w)-form RMSNorm + NVFP4 quantize producer. Computes the +// host norm exactly — y = x * rsqrt(mean(x^2) + eps) * (1 + w), all in +// fp32, one bf16 round at the write — and quantizes the same values +// with the production per-16-block amax/6 scale selection, e2m1 +// rounding table and 128x64 SFA layout. Emits BOTH the normed bf16 row +// (for hosts that still read it) and the packed FP4 + SFA the bound +// projections consume, off one read of x. Additive. +#pragma once +#include + +namespace flash_rt { +namespace fp4 { + +// x: (N, D) bf16 row-major; w: (D) bf16 (the residual-form weight, the +// kernel applies 1+w). normed: (N, D) bf16. packed: (N, D/2) bytes. +// sfa: ((N+127)/128)*((D+63)/64)*512 bytes. D%16==0, D<=8192. +int rms_norm_quantize_fp4_sfa_bf16( + const void* x_bf16, const void* w_bf16, float eps, void* normed_bf16, + void* dst_packed, void* dst_sfa, int N, int D, cudaStream_t stream); + +} // namespace fp4 +} // namespace flash_rt diff --git a/csrc/kernels/silu_mul_quantize_fp4_sfa_bf16.cu b/csrc/kernels/silu_mul_quantize_fp4_sfa_bf16.cu new file mode 100644 index 000000000..7241d492e --- /dev/null +++ b/csrc/kernels/silu_mul_quantize_fp4_sfa_bf16.cu @@ -0,0 +1,133 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Fused SwiGLU activation + NVFP4 quantize producer. See header for the +// contract. The quantize stage is a verbatim transcription of the +// production bf16 quantize kernel (same scale selection, rounding table +// and SFA layout); the activation stage rounds its fp32 product through +// bf16 so the quantizer sees the same value class the split chain fed +// it — measured bit-exact against (silu·mul kernel -> quantize kernel) +// across M in {1, 7, 2044} at H=17408. +#include "kernels/silu_mul_quantize_fp4_sfa_bf16.cuh" + +#include +#include +#include + +namespace flash_rt { +namespace fp4 { +namespace { + +__device__ __forceinline__ int smq_sfa_offset_128x64( + int row, int k, int dim) { + const int row_block = row >> 7; + const int row_in_block = row & 127; + const int k_block = k >> 6; + const int k_in_block = k & 63; + const int k_blocks = (dim + 63) >> 6; + return row_block * k_blocks * 512 + k_block * 512 + + (row_in_block & 31) * 16 + (row_in_block >> 5) * 4 + + (k_in_block >> 4); +} + +__device__ __forceinline__ uint8_t smq_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; +} + +__global__ void kernel_silu_mul_quantize_fp4_sfa_bf16( + const int4* __restrict__ src, // bf16 (N, 2H) as int4 chunks + uint2* __restrict__ dst_packed, // (N, H/2) bytes as uint2 + uint8_t* __restrict__ dst_sfa, + int N, int H8) { // H8 = H/8 int4 chunks per half + const int block_idx = blockIdx.x * blockDim.x + threadIdx.x; + const int row = blockIdx.y; + const int n_blocks = H8 >> 1; // 16 output elements per block + if (row >= N || block_idx >= n_blocks) return; + + const long row_off = (long)row * (H8 << 1); + const int4 g0 = src[row_off + 2 * block_idx]; + const int4 g1 = src[row_off + 2 * block_idx + 1]; + const int4 u0 = src[row_off + H8 + 2 * block_idx]; + const int4 u1 = src[row_off + H8 + 2 * block_idx + 1]; + const __nv_bfloat16* gh0 = reinterpret_cast(&g0); + const __nv_bfloat16* gh1 = reinterpret_cast(&g1); + const __nv_bfloat16* uh0 = reinterpret_cast(&u0); + const __nv_bfloat16* uh1 = reinterpret_cast(&u1); + + float vals[16]; + #pragma unroll + for (int i = 0; i < 8; ++i) { + const float g = __bfloat162float(gh0[i]); + const float u = __bfloat162float(uh0[i]); + vals[i] = __bfloat162float( + __float2bfloat16((g / (1.f + expf(-g))) * u)); + } + #pragma unroll + for (int i = 0; i < 8; ++i) { + const float g = __bfloat162float(gh1[i]); + const float u = __bfloat162float(uh1[i]); + vals[8 + i] = __bfloat162float( + __float2bfloat16((g / (1.f + expf(-g))) * u)); + } + + 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(fmaxf(desired, 0.f)); + const float bs_dq = static_cast(bs_q); + + const int H = H8 << 3; + dst_sfa[smq_sfa_offset_128x64(row, block_idx * 16, H)] = + *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 = smq_fp32_to_e2m1(vals[2 * p] * inv_bs); + const uint8_t hi = smq_fp32_to_e2m1(vals[2 * p + 1] * inv_bs); + ob[p] = static_cast(lo | (hi << 4)); + } + dst_packed[(long)row * n_blocks + block_idx] = out; +} + +} // namespace + +int silu_mul_quantize_fp4_sfa_bf16( + const void* merged_bf16, void* dst_packed, void* dst_sfa, + int N, int H, cudaStream_t stream) { + if (!merged_bf16 || !dst_packed || !dst_sfa) return 1; + if (N <= 0 || H <= 0 || (H % 16) != 0) return 2; + if ((reinterpret_cast(merged_bf16) & 15) || + (reinterpret_cast(dst_packed) & 7)) return 3; + const int n_blocks = H / 16; + const int threads = 128; + dim3 grid((n_blocks + threads - 1) / threads, N); + kernel_silu_mul_quantize_fp4_sfa_bf16<<>>( + reinterpret_cast(merged_bf16), + reinterpret_cast(dst_packed), + reinterpret_cast(dst_sfa), + N, H / 8); + const cudaError_t e = cudaGetLastError(); + return (e == cudaSuccess) ? 0 : -static_cast(e); +} + +} // namespace fp4 +} // namespace flash_rt diff --git a/csrc/kernels/silu_mul_quantize_fp4_sfa_bf16.cuh b/csrc/kernels/silu_mul_quantize_fp4_sfa_bf16.cuh new file mode 100644 index 000000000..006d2b840 --- /dev/null +++ b/csrc/kernels/silu_mul_quantize_fp4_sfa_bf16.cuh @@ -0,0 +1,23 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Fused SwiGLU activation + NVFP4 quantize producer: silu(gate)·up +// computed in fp32, rounded through bf16 (mirroring the elementwise +// producer it replaces), then quantized with the exact per-16-block +// amax/6 scale selection, e2m1 rounding table and 128x64 SFA layout of +// the production quantize kernel — bit-exact against that chain by +// construction. Additive. +#pragma once +#include + +namespace flash_rt { +namespace fp4 { + +// merged_bf16: (N, 2H) row-major, halves ordered [gate | up]; H%16==0. +// dst_packed: (N, H/2) bytes. dst_sfa: the 128x64-atom SFA block for +// (N, H), ((N+127)/128)*((H+63)/64)*512 bytes. Returns 0 on success. +int silu_mul_quantize_fp4_sfa_bf16( + const void* merged_bf16, void* dst_packed, void* dst_sfa, + int N, int H, cudaStream_t stream); + +} // namespace fp4 +} // namespace flash_rt diff --git a/docs/serving_engine_attach.md b/docs/serving_engine_attach.md new file mode 100644 index 000000000..876cc7bc5 --- /dev/null +++ b/docs/serving_engine_attach.md @@ -0,0 +1,314 @@ +# Attaching structures to a serving engine (vLLM / SGLang) + +The engine adapters seat this repo's kernels inside a running serving +engine — no fork, no model conversion, one hook installed before the +engine loads weights. Everything below is a measured configuration: +each command is the exact form its receipts were produced with. + +> **How to read these numbers — a fairness note, not a benchmark +> claim.** The vLLM/SGLang configurations measured here are **not an +> exhaustive tuning of either engine**, so none of the numbers should +> be read as a community performance comparison. The baseline arms +> use each engine's documented serving recipe for this model plus the +> minimum settings the protocol needed; both engines have many knobs +> outside my sweep — the "maximum context" figures in particular +> reflect the engines' default memory accounting under these settings +> and may move under other configurations. What this document and the +> adapters demonstrate is a **hot-pluggable, stackable optimization +> path**: what becomes possible when the structures layer attaches to +> an engine *as configured*, with everything reverting to the +> untouched host on detach or refusal. +> +> This line is a single-author project, and it exists because vLLM +> and SGLang are excellent hosts — everything here runs *through* +> their serving stacks, not around them. If a configuration I did not +> cover serves these workloads better, I would genuinely like to hear +> about it and will update the numbers; the tuning attempts I did +> make are recorded below so they can be checked and improved on. + +## Requirements + +These determine whether the seats bind at all — check them before +anything else. + +**Hardware** +- An SM120 GPU (RTX 5090 class). The W4A4 GEMV, small-M Marlin, and + paged FP8-KV attention tiers are `sm120a` builds; on other + architectures those seats refuse cleanly and the host runs its own + kernels (the adapter never blocks engine startup). +- The receipts below are from a single 32 GB card, `max_num_seqs=1` + (single-stream serving). Batch>1 is untested on this line. + +**Software** +- vLLM **0.27.x** (the adapter patches `GPUModelRunner.load_model` and + reads the FlashInfer backend's metadata layout of that series). +- SGLang **0.5.x** with the FlashInfer attention backend, run from its + official container image. +- torch **2.13 + cu130**: the published kernel build variant consumed + here is `torch213-cxx11-cu130-x86_64-linux`. A different torch/CUDA + pair needs the matching build variant on the hub (or mounted via + `FRT_KERNEL_DIR_*`, see the SGLang section). +- `kernels` (huggingface) library for hub resolution on the vLLM host. + Air-gapped SGLang containers skip it entirely through + `FRT_KERNEL_DIR_*`. + +**Model / checkpoint** +- Measured host: **`RadixArk/Qwen3.8-27B-NVFP4`** (Hugging Face) — the + NVIDIA Model Optimizer mixed-precision quantization of + [`Qwen/Qwen3.8-27B`](https://huggingface.co/Qwen/Qwen3.8-27B) + (Qwen3.5 backbone: 48 gated-delta + 16 full-attention layers). + Format the adopt tiers read: NVFP4-packed projections (uint8 nibble + pairs + FP8 block-16 scales + a per-tensor global scale), FP8 rows + on the gated-delta projections, a W4A16 head, and FP8 KV cache + (`kv_cache_dtype fp8_e4m3` resolved from the checkpoint config). +- Speculative draft for SGLang: + **`RadixArk/Qwen3.8-27B-DSpark`** (Hugging Face) — the DSpark + speculator for this model family, served with SGLang's + `--speculative-algorithm DSPARK`. vLLM's MTP path needs no separate + download (the MTP head ships inside the target checkpoint). +- The `auto`/`mirror` tiers *adopt* these packs, so this checkpoint + family is what they expect; a dense BF16 checkpoint works through + the re-grid tiers (`nvfp4`, `w8`) instead. +- Speculative arms: vLLM MTP (`method qwen3_5_mtp`, + `num_speculative_tokens=6` — measured optimum for this model; K≥7 + loses throughput at every context) / SGLang DSpark with its + published draft model. + +Real code/text prompts throughout — repeated/synthetic prompts +inflate speculative acceptance and void any spec-arm number. + +## vLLM + +### Attach + +```python +from flash_rt.structures.adapters import vllm_engine + +vllm_engine.install_load_hook( + seats=vllm_engine.DENSE_SEAT_SUFFIXES, # projection positions + precision="auto", # adopt the host's packs; carry FP8 rows to W4 + consume=True, # release replaced host weights (KV pool grows) + head=True, # LM head relay (see spec-decode note below) + fused_mlp=True) + +from vllm import LLM +llm = LLM(model=..., ...) # boot normally; seats install during load +``` + +The hook patches the model runner's `load_model` and attaches between +weight load and the engine's first trace — the only window where a +compiled vLLM host accepts a module swap. + +### Stock CLI serving (`vllm serve`), zero code changes + +A stock OpenAI-compatible server attaches the same way SGLang does: +a `sitecustomize.py` rides `PYTHONPATH` into the server's processes +and installs the load hook when its env gate is set. + +```python +# /sitecustomize.py +import os +if os.environ.get("FRT_VLLM_ATTACH") == "1": + import sys + p = os.environ.get("FRT_VLLM_STRUCTURES_PATH") + if p and p not in sys.path: + sys.path.insert(0, p) + from flash_rt.structures.adapters import vllm_engine + vllm_engine.install_load_hook( + seats=vllm_engine.DENSE_SEAT_SUFFIXES, + precision=os.environ.get("FRT_VLLM_PRECISION", "auto"), + consume=os.environ.get("FRT_VLLM_CONSUME", "1") == "1", + seat_draft=False, + head=os.environ.get("FRT_VLLM_HEAD", "1") == "1", + fused_mlp=True) +``` + +```bash +PYTHONPATH= FRT_VLLM_ATTACH=1 FRT_VLLM_STRUCTURES_PATH= vllm serve --trust-remote-code ... +``` + +Measured on the running server (OpenAI completions API, MTP K=6, 2K +real prompts, same client and boots back to back): + +| | stock serve | attached serve | +|---|---|---| +| code decode-only | 195.6 tok/s | **210.2 (+7.5%)** | +| text decode-only | 129.3 tok/s | **141.0 (+9.0%)** | +| TTFT | ~141 ms | **~123 ms (−13%)** | + +The spec-decode head rule applies unchanged (the server carries a +speculative config, so the Marlin head relay stands aside on its +own). + +### Precision tiers + +| `precision=` | behavior | +|---|---| +| `"auto"` | adopt NVFP4 packs zero-copy; re-grid the checkpoint's FP8 rows to W4 (a precision change, reported as such) | +| `"mirror"` | faithful per-position: NVFP4 adopted, FP8 stays FP8 | +| `"nvfp4"` | re-grid everything to W4 from dense rows | +| `"w8"` | per-channel weight-only INT8 from dense rows | + +### Environment knobs + +| env | default | meaning | +|---|---|---| +| `FRT_HEAD_MARLIN` | `1` (auto-`0` under spec decode) | relay the checkpoint's W4A16 head into the small-M Marlin layout | +| `FRT_ATTN_XQA` | `0` | route spec-verify attention through the paged BF16-Q/FP8-KV kernel | +| `FRT_MOE_BAND_T` | `16` | rows above this go to the host expert path | + +### Speculative decode: the head relay steps aside + +vLLM's MTP draft shares the target's `lm_head`, and acceptance is an +agreement test between draft and target. Relaying only the target's +head to a different (even more accurate) numeric path breaks that +agreement: judged paired over ten prompts, acceptance length landed +lower on all ten. With a `speculative_config` present the relay +therefore defaults off; set `FRT_HEAD_MARLIN=1` explicitly to force +it. Without spec decode it stays on — pure step-rate, no acceptance +to protect. + +### Long-context serving (the released memory becomes KV) + +With `consume=True` the replaced host weights are released and the KV +pool grows by that amount. Measured on the 27B host: stock vLLM tops +out near a 102K context on the 32 GB card; the attached engine boots +the model's native 262144 and generates at 200K. + +Two settings matter at the top of that range: + +```python +LLM(..., + max_model_len=262144, + kv_cache_memory_bytes=10_400_000_000, # explicit KV budget + max_num_batched_tokens=4096) # halves prefill activation peak +``` + +Sizing KV implicitly through `gpu_memory_utilization` leaves too +little headroom for long-prefill activation transients (measured OOM +at 200K with utilization-based sizing); an explicit KV budget with +1.5 GB+ left free is the configuration that survives. + +The same effect on SGLang (DSpark serving, real code prompts; both +arms carry the draft model card's own flags — `dspark-block-size 7`, +draft `unquant`, `mamba-radix-cache-strategy extra_buffer`): + +| context | stock server (mem 0.92) | attached server (mem 0.92) | +|---|---|---| +| KV pool (`max_total_num_tokens`) | 34,659 | **87,128 (2.51x)** | +| 32K decode | 205.4 tok/s | 200.6 tok/s | +| 60K request | **refused** (exceeds pool) | **147.7 tok/s** (AL 3.05) | +| 80K request | refused | **210.8 tok/s** (AL 4.49) | + +The stock arm was tuned before this table was written; the sweep is +recorded so it can be checked and improved on: + +| stock configuration attempted | pool (tokens) | long request | +|---|---|---| +| mem 0.85 (runner default) | — | refuses to boot with the draft | +| mem 0.92 + draft card's flags | 32,661–34,659 | 32K serves; 60K refused | +| + CUDA-graph trim (decode bs cap, no prefill graphs) | 32,661 | unchanged | +| + `context-length` 147456 | 55,822* | — | +| + `mamba-track-interval` 1024 | 55,334* | — | +| mem 0.95 | 55,334* | **48K fails server-side** (chunk 2048 and 1024) | +| mem 0.97 | — | fails during graph capture | + +\* paper pool only: at 0.95 the added fraction is exactly the +runtime headroom long prefills need, and the attached arm cannot +boot there either — **0.92 is the stable envelope for both arms**, +which is what the comparison table uses. The hybrid line's 2.25 GB +intermediate state cache is insensitive to the graph, context-length, +and track-interval knobs. Higher +single-server context figures published for this model family come +from multi-GPU serving (the model card's own recipe is `tp-size 4`, +which divides weight memory per GPU); on one 32 GB card, within the +envelope we covered, the released weight memory is the working lever +— and it multiplies the stock pool by ~2.5x. Decode at the new +lengths rides acceptance-length content variance like every +speculative number in this document; the receipt is that the band +exists at full speed at all, where the stock server refuses the +request. + +### Judging protocol (what the receipts require) + +- Speculative arms drift to different greedy continuations per arm, and + acceptance length rides content. Cross-arm AL comparisons are only + valid **paired by prompt** (same context length list, per-point + pairing); single-prompt A/B buries systematic effects in content + variance. +- At long contexts (32K+), decode and AL swing across boots; judge on + medians of three or more boots, and prefer the step-rate column + (decode ÷ AL) where continuations differ. +- Long-context prompts must be real text end to end. Tiled/repeated + prompts inflate draft acceptance far above honest values. Measure + decode as a difference of two long generations + (`(n₂-n₁)/(t₂-t₁)`, `ignore_eos=True`) so prefill jitter and early + stops cannot pollute the number. + +## SGLang + +The scheduler spawns worker processes, so the hook travels as a +`sitecustomize.py` on `PYTHONPATH` and activates through env vars. +Seats are the linear-attention projections (the fused-MLP and LM-head +surfaces differ from vLLM's and are not seated). Inside an air-gapped +container the kernels load from mounted snapshot directories through +`FRT_KERNEL_DIR_` — no hub access needed at serve time. + +```bash +C=/root/.cache/huggingface/hub +docker run -d --name sgl-attach --gpus all --shm-size 32g --ipc=host \ + --network host \ + -v :/models \ + -v :/frt:ro \ + -v :/frt-hook:ro \ + -v :/root/.cache/huggingface \ + -e PYTHONPATH=/frt-hook \ + -e FRT_SGLANG_ATTACH=1 \ + -e FRT_SGLANG_STRUCTURES_PATH=/frt \ + -e FRT_SGLANG_PRECISION=auto \ + -e FRT_SGLANG_RELEASE=1 \ + -e "FRT_SGLANG_SEATS=linear_attn.out_proj,linear_attn.in_proj_qkvz" \ + -e FRT_KERNEL_DIR_FLASHRT_FP4_GEMM=$C/kernels--flashrt--fp4-gemm/snapshots//build/ \ + -e FRT_KERNEL_DIR_FLASHRT_FP4_FUSED_OPS=$C/kernels--flashrt--fp4-fused-ops/snapshots//build/ \ + -e FRT_KERNEL_DIR_FLASHRT_FLASHRT_FP8_FFN=$C/kernels--flashrt--flashrt-fp8-ffn/snapshots//build/ \ + -e FRT_KERNEL_DIR_FLASHRT_FLASHRT_GEMM_EPILOGUES=$C/kernels--flashrt--flashrt-gemm-epilogues/snapshots//build/ \ + \ + sglang serve --model-path /models/ --trust-remote-code \ + --mem-fraction-static 0.85 --attention-backend flashinfer \ + --chunked-prefill-size 2048 --disable-radix-cache \ + --max-running-requests 1 --host 0.0.0.0 --port 30000 +``` + +`` contains the bridge's `sitecustomize.py` (see +`flash_rt/structures/adapters/sglang_engine.py:install`, which writes +it). `` is the build directory for the serving container's +torch/CUDA pair, e.g. `torch213-cxx11-cu130-x86_64-linux`. + +**Speculative serving (DSpark)**: add + +``` +--speculative-algorithm DSPARK \ +--speculative-draft-model-path /models/ +``` + +and raise `--mem-fraction-static` to **0.90–0.93** — at 0.85 the +draft weights leave no room for the KV pool and the server refuses to +boot (measured; 0.92 is the configuration the receipts used). + +## Measured receipts (fresh paired baselines, single 5090) + +Decode tok/s, real code/text prompts, greedy; spec = MTP K=6 (vLLM) / +DSpark (SGLang). All baselines re-measured in the same window as the +attached arms. + +| | vLLM 2K | vLLM 32K | SGLang 2K | SGLang 32K | +|---|---|---|---|---| +| attach vs base (no spec) | +17–27% | +13–17% | +11% | +9–10% | +| attach vs base (spec) | +2–6% (paired) | code +52%, text −5–15% | +9–14% | code −2%, text +29% | +| TTFT | −11–21% | −11% | −14% | −8% | +| max context | — | 102K → 262144 native | — | — | +| 200K decode | — | 178 (code) / 175 (text) tok/s | — | — | + +Spec-arm decode columns ride acceptance-length content variance (see +judging protocol); the step-rate column is uniformly +9–13% for the +attached arms. diff --git a/examples/serving_attach_demo/README.md b/examples/serving_attach_demo/README.md new file mode 100644 index 000000000..f48819db1 --- /dev/null +++ b/examples/serving_attach_demo/README.md @@ -0,0 +1,78 @@ +# Demo: one GPU, native-length context + +Two terminals, the same `vllm serve` command. One refuses the model's +native 262144 context on a 32 GB card; the other — three environment +variables later — boots it and answers questions about a 200K-token +codebase at full speed. + +Requirements are the engine-attach guide's +(`docs/serving_engine_attach.md`): an SM120 GPU, vLLM 0.27.x, and the +`RadixArk/Qwen3.8-27B-NVFP4` checkpoint. + +## Scene 1 — the ceiling + +```bash +MODEL= ./serve_stock.sh +``` + +The stock server refuses at boot: the KV cache the native context +needs does not fit next to the weights. The error message is the +demo's opening shot — the engine itself saying 262144 is out of +reach, and estimating ~102K as the best it could do. + +## Scene 2 — three env vars + +```bash +MODEL= FRT_REPO= ./serve_attached.sh +``` + +Same command underneath; the hook attaches during load, the replaced +weights are released, and the freed memory becomes the KV pool the +native context needs. The boot log shows the seats installing and the +server comes up at `max_model_len 262144`. + +## Scene 3 — ask the codebase + +```bash +python ask.py --corpus --ctx 200000 \ + --question "Summarize the architecture of this codebase." +``` + +The client streams the answer with a live decode meter. Measured on +one RTX 5090: 200K-token prompts answer at ~170-180 tok/s decode +after a ~68 s prefill; at 2K the same server does ~210 tok/s on code +continuations with TTFT around 120 ms. + +The same attach also serves shorter contexts faster than the stock +boot (see the guide's tables) — the demo's closing line: nothing was +converted, nothing forked; detach and the host is untouched. + + +## Scene 4 — the context race (the watchable cut) + +A single growing conversation instead of one giant prompt: each turn +appends the next slice of the codebase and asks for the running +summary to be updated. With prefix caching on, every turn prefills +only its increment — the pace stays conversational (~2 s a turn at +the start), and the context meter just climbs. + +```bash +# arm A: the best context the stock server can boot (~102K here) +MODEL= ./serve_stock_race.sh +python race.py --arm stock --corpus + +# arm B: the attached server at the native 262144 +MODEL= FRT_REPO= ./serve_attached_race.sh +python race.py --arm attach --corpus +``` + +The stock arm walks its meter up and dies at its ceiling with the +server's own 400 in the table — that row is the money shot. The +attached arm walks the same turns past it to the native maximum. One +32 GB card cannot hold two copies of the weights, so the arms are +recorded separately and cut side by side; the client prints the same +table either way, so the timelines align turn for turn. + +The race form runs without speculative decode (this vLLM series +disables prefix caching under it); the one-shot scenes above carry +the speculative form and its speed numbers. diff --git a/examples/serving_attach_demo/ask.py b/examples/serving_attach_demo/ask.py new file mode 100644 index 000000000..58f4c344b --- /dev/null +++ b/examples/serving_attach_demo/ask.py @@ -0,0 +1,84 @@ +"""Streaming demo client with a live tok/s meter. + +Feeds the server a real corpus (a directory of source files) up to +--ctx tokens, asks one question about it, and streams the answer with +a running decode-rate readout — the visible half of the demo. + + python ask.py --corpus --ctx 200000 \ + --question "Summarize the architecture of this codebase." +""" + +import argparse +import json +import pathlib +import time +import urllib.request + + +def build_prompt(corpus: str, ctx_chars: int) -> str: + parts, total = [], 0 + for p in sorted(pathlib.Path(corpus).rglob("*")): + if p.suffix not in (".py", ".md", ".h", ".cu", ".cpp", ".txt"): + continue + try: + t = p.read_text() + except Exception: # noqa: BLE001 + continue + if not t.strip(): + continue + parts.append(f"\n===== {p.name} =====\n{t}") + total += len(t) + if total >= ctx_chars: + break + return "".join(parts)[:ctx_chars] + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--corpus", required=True) + ap.add_argument("--ctx", type=int, default=200000, + help="approximate prompt size in tokens") + ap.add_argument("--question", + default="Summarize the architecture of this " + "codebase in ten bullet points.") + ap.add_argument("--url", default="http://127.0.0.1:8000") + ap.add_argument("--max-tokens", type=int, default=512) + args = ap.parse_args() + + body = build_prompt(args.corpus, args.ctx * 4) # ~4 chars/token + prompt = (f"Here is a codebase:\n{body}\n\n{args.question}\n") + payload = {"model": "demo", "prompt": prompt, "stream": True, + "max_tokens": args.max_tokens, "temperature": 0} + req = urllib.request.Request( + args.url + "/v1/completions", json.dumps(payload).encode(), + {"Content-Type": "application/json"}) + t0 = time.perf_counter() + first = None + n = 0 + with urllib.request.urlopen(req, timeout=3600) as r: + for line in r: + line = line.decode().strip() + if not line.startswith("data:") or line == "data: [DONE]": + continue + chunk = json.loads(line[5:]) + txt = chunk["choices"][0].get("text", "") + if not txt: + continue + now = time.perf_counter() + if first is None: + first = now + print(f"\n[TTFT {first - t0:.1f}s]\n", flush=True) + n += 1 + rate = n / (now - first) if now > first else 0.0 + print(txt, end="", flush=True) + if n % 32 == 0: + print(f" \033[36m[{rate:.0f} tok/s]\033[0m", + end="", flush=True) + if first is not None: + rate = n / (time.perf_counter() - first) + print(f"\n\n[done: {n} tokens, decode {rate:.1f} tok/s, " + f"TTFT {first - t0:.1f}s]", flush=True) + + +if __name__ == "__main__": + main() diff --git a/examples/serving_attach_demo/race.py b/examples/serving_attach_demo/race.py new file mode 100644 index 000000000..9be1f6bf0 --- /dev/null +++ b/examples/serving_attach_demo/race.py @@ -0,0 +1,96 @@ +"""The context race: one growing conversation until the server's +ceiling — run once against each arm and cut the recordings side by +side (one 32 GB card cannot hold two copies of the weights at once). + +Each turn appends the next slice of a real codebase and asks for the +running summary to be updated; with prefix caching on, every turn +prefills only its increment, so the pace stays conversational. The +stock arm dies at its ceiling with the server's own error — that row +prints as the crash line and the run stops; the attached arm walks on +to the model's native maximum. + + python race.py --arm stock --corpus + python race.py --arm attach --corpus +""" + +import argparse +import json +import pathlib +import time +import urllib.error +import urllib.request + + +def corpus_slices(corpus, chars_per_step): + buf = [] + for p in sorted(pathlib.Path(corpus).rglob("*")): + if p.suffix not in (".py", ".md", ".h", ".cu", ".cpp", ".txt"): + continue + try: + t = p.read_text() + except Exception: # noqa: BLE001 + continue + if t.strip(): + buf.append(f"\n===== {p.name} =====\n{t}") + text = "".join(buf) + for i in range(0, len(text), chars_per_step): + yield text[i:i + chars_per_step] + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--arm", choices=("stock", "attach"), + required=True) + ap.add_argument("--corpus", required=True) + ap.add_argument("--url", default="http://127.0.0.1:8000") + ap.add_argument("--step-tokens", type=int, default=8000) + ap.add_argument("--answer-tokens", type=int, default=64) + ap.add_argument("--max-ctx", type=int, default=262144) + args = ap.parse_args() + + slices = corpus_slices(args.corpus, args.step_tokens * 4) + history = "" + approx = 0 + print(f"{'turn':>4} {'~context':>9} {'TTFT':>6} {'rate':>9} " + f"status", flush=True) + turn = 0 + while approx < args.max_ctx - args.step_tokens: + try: + history += next(slices) + except StopIteration: + break + turn += 1 + approx += args.step_tokens + prompt = (f"Codebase so far:\n{history}\n\nUpdate your " + f"one-paragraph running summary of everything " + f"above.\n") + payload = {"model": "demo", "prompt": prompt, + "max_tokens": args.answer_tokens, + "min_tokens": max(16, args.answer_tokens // 2), + "temperature": 0} + req = urllib.request.Request( + args.url + "/v1/completions", + json.dumps(payload).encode(), + {"Content-Type": "application/json"}) + t0 = time.perf_counter() + try: + with urllib.request.urlopen(req, timeout=1800) as r: + out = json.loads(r.read()) + except urllib.error.HTTPError as e: + msg = e.read().decode()[:120] + print(f"{turn:>4} {approx:>9} {'—':>6} {'—':>9} " + f"\033[31m💥 {e.code}: {msg}\033[0m", flush=True) + print(f"\n[{args.arm}] ceiling at ~{approx} tokens.", + flush=True) + return + dt = time.perf_counter() - t0 + n = out["usage"]["completion_tokens"] + rate = n / dt if dt > 0 else 0.0 + print(f"{turn:>4} {approx:>9} {dt:>5.1f}s " + f"{rate:>6.1f}t/s \033[32mok\033[0m", flush=True) + print(f"\n[{args.arm}] reached ~{approx} tokens — native window " + f"served end to end.", flush=True) + + +if __name__ == "__main__": + main() diff --git a/examples/serving_attach_demo/serve_attached.sh b/examples/serving_attach_demo/serve_attached.sh new file mode 100755 index 000000000..27badc230 --- /dev/null +++ b/examples/serving_attach_demo/serve_attached.sh @@ -0,0 +1,15 @@ +#!/bin/bash +# Same server, three env vars more. FRT_REPO points at the FlashRT +# checkout. The long-context boot adds the explicit KV budget the +# docs describe (utilization-based sizing leaves no activation +# headroom at 200K). +MODEL=${MODEL:?set MODEL to the checkpoint path or hub id} +FRT_REPO=${FRT_REPO:?set FRT_REPO to the FlashRT checkout} +CTX=${CTX:-262144} +HERE=$(cd "$(dirname "$0")" && pwd) +PYTHONPATH="$HERE" FRT_VLLM_ATTACH=1 FRT_VLLM_STRUCTURES_PATH="$FRT_REPO" \ +vllm serve "$MODEL" --served-model-name demo --port 8000 \ + --trust-remote-code --max-model-len "$CTX" \ + --gpu-memory-utilization 0.90 --max-num-seqs 1 \ + --kv-cache-memory 10400000000 --max-num-batched-tokens 4096 \ + --speculative-config '{"method":"qwen3_5_mtp","num_speculative_tokens":6}' diff --git a/examples/serving_attach_demo/serve_attached_race.sh b/examples/serving_attach_demo/serve_attached_race.sh new file mode 100755 index 000000000..b478e0c43 --- /dev/null +++ b/examples/serving_attach_demo/serve_attached_race.sh @@ -0,0 +1,13 @@ +#!/bin/bash +# Race arm B: the attached server at the model's native maximum, +# same flags otherwise. Three env vars are the whole difference. +MODEL=${MODEL:?set MODEL to the checkpoint path or hub id} +FRT_REPO=${FRT_REPO:?set FRT_REPO to the FlashRT checkout} +CTX=${CTX:-262144} +HERE=$(cd "$(dirname "$0")" && pwd) +PYTHONPATH="$HERE" FRT_VLLM_ATTACH=1 FRT_VLLM_STRUCTURES_PATH="$FRT_REPO" \ +vllm serve "$MODEL" --served-model-name demo --port 8000 \ + --trust-remote-code --max-model-len "$CTX" \ + --gpu-memory-utilization 0.90 --max-num-seqs 1 \ + --kv-cache-memory 10400000000 --max-num-batched-tokens 4096 \ + --enable-prefix-caching diff --git a/examples/serving_attach_demo/serve_stock.sh b/examples/serving_attach_demo/serve_stock.sh new file mode 100755 index 000000000..6c94c6ad0 --- /dev/null +++ b/examples/serving_attach_demo/serve_stock.sh @@ -0,0 +1,10 @@ +#!/bin/bash +# Stock vLLM server. MODEL points at the checkpoint directory or hub id. +# At CTX=262144 on a 32 GB card this boot REFUSES — that refusal is +# scene one of the demo. +MODEL=${MODEL:?set MODEL to the checkpoint path or hub id} +CTX=${CTX:-262144} +vllm serve "$MODEL" --served-model-name demo --port 8000 \ + --trust-remote-code --max-model-len "$CTX" \ + --gpu-memory-utilization 0.90 --max-num-seqs 1 \ + --speculative-config '{"method":"qwen3_5_mtp","num_speculative_tokens":6}' diff --git a/examples/serving_attach_demo/serve_stock_race.sh b/examples/serving_attach_demo/serve_stock_race.sh new file mode 100755 index 000000000..23f55bf5b --- /dev/null +++ b/examples/serving_attach_demo/serve_stock_race.sh @@ -0,0 +1,12 @@ +#!/bin/bash +# Race arm A: the stock server at the best context it can actually +# boot on this card (its own error at 262144 estimates ~102K), with +# prefix caching on so each turn prefills only its increment. The +# race form runs without speculative decode (this vLLM series +# disables prefix caching under it). +MODEL=${MODEL:?set MODEL to the checkpoint path or hub id} +CTX=${CTX:-102400} +vllm serve "$MODEL" --served-model-name demo --port 8000 \ + --trust-remote-code --max-model-len "$CTX" \ + --gpu-memory-utilization 0.90 --max-num-seqs 1 \ + --enable-prefix-caching diff --git a/examples/serving_attach_demo/sitecustomize.py b/examples/serving_attach_demo/sitecustomize.py new file mode 100644 index 000000000..3b62c05a0 --- /dev/null +++ b/examples/serving_attach_demo/sitecustomize.py @@ -0,0 +1,26 @@ +"""Attach hook for a stock ``vllm serve`` process, env-gated. + +Put this file's directory on PYTHONPATH and set FRT_VLLM_ATTACH=1; +the server's processes install the load hook on their own. Without +the gate the file is inert. +""" +import os + +if os.environ.get("FRT_VLLM_ATTACH") == "1": + import sys + p = os.environ.get("FRT_VLLM_STRUCTURES_PATH") + if p and p not in sys.path: + sys.path.insert(0, p) + try: + from flash_rt.structures.adapters import vllm_engine + vllm_engine.install_load_hook( + verbose=True, + seats=vllm_engine.DENSE_SEAT_SUFFIXES, + precision=os.environ.get("FRT_VLLM_PRECISION", "auto"), + consume=os.environ.get("FRT_VLLM_CONSUME", "1") == "1", + seat_draft=False, + head=os.environ.get("FRT_VLLM_HEAD", "1") == "1", + fused_mlp=True) + print("[flash_rt] vllm serve attach hook installed", flush=True) + except Exception as e: # noqa: BLE001 — the server must still boot + print(f"[flash_rt] attach hook failed: {e!r}", flush=True) diff --git a/flash_rt/structures/__init__.py b/flash_rt/structures/__init__.py index fce219b35..a9a66d70f 100644 --- a/flash_rt/structures/__init__.py +++ b/flash_rt/structures/__init__.py @@ -106,7 +106,7 @@ def explain(plan): def decode_loop(model, *, max_len, compile_step=True, - compile_prefill=True, kv_band=None): + compile_prefill=True, kv_band=None, prefill_chunk=None): """Serving door: the whole-loop decode form (static cache + compiled step + whole-step CUDA graph) over whatever structures are attached. @@ -118,7 +118,8 @@ def decode_loop(model, *, max_len, compile_step=True, return build_decode_loop(model, max_len=max_len, compile_step=compile_step, compile_prefill=compile_prefill, - kv_band=kv_band) + kv_band=kv_band, + prefill_chunk=prefill_chunk) def aot_package(module, args=(), kwargs=None, diff --git a/flash_rt/structures/adapters/sglang_engine.py b/flash_rt/structures/adapters/sglang_engine.py index 9dcc9b8ea..904697c0d 100644 --- a/flash_rt/structures/adapters/sglang_engine.py +++ b/flash_rt/structures/adapters/sglang_engine.py @@ -159,9 +159,25 @@ def load_model(self, *a, **kw): orig(self, *a, **kw) seats = tuple(s for s in os.environ.get(_SEATS_VAR, "").split(",") if s) or DENSE_SEAT_SUFFIXES + precision = os.environ.get("FRT_SGLANG_PRECISION") try: - attach_engine(self.model, seats=seats, - release=os.environ.get("FRT_SGLANG_RELEASE") == "1") + if precision: + # this engine's layers share vLLM's module lineage down + # to the quant-method attribute names, so the precision + # tiers assemble here unchanged: positions already in a + # seam's width adopt the host's own tensors, and the + # auto tier's W8->W4 carry sources from the host's FP8 + # rows. The engine-specific part of this adapter stays + # what it was: the process-boundary carrier. + from .vllm_engine import attach_engine as _tiered + _tiered(self.model, seats=seats, precision=precision, + head=False, tag="sglang", + consume=os.environ.get( + "FRT_SGLANG_RELEASE") == "1") + else: + attach_engine( + self.model, seats=seats, + release=os.environ.get("FRT_SGLANG_RELEASE") == "1") except Exception as e: print(f"[structures.sglang] attach refused: {e!r}", flush=True) mr.ModelRunner.load_model = load_model diff --git a/flash_rt/structures/adapters/vllm_engine.py b/flash_rt/structures/adapters/vllm_engine.py index a529b5a3b..b7c303c77 100644 --- a/flash_rt/structures/adapters/vllm_engine.py +++ b/flash_rt/structures/adapters/vllm_engine.py @@ -75,15 +75,358 @@ def _(hidden, router_logits, top_idx, top_w, idx): return torch.empty_like(hidden) +_E2M1_LUT = None +_SRC_INDEX = None + +#: engine fused-projection composition, in the engine's own concat +#: order (the split sites in its forwards are the receipts) +_FUSE = { + "in_proj_qkvz": ("in_proj_qkv", "in_proj_z"), + "qkv_proj": ("q_proj", "k_proj", "v_proj"), + "gate_up_proj": ("gate_proj", "up_proj"), + "in_proj_ba": ("in_proj_b", "in_proj_a"), +} + + +def _source_ckpt_weight(seat_name): + """Full-precision rows for a seat, read from ``FRT_SOURCE_CKPT``. + + ``seat_name`` is the engine module path; the checkpoint key drops + the engine's leading ``language_model.`` and prefixes ``model.``. + Fused engine projections concat their checkpoint constituents in + the engine's own order. Returns None when the env is unset or the + key cannot be resolved — the caller falls through to pack-specific + dequant, and its bind probe stays the last word either way. + """ + import json + import os + + root = os.environ.get("FRT_SOURCE_CKPT") + if not root or not seat_name: + return None + global _SRC_INDEX + if _SRC_INDEX is None: + from safetensors import safe_open + idx = json.load(open(os.path.join( + root, "model.safetensors.index.json"))) + _SRC_INDEX = (idx["weight_map"], {}, root, safe_open) + wmap, handles, root, safe_open = _SRC_INDEX + + def read(key): + fn = wmap.get(key) + if fn is None: + return None + if fn not in handles: + handles[fn] = safe_open(os.path.join(root, fn), + framework="pt", device="cpu") + return handles[fn].get_tensor(key) + + cands = [seat_name] + if seat_name.startswith("language_model.model."): + cands.append("model.language_model." + + seat_name[len("language_model.model."):]) + if seat_name.startswith("model."): + cands.append("model.language_model." + + seat_name[len("model."):]) + cands.append("model." + seat_name) + for path in cands: + leaf = path.rsplit(".", 1)[-1] + parts = _FUSE.get(leaf) + if parts is None: + t = read(path + ".weight") + if t is not None: + return t.to("cuda", torch.bfloat16) + continue + base = path.rsplit(".", 1)[0] + pieces = [read(f"{base}.{p}.weight") for p in parts] + if all(p is not None for p in pieces): + return torch.cat( + [p.to("cuda", torch.bfloat16) for p in pieces], + dim=0).contiguous() + return None + + +#: FP8 weight dtypes a checkpoint may store a projection in +_FP8_DTYPES = tuple( + d for d in (getattr(torch, "float8_e4m3fn", None), + getattr(torch, "float8_e5m2", None)) if d is not None) + + +def _host_precision(mod) -> str: + """Which precision the *checkpoint* put this projection in. + + A mixed-precision checkpoint is a per-projection decision that was + made with calibration data — attention projections held at FP8 while + the FFN goes to NVFP4, a draft head excluded from quantization + entirely. Seating every position at the adapter's favourite width + overrides that decision silently: the weight stream shrinks, the + step gets faster, and the accuracy it cost shows up somewhere no + throughput number looks. Read the choice off what the host is + actually holding rather than off module names, which is the same + discipline the rest of this adapter uses for seam recognition. + """ + w = getattr(mod, "weight", None) + if w is None: + return "none" + dt = w.data.dtype + if dt is torch.uint8: + return "nvfp4" + if dt in _FP8_DTYPES: + return "fp8" + return "unquantized" + + +def _projection_weight(mod) -> torch.Tensor: + """The projection's dense rows, whatever the checkpoint stored. + + A bf16/fp16 weight is the rows. A uint8 weight is a modelopt NVFP4 + pack (e2m1 nibble pairs + fp8 block-16 scales + a global scalar): + dequantize it here so the seam re-grids from real values. Nibble + order inside the byte is resolved by the caller's bind probe — this + helper emits the low-nibble-first convention, and a probe failure + is a refusal, never silent garbage. + """ + w = mod.weight.data + if w.dtype in (torch.bfloat16, torch.float16, torch.float32): + return w + if w.dtype in _FP8_DTYPES: + # an FP8 host needs no source checkpoint: every FP8 code is + # exact in BF16, so the host's own tensors are the rows — + # dequantization here loses nothing. The engine stores this + # weight transposed ([K, N], its cutlass B), so orientation + # comes from the layer's declaration, same as the adopt path. + ws = getattr(mod, "weight_scale", None) + if ws is None: + raise ValueError("fp8 weight without weight_scale") + n = int(getattr(mod, "output_size_per_partition", 0) or 0) + k = int(getattr(mod, "input_size_per_partition", 0) or 0) + rows = w.to(torch.float32) * ws.data.reshape(-1)[0].float() + if n and k and tuple(w.shape) == (k, n): + rows = rows.t() + return rows.to(torch.bfloat16).contiguous() + src = _source_ckpt_weight(getattr(mod, "_frt_seat_name", None)) + if src is not None: + # a quantized runtime weight, but the caller pointed + # FRT_SOURCE_CKPT at the full-precision checkpoint: re-grid + # from the real rows (the adopt door's own discipline) instead + # of dequantizing whatever runtime pack the engine chose + return src + if w.dtype is not torch.uint8: + raise ValueError(f"unrecognised weight dtype {w.dtype}") + global _E2M1_LUT + if _E2M1_LUT is None: + _E2M1_LUT = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, + -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], + device=w.device, dtype=torch.float32) + ws = mod.weight_scale.data.to(torch.float32) # [N, K/16] + ws2 = getattr(mod, "weight_scale_2", None) + g = (float(ws2.data.reshape(-1)[0]) if ws2 is not None else 1.0) + n = w.shape[0] + lo = (w & 0xF).long() + hi = (w >> 4).long() + codes = torch.stack([lo, hi], dim=-1).reshape(n, -1) + vals = _E2M1_LUT[codes] # [N, K] + scale = ws.repeat_interleave(16, dim=1) * g + return (vals * scale).to(torch.bfloat16).contiguous() + + +def _adopt_fp8_pack(mod, rows_hint): + """A seat on the host's own FP8 tensors, or None. + + Same hot-plug contract as the NVFP4 adopt: the engine already holds + the weight as ``float8_e4m3fn`` with the checkpoint's calibrated + per-tensor scales, and that per-tensor W8A8 scheme is exactly what + ``fp8_static`` executes — so the seat adopts the storage by + reference and changes only which kernel reads it. No dequantize, no + re-quantize, no second copy on the card. + """ + from ..impls.linear_proj import fp8_static as _fp8 + + w = getattr(mod, "weight", None) + ws = getattr(mod, "weight_scale", None) + xs = getattr(mod, "input_scale", None) + if w is None or ws is None or xs is None: + return None + if w.data.dtype not in _FP8_DTYPES: + return None + # the engine stores this weight transposed — [K, N], a column-major + # B for its cutlass entry — so the projection's dims must come from + # the layer's own declaration, not the storage order + n = int(getattr(mod, "output_size_per_partition", 0) or 0) + k = int(getattr(mod, "input_size_per_partition", 0) or 0) + if not n or not k: + return None + for name, dim in (("K", k), ("N", n)): + lo = _fp8.SUPPORT[name]["min"] + hi = _fp8.SUPPORT[name]["max"] + if not lo <= dim <= hi: + raise ValueError( + f"fp8 adopt: {name}={dim} outside support envelope") + if tuple(w.data.shape) == (k, n): + # one transposed copy, then the original storage goes: after + # seating, this seam is the weight's only consumer, so total + # memory is unchanged and the transient peak is one projection. + # The host module keeps a [K, N] *view* of the new storage — + # same logical content, so detach still reads correct values. + w_nk = w.data.t().contiguous() + mod.weight.data = w_nk.t() + elif tuple(w.data.shape) == (n, k): + w_nk = w.data + else: + raise ValueError( + f"fp8 adopt: storage {tuple(w.data.shape)} matches neither " + f"orientation of N={n} K={k}") + ms = sorted(int(m) for m in rows_hint) + form = _fp8._form_for(None, "bf16", + float(ms[len(ms) // 2]) * n * k) + # the form bands were measured against a BF16-Linear host; against + # an engine whose FP8 path fuses its own quantize they are not + # gospel, so an explicit override stays available for A/B + form = os.environ.get("FRT_FP8_FORM", form) + bias = torch.zeros(n, device=w_nk.device, dtype=torch.bfloat16) + return _fp8.FusedLinearProj( + w_nk, bias, + xs.data.reshape(-1)[:1].to(torch.float32), + ws.data.reshape(-1)[:1].to(torch.float32), + original=None, form=form) + + +def _adopt_nvfp4_pack(mod): + """Build a seam on the host's own packed weights, or return None. + + This is the hot-plug form: the engine already holds NVFP4, so the + seat changes which kernel reads it and nothing else. Numerics stay + the checkpoint's — which for a speculative host means acceptance + stays the checkpoint's too, and every millisecond the seat saves is + kept rather than paid back as a lower accept rate. + + Returns None when this host's pack is not in a shape the seam can + adopt, so the caller falls back to regridding and the difference is + a reported choice rather than a silent one. + """ + w = getattr(mod, "weight", None) + ws = getattr(mod, "weight_scale", None) + if w is None or ws is None or w.data.dtype is not torch.uint8: + return None + w = w.data + # the engine's own kernel may pad the packed columns for its tile + # shape; those columns are not part of the projection + pad = int(getattr(mod, "weights_padding_cols", 0) or 0) + if pad: + w = w[:, :w.shape[1] - pad] + gs = None + for attr in ("weight_global_scale", "weight_scale_2"): + v = getattr(mod, attr, None) + if v is not None: + flat = v.data.reshape(-1).float() + if flat.numel() > 1 and not bool( + (flat == flat[0]).all()): + # a merged projection may carry one global factor per + # constituent; a single alpha can only stand in for + # them when they agree, and pretending otherwise would + # scale one half by the other's factor + raise ValueError( + "per-part global scales differ; cannot adopt " + "under one alpha") + gs = float(flat[0]) + break + n, k = w.shape[0], w.shape[1] * 2 + return _linear.bind_proj_seam_packed( + w.contiguous() if pad else w, ws.data, n, k, global_scale=gs) + + +def _bind_fp8_seam(mod, w, rows_hint): + """An FP8 seat for a projection the checkpoint quantized to FP8. + + The input scale is the checkpoint's own calibrated one where the + host carries it; a projection whose activations the checkpoint + scaled dynamically has none, and the amax of its bind probe would be + a statistic invented here rather than one measured on data, so that + case is refused instead of guessed. + """ + from ..impls.linear_proj import fp8_static as _fp8 + + s = getattr(mod, "input_scale", None) + if s is None: + raise ValueError( + "FP8 seat needs the checkpoint's calibrated input scale; " + "this projection carries none") + return _fp8.bind_proj_seam( + {"w": w}, input_scale=float(s.data.reshape(-1)[0]), + row_profile=list(rows_hint), original=mod) + + class _ProjSeat(nn.Module): - """Preserves the engine's ``(out, bias)`` projection contract.""" + """Preserves the engine's ``(out, bias)`` projection contract. + + Attribute reads fall through to the replaced module: an engine that + planned a fusion around this projection reads its scale attributes + off the module object itself, and a seat that answers only + ``forward`` turns that read into a startup crash far from here. + """ - def __init__(self, seam): + def __init__(self, seam, host=None): super().__init__() self.seam = seam + # the engine's linear layers carry a return_bias switch, and a + # call site built for return_bias=False receives a bare tensor + # — handing it the tuple anyway detonates at the next op, at + # trace time, far from here + object.__setattr__( + self, "_frt_tuple_out", + getattr(host, "return_bias", True) is not False) + if host is not None: + object.__setattr__(self, "_frt_host", host) + + def __getattr__(self, name): + try: + return super().__getattr__(name) + except AttributeError: + host = self.__dict__.get("_frt_host") + if host is None: + raise + return getattr(host, name) def forward(self, x, *args, **kwargs): - return self.seam(x), None + y = self.seam(x) + return (y, None) if self._frt_tuple_out else y + + +class _FusedMlpSeat(nn.Module): + """Stands where the host's SwiGLU MLP stood. + + The host's own dataflow is ``gate_up -> SiLU·mul -> down``, and it + already merges gate and up into one projection — so once both + projections carry seats, the activation between them is the only + step that still round-trips through BF16 and re-quantizes for the + down projection. This seat runs the merged projection through its + seam, collapses activation + quantization into the fused producer, + and hands the packed rows straight to the down seam. Anything the + host did around that (an expert gate) is kept by delegating to the + retained module for the parts this seat does not own. + """ + + def __init__(self, host, gate_up_seam, down_seam, silu_mul): + super().__init__() + self.host_mlp = host + self.gate_up_seam = gate_up_seam + self.down_seam = down_seam + self._silu_mul = silu_mul + + def forward(self, x): + shape = x.shape + flat = x.reshape(-1, shape[-1]) + a_packed, a_sfa = _linear._quantize_activation( + self.gate_up_seam._kern, flat) + merged = self.gate_up_seam._mm_packed(a_packed, a_sfa) + p2, s2 = self._silu_mul(merged.contiguous()) + out = self.down_seam._mm_packed(p2, s2) + out = out.reshape(*shape[:-1], self.down_seam._n).type_as(x) + gate = getattr(self.host_mlp, "expert_gate", None) + if gate is not None: + out = torch.sigmoid(gate(x)[0]) * out + return out class _MoESeat(nn.Module): @@ -144,7 +487,14 @@ def forward(self, hidden_states, router_logits): class _SlabbedHeadMethod: """Stands in for the LM head's quant method: the engine computes - logits through ``quant_method.apply``, never module forward.""" + logits through ``quant_method.apply``, never module forward. + + The vocabulary projection stays weight-only INT8 rather than joining + the FP4 band: it is the logits family, where the decision on this + model line has been W8 from the start. One seam covers the whole + vocabulary when the entry's row support reaches it; otherwise the + rows split into slabs and concatenate. + """ def __init__(self, seams, orig): self.seams = seams @@ -152,7 +502,8 @@ def __init__(self, seams, orig): def apply(self, layer, x, bias=None): xb = x.to(torch.bfloat16) - y = torch.cat([s(xb) for s in self.seams], dim=-1) + y = (self.seams[0](xb) if len(self.seams) == 1 + else torch.cat([s(xb) for s in self.seams], dim=-1)) if bias is not None: y = y + bias return y.to(x.dtype) @@ -161,12 +512,63 @@ def __getattr__(self, name): return getattr(self.orig, name) +class _MarlinHeadMethod: + """The checkpoint's own W4A16 head, re-laid for small M. + + Same NVFP4 grid, same dequantized values — only the storage layout + changes (bind-time Marlin repack) so the verify-step logits calls + (M ≤ 16) leave the host's tiled kernel for one built for that band. + The logits family keeps its precision class; anything outside the + band, and any capture pass, goes to the host method untouched. + """ + + def __init__(self, kernel, packs, rows, orig): + self.kernel = kernel + self.packs = packs # (weight, scale, global, lock) + self.rows = rows # vocabulary rows of the output + self.orig = orig + + def apply(self, layer, x, bias=None): + if (1 <= x.shape[0] <= 16 and x.dtype == torch.bfloat16 + and x.dim() == 2 + and not torch.cuda.is_current_stream_capturing()): + wm, wsm, wgs, lock = self.packs + out = torch.empty(x.shape[0], self.rows, + dtype=x.dtype, device=x.device) + self.kernel.nvfp4_w4a16_marlin_bf16( + x, wm, wsm, wgs, workspace=lock, out=out) + if bias is not None: + out = out + bias + return out + return self.orig.apply(layer, x, bias) + + def __getattr__(self, name): + return getattr(self.orig, name) + + def _is_projection(module) -> bool: w = getattr(module, "weight", None) return (isinstance(w, torch.Tensor) and w.dim() == 2 and hasattr(module, "quant_method")) +def _arm_runtime_dispatch(seam) -> None: + """Move this seam's tier dispatch to call time. + + This engine compiles a seam forward once per shape *range* and + replays it without re-evaluating shape guards, so the row count a + trace observed is not the row count a replay carries: a Python-level + ``if m >= N`` freezes the tracing sample's choice and then runs it + for every M (measured as a hard refusal from the M256 tier at engine + start, and — quieter and more expensive — as every decode row taking + the tiled GEMM instead of the GEMV). Registering the seam for + runtime dispatch puts the branch inside a custom op, so the trace + sees one opaque call and the tier is chosen on the shape the call + actually receives. Every tier stays available. + """ + _linear.register_runtime_dispatch(seam) + + def _expert_holder(module): for _, child in module.named_modules(): if torch.is_tensor(getattr(child, "w13_weight", None)): @@ -174,6 +576,206 @@ def _expert_holder(module): return None +#: install-once state for the verify-attention seat; keyed data that +#: the wrapped ``forward`` reads on every call lives here so a second +#: ``attach_verify_attention`` call is a no-op rather than a re-wrap. +_XQA_STATE: dict[str, Any] = {} + + +def attach_verify_attention(*, verbose: bool = True) -> dict | None: + """Route spec-verify attention through the paged BF16-Q/FP8-KV XQA + kernel; host FlashInfer keeps everything else. + + A speculative verify step is a handful of query rows against the + whole paged cache, and the host runs it on its chunked-prefill + kernel — at long context the most off-roofline structure in the + decode chain. This seat wraps the FlashInfer backend's ``forward`` + and takes only calls whose metadata is exactly a verify batch + (single request, no decode rows, causal, fp8 cache) on a layer + geometry the kernel supports; anything else, and any error inside + the seat's own path, falls through to the host untouched. + + Two engine facts carried here so callers do not rediscover them: + + - The kv-cache tensor arriving at ``forward`` is a permuted view. + Flattening it (``reshape``/``flatten``) silently copies the whole + pool per call; the zero-copy road is ``as_strided`` claims over + the raw storage, where the view's own stride metadata *is* the + address walk (page/token/head strides read off the view, V half + at ``+head_dim``). + - The prefill wrapper's ``_paged_kv_indices_buf`` is the kernel's + page table verbatim: kernel pages equal engine blocks, so prefix + reuse and non-contiguous block runs need no translation and no + contiguity guard. + """ + if _XQA_STATE.get("installed"): + return _XQA_STATE + try: + import vllm.v1.attention.backends.flashinfer as fi + except Exception as e: # noqa: BLE001 — host without this backend + if verbose: + print(f"[flash_rt] verify-attention seat refused: " + f"no FlashInfer backend ({e!r})") + return None + from ..impls import hub_kernel + try: + kx = hub_kernel("flashrt/fp8-kv-attention", "4") + except Exception as e: # noqa: BLE001 — no build for this host + if verbose: + print(f"[flash_rt] verify-attention seat refused: {e!r}") + return None + page = getattr(kx, "PAGE_SIZE", None) + if page != 32: + if verbose: + print(f"[flash_rt] verify-attention seat refused: kernel " + f"page size {page} != engine block size 32") + return None + supported = {tuple(c) for c in getattr(kx, "SUPPORTED_CONFIGS", ())} + + st = _XQA_STATE + st.update(installed=True, kernel=kx, taken=0, fell_through=0, + masks={}, sems=None, scratch=None, sl=None, sl_val=-1) + orig = fi.FlashInferImpl.forward + fp8 = torch.float8_e4m3fn + + def _stash_plan(wrap): + # The page count and last-page length the forward needs are on + # the device buffers, and reading them there is a host sync — + # sixteen of those per verify step erase the kernel's win by + # stalling the launch pipeline. The builder hands ``plan`` the + # same values as CPU tensors, so a wrap of ``plan`` stashes + # them as plain ints and the forward path never syncs. + orig_plan = wrap.plan + + def plan(*a, **kw): + try: + ind = kw.get("paged_kv_indptr") + lpl = kw.get("paged_kv_last_page_len") + if (ind is not None and lpl is not None + and not ind.is_cuda and ind.numel() == 2): + wrap._frt_plan_np = (int(ind[1]), int(lpl[0])) + else: + wrap._frt_plan_np = None + except Exception: # noqa: BLE001 + wrap._frt_plan_np = None + return orig_plan(*a, **kw) + + wrap.plan = plan + wrap._frt_plan_hooked = True + + def forward(self, layer, query, key, value, kv_cache, attn_metadata, + output, *args, **kwargs): + m = attn_metadata + take = ( + m is not None + # no fused output quant on the seat path + and not any(a is not None for a in args) + and kwargs.get("output_scale") is None + and kwargs.get("output_block_scale") is None + and hasattr(self, "do_kv_cache_update") + and getattr(m, "num_decode_tokens", 1) == 0 + and getattr(m, "num_prefills", 0) == 1 + and getattr(m, "causal", False) + and 2 <= getattr(m, "num_actual_tokens", 0) <= 32 + and m.prefill is not None + and query.dim() == 3 + and query.dtype == torch.bfloat16 + and kv_cache.dim() == 4 + and (query.shape[1], kv_cache.shape[1], + query.shape[2]) in supported + and kv_cache.shape[2] == 32 + and kv_cache.shape[3] == 2 * query.shape[2] + and getattr(self, "sinks", None) is None + and getattr(self, "window_left", -1) == -1 + and not getattr(self, "logits_soft_cap", None) + and abs(self.scale - query.shape[2] ** -0.5) < 1e-9 + and layer._k_scale_float == layer._v_scale_float + # capture bakes this call's shapes (page count, seq) into + # the replayed graph while the Python that refreshes them + # never runs again — the seat serves eager calls only + and not torch.cuda.is_current_stream_capturing() + ) + if take: + try: + wrap = m.prefill.wrapper + idx = getattr(wrap, "_paged_kv_indices_buf", None) + if idx is None: + raise LookupError("wrapper plan buffers absent") + meta_np = getattr(wrap, "_frt_plan_np", None) + if meta_np is None: + # first sighting of this wrapper: hook its plan so + # every later step has the counts sync-free; this + # step goes to the host + if not getattr(wrap, "_frt_plan_hooked", False): + _stash_plan(wrap) + raise LookupError("plan counts not stashed yet") + n_pages, last = meta_np + qs = int(m.num_actual_tokens) + # the host forward is also the cache writer: the new + # tokens' K/V must land in the paged pool before any + # attention (this step reads them, every later step + # reads them as prefix) — skipping it corrupts the + # sequence forever, not just this call + self.do_kv_cache_update(layer, key, value, kv_cache, + m.slot_mapping) + seq = (n_pages - 1) * 32 + last + if n_pages <= 0 or not (0 < last <= 32) or seq < qs: + raise ValueError("verify plan out of shape") + nq, hd = query.shape[1], query.shape[2] + kh = kv_cache.shape[1] + kv8 = (kv_cache if kv_cache.dtype == fp8 + else kv_cache.view(fp8)) + elems = 32 * kh * hd + + def claim(off): + return kv8.as_strided( + (n_pages, 32, kh, hd), + (elems, kh * hd, hd, 1), storage_offset=off) + + mask = st["masks"].get(qs) + if mask is None: + mask = kx.causal_spec_mask(qs, device=query.device) + st["masks"][qs] = mask + if st["sems"] is None: + st["sems"], st["scratch"] = kx.allocate_workspace( + q_seq=32, num_q_heads=nq, num_kv_heads=kh, + device=query.device) + st["sl"] = torch.zeros( + 1, 1, device=query.device, dtype=torch.int32) + if st["sl_val"] != seq: + st["sl"].fill_(seq) + st["sl_val"] = seq + kx.ops.xqa_bf16_fp8kv( + query[:qs], claim(0), claim(hd), + idx[:n_pages].view(1, -1), st["sl"], mask, + output[:qs], st["sems"], st["scratch"], + n_pages * 32, 1.0, float(layer._k_scale_float), + # PDL off: with dependent launch the next layer's + # call can overlap this one, and the workspace + # (semaphores, split scratch) is shared state + False, 0, + int(kv_cache.stride(0)), int(kv_cache.stride(2)), + int(kv_cache.stride(1))) + st["taken"] += 1 + if st["taken"] == 1 or st["taken"] % 2000 == 0: + print(f"[flash_rt] verify-attention xqa " + f"taken={st['taken']} " + f"fell_through={st['fell_through']}", + flush=True) + return output + except Exception: # noqa: BLE001 — host path must survive + st["fell_through"] += 1 + return orig(self, layer, query, key, value, kv_cache, + attn_metadata, output, *args, **kwargs) + + fi.FlashInferImpl.forward = forward + st["revert"] = lambda: setattr(fi.FlashInferImpl, "forward", orig) + if verbose: + print("[flash_rt] verify-attention seat installed " + "(fp8-kv-attention v4, page 32)") + return st + + class _NoSeats: """The handle shape for a host where nothing could be seated. @@ -206,7 +808,10 @@ def summary(self): def attach_engine(model, *, seats=DENSE_SEAT_SUFFIXES, experts=True, - head=True, use_gemv=None, verbose=True, strict=False): + head=True, use_gemv=None, verbose=True, strict=False, + fused_mlp=True, consume=False, ckpt_rename=None, + tag="model", precision="nvfp4", rows_hint=(1, 8), + adopt_pack=True): """Seat a vLLM model: dense projections, expert banks, LM head. Call between weight load and the engine's first trace (see @@ -228,8 +833,173 @@ def _init(self, *a, **kw): swaps: dict[str, nn.Module] = {} reverts: list = [] refused: list = [] + # opt-in verify-attention seat: a backend-level wrap, independent of + # the module seats below, reverted with the same handle + if os.environ.get("FRT_ATTN_XQA") == "1": + xst = attach_verify_attention(verbose=verbose) + if xst is not None and "revert" in xst: + reverts.append(xst["revert"]) + parked_tiers: dict[str, int] = {} + fused_mlps = 0 + staged = 0 + probe_cos: list[tuple[float, str]] = [] + #: the packer's own pack-and-unpack relative L2, which ``bind_proj_seam`` + #: returns and callers have been discarding. It is the grid's quality + #: with no activation and no host in the way — the one number that + #: says whether a seat reproduces the checkpoint's precision or + #: quietly lowers it. + pack_rels: list[float] = [] + kinds: dict[str, int] = {} + adopted = 0 + requant = 0 + skipped: list[str] = [] modules = dict(model.named_modules()) + # The head binds first. Its quantize needs an fp32 transient the + # width of the vocabulary, and by the time the dense seats are in + # place this process holds both the engine's original weights and + # the seats' packed copies — the transient is exactly what is no + # longer there. Bound first, it runs while the engine's weights are + # the only thing resident. + head_slabs = 0 + if head: + lm = next((m for n, m in modules.items() + if n.endswith("lm_head") + and isinstance(getattr(m, "weight", None), torch.Tensor)), + None) + # a head the checkpoint already packed as NVFP4 can be adopted + # zero-copy — but the seam's row tiers quantize activations to + # FP4, and the head is the logits family: measured on real + # streams, A4 logits moved speculative acceptance well outside + # the noise band while the kernel win was microseconds. The + # host's own W4A16 head is both exact-class and near its + # weight-stream floor, so adoption is explicit opt-in only. + if (lm is not None and adopt_pack + and os.environ.get("FRT_HEAD_W4A4") == "1" + and precision in ("mirror", "auto") + and _host_precision(lm) == "nvfp4"): + try: + lm._frt_seat_name = "lm_head" + seam = _adopt_nvfp4_pack(lm) + if seam is not None: + orig_method = lm.quant_method + lm.quant_method = _SlabbedHeadMethod( + [seam], orig_method) + reverts.append( + lambda lm=lm, m=orig_method: setattr( + lm, "quant_method", m)) + head_slabs = 1 + adopted += 1 + _arm_runtime_dispatch(seam) + lm = None # claimed; skip the W8 slab path + except Exception as e: # noqa: BLE001 — W8 path may still work + if verbose: + print(f"[structures.vllm] lm_head adopt fell back: " + f"{repr(e)[:160]}", flush=True) + # the checkpoint's own W4A16 head re-laid for small M: identical + # NVFP4 grid and dequant values, Marlin layout, so verify-step + # logits calls run a kernel built for M<=16 instead of the tiled + # path. Precision class unchanged — default on, env to disable. + if (lm is not None and adopt_pack + and os.environ.get("FRT_HEAD_MARLIN", "1") == "1" + and precision in ("mirror", "auto")): + try: + kxg = _linear._kernel() + w = getattr(lm, "weight", None) + ws = getattr(lm, "weight_scale", None) + # modelopt spells the per-tensor scale two ways across + # its linear methods + ws2 = getattr(lm, "weight_scale_2", None) + if ws2 is None: + ws2 = getattr(lm, "weight_global_scale", None) + if (hasattr(kxg, "adopt_nvfp4_w4a16_marlin") + and isinstance(w, torch.Tensor) and w.dim() == 2 + and w.dtype == torch.uint8 + and isinstance(ws, torch.Tensor) and ws.dim() == 2 + and ws.shape[0] == w.shape[0] + and isinstance(ws2, torch.Tensor) + and ws2.numel() == 1): + packs = kxg.adopt_nvfp4_w4a16_marlin( + w.contiguous(), + ws.view(torch.float8_e4m3fn).contiguous(), + ws2.float().reshape(())) + orig_method = lm.quant_method + mh = _MarlinHeadMethod(kxg, packs, w.shape[0], + orig_method) + # bind receipt: the relaid pack must reproduce the + # checkpoint's own dequantized rows. The host's + # apply is not the reference — its activation + # quantization adds noise of its own (measured cos + # ~0.92 on an exact relay); the E2M1 table decode + # of a row slice is the ground truth. + rows_p = 4096 + tbl = torch.tensor( + [0, .5, 1, 1.5, 2, 3, 4, 6, + -0., -.5, -1, -1.5, -2, -3, -4, -6], + device=w.device) + pw = w[:rows_p].to(torch.int32) + q = torch.stack([tbl[pw & 0xF], tbl[pw >> 4]], + dim=-1).reshape(rows_p, -1) + sc = (ws.view(torch.float8_e4m3fn)[:rows_p].float() + .repeat_interleave(16, dim=1)) + wd = q * sc * float(ws2) + xp = torch.randn(8, w.shape[1] * 2, device=w.device, + dtype=torch.bfloat16) + ours = mh.apply(lm, xp)[:, :rows_p].float() + ref = xp.float() @ wd.T + cos = torch.nn.functional.cosine_similarity( + ours.reshape(-1), ref.reshape(-1), dim=0).item() + del pw, q, sc, wd, ref, xp + if cos < 0.999: + raise RuntimeError( + f"marlin head probe cos {cos:.6f} vs " + f"checkpoint dequant") + lm.quant_method = mh + reverts.append( + lambda lm=lm, m=orig_method: setattr( + lm, "quant_method", m)) + head_slabs = 1 + adopted += 1 + if verbose: + print(f"[structures.vllm] lm_head marlin relay: " + f"probe cos {cos:.6f}", flush=True) + lm = None # claimed; skip the W8 slab path + except Exception as e: # noqa: BLE001 — host head still fine + if verbose: + print(f"[structures.vllm] lm_head marlin fell back: " + f"{repr(e)[:160]}", flush=True) + if lm is not None: + try: + from ..impls.linear_proj import w8a16_static as _w8 + rows = lm.weight.shape[0] + # name it so a quantized runtime head can be re-gridded + # from the checkpoint's own rows rather than unpacked + lm._frt_seat_name = "lm_head" + w = _projection_weight(lm) + # slabs, even where the entry's row support would take + # the whole vocabulary in one bind: the engine has + # already claimed its memory fraction by this point, so + # the transient a whole-vocabulary quantize needs is + # exactly what is not there. Four slabs put the peak + # inside what the engine leaves behind. + cap = min(_w8.SUPPORT["N"]["max"], -(-rows // 4)) + slab = -(-cap // 64) * 64 + seams = [_w8.bind_proj_seam({"w": w[lo:lo + slab]}) + for lo in range(0, rows, slab)] + orig_method = lm.quant_method + lm.quant_method = _SlabbedHeadMethod(seams, orig_method) + reverts.append( + lambda lm=lm, m=orig_method: setattr( + lm, "quant_method", m)) + head_slabs = len(seams) + except Exception as e: + refused.append(("lm_head", repr(e)[:200])) + if verbose: + print(f"[structures.vllm] lm_head refused: " + f"{repr(e)[:200]}", flush=True) + + + # dense projections, smallest first: on tight cards early frees # make room for the big binds targets = [(n, m) for n, m in modules.items() @@ -237,10 +1007,158 @@ def _init(self, *a, **kw): targets.sort(key=lambda t: t[1].weight.numel()) for name, mod in targets: try: - seam, _ = _linear.bind_proj_seam({"w": mod.weight.data}) - swaps[name] = _ProjSeat(seam) + mod._frt_seat_name = (ckpt_rename(name) if ckpt_rename + else name) + kind = (_host_precision(mod) + if precision in ("mirror", "auto") + else "w8" if precision == "w8" else "nvfp4") + if precision == "auto" and kind == "fp8": + # the auto tier's one opinion: a W8 position is carried + # to W4 (the measured arbitrage: smaller weight stream, + # acceptance unmoved on real streams), sourced from the + # host's own FP8 rows — no external checkpoint. This is + # a precision change and the tier says so in its report. + kind = "nvfp4" + requant += 1 + if kind == "unquantized": + # the checkpoint held this projection out of its own + # quantization; a seat here is not an acceleration of + # the host's decision, it is a replacement of it + skipped.append(name) + continue + seam, seam_shares_host, k_in = None, False, None + if kind == "fp8" and adopt_pack: + seam = _adopt_fp8_pack(mod, rows_hint) + if seam is not None: + adopted += 1 + seam_shares_host = True + # the engine stores FP8 as [K, N] (its cutlass B), + # so neither storage axis can be assumed; the seam's + # own [N, K] weight is the one orientation-safe place + # to read the projection's input width + k_in = int(seam._w_fp8.shape[1]) + if kind == "nvfp4" and adopt_pack: + # try the host's own pack first: regridding is for a + # host holding dense rows, and paying it here would + # substitute our packer's rounding for the checkpoint's + seam = _adopt_nvfp4_pack(mod) + if seam is not None: + adopted += 1 + # the seam holds the engine's own tensors, so the + # engine's copy is the seam's copy: releasing it + # would release the weights the seat executes + seam_shares_host = True + if seam is None: + w_bind = _projection_weight(mod) + if kind == "fp8": + seam = _bind_fp8_seam(mod, w_bind, rows_hint) + elif kind == "w8": + # the logits-family middle tier: half the weight + # stream of BF16 at a fraction of W4's rounding — + # the draft-precision experiment band + from ..impls.linear_proj import w8a16_static as _w8s + seam = _w8s.bind_proj_seam({"w": w_bind}) + else: + seam, pack_rel = _linear.bind_proj_seam( + {"w": w_bind}) + pack_rels.append(pack_rel) + k_in = w_bind.shape[1] + elif k_in is None: + # the probe width is the projection's K, and the adopt + # path has no dense rows to read it off — the NVFP4 host + # holds packed bytes, whose second dim is K/2. A probe + # built at that width fails inside the *host's* forward, + # which reads as the seam being rejected when nothing + # about the seam was tested at all. + k_in = int(seam._k) + # bind acceptance: the seam must reproduce the host module + # on a live probe — this is what catches a wrong weight + # layout (a packed checkpoint mistaken for dense rows) at + # bind time instead of as garbage tokens later + probe = torch.randn(4, k_in, device="cuda", + dtype=torch.bfloat16) + with torch.no_grad(): + host_out = mod(probe) + host_out = (host_out[0] if isinstance(host_out, tuple) + else host_out) + cos = torch.nn.functional.cosine_similarity( + seam(probe).float().reshape(-1), + host_out.float().reshape(-1), dim=0) + if float(cos) < 0.98: + raise ValueError( + f"bind probe cos {float(cos):.4f} < 0.98") + # the threshold is an admission gate, not a verdict: a tree + # of seats all sitting just above it is a different model + # from one sitting at 0.9999, and only the distribution says + # which this is + probe_cos.append((float(cos), name)) + # the dense rows were only ever the seam's input; holding + # them across the next bind doubles the transient peak + del probe, host_out + if kind == "nvfp4": + _arm_runtime_dispatch(seam) + kinds[kind] = kinds.get(kind, 0) + 1 + swaps[name] = _ProjSeat(seam, host=mod) + if consume and kind == "nvfp4" and not seam_shares_host: + # nvfp4 only: the FP8 seam retains the host module for + # its own fallback form, and releasing rows it may still + # execute would turn a fallback into a device mismatch. + # Release the engine's copy of these rows now, not after + # the whole tree is seated: the seat owns the packed + # form from here on, and holding both copies is exactly + # what makes the *next* bind fail. Staged to host memory + # so the handle's restore path still has them. + dev = mod.weight.data.device + cpu_rows = mod.weight.data.to("cpu") + # the revert closure may capture only the host copy: a + # default argument holding the device tensor pins the + # very allocation this is releasing, and the release + # then measures as a no-op. + mod.weight.data = cpu_rows + reverts.append( + lambda m=mod, w=cpu_rows, d=dev: setattr( + m.weight, "data", w.to(d))) + del cpu_rows + staged += 1 + if staged % 16 == 0: + torch.cuda.empty_cache() + if verbose and staged % 48 == 0: + free_b, _ = torch.cuda.mem_get_info() + print(f"[structures.vllm] staged {staged} seats, " + f"{free_b / 2**30:.2f} GiB free", flush=True) except Exception as e: - refused.append((name, repr(e)[:120])) + # a bare exception type carries no diagnosis; the first few + # refusals keep their frames so "refused" names a line + if len(refused) < 3: + import traceback + refused.append((name, repr(e)[:80] + " | " + + traceback.format_exc()[-1600:])) + else: + refused.append((name, repr(e)[:120])) + + if fused_mlp: + from ..impls.decoder_ffn import nvfp4_fused as _ffn + + silu_mul = _ffn._native_silu_mul() + if silu_mul is not None: + for name, mod in modules.items(): + gu = swaps.get(f"{name}.gate_up_proj") + dn = swaps.get(f"{name}.down_proj") + if gu is None or dn is None: + continue + if not hasattr(mod, "act_fn"): + continue + # the fused producer emits packed FP4 for the down seam: + # a pair the checkpoint put in different widths has no + # such handoff + if not all(isinstance(s.seam, _linear.LinearProjNvfp4Dynamic) + for s in (gu, dn)): + continue + swaps[name] = _FusedMlpSeat(mod, gu.seam, dn.seam, + silu_mul) + swaps.pop(f"{name}.gate_up_proj") + swaps.pop(f"{name}.down_proj") + fused_mlps += 1 if experts: impl = (_experts_w4a4 if use_gemv else _experts_w4a16) @@ -267,29 +1185,6 @@ def _init(self, *a, **kw): except Exception as e: refused.append((name, repr(e)[:120])) - head_slabs = 0 - if head: - lm = next((m for n, m in modules.items() - if n.endswith("lm_head") - and isinstance(getattr(m, "weight", None), torch.Tensor)), - None) - if lm is not None: - try: - rows = lm.weight.shape[0] - slab = -(-rows // 4) // 64 * 64 - seams = [ - _linear.bind_proj_seam( - {"w": lm.weight.data[lo:lo + slab]})[0] - for lo in range(0, rows, slab)] - orig_method = lm.quant_method - lm.quant_method = _SlabbedHeadMethod(seams, orig_method) - reverts.append( - lambda lm=lm, m=orig_method: setattr( - lm, "quant_method", m)) - head_slabs = len(seams) - except Exception as e: - refused.append(("lm_head", repr(e)[:120])) - model.eval() if not swaps: if strict: @@ -307,15 +1202,79 @@ def _init(self, *a, **kw): handle.notes["refused"] = refused return handle handle = _swap.attach(model, swaps, revert=reverts) + if consume: + # Until this runs the engine's own weights and the seats' packed + # copies are both resident, and on a card sized for the model + # alone that doubling is what makes the later binds fail — the + # measured refusals were 80 MiB allocations against a full card. + # Consuming moves each replaced module's truth to the weight + # store; fallback and detach survive as restore-from-store. + freed = handle.consume().get("freed_bytes", 0) + if verbose: + print(f"[structures.vllm] consumed host weights: " + f"{freed / 2**30:.2f} GiB freed", flush=True) if verbose: - print(f"[structures.vllm] {len(swaps)} seats " - f"({head_slabs} head slabs), {len(refused)} refused", - flush=True) - handle.notes = {"refused": refused, "head_slabs": head_slabs} + parked = "" + fused = f", {fused_mlps} fused MLPs" if fused_mlps else "" + by = (", ".join(f"{k}x{v}" for k, v in sorted(kinds.items())) + if precision == "mirror" else "") + skip = (f", {len(skipped)} left unquantized (checkpoint's own " + f"exclusion)" if skipped else "") + print(f"[structures.vllm] {tag}: {len(swaps)} seats " + f"({head_slabs} head slabs), {len(refused)} refused" + f"{parked}{fused}" + f"{(' [' + by + ']') if by else ''}{skip}", flush=True) + if adopted: + print(f"[structures.vllm] {tag}: {adopted} seats adopted " + f"the host's own pack (no regrid, no extra weight " + f"memory)", flush=True) + if requant: + print(f"[structures.vllm] {tag}: {requant} seats carried " + f"W8->W4 (auto tier: precision change, task-level " + f"validation is the caller's gate)", flush=True) + if pack_rels: + pr = sorted(pack_rels) + print(f"[structures.vllm] {tag}: pack relL2 median " + f"{pr[len(pr) // 2]:.5f}, worst {pr[-1]:.5f}", + flush=True) + if probe_cos: + cs = sorted(probe_cos) + print(f"[structures.vllm] {tag}: bind cos min " + f"{cs[0][0]:.5f} ({cs[0][1].rsplit('.', 2)[-2:] and '.'.join(cs[0][1].split('.')[-3:])}), " + f"p10 {cs[len(cs) // 10][0]:.5f}, " + f"median {cs[len(cs) // 2][0]:.5f}", flush=True) + seen = set() + for nm, why in refused: + key = why[:60] + if key in seen: + continue + seen.add(key) + print(f"[structures.vllm] refused {nm}: {why[:1800]}", + flush=True) + handle.notes = {"refused": refused, "head_slabs": head_slabs, + "parked_tiers": parked_tiers, + "fused_mlps": fused_mlps} return handle -def install_load_hook(*, on_attached=None, **attach_kwargs): +def _draft_ckpt_rename(name: str) -> str: + """Draft module path -> its key in the full-precision checkpoint. + + The proposer holds its own model, so its paths restart at ``model.`` + and collide with the target's layer 0 — a source lookup on the raw + path silently reads the *target's* rows, and only the bind probe + stands between that and a wrong seat. The checkpoint files the draft + under its own ``mtp.`` subtree; naming it that way is what makes the + lookup mean what it says. + """ + for lead in ("model.model.", "model."): + if name.startswith(lead): + return "mtp." + name[len(lead):] + return "mtp." + name + + +def install_load_hook(*, on_attached=None, seat_draft=True, + draft_precision=None, **attach_kwargs): """Patch every importable vLLM model-runner so :func:`attach_engine` runs after weights load and before the engine's first trace. Set ``VLLM_DISABLE_COMPILE_CACHE=1``: the engine's compile cache key @@ -340,7 +1299,51 @@ def install_load_hook(*, on_attached=None, **attach_kwargs): def load_model(self, *a, __orig=orig, **kw): __orig(self, *a, **kw) + # A speculative engine judges drafts against target logits, + # and vLLM's MTP draft shares the target's lm_head. The + # Marlin head relay changes the target's head numerics + # only, so draft and target stop agreeing at the head — + # measured as a systematic acceptance drop (paired over + # ten prompts, lower on all ten). Under spec decode the + # relay steps aside unless the caller forces it. + _spec = getattr(getattr(self, "vllm_config", None), + "speculative_config", None) + if _spec is not None and "FRT_HEAD_MARLIN" not in os.environ: + os.environ["FRT_HEAD_MARLIN"] = "0" handle = attach_engine(self.model, **attach_kwargs) + # A speculative engine is two models, and its throughput is + # the product of acceptance and step rate. Seating only the + # target re-grids one side of the agreement test: the draft + # still proposes from the checkpoint's own grid, the target + # now judges from ours, and the measured cost is acceptance + # — a loss no kernel win can pay back. The draft's seats are + # bound from the same source rows for that reason first, and + # for its own speed second. + draft = getattr(getattr(self, "drafter", None), "model", None) + if draft is not None and isinstance(draft, nn.Module) and seat_draft: + kw2 = dict(attach_kwargs) + kw2["head"] = False # the draft shares the target's + kw2["ckpt_rename"] = _draft_ckpt_rename + kw2["tag"] = "draft" + # a draft the checkpoint left in BF16 contributes no + # seats under the adopt tiers; an explicit draft + # precision re-grids it (proposal quality is the only + # exposure — the target still judges every token — and + # the caller gates it on measured acceptance) + if draft_precision: + kw2["precision"] = draft_precision + kw2["seats"] = tuple( + kw2.get("seats", DENSE_SEAT_SUFFIXES) + ) + ("model.fc",) + try: + dh = attach_engine(draft, **kw2) + # one detach undoes both models: the draft handle's + # undo list rides on the target's + handle._revert.append(dh.detach) + except Exception as e: # noqa: BLE001 + if attach_kwargs.get("verbose", True): + print(f"[structures.vllm] draft not seated: " + f"{repr(e)[:160]}", flush=True) if on_attached is not None: on_attached(handle) runner.load_model = load_model diff --git a/flash_rt/structures/impls/__init__.py b/flash_rt/structures/impls/__init__.py index e85ab2154..853cff010 100644 --- a/flash_rt/structures/impls/__init__.py +++ b/flash_rt/structures/impls/__init__.py @@ -144,12 +144,99 @@ def _check_arch(repo: str, module) -> None: #: refusal path must not manufacture that error on retry) _LOADED: dict[tuple[str, str], object] = {} +#: packages this host cannot resolve, by the same key. Absence is as +#: cacheable as presence and far more expensive to re-establish. +_UNRESOLVED: dict[tuple[str, str], BaseException] = {} + +#: highest version tag the resolution fallback searches downward from. +#: Only an upper bound for the walk — a repo that has not published +#: that many releases simply misses those revisions and continues. +_TAG_SEARCH_TOP = 32 + + +def _newest_loadable(get_kernel, repo, version, kw, first=None, + reported=None): + """The newest release at or above ``version``'s floor that loads. + + A resolver that only accepts an exact major still reports which + majors exist when it rejects a range, so the published set comes + from that report rather than a blind walk; ``_TAG_SEARCH_TOP`` is + only the fallback ceiling for a resolver that says nothing. + + ``first`` is an optional resolution to try before the search (the + pre-semver library's repo default, which is usually right and + costs nothing to attempt). + """ + if first is not None: + try: + return first() + except FileNotFoundError as no_variant: + if "build variants" not in str(no_variant): + raise + floor = re.match(r"^\s*>=\s*v?(\d+)", str(version)) + lo = int(floor.group(1)) if floor else 1 + published = sorted( + {int(n) for n in re.findall(r"\d+", reported or "") + if int(n) >= lo}, reverse=True) if reported else [] + last = None + for major in published: + try: + return get_kernel(repo, version=major, **kw) + except Exception as e: # noqa: BLE001 — try the next release + last = e + for major in range(_TAG_SEARCH_TOP, lo - 1, -1): + try: + return get_kernel(repo, revision=f"v{major}", **kw) + except Exception as e: # noqa: BLE001 — try the next tag + last = e + raise last if last is not None else RuntimeError( + f"no loadable release for {repo!r} at {version}") + @lru_cache(maxsize=None) def hub_kernel(repo: str, version: str): from kernels import get_kernel + # A repo with no build variant for this host resolves by walking + # releases, and every step of that walk is a network round trip. Left + # unbounded those trips can hang far longer than the absence they are + # establishing is worth — a serving engine calling this during model + # load stalls with no output at all. Bound them so "this host has no + # variant" arrives as a refusal in seconds. Defaults only: a caller + # that set its own timeout meant it. + os.environ.setdefault("HF_HUB_ETAG_TIMEOUT", "10") + os.environ.setdefault("HF_HUB_DOWNLOAD_TIMEOUT", "60") + key = (repo, version) + slug = re.sub(r"[^A-Za-z0-9]", "_", repo).upper() + var_dir = os.environ.get("FRT_KERNEL_DIR_" + slug) + if var_dir and key not in _LOADED: + # filesystem-direct load for hosts whose process cannot reach + # the hub at all (air-gapped serving containers): the variant + # directory is the package — importing it here is exactly what + # the resolver would do after its network walk, minus the walk. + # The arch check below still runs; a wrong variant refuses the + # same way a resolved one would. + import importlib.util + import sys as _sys + init = os.path.join(var_dir, "__init__.py") + if not os.path.isfile(init): + raise KernelUnavailable( + f"kernel package {repo!r}: FRT_KERNEL_DIR_{slug} does " + f"not point at a build variant (no __init__.py in " + f"{var_dir!r})") + mod_name = "_frt_kernel_" + slug.lower() + spec = importlib.util.spec_from_file_location( + mod_name, init, submodule_search_locations=[var_dir]) + mod = importlib.util.module_from_spec(spec) + _sys.modules[mod_name] = mod + spec.loader.exec_module(mod) + _LOADED[key] = mod + if key in _UNRESOLVED: + prior = _UNRESOLVED[key] + raise KernelUnavailable( + f"kernel package {repo!r} ({version}) is unavailable on " + f"this host: {type(prior).__name__}: {prior}") from prior if key not in _LOADED: # author pin for artifact bisection: an exact hub revision # outranks version resolution for this repo only. A perf or @@ -166,33 +253,45 @@ def hub_kernel(repo: str, version: str): # the trust gate arrived with newer kernels; our own # first-party artifacts are the explicit trust set _kw["trust_remote_code"] = True - try: + if rev: + _LOADED[key] = get_kernel(repo, revision=rev, **_kw) + else: try: - _LOADED[key] = (get_kernel(repo, revision=rev, - **_kw) - if rev - else get_kernel(repo, - version=version, - **_kw)) - except ValueError as ve: - # newer kernels resolve an exact integer version - # where older ones accepted a range string; the - # range's floor is the same request in both bands - m = re.match(r"^\s*>=\s*v?(\d+)", str(version)) - if not (m and "available versions" in str(ve)): - raise - _LOADED[key] = get_kernel( - repo, version=int(m.group(1)), **_kw) - except TypeError: - # kernels<0.13 — the band transformers pins — has no - # semver resolution kwarg; the default revision is - # exactly what that library resolved before semver - # tags existed. Widest-band compat: 0.12 through 0.16 - # serve the same call site. - _LOADED[key] = (get_kernel(repo, revision=rev) if rev - else get_kernel(repo)) # pre-semver band + _LOADED[key] = get_kernel(repo, version=version, + **_kw) + except TypeError: + # kernels<0.13 — the band transformers pins — has + # no semver kwarg; it resolves the repo default + _LOADED[key] = _newest_loadable( + get_kernel, repo, version, _kw, + first=lambda: get_kernel(repo)) + except (ValueError, FileNotFoundError) as unresolved: + # Two ways a version range fails to land on a + # usable artifact, both routine: the newer library + # resolves only an exact major (a range string is + # rejected outright), and any resolved release may + # have no build for this host's torch/CUDA pair — + # publishers add variants release by release, so a + # newer release can drop a pair an older one + # carried. Both are answered the same way: take + # the newest release at or above the floor that + # this host can actually load. That is what the + # range in the dependency spec asks for. + _LOADED[key] = _newest_loadable( + get_kernel, repo, version, _kw, + reported=(str(unresolved).split( + "available versions:")[-1] + if "available versions:" in str(unresolved) + else None)) except (OSError, RuntimeError, ValueError) as unavailable: _record_unavailable(repo, version, unavailable) + # remember the absence, not just the presence: resolution + # walks releases over the network, and a caller binding a + # whole model asks for the same package once per seat. Without + # this, one unpublished build variant costs that walk hundreds + # of times over — the failure is a property of this host, and + # it cannot change inside one process. + _UNRESOLVED[key] = unavailable raise KernelUnavailable( f"kernel package {repo!r} ({version}) is unavailable on " f"this host: {type(unavailable).__name__}: " diff --git a/flash_rt/structures/impls/decode_loop/dspark_block.py b/flash_rt/structures/impls/decode_loop/dspark_block.py new file mode 100644 index 000000000..786e6434a --- /dev/null +++ b/flash_rt/structures/impls/decode_loop/dspark_block.py @@ -0,0 +1,589 @@ +"""Block-draft speculative decoding for the whole-step loop (DFlash/DSpark). + +The draft is a small full-attention transformer that infills a block of +masked positions in one forward, conditioned on auxiliary features tapped +from a handful of the target's layers. One round costs one draft forward +plus one multi-token verify — against the autoregressive MTP draft chain +this removes the per-token draft steps that dominate that family's round. + +Serving geometry (learned the hard way): the draft emits one logit row +per noise slot *including the seed slot* — the seed's own row, fed the +real token embedding, is the strongest draft of the block, and dropping +it (as the checkpoint's reference implementation does) halves the +acceptance length. gamma drafts come from gamma noise slots +``[seed, mask x (gamma-1)]``, the Markov bigram head is applied +*sequentially* (each slot's bias conditioned on the previously sampled +token, anchored at the seed), and the verify window is gamma+1 tokens. + +The round is fully captured, following the MTP member's graph families: + +- the draft forward is one graph — static context K/V buffers take the + block's own keys by ``index_copy_`` at the block positions (the block + region ``[start, start+gamma)`` always sits right past the appended + context, so one masked SDPA over the static window covers context and + block siblings alike), and the serial Markov head runs unrolled on + device, so proposing costs zero host syncs; +- the gamma+1 verify is one graph with the gated-delta state snapshot + captured at its head (the MTP draft graph's ``snap_states`` idiom); +- re-advance is gated-delta-only: the verify pass already produced + everything else a rejected round needs — the bonus token is the + verify logit at the cut (its prefix rows saw exactly the accepted + tokens), and the KV rows for the accepted region were committed by + the verify itself (same tokens, same positions, FP8 pages included). + The only state a rejection actually loses is the gated-delta + recurrence past the cut, so each rejected-prefix length gets a graph + that restores the snapshot and re-drives just the gated-delta + sublayers over their stashed verify inputs — forty-eight state + layers instead of the whole model. A fully-accepted round skips + rollback outright: the verify pass committed exactly the accepted + stream; +- the arbiter costs the round's single host sync, same as the MTP loop. + +Aux features are the *inputs* of the configured target layers (the +serving-side capture point for this host family), collected by forward +pre-hooks. Hooks only execute during eager passes and graph capture — +the tensors they see during the verify capture live in the graph's +private pool, which every replay rewrites in place, so the captured +references stay valid round after round and the accepted rows' features +read straight out of them. + +Weights load from the draft checkpoint's single safetensors file; the +draft shares the target's embedding and lm_head (the draft was trained +against the target's own head). The confidence head ships in the +checkpoint but is not consumed here — the block length stays fixed. +""" + +from __future__ import annotations + +import json +import os +import pathlib + +import torch +import torch.nn.functional as F + +__all__ = ["DSparkBlockDraft", "DSparkRunner"] + + +@torch.library.custom_op("flashrt_native::dspark_draft_attn", + mutates_args=()) +def _draft_attn(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + seqused: torch.Tensor, scale: float) -> torch.Tensor: + """Non-causal fa2 over the draft's static context window. + + ``seqused`` is the device-side valid-K length (context + block): + the kernel reads it at run time, so the captured launch replays + with whatever the round wrote there — no mask, no materialised + attention math. q [S, Hq, D]; k/v [max, Hkv, D].""" + from flash_rt import flash_rt_fa2 as fa2 + + S, Hq, D = q.shape + o = torch.empty_like(q) + lse = torch.empty(1, Hq, S, device=q.device, dtype=torch.float32) + # strides are (batch, token, head) with a unit dim stride — the + # production call sites' convention for the vendored fa2 surface + mk, hk = k.shape[0], k.shape[1] + fa2.fwd_bf16_seqused( + q.data_ptr(), k.data_ptr(), v.data_ptr(), o.data_ptr(), + lse.data_ptr(), seqused.data_ptr(), 1, S, mk, Hq, hk, D, + (S * Hq * D, Hq * D, D), + (mk * hk * D, hk * D, D), + (mk * hk * D, hk * D, D), + (S * Hq * D, Hq * D, D), scale, 0, + torch.cuda.current_stream().cuda_stream) + return o + + +@_draft_attn.register_fake +def _(q, k, v, seqused, scale): + return torch.empty_like(q) + + +def _rms(x, w, eps=1e-6): + v = x.float() + v = v * torch.rsqrt(v.pow(2).mean(-1, keepdim=True) + eps) + return (v * w.float()).to(x.dtype) + + +class DSparkBlockDraft: + """The draft network, functional form over checkpoint tensors. + + ``propose_`` is capture-ready: every input rides a persistent device + buffer (seed token, block start), every write lands in a persistent + buffer (the draft tokens), and no host value is consulted. + """ + + def __init__(self, draft_dir, embed, lm_head, max_len, device="cuda"): + from safetensors.torch import load_file + + d = pathlib.Path(str(draft_dir)) + cfg = json.loads((d / "config.json").read_text()) + dc = cfg.get("dflash_config") or {} + self.taps = list(dc["target_layer_ids"]) + self.mask_id = int(dc["mask_token_id"]) + #: gamma = draft tokens per round = noise slots (seed included) + self.gamma = int(cfg["block_size"]) + self.n_layers = int(cfg["num_hidden_layers"]) + self.n_q = int(cfg["num_attention_heads"]) + self.n_kv = int(cfg["num_key_value_heads"]) + self.hd = int(cfg["head_dim"]) + self._embed = embed + self._head = lm_head + self._dev = device + self._max = int(max_len) + + t = {k: v.to(device, torch.bfloat16) + for k, v in load_file(str(d / "model.safetensors")).items()} + self._t = t + # draft projection precision: the verify pass anchors the + # output stream regardless, so acceptance length alone judges + # the draft's precision. Measured on the real-source coding + # stream: W4A4 costs it (AL 2.93 -> 2.79, net negative + # despite the halved weight stream) - the bf16 default + # stands, the arms stay for other streams. FP8 (W8A8, + # per-tensor scales) is the middle arm. + mode = os.environ.get("FRT_DSPARK_DRAFT_QUANT", "bf16") + self._proj = {} + self._proj8 = {} + names = [f"layers.{li}.self_attn.{nm}" + for li in range(self.n_layers) + for nm in ("q_proj", "k_proj", "v_proj", "o_proj")] + names += [f"layers.{li}.mlp.{nm}" + for li in range(self.n_layers) + for nm in ("gate_proj", "up_proj", "down_proj")] + if mode == "fp8": + for nm in names: + key = nm + ".weight" + w = t[key].float() + sw = (w.abs().amax() / 448.0).clamp_min(1e-12) + self._proj8[nm] = ( + (w / sw).to(torch.float8_e4m3fn).contiguous(), + sw.to(torch.float32)) + del t[key] + torch.cuda.empty_cache() + elif mode == "fp4": + from ..linear_proj import nvfp4_dynamic + for nm in names: + key = nm + ".weight" + seam, _rel = nvfp4_dynamic.bind_proj_seam( + {"w": t[key]}) + self._proj[nm] = seam + del t[key] + torch.cuda.empty_cache() + + # the checkpoint's own rope parameters through the host library's + # rope init - yarn attention scaling included in cos/sin + from transformers import Qwen3Config + from transformers.models.qwen3.modeling_qwen3 import ( + Qwen3RotaryEmbedding) + rp = dict(cfg["rope_parameters"]) + qc = Qwen3Config( + hidden_size=cfg["hidden_size"], head_dim=self.hd, + num_attention_heads=self.n_q, + num_key_value_heads=self.n_kv, + max_position_embeddings=cfg["max_position_embeddings"], + rope_scaling=rp, rope_theta=rp.get("rope_theta", 1e7)) + rot = Qwen3RotaryEmbedding(qc).to(device) + pos = torch.arange(self._max, device=device)[None] + cos, sin = rot(torch.empty(1, 1, device=device), pos) + self._cos = cos[0].to(torch.bfloat16) # [max_len, hd] + self._sin = sin[0].to(torch.bfloat16) + + # per-layer context K/V (roped keys). Appended rows carry the + # fc-projected target features; the block region past ``_len`` + # holds the propose pass's own keys until the next append + # overwrites it — every stale row dies before it is read. + self._ck = [torch.zeros(self._max, self.n_kv, self.hd, + device=device, dtype=torch.bfloat16) + for _ in range(self.n_layers)] + self._cv = [torch.zeros(self._max, self.n_kv, self.hd, + device=device, dtype=torch.bfloat16) + for _ in range(self.n_layers)] + self._len = 0 + + g = self.gamma + self._ids = torch.full((g,), self.mask_id, dtype=torch.long, + device=device) + self._start = torch.zeros(1, dtype=torch.long, device=device) + self._drafts = torch.zeros(g, dtype=torch.long, device=device) + self._arg = torch.arange(g, device=device) + self._armax = torch.arange(self._max, device=device) + self._seqk = torch.zeros(1, dtype=torch.int32, device=device) + + def reset(self): + self._len = 0 + + def _mm(self, x, name): + """Projection router: quantized arm when bound, bf16 otherwise.""" + seam = self._proj.get(name) + if seam is not None: + return seam(x) + w8 = self._proj8.get(name) + if w8 is not None: + wq, sw = w8 + sx = (x.float().abs().amax() / 448.0).clamp_min(1e-12) + xq = (x.float() / sx).to(torch.float8_e4m3fn) + return torch._scaled_mm( + xq, wq.t(), scale_a=sx, scale_b=sw, + out_dtype=torch.bfloat16) + return x @ self._t[name + ".weight"].T + + def _rope(self, x, pos): + # x [S, H, hd]; standard interleaved-half rotation + cos = self._cos[pos].unsqueeze(1) + sin = self._sin[pos].unsqueeze(1) + h = self.hd // 2 + x1, x2 = x[..., :h], x[..., h:] + rot = torch.cat((-x2, x1), dim=-1) + return x * cos + rot * sin + + @torch.no_grad() + def append_ctx(self, feats, pos): + """feats [T, taps*hidden] raw tap concat; pos [T] absolute.""" + t = self._t + tgt = _rms(feats @ t["fc.weight"].T, t["hidden_norm.weight"]) + n = feats.shape[0] + for li in range(self.n_layers): + p = f"layers.{li}.self_attn." + k = self._mm(tgt, p + "k_proj").view(n, self.n_kv, self.hd) + k = self._rope(_rms(k, t[p + "k_norm.weight"]), pos) + v = self._mm(tgt, p + "v_proj").view(n, self.n_kv, self.hd) + self._ck[li].index_copy_(0, pos, k) + self._cv[li].index_copy_(0, pos, v) + self._len += n + + @torch.no_grad() + def propose_(self): + """One block forward, buffers in, buffers out (capture-ready). + + Reads ``_ids`` (slot 0 = seed) and ``_start``; writes the gamma + draft tokens into ``_drafts``. The block's keys land in the + context buffers at the block positions, so a single masked SDPA + over the static window serves the dual-source attention: every + noise slot sees all context and every noise sibling. + """ + t = self._t + g = self.gamma + pos = self._start + self._arg + h = self._embed(self._ids.unsqueeze(0))[0] # [g, hidden] + # device-side valid-K length: context + this block + self._seqk.copy_((self._start + g).to(torch.int32)) + scale = 1.0 / (self.hd ** 0.5) + for li in range(self.n_layers): + p = f"layers.{li}." + a = f"{p}self_attn." + x = _rms(h, t[p + "input_layernorm.weight"]) + q = self._mm(x, a + "q_proj").view(g, self.n_q, self.hd) + q = self._rope(_rms(q, t[a + "q_norm.weight"]), pos) + kn = self._mm(x, a + "k_proj").view(g, self.n_kv, self.hd) + kn = self._rope(_rms(kn, t[a + "k_norm.weight"]), pos) + vn = self._mm(x, a + "v_proj").view(g, self.n_kv, self.hd) + self._ck[li].index_copy_(0, pos, kn) + self._cv[li].index_copy_(0, pos, vn) + o = _draft_attn(q.contiguous(), self._ck[li], self._cv[li], + self._seqk, scale).reshape(g, -1) + h = h + self._mm(o, a + "o_proj") + x = _rms(h, t[p + "post_attention_layernorm.weight"]) + m = f"{p}mlp." + h = h + self._mm( + F.silu(self._mm(x, m + "gate_proj")) + * self._mm(x, m + "up_proj"), m + "down_proj") + logits = self._head(_rms(h, t["norm.weight"]).unsqueeze(0))[0] + # serial markov, unrolled on device: each slot biased by the + # previously sampled token, anchored at the seed — no host + # round-trips inside the chain. The bias GEMV reads the [vocab, + # r] factor once per slot; it stays bf16 (fp32 accumulate) so + # the chain reads half the bytes + w1 = t["markov_head.markov_w1.weight"] + w2 = t["markov_head.markov_w2.weight"] + prev = self._ids[0:1] + for i in range(g): + bias = w1.index_select(0, prev) @ w2.T + tok = (logits[i].float() + bias[0].float()).argmax().view(1) + self._drafts[i:i + 1].copy_(tok) + prev = tok + + +class DSparkRunner: + """Captured-round driver over a built whole-step loop.""" + + def __init__(self, loop, draft_dir, model=None): + self._loop = loop + del model # reserved: host handle for cross-vehicle diagnostics + self._draft = DSparkBlockDraft( + draft_dir, loop._embed, loop._head, loop._max, + device=loop._embed.weight.device) + self._tap_in = {} + lm_layers = loop._layers + self._hooks = [] + for slot, li in enumerate(self._draft.taps): + self._hooks.append(lm_layers[li].register_forward_pre_hook( + self._make_hook(slot))) + # the gated-delta sublayers and their inputs: the re-advance + # graphs re-drive exactly these modules over the rows the + # verify pass showed them (input = the decoder layer's normed + # hidden, which is what the sublayer consumes again) + self._gdn_mods = {} + self._gdn_in = {} + for i, lyr in enumerate(lm_layers): + if i in loop._full_set: + continue + for _name, mod in lyr.named_children(): + if hasattr(mod, "conv1d") and hasattr(mod, "A_log"): + self._gdn_mods[i] = mod + self._hooks.append(mod.register_forward_pre_hook( + self._make_gdn_hook(i), with_kwargs=True)) + break + g = self._draft.gamma + dev = self._draft._dev + self._vblk = torch.zeros(1, g + 1, dtype=torch.long, device=dev) + self._vpos = torch.zeros(g + 1, dtype=torch.long, device=dev) + self._ag1 = torch.arange(g + 1, device=dev) + self._outbuf = torch.zeros(loop._max, dtype=torch.long, + device=dev) + self._gp = None + self.last_acceptance = 0.0 + + def _make_hook(self, slot): + def hook(module, args): + self._tap_in[slot] = args[0] + return hook + + def _make_gdn_hook(self, idx): + def hook(module, args, kwargs): + self._gdn_in[idx] = args[0] if args else \ + kwargs["hidden_states"] + return hook + + def _taps_rows(self, taps, rows): + return torch.cat([taps[s][0, :rows] for s in + range(len(self._draft.taps))], dim=-1) + + def _capture(self): + """Warm every round shape, then capture the graph families. + + Runs once, right after the first prompt pass, when every buffer + holds live values. The warmups execute on a side stream (the + capture-stream discipline the loop's own capture established); + the gated-delta states they advance are rolled back before + capture so the first replay starts from the real prompt state. + """ + loop, draft = self._loop, self._draft + g = draft.gamma + if loop._kv_band is not None: + # every round shape's fp8 mask exists before capture — a + # lazy fill between warmup and capture flips state mid-graph + loop._kv_band.prewarm(range(1, g + 2)) + # compile-then-capture, the MTP verify recipe: inductor's + # elementwise fusion buys the multi-row passes the same third + # it buys the plain step. Each round shape specialises once + # during warmup; the loop's init already raised the recompile + # budget for per-layer specialisation. + if loop._use_compile: + if getattr(loop, "_fwd_full_c", None) is None: + loop._fwd_full_c = torch.compile(loop._fwd_full, + dynamic=False) + fwd = loop._fwd_full_c + if getattr(self, "_propose_c", None) is None: + self._propose_c = torch.compile(draft.propose_, + dynamic=False) + prop = self._propose_c + else: + fwd = loop._fwd_full + prop = draft.propose_ + self._snap = { + i: (torch.empty_like(loop.cache.conv_states[i]), + torch.empty_like(loop.cache.recurrent_states[i])) + for i in loop._gdn_slots()} + for i, (cb, rb) in self._snap.items(): + cb.copy_(loop.cache.conv_states[i]) + rb.copy_(loop.cache.recurrent_states[i]) + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + prop() + self._vblk[0, 1:].copy_(draft._drafts) + fwd(self._vblk, self._vpos) + for m in range(1, g + 1): + fwd(self._vblk[:, :m], self._vpos[:m]) + torch.cuda.current_stream().wait_stream(side) + for i, (cb, rb) in self._snap.items(): + loop.cache.conv_states[i].copy_(cb) + loop.cache.recurrent_states[i].copy_(rb) + + self._gp = torch.cuda.CUDAGraph() + with torch.cuda.graph(self._gp): + # the verify inputs assemble in-graph: seed and start + # already ride draft buffers, so the whole block-and- + # positions setup costs zero host launches per round + prop() + self._vblk[0, 0].copy_(draft._ids[0]) + self._vblk[0, 1:].copy_(draft._drafts) + self._vpos.copy_(self._ag1) + self._vpos.add_(draft._start) + # the pre-round state snapshot rides at the verify graph's head: + # ninety-six host-launched little copies a round otherwise sit + # squarely inside the per-round sync window + self._gv = torch.cuda.CUDAGraph() + with torch.cuda.graph(self._gv): + for i, (cb, rb) in self._snap.items(): + cb.copy_(loop.cache.conv_states[i]) + rb.copy_(loop.cache.recurrent_states[i]) + self._vlg, self._vhn = fwd(self._vblk, self._vpos) + # the verify capture just ran the hooks: these references live + # in the verify graph's pool and every replay rewrites them in + # place. Saved now — later captures would overwrite the hook + # dicts with their own pools. + self._vtaps = [self._tap_in[s] + for s in range(len(draft.taps))] + gdn_v = dict(self._gdn_in) + # gated-delta-only re-advance: restore the snapshot, then + # re-drive just the state sublayers over the rows the verify + # showed them. Everything else a rejected round needs already + # exists — bonus logits sit in the verify output at the cut, + # and the accepted region's KV rows were the verify's own + # writes (same tokens, same positions). + self._ra = {} + for m in range(1, g + 1): + gr = torch.cuda.CUDAGraph() + with torch.cuda.graph(gr): + if getattr(self, "_stash_ok", False): + # rollback by selection: the verify pass stashed + # the carried state after every row (bit-equal to + # a re-advance ending there), and the conv window + # rebuilds from the snapshot tail plus the stashed + # raw rows — the epilogue's own semantics + for i, mod in self._gdn_mods.items(): + rec = loop.cache.recurrent_states[i] + rec.copy_(mod._stash_rec[m - 1].view_as(rec)) + cslot = loop.cache.conv_states[i] + kk = cslot.shape[-1] + take = min(kk, m) + if m < kk: + cslot[:, :, :kk - m].copy_( + self._snap[i][0][:, :, m:]) + cslot[0, :, kk - take:].copy_( + mod._stash_mixed[m - take:m].t()) + else: + for i, (cb, rb) in self._snap.items(): + loop.cache.conv_states[i].copy_(cb) + loop.cache.recurrent_states[i].copy_(rb) + for i, mod in self._gdn_mods.items(): + mod(gdn_v[i][:, :m], loop.cache, None) + self._ra[m] = gr + + @torch.no_grad() + def generate(self, input_ids, max_new_tokens): + loop, draft = self._loop, self._draft + dev = input_ids.device + L = int(input_ids.shape[1]) + g = draft.gamma + if L + max_new_tokens + g + 2 > loop._max: + raise ValueError("refused: window exceeds the static max") + draft.reset() + self._round_log = [] + loop.cache.frt_continue = False + loop._rope_delta.zero_() + if loop._kv_band is not None: + loop._kv_band.reset() + # canonical window, same as the loop's own generate: rows past + # the prompt zero before any produced token + for i in loop._full: + loop.cache.key_cache[i][:, :, L:].zero_() + loop.cache.value_cache[i][:, :, L:].zero_() + ck = getattr(loop, "_prefill_chunk", None) + if ck and L > ck: + # deep-window prompt: same sliced pass as the loop's own + # generate; the tap hooks fire per slice, so the draft's + # context features append slice by slice — the features of + # a row do not depend on how the prompt was sliced (they + # are that row's layer inputs) + for s in range(0, L, ck): + e = min(s + ck, L) + if loop._kv_band is not None: + # the deep flex arm keys its block mask on the + # slice span; host ints, eager path only + loop._kv_band._prompt_span = (s, e) + # _fwd, not _fwd_full: the prompt path's head reads the + # last row only — the full-rows head is a whole-vocab + # [rows, 248K] slab per slice (the measured ~1GB OOM at + # depth) and nothing in a prompt slice consumes it. The + # tap hooks ride the layers either way. + logits = loop._fwd( + input_ids[:, s:e], torch.arange(s, e, device=dev)) + draft.append_ctx( + self._taps_rows(self._tap_in, e - s), + torch.arange(s, e, device=dev)) + loop.cache.frt_continue = True + if loop._kv_band is not None: + loop._kv_band._prompt_span = None + else: + logits, _ = loop._fwd_full(input_ids, + torch.arange(L, device=dev)) + draft.append_ctx( + self._taps_rows(self._tap_in, L), + torch.arange(L, device=dev)) + loop.cache.frt_continue = True + seed = logits[0, -1].float().argmax().view(1) + self._outbuf[0].copy_(seed[0]) + produced = 1 + start = L + rounds = 0 + accepted_total = 0 + if self._gp is None: + # arm the per-row state stash on every gated-delta sublayer + # (eagerly, before any compiled pass traces): a rejected + # round then rolls back by selecting a stash row instead of + # re-driving the sublayers. Falls back to the re-drive form + # when the native build lacks the stash kernel. + self._stash_ok = bool(self._gdn_mods) and all( + getattr(mod, "frt_enable_stash", lambda *_: False)( + g + 1, loop.cache) + for mod in self._gdn_mods.values()) + self._vpos.copy_(self._ag1 + start) + draft._ids[0].copy_(seed[0]) + draft._start.fill_(start) + self._vblk[0, 0].copy_(seed[0]) + self._capture() + while produced < max_new_tokens: + draft._ids[0].copy_(seed[0]) + draft._start.fill_(start) + self._gp.replay() + self._gv.replay() + post = self._vlg[0].float().argmax(-1) # [g+1] + match = (self._vblk[0, 1:] == post[:-1]).long() + # device-side prefix match: the round's single host sync + a = int(match.cumprod(0).sum().item()) + # the bonus is the verify logit at the cut either way: row + # a's prefix saw exactly the accepted tokens + bonus = post[a:a + 1] + if a < g: + # rejection loses only the gated-delta recurrence past + # the cut: restore and re-drive the state sublayers + self._ra[a + 1].replay() + rows = a + 1 + draft.append_ctx( + self._taps_rows(self._vtaps, rows), + torch.arange(start, start + rows, device=dev)) + if a > 0: + self._outbuf[produced:produced + a].copy_( + self._vblk[0, 1:a + 1]) + self._outbuf[produced + a].copy_(bonus[0]) + seed = bonus + produced += rows + start += rows + accepted_total += rows + rounds += 1 + if rounds <= 40: + self._round_log.append(rows) + self.last_acceptance = accepted_total / max(rounds, 1) + self.last_rounds = self._round_log + loop.cache.frt_continue = False + return torch.cat( + [input_ids, self._outbuf[:max_new_tokens].view(1, -1)], dim=1) + + def detach(self): + for h in self._hooks: + h.remove() + self._hooks = [] diff --git a/flash_rt/structures/impls/decode_loop/fp8_kv.py b/flash_rt/structures/impls/decode_loop/fp8_kv.py index ba2e6fd98..0f31de275 100644 --- a/flash_rt/structures/impls/decode_loop/fp8_kv.py +++ b/flash_rt/structures/impls/decode_loop/fp8_kv.py @@ -58,6 +58,14 @@ def __init__(self, attn_layers, max_len, device): kern = _kernel() self._kern = kern self._max = int(max_len) + #: set by a loop built with ``prefill_chunk``: routes its prompt + #: slices through the GQA-native deep arm. Off by default so + #: every established short-window form keeps its exact path. + self._deep_prompt = False + #: (start, end) of the prompt slice in flight, host ints, set by + #: the eager chunked prompt drivers; None outside a slice + self._prompt_span = None + self._blockmasks = {} pages = (self._max + _PAGE - 1) // _PAGE # the kernel's max_seq_len speaks in whole pages self._max_paged = pages * _PAGE @@ -149,6 +157,35 @@ def attend(self, layer_idx, q): max_seq_len=self._max_paged) return out.view(1, s, _QH, _HD) + def prompt_attend(self, q, k, v, scaling): + """One prompt slice's attention: offset-causal flex over the + static window. The block mask depends only on (slice length, + slice start), so a prompt's handful of masks build once and + replay across layers, chunks, and repeat runs. The flex entry + itself is compiled once — eager flex is the math fallback and + materialises the score matrix, the very allocation this arm + exists to avoid; the block mask rides as a tensor input, so + chunk starts do not retrace.""" + from torch.nn.attention.flex_attention import ( + create_block_mask, flex_attention) + + s0, _e = self._prompt_span + S = q.shape[2] + bm = self._blockmasks.get((S, s0)) + if bm is None: + def causal_off(b, h, qi, ki): + return ki <= qi + s0 + + bm = create_block_mask(causal_off, 1, 1, S, self._max, + device=q.device) + self._blockmasks[(S, s0)] = bm + fx = getattr(self, "_flex_c", None) + if fx is None: + fx = torch.compile(flex_attention, dynamic=False) + self._flex_c = fx + return fx(q, k, v, block_mask=bm, scale=scaling, + enable_gqa=True) + def _interface(module, q, k, v, attention_mask, scaling=None, **kwargs): band = getattr(module, "_frt_fp8_band", None) @@ -161,6 +198,18 @@ def _interface(module, q, k, v, attention_mask, scaling=None, **kwargs): and q.shape[2] <= _XQA_MAX_Q \ and k.shape[2] == band._max: return band.attend(module.layer_idx, q), None + if band is not None and getattr(band, "_deep_prompt", False) \ + and band._prompt_span is not None and q.shape[0] == 1 \ + and q.shape[2] > _XQA_MAX_Q and k.shape[2] == band._max: + # deep-window prompt slice: flex attention over the offset + # causal block mask. The host's sdpa interface repeat_kv-expands + # 4 KV heads to 24 across the whole window (gigabytes per layer + # at 32K+), and masked SDPA with GQA falls to the math backend, + # which materialises the [heads, S, window] score matrix — both + # were measured as the deep-prompt OOM. Flex reads the same + # rows fused, GQA-native, nothing materialised. + o = band.prompt_attend(q, k, v, scaling) + return o.transpose(1, 2).contiguous(), None if attention_mask is None and q.shape[2] > 1 \ and k.shape[2] > q.shape[2]: # maskless prompt rows over the full static window: causal diff --git a/flash_rt/structures/impls/decode_loop/whole_step.py b/flash_rt/structures/impls/decode_loop/whole_step.py index 2f9b378e2..0670d754a 100644 --- a/flash_rt/structures/impls/decode_loop/whole_step.py +++ b/flash_rt/structures/impls/decode_loop/whole_step.py @@ -95,11 +95,17 @@ class WholeStepDecodeLoop: """Compiled, graph-captured greedy decode over the attached model.""" def __init__(self, model, *, max_len, compile_step=True, - compile_prefill=True, kv_band=None): + compile_prefill=True, kv_band=None, prefill_chunk=None): lm = _find_stack(model) self._model = model self._lm = lm self._compile_prefill = bool(compile_prefill) + #: deep-window prompt form: slices of this many rows drive the + #: same offset-mask forward, gated-delta state carrying across + #: slices through the continuation branch the verify batches + #: already exercise. Bounds the prompt-pass transients to one + #: slice — what admits 16K+ windows on this card's budget. + self._prefill_chunk = int(prefill_chunk) if prefill_chunk else None head = getattr(model, "lm_head", None) if head is None: raise ValueError("refused: host carries no lm_head") @@ -145,6 +151,10 @@ def __init__(self, model, *, max_len, compile_step=True, raise ValueError( "refused: fp8 kv band serves the kernel's v1 head " "profile (24/4/256); this host keeps BF16 KV") + if self._prefill_chunk: + # deep-window prompt slices take the GQA-native SDPA + # arm; established short-window forms keep their path + self._kv_band._deep_prompt = True elif kv_band is not None: raise ValueError(f"refused: unknown kv band {kv_band!r}") # no quadratic causal table: the decode mask is one static row @@ -260,8 +270,26 @@ def generate(self, input_ids, max_new_tokens): self.cache.key_cache[i][:, :, L:].zero_() self.cache.value_cache[i][:, :, L:].zero_() pf = self._aot_pf or self._prefill_callable() - logits = pf(input_ids, - torch.arange(L, device=input_ids.device)) + ck = self._prefill_chunk + if ck and L > ck: + # sliced prompt pass. Not bitwise-identical to the one-shot + # form: two runs of the FP4 pipeline decorrelate at low + # mantissa bits and the quantization grid amplifies them + # (measured: chunked-vs-oneshot logits cos ~= the arbiter + # band squared — two independent samples of the same + # quantization ball). The equivalence contract is the + # arbiter gate plus continuation identity, not bit equality. + dev = input_ids.device + logits = None + for s in range(0, L, ck): + e = min(s + ck, L) + logits = pf(input_ids[:, s:e], + torch.arange(s, e, device=dev)) + self.cache.frt_continue = True + self.cache.frt_continue = False + else: + logits = pf(input_ids, + torch.arange(L, device=input_ids.device)) self._cur.copy_(logits.float().argmax(-1)) self._pos.fill_(L) toks = [self._cur.clone()] @@ -456,6 +484,21 @@ def enable_mtp(self, ckpt_dir=None, head=None, self._verify_capture = bool(verify_capture) return self._mtp + @torch.no_grad() + def enable_dspark(self, draft_dir): + """Attach a block-draft (DFlash/DSpark) speculative runner. + + The draft checkpoint's own config supplies the tap layers, + block size, and mask token; the runner rides this loop's graph + families (captured propose/verify, rollback by stash selection + where the build carries the stash kernel). Returns the runner; + its ``generate`` replaces ``loop.generate`` for speculative + decoding and anchors exactness on this loop's own verify. + """ + from .dspark_block import DSparkRunner + + return DSparkRunner(self, draft_dir) + def _fwd_full(self, tok_ids, pos_t): h = self._embed(tok_ids) pe = self._rotary(h, pos_t.view(1, -1)) @@ -844,9 +887,11 @@ def forward(self, tok_ids, pos_t): def build_decode_loop(model, *, max_len, compile_step=True, - compile_prefill=True, kv_band=None): + compile_prefill=True, kv_band=None, + prefill_chunk=None): """Build the whole-loop form over whatever is attached to ``model``.""" return WholeStepDecodeLoop(model, max_len=max_len, compile_step=compile_step, compile_prefill=compile_prefill, - kv_band=kv_band) + kv_band=kv_band, + prefill_chunk=prefill_chunk) diff --git a/flash_rt/structures/impls/decoder_ffn/nvfp4_fused.py b/flash_rt/structures/impls/decoder_ffn/nvfp4_fused.py new file mode 100644 index 000000000..5b5f538f5 --- /dev/null +++ b/flash_rt/structures/impls/decoder_ffn/nvfp4_fused.py @@ -0,0 +1,200 @@ +"""NVFP4 fused-GLU implementation of the ``decoder_ffn`` structure. + +Composes the whole SwiGLU block behind one seam over already-bound +``linear_proj_nvfp4`` projections: the gate and up weights concatenate +into a single merged seam (one weight stream, one launch per tier), the +activation between them collapses into the ``flashrt/fp4-fused-ops`` +silu-mul producer that emits the down projection's packed FP4 input +directly — no BF16 round-trip, no standalone quantize launch, no +elementwise-mul kernel. + +Adoption is a post-pass over a model whose projections the +``linear_proj_nvfp4`` scheme has already bound: the three bound seams +are fused in place, reusing their packed weights (the concat is layout +sound because the SFB blocks are N-major and both halves are multiples +of the 128-row block). The host MLP is retained whole for fallback and +introspection, same contract as the other ``decoder_ffn`` impls. +""" + +from __future__ import annotations + +from functools import lru_cache + +import torch + +from ...guard import CAST_OK, PROCEED, GuardedSeam +from ..linear_proj.nvfp4_dynamic import (LinearProjNvfp4Dynamic, + _quantize_activation) + +FUSED_DEP = {"provider": "hf", "repo": "flashrt/fp4-fused-ops", + "version": ">=1"} + + +@lru_cache(maxsize=1) +def _native_silu_mul(): + """The fused SwiGLU + NVFP4 quantize producer. + + Bit-exact against the split (elementwise mul kernel -> production + quantize kernel) chain by construction, so adopting it moves no + numerics. + + Hub artifact first: its entry is already a torch op with a fake, so + a host that compiles this call traces it without help, and a host + process that cannot load our native extension (a serving engine on + a different torch/CUDA pair) still gets the fused producer. The + local build follows, wrapped as a custom op for the same reason. + """ + from flash_rt.structures.impls import hub_kernel + + try: + hub = hub_kernel(FUSED_DEP["repo"], FUSED_DEP["version"]) + except Exception: # noqa: BLE001 — absence is not a refusal + hub = None + hub_fn = getattr(hub, "silu_mul_quantize_fp4_sfa_bf16", None) + if hub_fn is not None: + def _hub_entry(merged): + # allocate the outputs here rather than letting the wrapper + # size them: its size helper is a custom op without a fake, + # so a host tracing this call on Meta tensors dies inside + # it. The layout is the packaged quantizer's own and stated + # in its contract, so computing it here is reading the + # contract, not guessing at it. + m, two_h = merged.shape + h = two_h // 2 + packed = merged.new_empty((m, h // 2), dtype=torch.uint8) + sfa = merged.new_zeros( + (((m + 127) // 128) * ((h + 63) // 64) * 512,), + dtype=torch.uint8) + return hub_fn(merged, packed=packed, sfa=sfa) + + return _hub_entry + try: + from flash_rt import flash_rt_kernels as _fk + except ImportError: + return None + fn = getattr(_fk, "silu_mul_quantize_fp4_sfa_bf16", None) + if fn is None: + return None + + from torch.library import custom_op + + @custom_op("flashrt_native::silu_mul_quant_fp4", mutates_args=()) + def _op(merged: torch.Tensor) -> list[torch.Tensor]: + m, two_h = merged.shape + h = two_h // 2 + packed = torch.empty(m, h // 2, device=merged.device, + dtype=torch.uint8) + # zero-filled: a partial 128-row scale atom leaves tail rows + # unwritten, and run-to-run garbage there is a nondeterminism + # channel for any consumer that loads whole atoms + sfa = torch.zeros( + ((m + 127) // 128) * ((h + 63) // 64) * 512, + device=merged.device, dtype=torch.uint8) + rc = fn(merged.data_ptr(), packed.data_ptr(), sfa.data_ptr(), + m, h, torch.cuda.current_stream().cuda_stream) + if rc != 0: + raise RuntimeError( + f"silu_mul_quant_fp4 refused rc={rc} for " + f"M={m} H={h}") + return [packed, sfa] + + @_op.register_fake + def _(merged): + m, two_h = merged.shape + h = two_h // 2 + return [merged.new_empty((m, h // 2), dtype=torch.uint8), + merged.new_empty( + (((m + 127) // 128) * ((h + 63) // 64) * 512,), + dtype=torch.uint8)] + + return _op + + +class FusedGluMlpNvfp4(GuardedSeam, torch.nn.Module): + """MLP seam: quantize once -> merged gate|up -> silu-mul-quant -> down. + + No host retention: like the bound projection seams it fuses, the + packed weights are the only copy — retaining the pre-fusion seams + would hold a second full FFN weight set on device. + """ + + def __init__(self, gate_up: LinearProjNvfp4Dynamic, + down: LinearProjNvfp4Dynamic, silu_mul): + super().__init__() + self.gate_up = gate_up + self.down = down + self._silu_mul = silu_mul + self._d = int(down._n) + self._frt_arm(dtypes=CAST_OK, + device=gate_up._w_packed.device, + k=int(gate_up._k)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + admitted = self._frt_admit(x) + if admitted is not PROCEED: + return admitted + shape = x.shape + cell = getattr(self, "_share", None) + if cell is not None and cell.x is x: + a_packed, a_sfa = cell.a, cell.sfa + else: + flat = x.reshape(-1, shape[-1]) + a_packed, a_sfa = _quantize_activation(self.gate_up._kern, + flat) + if cell is not None: + cell.x, cell.a, cell.sfa = x, a_packed, a_sfa + merged = self.gate_up._mm_packed(a_packed, a_sfa) + p2, s2 = self._silu_mul(merged.contiguous()) + y = self.down._mm_packed(p2, s2) + return y.reshape(*shape[:-1], self._d).type_as(x) + + +@torch.no_grad() +def fuse_bound_mlp(mlp: torch.nn.Module) -> FusedGluMlpNvfp4 | None: + """Fuse one MLP module whose projections are already bound seams. + + Returns None (leave the host untouched) when any projection is not + a bound ``linear_proj_nvfp4`` seam, carries a bias, or the merged + halves would break the SFB block layout (either N not a multiple of + the 128-row scale block). + """ + g = getattr(mlp, "gate_proj", None) + u = getattr(mlp, "up_proj", None) + d = getattr(mlp, "down_proj", None) + for seam in (g, u, d): + if not isinstance(seam, LinearProjNvfp4Dynamic): + return None + if seam._bias is not None: + return None + if g._n % 128 or u._n % 128 or g._n != u._n or g._k != u._k: + return None + silu_mul = _native_silu_mul() + if silu_mul is None: + return None + wp = torch.cat([g._w_packed, u._w_packed], dim=0).contiguous() + ws = torch.cat([g._w_sfb, u._w_sfb], dim=0).contiguous() + merged = LinearProjNvfp4Dynamic(wp, ws, None, 2 * int(g._n), + int(g._k)) + return FusedGluMlpNvfp4(merged, d, silu_mul) + + +@torch.no_grad() +def adopt_fused_glu(root: torch.nn.Module, verbose: bool = False) -> int: + """Post-pass: fuse every bound SwiGLU MLP under ``root`` in place.""" + count = 0 + for name, mod in list(root.named_modules()): + for child_name, child in list(mod.named_children()): + if not (hasattr(child, "gate_proj") + and hasattr(child, "up_proj") + and hasattr(child, "down_proj")): + continue + fused = fuse_bound_mlp(child) + if fused is None: + continue + setattr(mod, child_name, fused) + count += 1 + if verbose: + print(f"[decoder_ffn.nvfp4_fused] {name}.{child_name} " + f"fused (2x{fused.gate_up._n // 2} -> " + f"{fused._d})") + return count diff --git a/flash_rt/structures/impls/gated_delta_core/fused_layer.py b/flash_rt/structures/impls/gated_delta_core/fused_layer.py index 4ae2ba222..9f1de001d 100644 --- a/flash_rt/structures/impls/gated_delta_core/fused_layer.py +++ b/flash_rt/structures/impls/gated_delta_core/fused_layer.py @@ -35,6 +35,10 @@ from ...guard import CAST_OK, PROCEED, GuardedSeam +#: a range, not a pin: the resolver picks the newest release whose +#: build matrix covers the caller's torch/CUDA pair, and the entries +#: this structure calls have been stable since 3. Pinning the newest +#: release strands any host whose variant that release did not build. GDA_DEP = {"provider": "hf", "repo": "flashrt/gated-delta-attention", "version": ">=3"} CONV_DEP = {"provider": "hf", "repo": "flashrt/causal-conv1d-state", @@ -44,6 +48,343 @@ @lru_cache(maxsize=1) +def _native_ltri_inv(): + """The native batched 64x64 unit-lower-triangular inverse. + + Same forward-substitution recurrence class as the batched cuBLAS + solve it replaces (fp32 error band matches), with the identity/ + tril preparation folded in — the eye/expand/tril materializations + disappear with the solve. Registered as a torch custom op with a + fake so the compiled prefill traces through it. Absence is not a + refusal: the cuBLAS solve path keeps serving. + """ + try: + from flash_rt import flash_rt_kernels as _fk + except ImportError: + return None + fn = getattr(_fk, "batched_unit_ltri_inv64_f32", None) + if fn is None: + return None + + from torch.library import custom_op + + @custom_op("flashrt_native::unit_ltri_inv64", mutates_args=()) + def _op(big_a: torch.Tensor) -> torch.Tensor: + flat = big_a.reshape(-1, 64, 64).contiguous() + x = torch.empty_like(flat) + rc = fn(flat.data_ptr(), x.data_ptr(), flat.shape[0], + torch.cuda.current_stream().cuda_stream) + if rc != 0: + raise RuntimeError( + f"unit_ltri_inv64 refused rc={rc} for " + f"B={flat.shape[0]}") + return x.reshape(big_a.shape) + + @_op.register_fake + def _(big_a): + return torch.empty_like(big_a) + + return _op + + +import os as _os + +_NCP_ON = _os.environ.get("FRT_WY_NCP_V2", "1") != "0" + + +@lru_cache(maxsize=1) +def _native_gated_norm_quant(): + """The gated norm that also emits its consumer's NVFP4 input. + + The output projection is the norm's only consumer and quantizes + what it receives; the quantizer's blocks tile a head's lanes + exactly, so the whole step fits in the block that produced the + row. Measured bit-identical (normed, packed and scales alike) + against the packaged norm followed by the production quantize, + at 2.9x of the pair. + """ + if _os.environ.get("FRT_GDN_NORMQUANT", "1") == "0": + return None + # hub artifact first: its entry is a torch op with a fake, so a + # host that compiles this call traces it unaided and a process + # that cannot load our native extension still gets the producer + from flash_rt.structures.impls import hub_kernel + try: + hub = hub_kernel(FUSED_DEP["repo"], FUSED_DEP["version"]) + except Exception: # noqa: BLE001 — absence is not a refusal + hub = None + hub_fn = getattr(hub, "rms_norm_gated_silu_quant_fp4_bf16", None) + if hub_fn is not None: + def _hub_entry(x, gate, weight, out, packed, sfa, eps): + hub_fn(x, gate, weight, eps=eps, out=out, packed=packed, + sfa=sfa) + return _hub_entry + try: + from flash_rt import flash_rt_kernels as _fk + except ImportError: + return None + fn = getattr(_fk, "rms_norm_gated_silu_quant_fp4_bf16", None) + if fn is None: + return None + + from torch.library import custom_op + + @custom_op("flashrt_native::gated_norm_quant", + mutates_args=("out", "packed", "sfa")) + def _op(x: torch.Tensor, gate: torch.Tensor, weight: torch.Tensor, + out: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, + eps: float) -> None: + m, d = int(x.shape[0]), int(x.shape[1]) + rc = fn(x.data_ptr(), gate.data_ptr(), weight.data_ptr(), + out.data_ptr(), packed.data_ptr(), sfa.data_ptr(), + m, d, float(eps), + torch.cuda.current_stream().cuda_stream) + if rc != 0: + raise RuntimeError( + f"gated_norm_quant refused rc={rc} M={m} D={d}") + + @_op.register_fake + def _(x, gate, weight, out, packed, sfa, eps): + return None + + return _op + + +@lru_cache(maxsize=1) +def _native_recurrent_stream(): + """The streaming-column form of the gated-delta recurrent step. + + Bit-identical to the packaged kernel — same block reduction, same + per-column arithmetic in the same order. The column simply streams + in two passes instead of living in a 128-entry per-thread array + that spills to local memory, which is where the step was paying + DRAM for state it believed was resident. Measured 1.5x. + """ + if _os.environ.get("FRT_GDN_STREAM", "1") == "0": + return None + from flash_rt.structures.impls import hub_kernel + try: + hub = hub_kernel(GDA_DEP["repo"], GDA_DEP["version"]) + except Exception: # noqa: BLE001 — absence is not a refusal + hub = None + hub_fn = getattr(hub, "gdn_recurrent_inout_stream_bf16", None) + if hub_fn is not None: + def _hub_entry(q, k, v, g, beta, state_in, state_out, out): + return hub_fn(q, k, v, g, beta, state_in, + use_qk_l2norm=True, state_out=state_out, + out=out) + return _hub_entry + try: + from flash_rt import flash_rt_kernels as _fk + except ImportError: + return None + fn = getattr(_fk, "gdn_recurrent_inout_stream_bf16", None) + if fn is None: + return None + + from torch.library import custom_op + + @custom_op("flashrt_native::gdn_recurrent_stream", + mutates_args=("state_out", "out")) + def _op(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + g: torch.Tensor, beta: torch.Tensor, + state_in: torch.Tensor, state_out: torch.Tensor, + out: torch.Tensor) -> None: + h, d = int(q.shape[1]), int(q.shape[2]) + rc = fn(q.data_ptr(), k.data_ptr(), v.data_ptr(), g.data_ptr(), + beta.data_ptr(), state_in.data_ptr(), + state_out.data_ptr(), out.data_ptr(), 1, h, d, True, + torch.cuda.current_stream().cuda_stream) + if rc != 0: + raise RuntimeError( + f"gdn_recurrent_stream refused rc={rc} H={h} D={d}") + + @_op.register_fake + def _(q, k, v, g, beta, state_in, state_out, out): + return None + + def _entry(q, k, v, g, beta, state_in, state_out, out): + _op(q.contiguous(), k.contiguous(), v.contiguous(), + g.contiguous(), beta.contiguous(), state_in, state_out, + out) + return out, state_out + + return _entry + + +@lru_cache(maxsize=1) +def _native_conv_steps_gqa(): + """The step-batched conv1d update with GQA split outputs. + + The packaged chunk-parallel conv reads every input element K times + from DRAM (one token per thread); this arm rolls the taps through + registers over 8 consecutive tokens with the packaged tap order and + fma chain — bit-exact — and writes the q/k/v splits directly. + Fixed 2048/2048/6144, K=4 family only. + """ + try: + from flash_rt import flash_rt_kernels as _fk + except ImportError: + return None + fn = getattr(_fk, "causal_conv1d_update_steps_gqa_bf16", None) + if fn is None: + return None + + from torch.library import custom_op + + @custom_op("flashrt_native::conv1d_steps_gqa", mutates_args=()) + def _op(x: torch.Tensor, w: torch.Tensor, bias: torch.Tensor, + state: torch.Tensor) -> list[torch.Tensor]: + # bias is required: the Optional-tensor marshalling was + # measured to corrupt the following pointer operands; a + # bias-free host passes an explicit zeros row (acc starts at + # 0.0 either way - bit-identical to the null-bias path) + s = x.shape[0] + q16 = torch.empty((s, 16, 128), device=x.device, + dtype=torch.bfloat16) + k16 = torch.empty((s, 16, 128), device=x.device, + dtype=torch.bfloat16) + v48 = torch.empty((s, 48, 128), device=x.device, + dtype=torch.bfloat16) + rc = fn(x.data_ptr(), w.data_ptr(), bias.data_ptr(), + state.data_ptr(), q16.data_ptr(), k16.data_ptr(), + v48.data_ptr(), s, True, + torch.cuda.current_stream().cuda_stream) + if rc != 0: + raise RuntimeError(f"conv1d_steps_gqa refused rc={rc}") + return [q16, k16, v48] + + @_op.register_fake + def _(x, w, bias, state): + s = x.shape[0] + return [x.new_empty((s, 16, 128)), x.new_empty((s, 16, 128)), + x.new_empty((s, 48, 128))] + + return _op + + +@lru_cache(maxsize=1) +def _native_norm_cumsum_pack(): + """The native v2 launch of the WY norm/pack + gate-cumsum pair. + + Math transcribed verbatim from the packaged fast arm; the gate + cumsum is parallelized over the independent (chunk, head) pairs + (the packaged kernel walks the whole prompt from one 64-thread + block). Bit-exact against the packaged pair. Fixed 16/48/128/64 + family only — absence or another family keeps the packaged op. + """ + try: + from flash_rt import flash_rt_kernels as _fk + except ImportError: + return None + fn = getattr(_fk, "gdn_wy_norm_cumsum_pack_qk_v2_bf16", None) + if fn is None: + return None + + from torch.library import custom_op + + @custom_op("flashrt_native::gdn_wy_ncp_v2", mutates_args=()) + def _op(q16: torch.Tensor, k16: torch.Tensor, + g: torch.Tensor) -> list[torch.Tensor]: + s = q16.shape[0] + c = (s + 63) // 64 + q16_l2 = torch.empty_like(q16) + k16_l2 = torch.empty_like(k16) + q_pack_hv = torch.empty((c, 48, 64, 128), device=q16.device, + dtype=q16.dtype) + k_pack_hk = torch.empty((c, 16, 64, 128), device=q16.device, + dtype=q16.dtype) + g_cumsum = torch.empty_like(g) + rc = fn(q16.data_ptr(), k16.data_ptr(), g.data_ptr(), + q16_l2.data_ptr(), k16_l2.data_ptr(), + q_pack_hv.data_ptr(), k_pack_hk.data_ptr(), + g_cumsum.data_ptr(), s, + torch.cuda.current_stream().cuda_stream) + if rc != 0: + raise RuntimeError(f"gdn_wy_ncp_v2 refused rc={rc} S={s}") + return [q16_l2, k16_l2, q_pack_hv, k_pack_hk, g_cumsum] + + @_op.register_fake + def _(q16, k16, g): + s = q16.shape[0] + c = (s + 63) // 64 + return [torch.empty_like(q16), torch.empty_like(k16), + q16.new_empty((c, 48, 64, 128)), + q16.new_empty((c, 16, 64, 128)), + torch.empty_like(g)] + + def _entry(q16, k16, g): + if (q16.shape[1:] != (16, 128) + or g.shape[-1] != 48): + return None + return _op(q16.contiguous(), k16.contiguous(), g.contiguous()) + + return _entry + + +def _wy_ai_inverse(big_a: torch.Tensor) -> torch.Tensor: + """inv(I + strict_tril(A)) for the WY chain, batched 64x64 fp32.""" + import os + if (big_a.dtype is torch.float32 + and big_a.shape[-2:] == (64, 64) + and os.environ.get("FRT_TRSM_NATIVE", "1") != "0"): + inv = _native_ltri_inv() + if inv is not None: + return inv(big_a) + eye = torch.eye(64, device=big_a.device, + dtype=big_a.dtype).expand_as(big_a).contiguous() + return torch.linalg.solve_triangular( + eye + torch.tril(big_a, -1), eye, upper=False) + + +@lru_cache(maxsize=1) +def _native_stash_op(): + """The native per-row-stash arm of the from-conv chunk core. + + Registered lazily as a mutating torch custom op so the compiled + and captured spec-verify passes trace through it. Absence of the + native build is not a refusal — the plain hub chunk kernel keeps + serving, and rejected rounds re-drive the state sublayers instead + of selecting a stash row. + """ + global _STASH_OP + if _STASH_OP is not None: + return _STASH_OP if _STASH_OP is not False else None + try: + from flash_rt import flash_rt_kernels as _fk + fn = getattr(_fk, "gdn_chunk_from_conv_smem_h_stash_bf16", None) + except ImportError: + fn = None + if fn is None: + _STASH_OP = False + return None + + @torch.library.custom_op( + "flashrt_native::gdn_chunk_stash", + mutates_args=("state", "out", "stash")) + def _op(conv_out: torch.Tensor, a: torch.Tensor, b: torch.Tensor, + neg_exp_a: torch.Tensor, dt_bias: torch.Tensor, + state: torch.Tensor, out: torch.Tensor, + stash: torch.Tensor, num_v_heads: int, num_k_heads: int, + head_dim: int) -> None: + fn(conv_out.data_ptr(), a.data_ptr(), b.data_ptr(), + neg_exp_a.data_ptr(), dt_bias.data_ptr(), state.data_ptr(), + out.data_ptr(), stash.data_ptr(), conv_out.shape[0], + num_v_heads, num_k_heads, head_dim, a.stride(0), b.stride(0), + True, torch.cuda.current_stream().cuda_stream) + + @_op.register_fake + def _(conv_out, a, b, neg_exp_a, dt_bias, state, out, stash, + num_v_heads, num_k_heads, head_dim): + return None + + _STASH_OP = _op + return _op + + +_STASH_OP = None + + def _packages(): from flash_rt.structures.impls import hub_kernel @@ -139,6 +480,11 @@ def _gate(a, b, neg_exp_a, dt_bias): self._conv_w = host.conv1d.weight.detach().squeeze(1).contiguous() self._conv_b = (host.conv1d.bias.detach().contiguous() if host.conv1d.bias is not None else None) + # dense bias operand for the step-batched conv arm (required + # tensor; zeros reproduce the null-bias accumulator exactly) + self._conv_b_dense = (self._conv_b if self._conv_b is not None + else torch.zeros_like(self._conv_w[:, 0]) + .contiguous()) self._neg_exp_a = (-host.A_log.detach().float().exp()).contiguous() self._dt_bias = host.dt_bias.detach().float().contiguous() self._eps = float(getattr(host.norm, "variance_epsilon", @@ -161,6 +507,25 @@ def _gate(a, b, neg_exp_a, dt_bias): self._proj_rel = (rel_in, rel_out) except ValueError: self._proj_in = self._proj_out = None + elif projection_format == "nvfp4_balance": + # balanced fold: same FP4 band, with the per-input-channel + # balance fitted on calibrated activation amax — the lever + # that cut the single-seat tail error 2.5x on this host. + # Calibration is a precondition, not a default: a host layer + # without the attached amax keeps the BF16 band (counted as + # a refusal, never raised), and the amax only ever comes + # from calibrate_gdn_channel_amax's real-forward statistics. + from ..linear_proj import nvfp4_balance + amax = getattr(host, "_frt_gdn_channel_amax", None) + if amax is not None: + try: + self._proj_in = nvfp4_balance.bind_proj_seam( + {"w": self._packed_w}, channel_amax=amax["in"]) + self._proj_out = nvfp4_balance.bind_proj_seam( + {"w": host.out_proj.weight.detach()}, + channel_amax=amax["out"]) + except ValueError: + self._proj_in = self._proj_out = None elif projection_format is not None: raise ValueError( f"refused: unknown gdn projection format " @@ -169,6 +534,8 @@ def _gate(a, b, neg_exp_a, dt_bias): device=dev, dtype=torch.bfloat16) self._core_out = torch.empty(1, self._hv, self._d, device=dev, dtype=torch.bfloat16) + self._gn_packed = self._gn_sfa = self._gn_normed = None + self._norm_w = None # prefill chain needs the chunk entries; their absence is not a # bind refusal — prompts simply keep the host form self._chunk_ok = ( @@ -221,6 +588,38 @@ def _host_form(self, *args, **kwargs): guard.notes["host_form_calls"] += 1 return self.host_layer(*args, **kwargs) + @torch.no_grad() + def frt_enable_stash(self, rows, cache_params): + """Arm the per-row state stash for spec-verify passes. + + Called eagerly by the speculative runner before any compiled + pass traces: the buffers exist up front, so the compiled chunk + branch is straight-line. Returns False (and arms nothing) when + the native build does not carry the stash kernel. + """ + hub = getattr(self._gda, "gdn_chunk_from_conv_smem_stash_bf16", + None) + self._stash_hub = hub + if hub is None and _native_stash_op() is None: + return False + if not self._chunk_ok or self._chunk_name not in ( + "gdn_chunk_from_conv_smem_bf16", + "gdn_chunk_from_conv_smem_h_bf16"): + return False + dev = self._conv_w.device + self._stash_rec = torch.empty( + rows, self._hv, self._d, self._d, device=dev, + dtype=torch.bfloat16) + self._stash_mixed = torch.empty( + rows, self._conv_w.shape[0], device=dev, + dtype=torch.bfloat16) + reg = getattr(cache_params, "frt_stash_layers", None) + if reg is None: + reg = {} + cache_params.frt_stash_layers = reg + reg[self._idx] = self + return True + def _prefill_chain(self, hidden_states, cache_params): """Whole-prompt form: conv chunk + fused gating/split/recurrent. @@ -235,9 +634,13 @@ def _prefill_chain(self, hidden_states, cache_params): """ host = self.host_layer S = hidden_states.shape[1] - x = hidden_states.view(S, -1) - allp = (self._proj_in(x) if self._proj_in is not None - else torch.nn.functional.linear(x, self._packed_w)) + # hand the seam the caller's tensor, not a view of it: an + # upstream producer that pre-quantized this activation keys the + # handoff on tensor identity, and a .view here would break it + allp = (self._proj_in(hidden_states).view(S, -1) + if self._proj_in is not None + else torch.nn.functional.linear( + hidden_states.view(S, -1), self._packed_w)) (q0, q1), (z0, z1), (b0, b1), (a0, a1) = self._splits mixed = allp[:, q0:q1].contiguous() a_all = allp[:, a0:a1].contiguous() @@ -289,6 +692,34 @@ def _prefill_chain(self, hidden_states, cache_params): return self._prefill_epilogue( hidden_states, cache_params, allp, mixed, core_out, state, cont, old_slot, S) + stash_rec = getattr(self, "_stash_rec", None) + if stash_rec is not None and S <= stash_rec.shape[0]: + # spec-verify arm: same conv update, same chunk recurrence, + # plus the per-row state stash a rejected round selects + # from instead of re-driving this layer + self._stash_mixed[:S].copy_(mixed) + conv_out = self._conv.causal_conv1d_update_chunk_parallel_bf16( + mixed.view(1, S, -1), self._conv_w, conv_state, + self._conv_b, apply_silu=True) + core_out = torch.empty(S, self._hv, self._d, + device=mixed.device, + dtype=torch.bfloat16) + hub = getattr(self, "_stash_hub", None) + if hub is not None: + hub(conv_out.view(S, -1), a_all, b_all, + self._neg_exp_a, self._dt_bias, state, + self._stash_rec, num_v_heads=self._hv, + num_k_heads=self._hk, head_dim=self._d, + out=core_out) + else: + _native_stash_op()( + conv_out.view(S, -1), a_all, b_all, + self._neg_exp_a, self._dt_bias, state, + core_out.view(S, -1), self._stash_rec, self._hv, + self._hk, self._d) + return self._prefill_epilogue( + hidden_states, cache_params, allp, mixed, core_out, + state, cont, old_slot, S) core_out = torch.empty(S, self._hv, self._d, device=mixed.device, dtype=torch.bfloat16) for s0 in range(0, S, 64): @@ -323,27 +754,58 @@ def _wy_core(self, mixed, a_all, b_all, conv_state, state, S): fallback core pays. """ gda = self._gda - conv_out = self._conv.causal_conv1d_update_chunk_parallel_bf16( - mixed.view(1, S, -1), self._conv_w, conv_state, - self._conv_b, apply_silu=True) - co = conv_out.view(S, -1) + # the GQA conv variant writes the q/k/v splits directly (same + # channel mapping and tap order as conv -> lin_split, bit-exact) + # and saves the full-width conv_out round trip; the head-generic + # arm keeps the plain conv + host-side slicing + conv_gqa = (None if self._wy_h + or _os.environ.get("FRT_CONV_GQA", "1") == "0" + else getattr( + self._conv, "causal_conv1d_update_chunk_parallel_gqa_bf16", + None)) + if conv_gqa is None: + conv_out = self._conv.causal_conv1d_update_chunk_parallel_bf16( + mixed.view(1, S, -1), self._conv_w, conv_state, + self._conv_b, apply_silu=True) + co = conv_out.view(S, -1) g, beta = self._gate_fn( a_all.view(S, self._hv), b_all.view(S, self._hv), self._neg_exp_a, self._dt_bias) if self._wy_h: return self._wy_core_h(gda, co, g, beta, state, S) - q16, k16, v48 = gda.lin_split_qkv_gqa_bf16(co) - q16_l2, k16_l2, q_pack_hv, _k_pack_hk, g_cumsum = \ - gda.gdn_wy_norm_cumsum_pack_qk_bf16(q16, k16, g) - big_a = gda.gdn_wy_kkt_b64_bf16(k16_l2, beta, g_cumsum) - # the packaged triangular solve walks its rows serially and is - # the measured 82% of this chain; the same inverse — semantics - # pinned numerically: inv(I + strict_tril(A)) — through the - # batched cuBLAS solve runs ~40x faster and stays deterministic - eye = torch.eye(64, device=big_a.device, - dtype=big_a.dtype).expand_as(big_a).contiguous() - ai = torch.linalg.solve_triangular( - eye + torch.tril(big_a, -1), eye, upper=False).contiguous() + if conv_gqa is not None: + steps = _native_conv_steps_gqa() + if steps is not None and mixed.shape[-1] == 10240: + q16, k16, v48 = steps( + mixed.view(S, -1), self._conv_w.view(-1, 4), + self._conv_b_dense, conv_state.view(-1, 3)) + else: + q16, k16, v48 = conv_gqa( + mixed.view(1, S, -1), self._conv_w, conv_state, + self._conv_b, apply_silu=True) + q16 = q16.view(S, 16, 128) + k16 = k16.view(S, 16, 128) + v48 = v48.view(S, 48, 128) + else: + q16, k16, v48 = gda.lin_split_qkv_gqa_bf16(co) + ncp = _native_norm_cumsum_pack() if _NCP_ON else None + packed_qk = ncp(q16, k16, g) if ncp is not None else None + if packed_qk is not None: + q16_l2, k16_l2, q_pack_hv, _k_pack_hk, g_cumsum = packed_qk + else: + q16_l2, k16_l2, q_pack_hv, _k_pack_hk, g_cumsum = \ + gda.gdn_wy_norm_cumsum_pack_qk_bf16(q16, k16, g) + # the wmma Gram tier replaces the scalar walk where the + # installed artifact carries it - same signature, same A + # layout, 32.8x on the measured long-prompt term + kkt = getattr(gda, "gdn_wy_kkt_b64_mma_bf16", None) \ + or gda.gdn_wy_kkt_b64_bf16 + big_a = kkt(k16_l2, beta, g_cumsum) + # inv(I + strict_tril(A)): the native fused inverse where the + # build carries it (same fp32 forward-substitution recurrence + # as the batched cuBLAS solve, eye/tril prep folded in), the + # cuBLAS solve otherwise + ai = _wy_ai_inverse(big_a).contiguous() ai_pack = gda.gdn_wy_cast_ai_f32_to_bf16(ai, S) w_pack, u_pack = gda.gdn_wy_recompute_wu_b64_mma_fla_bf16( k16_l2, v48, beta, g_cumsum, ai_pack) @@ -369,10 +831,7 @@ def _wy_core_h(self, gda, co, g, beta, state, S): q_l2, k_l2, q_pack_hv, _k_pack_hk, g_cumsum = \ gda.gdn_wy_norm_cumsum_pack_qk_h_bf16(q, k, g, **hp) big_a = gda.gdn_wy_kkt_b64_h_bf16(k_l2, beta, g_cumsum, **hp) - eye = torch.eye(64, device=big_a.device, - dtype=big_a.dtype).expand_as(big_a).contiguous() - ai = torch.linalg.solve_triangular( - eye + torch.tril(big_a, -1), eye, upper=False).contiguous() + ai = _wy_ai_inverse(big_a).contiguous() ai_pack = gda.gdn_wy_cast_ai_h_f32_to_bf16( ai, S, num_v_heads=self._hv) w_pack, u_pack = gda.gdn_wy_recompute_wu_b64_mma_fla_h_bf16( @@ -448,9 +907,11 @@ def forward(self, hidden_states, cache_params=None, def _decode_one(self, hidden_states, cache_params): host = self.host_layer - x = hidden_states.view(1, -1) - allp = (self._proj_in(x) if self._proj_in is not None - else torch.nn.functional.linear(x, self._packed_w)) + # same identity-preserving handoff as the prefill chain + allp = (self._proj_in(hidden_states).view(1, -1) + if self._proj_in is not None + else torch.nn.functional.linear( + hidden_states.view(1, -1), self._packed_w)) # column slices of a single-row output stay contiguous (q0, q1), (z0, z1), (b0, b1), (a0, a1) = self._splits mixed = allp[:, q0:q1] @@ -479,15 +940,42 @@ def _decode_one(self, hidden_states, cache_params): # never changes, which is what graph replay requires state_in = state_in.to(torch.bfloat16).contiguous() cache_params.recurrent_states[self._idx] = state_in - core_out, new_state = self._gda.gated_delta_recurrent_inout_bf16( - q.view(1, self._hv, self._d), k.view(1, self._hv, self._d), - v.view(1, self._hv, self._d), g, beta, - state_in, use_qk_l2norm=True, - state_out=self._state_a, out=self._core_out) + stream_rec = (_native_recurrent_stream() if self._d == 128 + else None) + if stream_rec is not None: + core_out, new_state = stream_rec( + q.view(1, self._hv, self._d), + k.view(1, self._hv, self._d), + v.view(1, self._hv, self._d), g, beta, state_in, + self._state_a, self._core_out) + else: + core_out, new_state = \ + self._gda.gated_delta_recurrent_inout_bf16( + q.view(1, self._hv, self._d), + k.view(1, self._hv, self._d), + v.view(1, self._hv, self._d), g, beta, + state_in, use_qk_l2norm=True, + state_out=self._state_a, out=self._core_out) # scratch -> slot copy keeps the slot pointer stable; the core # cannot write the slot it is reading within the same step state_in.copy_(new_state) + gnq = (_native_gated_norm_quant() + if self._proj_out is not None + and self._d == 128 else None) + if gnq is not None: + # the norm hands the projection packed rows directly: its + # only consumer would otherwise re-read the row to + # quantize it, one launch per layer + if self._gn_packed is None: + self._arm_gated_norm_quant() + gnq(core_out.view(self._hv, self._d), + z.view(self._hv, self._d), self._norm_w, + self._gn_normed, self._gn_packed, self._gn_sfa, + self._eps) + out = self._proj_out._mm_packed(self._gn_packed, + self._gn_sfa) + return out.view(1, 1, -1).to(hidden_states.dtype) normed = self._fused.rms_norm_gated_silu_bf16( core_out.view(self._hv, self._d), z.view(self._hv, self._d), host.norm.weight, eps=self._eps) @@ -497,6 +985,25 @@ def _decode_one(self, hidden_states, cache_params): host.out_proj.weight)) return out.view(1, 1, -1).to(hidden_states.dtype) + @torch.no_grad() + def _arm_gated_norm_quant(self): + """Stable buffers for the fused gated-norm producer. + + Allocated once, before any capture: the graph records these + addresses, and the norm weight is copied detached so the op + never sits on an autograd edge.""" + host = self.host_layer + dev = self._conv_w.device + n = self._hv * self._d + self._norm_w = host.norm.weight.detach().to( + dev, torch.bfloat16).contiguous().clone() + self._gn_normed = torch.empty(self._hv, self._d, device=dev, + dtype=torch.bfloat16) + self._gn_packed = torch.empty(1, n // 2, device=dev, + dtype=torch.uint8) + self._gn_sfa = torch.zeros(((n + 63) // 64) * 512, device=dev, + dtype=torch.uint8) + @torch.no_grad() def bind_fused_decode_layer(host, layer_idx: int, @@ -543,3 +1050,61 @@ class _Cache: if guard is not None: guard.notes["host_weights"] = "released (one-way)" return bound + + +@torch.no_grad() +def calibrate_gdn_channel_amax(lm, run_once, *, samples: int = 1, + percentile: float = 99.9, + verbose: bool = False) -> int: + """Attach calibrated per-input-channel amax to gated-delta hosts. + + The house calibration front door: the structures ``Collector`` + observes the two projection inputs (``in_proj_qkv``, ``out_proj``) + of every gated-delta host layer while ``run_once`` drives a real + forward — the statistics only ever come from the host's own data + path, never from synthetic tensors. Each observed layer receives + ``_frt_gdn_channel_amax = {"in": [K], "out": [K2]}``, which is the + precondition the ``nvfp4_balance`` projection format checks at + bind. Returns the number of layers calibrated. + """ + from types import SimpleNamespace + + from ...points import Collector, Point + + mods = dict(lm.named_modules()) + hosts = [(name, mod) for name, mod in mods.items() + if hasattr(mod, "conv1d") and hasattr(mod, "A_log") + and hasattr(mod, "in_proj_qkv")] + points, request = [], {} + for name, _mod in hosts: + for attr in ("in_proj_qkv", "out_proj"): + path = f"{name}.{attr}" if name else attr + p = Point("x", path, "input") + points.append(p) + request[f"{p.path}|{p.name}"] = SimpleNamespace( + stat="amax", granularity="channel") + collector = Collector(points=points) + collector.request = request + handles = collector.hooks(lambda path: mods[path]) + try: + for _ in range(max(1, samples)): + run_once() + collector.end_sample() + finally: + for h in handles: + h.remove() + collector.reduce(percentile, verbose=verbose, + label="gdn_channel_amax") + calibrated = 0 + for name, mod in hosts: + pin = f"{name}.in_proj_qkv" if name else "in_proj_qkv" + pout = f"{name}.out_proj" if name else "out_proj" + a_in = collector.channel_amax(pin, "x") + a_out = collector.channel_amax(pout, "x") + if a_in is None or a_out is None: + continue + mod._frt_gdn_channel_amax = { + "in": torch.as_tensor(a_in, dtype=torch.float32), + "out": torch.as_tensor(a_out, dtype=torch.float32)} + calibrated += 1 + return calibrated diff --git a/flash_rt/structures/impls/linear_proj/fp8_static.py b/flash_rt/structures/impls/linear_proj/fp8_static.py index dfed873da..a6f8c0505 100644 --- a/flash_rt/structures/impls/linear_proj/fp8_static.py +++ b/flash_rt/structures/impls/linear_proj/fp8_static.py @@ -182,8 +182,13 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return y.reshape(*shape[:-1], y.shape[-1]) flat = flat.to(torch.bfloat16).contiguous() if self.form == "no_bias": - self._quantize(flat, self._chan, self._input_scale, out=x_fp8) - y = self._gemm(x_fp8, self._w_fp8, self._input_scale, + # bind the return, not the buffer: on the traced path the + # persistent out-buffer is None and the quantize's returned + # tensor is the only handle to the result — feeding the + # buffer forward hands the GEMM a None + x_q = self._quantize(flat, self._chan, self._input_scale, + out=x_fp8) + y = self._gemm(x_q, self._w_fp8, self._input_scale, self._weight_scale, out=out) else: y = self._fn(flat, self._w_fp8, self._bias, self._input_scale, diff --git a/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py b/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py index 98802ef61..3cd2e53bd 100644 --- a/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py +++ b/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py @@ -24,6 +24,7 @@ from __future__ import annotations +import os from collections.abc import Mapping from functools import lru_cache @@ -55,6 +56,56 @@ def _kernel(): return hub_kernel(KERNEL_DEP["repo"], KERNEL_DEP["version"]) +@lru_cache(maxsize=1) +def _native_mrows(): + """The locally built native extension's small-M warp-split GEMM. + + The 16x8x64 block-scaled MMA atom computes a full 16-row tile, so + M<=16 rows cost the same weight stream as one — the spec-verify + rows (draft block + 1, and the shorter re-advance prefixes) are + the customers. Absence is not a refusal: the tiled GEMM serves + every shape correctly, this tier just reads the weights once for + all rows where the build carries it. + + The pointer-style native entry registers as a torch custom op + (with a fake shim) so the compiled multi-row passes trace through + it instead of tripping over the raw stream handle. + """ + try: + from flash_rt import flash_rt_kernels as _fk + except ImportError: + return None + fn = getattr(_fk, "fp4_w4a4_mma_sm120_warpsplit_mrows_bf16out", + None) + if fn is None: + return None + + from torch.library import custom_op + + @custom_op("flashrt_native::warpsplit_mrows", mutates_args=()) + def _op(a_packed: torch.Tensor, w_packed: torch.Tensor, + a_sfa: torch.Tensor, w_sfb: torch.Tensor, n: int, k: int, + warps: int, stages: int) -> torch.Tensor: + m = a_packed.shape[0] + y = torch.empty(m, n, device=a_packed.device, + dtype=torch.bfloat16) + rc = fn(a_packed.data_ptr(), w_packed.data_ptr(), y.data_ptr(), + m, n, k, a_sfa.data_ptr(), w_sfb.data_ptr(), 1.0, + warps, stages, + torch.cuda.current_stream().cuda_stream) + if rc != 0: + raise RuntimeError( + f"warpsplit_mrows refused rc={rc} for M={m} N={n} K={k}") + return y + + @_op.register_fake + def _(a_packed, w_packed, a_sfa, w_sfb, n, k, warps, stages): + return a_packed.new_empty((a_packed.shape[0], n), + dtype=torch.bfloat16) + + return _op + + def _check(weights: Mapping[str, torch.Tensor]) -> tuple[int, int]: w = weights["w"] if w.dim() != 2: @@ -87,18 +138,185 @@ def _quantize_activation(kern, flat: torch.Tensor): flat.to(torch.float16).contiguous()) +#: seams whose tier dispatch runs at call time, by index. A host that +#: compiles one graph for a whole range of row counts cannot carry a +#: Python-level tier branch (it would freeze the tracing sample's +#: choice), so those seams register here and dispatch inside a custom +#: op instead: the trace sees one opaque call, and the branch runs on +#: the shape the call actually receives — at capture, that is the size +#: being captured; outside a capture, it is the live batch. +_RT_SEATS: dict[int, object] = {} + + +def register_runtime_dispatch(seam) -> int: + """Give ``seam`` a call-time dispatching entry; returns its index.""" + idx = len(_RT_SEATS) + _RT_SEATS[idx] = seam + seam._rt_idx = idx + return idx + + +# Registered at import, never on first use: schema inference inside a +# traced region graph-breaks the host's compiled forward, and the first +# call of a lazily registered op lands exactly there. +@torch.library.custom_op("flash_rt_structures::nvfp4_linear_rt", + mutates_args=()) +def _nvfp4_linear_rt(a_packed: torch.Tensor, a_sfa: torch.Tensor, + idx: int) -> torch.Tensor: + return _RT_SEATS[idx]._mm_packed_impl(a_packed, a_sfa) + + +@_nvfp4_linear_rt.register_fake +def _(a_packed, a_sfa, idx): + return a_packed.new_empty((a_packed.shape[0], _RT_SEATS[idx]._n), + dtype=torch.bfloat16) + + +def _runtime_dispatch(): + return _nvfp4_linear_rt + + +#: opt-in census of which tier each call actually lands in, keyed by +#: (row count, tier). Dispatch is data-dependent and lives inside a +#: custom op, so a host's own profile attributes every tier to the same +#: opaque call — reading the choice off the kernel names is exactly +#: what this makes unnecessary. +_TIER_CENSUS: dict | None = ({} if os.environ.get("FRT_TIER_CENSUS") + else None) + +#: whether an adopted pack's per-tensor factor rides each tier's alpha +#: instead of a separate pass over the result. The routes are equivalent +#: (2.3e-3 apart at the seam, both the same distance from BF16), and on +#: real input streams a speculative host's acceptance length is +#: indifferent between them — each numeric path lands somewhere in a +#: ±0.2 content-dependent band around the host's own, with no +#: systematically better draw. On degenerate repeated-sentence prompts +#: the same choice swings acceptance 14%, which is a fact about that +#: protocol, not about the numerics: never judge this switch (or any +#: speculative A/B) on synthetic repetition. Folding is the default +#: because it is free — the separate pass costs ~10ms of 2K TTFT. +_ALPHA_FOLD = os.environ.get("FRT_NVFP4_ALPHA_FOLD", "1") != "0" + + +def _tier_census_note(seam, m) -> None: + if not isinstance(m, int): + tier = "sym->gemm" + elif m == 1 and seam._gemv is not None: + tier = "gemv" + elif 2 <= m <= 16 and seam._mrows is not None: + tier = "mrows" + elif m >= 512 and seam._m256 is not None: + tier = "m256" + else: + tier = "gemm" + key = (m if isinstance(m, int) else -1, tier) + _TIER_CENSUS[key] = _TIER_CENSUS.get(key, 0) + 1 + + +def tier_census() -> dict: + """The census so far, or an empty mapping when it is not armed.""" + return dict(_TIER_CENSUS or {}) + + +if _TIER_CENSUS is not None: + import atexit + + @atexit.register + def _dump_tier_census(): + # printed from whichever process ran the seams: on a serving + # host that is the engine worker, not the caller + rows = sorted(_TIER_CENSUS.items(), key=lambda kv: -kv[1]) + print("[linear_proj.nvfp4] tier census (M, tier): calls", + flush=True) + for (m, tier), n in rows: + print(f"[linear_proj.nvfp4] M={m:<6} {tier:<10} {n}", + flush=True) + + +class _ShareCell: + """One activation-quantization seat shared by sibling projections. + + Reuse is keyed on tensor *identity*: a sibling consumes the stored + quantization only when its input is the very object that produced + it, so a fresh activation can never be served a stranger's data — + at worst the cell misses and the seam quantizes for itself. Inside + one traced graph the identity resolves at trace time, which bakes + the single-quantize dataflow into the compiled prefill and the + captured decode step alike; in eager it holds per call because the + host hands every sibling the same normed hidden. + + Contract: activations must be functional — a caller that mutates a + tensor *in place* and feeds the same object again would hit the + cell with stale contents. This host family never does (every step's + layernorm output is a fresh allocation, and the captured paths + re-run the quantize inside the graph), which is why the linking is + an explicit opt-in per adopted model rather than ambient behavior. + """ + + __slots__ = ("x", "a", "sfa") + + def __init__(self): + self.x = None + self.a = None + self.sfa = None + + +#: sibling groups that consume the same activation in this host family: +#: the attention trio reads the input layernorm's output, the MLP pair +#: reads the post-attention layernorm's output. down/o are not grouped — +#: their inputs are their own. +_SHARE_GROUPS = ( + ("self_attn", ("q_proj", "k_proj", "v_proj")), + ("mlp", ("gate_proj", "up_proj")), +) + + +def link_shared_producers(root: torch.nn.Module) -> int: + """Link sibling NVFP4 seams so each shared activation quantizes once. + + Walks ``root`` for the host family's sibling groups and hands every + group one :class:`_ShareCell`. Returns the number of groups linked. + Additive and reversible: seams without a cell behave exactly as + before, and a group only forms when at least two members are bound. + """ + n_groups = 0 + for name, mod in root.named_modules(): + for tag, members in _SHARE_GROUPS: + if not name.endswith(tag): + continue + seams = [getattr(mod, m, None) for m in members] + seams = [s for s in seams + if isinstance(s, LinearProjNvfp4Dynamic)] + if len(seams) < 2: + continue + cell = _ShareCell() + for s in seams: + s._share = cell + n_groups += 1 + return n_groups + + class LinearProjNvfp4Dynamic(GuardedSeam, torch.nn.Module): """Packed-weight projection: FP4 GEMM with runtime activation scales.""" _frt_can_fallback = False - def __init__(self, w_packed, w_sfb, bias, n, k): + def __init__(self, w_packed, w_sfb, bias, n, k, global_scale=None): super().__init__() self.register_buffer("_w_packed", w_packed) self.register_buffer("_w_sfb", w_sfb) self._bias = bias self._n = n self._k = k + #: a per-tensor factor sitting outside the block scales. Weights + #: this seam packs itself fold everything into the block scale + #: and leave this None; weights adopted from a checkpoint that + #: stores a separate global scale carry it here rather than + #: having the block scales rescaled to absorb it — rescaling + #: would re-round every E4M3 scale and put a lossy step into + #: what is otherwise a pure relayout. + self._w_gs = (None if global_scale is None + else float(global_scale)) kern = _kernel() self._kern = kern self._gemm = kern.fp4_w4a16_linear_bf16 @@ -110,6 +328,46 @@ def __init__(self, w_packed, w_sfb, bias, n, k): gemv = getattr(kern, "fp4_w4a4_gemv_warpsplit_bf16", None) self._gemv = (gemv if gemv is not None and n % 8 == 0 and k % (64 * 4) == 0 else None) + # measured per-shape launch config (512MB-rotation protocol, + # bench_v1_sweep): small-N underfilled shapes want max warps + # (+29%), long-K reads take (8,3) (+4%), the widest tall rows + # edge to (2,3); everything else stays the entry default (4,4). + # Each pick honors the kernel's (K/64)%warps==0 contract. + kt = k // 64 + if os.environ.get("FRT_GEMV_CFG", "1") == "0": + self._gemv_cfg = (4, 4) + elif n <= 2048 and kt % 8 == 0: + self._gemv_cfg = (8, 4) + elif k >= 16384 and kt % 8 == 0: + self._gemv_cfg = (8, 3) + elif n >= 17000 and kt % 2 == 0: + self._gemv_cfg = (2, 3) + else: + self._gemv_cfg = (4, 4) + # 2<=M<=16 rows route to the native multi-row warp-split tier + # where the local build carries it. Per-shape launch config: + # deeper stages hide the strided-B latency the extra A-row + # loads expose — the wide/long-K families take (2,6), the + # short-K o-class (2,4), tiny-N (4,4); each config's K + # divisibility is its own gate. + # hub artifact first (its ops are compile-safe torch ops), + # local native build second, tiled GEMM always the floor + hub_mrows = getattr(kern, "fp4_w4a4_gemm_warpsplit_mrows_bf16", + None) + mrows = hub_mrows if hub_mrows is not None else _native_mrows() + self._mrows_hub = hub_mrows is not None + self._mrows = None + if mrows is not None and n % 8 == 0: + if (k >= 8192 or n >= 8192) and k % 128 == 0: + self._mrows, self._mr_cfg = mrows, (2, 6) + elif n <= 2048 and k % 256 == 0: + self._mrows, self._mr_cfg = mrows, (4, 4) + elif k % 128 == 0: + self._mrows, self._mr_cfg = mrows, (2, 4) + # M>=512 prefill slabs take the cooperative 256-tile tier + # where the artifact carries it - wins every measured prefill + # family over the base tile (the wrapper owns its workspace) + self._m256 = getattr(kern, "nvfp4_gemm_m256_bf16", None) self._frt_arm(dtypes=CAST_OK, device=w_packed.device, k=int(k)) def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -117,16 +375,139 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: if admitted is not PROCEED: return admitted shape = x.shape - flat = x.reshape(-1, shape[-1]) - a_packed, a_sfa = _quantize_activation(self._kern, flat) - if flat.shape[0] == 1 and self._gemv is not None: - y = self._gemv(a_packed, self._w_packed, a_sfa, self._w_sfb) + cell = getattr(self, "_share", None) + if cell is not None and cell.x is x: + a_packed, a_sfa = cell.a, cell.sfa else: + flat = x.reshape(-1, shape[-1]) + a_packed, a_sfa = _quantize_activation(self._kern, flat) + if cell is not None: + cell.x, cell.a, cell.sfa = x, a_packed, a_sfa + y = self._mm_packed(a_packed, a_sfa) + return y.reshape(*shape[:-1], self._n).type_as(x) + + def _mm_packed(self, a_packed: torch.Tensor, + a_sfa: torch.Tensor) -> torch.Tensor: + """Tier-dispatched matmul, through the runtime op when armed.""" + if getattr(self, "_rt_idx", None) is not None: + # armed for a host that reuses one compiled graph across row + # counts: whether this trace shows a symbolic M or a + # specialized one, the branch it records is not the branch + # the replay needs, so every call goes through the op whose + # body runs on the shape actually received + return _runtime_dispatch()(a_packed, a_sfa, self._rt_idx) + return self._mm_packed_impl(a_packed, a_sfa) + + def _mm_packed_impl(self, a_packed: torch.Tensor, + a_sfa: torch.Tensor) -> torch.Tensor: + """Tier-dispatched matmul over a pre-quantized activation. + + The same dispatch the seam's own forward uses, exposed so a + producer that already holds the packed activation (a fused + silu-mul epilogue, a sibling seam's shared quantization) can + feed the weights without a decode round-trip through BF16. + Returns the (m, n) BF16 product with bias applied. + """ + m = a_packed.shape[0] + if _TIER_CENSUS is not None: + _tier_census_note(self, m) + # every tier's entry takes the output scale as its own alpha, so + # an adopted checkpoint's per-tensor factor rides the epilogue + # that is already writing the result. Applying it afterwards + # instead costs a full-size read-modify-write per projection — + # invisible at M=1, and the whole prefill regression at M=2048. + al = 1.0 if (self._w_gs is None or not _ALPHA_FOLD) else self._w_gs + if not isinstance(m, int): + # a symbolic row count with no runtime op armed: a tier + # branch here would freeze the tracing sample's choice, and + # the tiled GEMM is the one tier that serves every M + # this branch returns before the shared epilogue below, so + # it always folds — the switch exists to A/B the dispatched + # tiers, and leaving a path that drops the factor entirely + # would be a bug wearing an experiment's clothes y = self._gemm(a_packed, self._w_packed, a_sfa, self._w_sfb, + 1.0 if self._w_gs is None else self._w_gs, variant=2) + if self._bias is not None: + y = y + self._bias + return y + if m == 1 and self._gemv is not None: + gw, gs = self._gemv_cfg + y = self._gemv(a_packed, self._w_packed, a_sfa, self._w_sfb, + alpha=al, warps=gw, stages=gs) + elif 2 <= m <= 16 and self._mrows is not None: + w_, s_ = self._mr_cfg + if self._mrows_hub: + y = self._mrows(a_packed, self._w_packed, a_sfa, + self._w_sfb, alpha=al, warps=w_, + stages=s_) + else: + # the local native op fixes alpha at 1.0 in its schema; + # scaling after it is the only route on that build + y = self._mrows(a_packed, self._w_packed, a_sfa, + self._w_sfb, self._n, self._k, w_, s_) + if self._w_gs is not None and _ALPHA_FOLD: + y = y * self._w_gs + elif m >= 512 and self._m256 is not None: + y = self._m256(a_packed, self._w_packed, a_sfa, + self._w_sfb, alpha=al) + else: + y = self._gemm(a_packed, self._w_packed, a_sfa, self._w_sfb, + al, variant=2) + if self._w_gs is not None and not _ALPHA_FOLD: + y = y * self._w_gs if self._bias is not None: y = y + self._bias - return y.reshape(*shape[:-1], self._n).type_as(x) + return y + + +@torch.no_grad() +def bind_proj_seam_packed( + w_packed: torch.Tensor, + w_sfb: torch.Tensor, + n: int, + k: int, + *, + global_scale=None, + bias: torch.Tensor | None = None, +) -> LinearProjNvfp4Dynamic: + """Adopt a projection that is *already* NVFP4, without re-gridding. + + The regridding entry (:func:`bind_proj_seam`) exists for a host + holding dense rows. A host holding a packed checkpoint does not need + it, and should not pay it: dequantizing someone else's grid and + quantizing it again with ours replaces their calibration with our + packer's rounding, which is a change of model wearing the costume of + an acceleration. The block-scale layout is the only thing that + differs between the two conventions, and a relayout is a + permutation — it loses nothing. + + ``w_packed`` is ``[N, K/2]`` E2M1 nibble pairs; ``w_sfb`` is the + block scales already in this kernel's atom layout; ``global_scale`` + is the checkpoint's per-tensor factor, applied at the output. + Tensors are adopted by reference: the caller's copy *is* the seam's, + so seating a whole model costs no additional weight memory. + """ + if w_packed.dtype is not torch.uint8: + raise ValueError( + f"packed weight must be uint8 nibble pairs, got " + f"{w_packed.dtype}") + if w_packed.shape != (n, k // 2): + raise ValueError( + f"packed weight {tuple(w_packed.shape)} does not match " + f"N={n} K={k} (expected {(n, k // 2)})") + bound = LinearProjNvfp4Dynamic( + w_packed, w_sfb.view(torch.uint8).reshape(-1), + (None if bias is None else bias.detach().to(torch.bfloat16)), + n, k, global_scale=global_scale) + probe = bound(torch.zeros(1, k, device=w_packed.device, + dtype=torch.bfloat16)) + if probe.shape != (1, n) or not torch.isfinite(probe).all(): + raise ValueError( + f"refused: nvfp4 pack-adopt smoke produced shape " + f"{tuple(probe.shape)}, finite=" + f"{bool(torch.isfinite(probe).all())}") + return bound @torch.no_grad() diff --git a/flash_rt/structures/impls/norm_fused/nvfp4_producer.py b/flash_rt/structures/impls/norm_fused/nvfp4_producer.py new file mode 100644 index 000000000..9c9ee11f4 --- /dev/null +++ b/flash_rt/structures/impls/norm_fused/nvfp4_producer.py @@ -0,0 +1,170 @@ +"""A (1+w)-form RMSNorm that also emits its consumers' NVFP4 input. + +The pipeline fact this serves: a pre-norm decoder norm's output has +exactly one consumer group — the bound FP4 projections that read it — +and each group's first act is to quantize that activation. The fused +producer computes the host norm exactly (fp32, one bf16 round) and +quantizes the same values in the same kernel, publishing the packed +FP4 + SFA through the identity-keyed share cell the consumer group +already honors. The normed BF16 tensor still flows through the host +graph unchanged, so nothing off the calibrated path ever sees a packed +tensor: at worst a consumer misses the cell and quantizes for itself. +""" + +from __future__ import annotations + +import os +from functools import lru_cache + +import torch + +from ..linear_proj.nvfp4_dynamic import (LinearProjNvfp4Dynamic, + _ShareCell) + + +@lru_cache(maxsize=1) +def _native_norm_quant(): + """The local build's fused RMSNorm + NVFP4 quantize producer.""" + try: + from flash_rt import flash_rt_kernels as _fk + except ImportError: + return None + fn = getattr(_fk, "rms_norm_quantize_fp4_sfa_bf16", None) + if fn is None: + return None + + from torch.library import custom_op + + @custom_op("flashrt_native::rms_norm_quant_fp4", mutates_args=()) + def _op(x: torch.Tensor, w: torch.Tensor, + eps: float) -> list[torch.Tensor]: + flat = x.reshape(-1, x.shape[-1]).contiguous() + m, d = flat.shape + normed = torch.empty_like(flat) + packed = torch.empty(m, d // 2, device=x.device, + dtype=torch.uint8) + # zero-filled: the tail rows of the 128-row scale atom are + # never written for a partial block, and run-to-run garbage + # there is the one nondeterministic input a consumer could see + sfa = torch.zeros( + ((m + 127) // 128) * ((d + 63) // 64) * 512, + device=x.device, dtype=torch.uint8) + rc = fn(flat.data_ptr(), w.data_ptr(), float(eps), + normed.data_ptr(), packed.data_ptr(), sfa.data_ptr(), + m, d, torch.cuda.current_stream().cuda_stream) + if rc != 0: + raise RuntimeError( + f"rms_norm_quant_fp4 refused rc={rc} for M={m} D={d}") + return [normed, packed, sfa] + + @_op.register_fake + def _(x, w, eps): + flat_shape = (x.numel() // x.shape[-1], x.shape[-1]) + m, d = flat_shape + return [x.new_empty(flat_shape), + x.new_empty((m, d // 2), dtype=torch.uint8), + x.new_empty( + (((m + 127) // 128) * ((d + 63) // 64) * 512,), + dtype=torch.uint8)] + + return _op + + +class RMSNormQuantFp4Producer(torch.nn.Module): + """Drop-in for the host RMSNorm; feeds the consumer group's cell.""" + + def __init__(self, host_norm: torch.nn.Module, cell: _ShareCell): + super().__init__() + self.host_norm = host_norm + self._cell = cell + self._op = _native_norm_quant() + self._eps = float(getattr(host_norm, "variance_epsilon", + getattr(host_norm, "eps", 1e-6))) + self._min_m = int(os.environ.get("FRT_NORM_QUANT_MIN_M", "64")) + # a detached copy keeps autograd out of the custom op: the + # host weight is a Parameter, and a traced graph that sees it + # demands a backward formula this producer does not carry + self.register_buffer( + "w", host_norm.weight.detach().to(torch.bfloat16) + .contiguous().clone()) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self._op is None or x.dtype is not torch.bfloat16: + return self.host_norm(x) + # measured M-dispatch: a host norm the compiler folded into + # its neighbours can beat this producer's standalone launch at + # narrow widths, so the band where the fused pass pays is + # measured per host rather than assumed + if x.numel() // x.shape[-1] < self._min_m: + return self.host_norm(x) + # inference producer: detach severs the autograd edge a + # Parameter-derived input would otherwise demand a backward + # formula for (the op carries none) + normed, packed, sfa = self._op(x.detach(), self.w, self._eps) + out = normed.reshape(x.shape) + cell = self._cell + cell.x, cell.a, cell.sfa = out, packed, sfa + return out + + +def _consumer_cell(layer: torch.nn.Module, + which: str) -> _ShareCell | None: + """Locate/attach the share cell of a norm's consumer group.""" + if which == "attn": + attn = getattr(layer, "self_attn", None) + if attn is not None: + seams = [getattr(attn, n, None) + for n in ("q_proj", "k_proj", "v_proj")] + if all(isinstance(s, LinearProjNvfp4Dynamic) + for s in seams): + cell = getattr(seams[0], "_share", None) + if cell is None: + cell = _ShareCell() + for s in seams: + s._share = cell + return cell + gdn = getattr(layer, "linear_attn", None) + seam = getattr(gdn, "_proj_in", None) if gdn is not None else None + if isinstance(seam, LinearProjNvfp4Dynamic): + cell = getattr(seam, "_share", None) + if cell is None: + cell = _ShareCell() + seam._share = cell + return cell + return None + mlp = getattr(layer, "mlp", None) + if mlp is not None and hasattr(mlp, "gate_up"): + cell = getattr(mlp, "_share", None) + if cell is None: + cell = _ShareCell() + mlp._share = cell + return cell + return None + + +@torch.no_grad() +def adopt_norm_quant(root: torch.nn.Module, + verbose: bool = False) -> int: + """Swap each decoder norm whose consumer group is FP4-bound.""" + if _native_norm_quant() is None: + return 0 + count = 0 + layers = getattr(root, "layers", None) + if layers is None: + return 0 + for li, layer in enumerate(layers): + for norm_name, which in (("input_layernorm", "attn"), + ("post_attention_layernorm", "mlp")): + norm = getattr(layer, norm_name, None) + if norm is None or not hasattr(norm, "weight"): + continue + cell = _consumer_cell(layer, which) + if cell is None: + continue + setattr(layer, norm_name, + RMSNormQuantFp4Producer(norm, cell)) + count += 1 + if verbose: + print(f"[norm_fused.nvfp4_producer] layer {li} " + f"{norm_name} -> producer") + return count diff --git a/flash_rt/structures/prequantized.py b/flash_rt/structures/prequantized.py index 85771cdc7..92a98c326 100644 --- a/flash_rt/structures/prequantized.py +++ b/flash_rt/structures/prequantized.py @@ -50,6 +50,9 @@ class AdoptionReport: fmt: str replaced: list[str] = field(default_factory=list) conversion_rel_l2: dict[str, float] = field(default_factory=dict) + #: projections seen but deliberately not adopted, path -> reason + #: (a refusal is a routing event; the receipt keeps its reason) + retained: dict[str, str] = field(default_factory=dict) @property def worst_conversion(self) -> float: diff --git a/flash_rt/structures/quantize_on_adopt.py b/flash_rt/structures/quantize_on_adopt.py index bcea0b620..276bf23f6 100644 --- a/flash_rt/structures/quantize_on_adopt.py +++ b/flash_rt/structures/quantize_on_adopt.py @@ -8,6 +8,23 @@ into structure impls brings the whole model into card budget while the attention, norms, and router stay in the host's own precision. +The second format serves the dense sibling of that situation: a +full-precision checkpoint whose weight mass sits in ordinary 2-D linear +projections (attention q/k/v/o and MLP gate/up/down). Each projection +packs to the Hub NVFP4 layout through the same seam binder the +pre-quantized door uses, so both doors produce the same executable +form and everything downstream of adoption is shared. Two families are +deliberately not adopted here, each because it carries its own door: +projections living inside a gated-delta state layer (recognised by the +``conv1d`` + ``A_log`` profile) belong to the ``gated_delta_core`` +scheme, which packs them together with the fused layer and owns their +release semantics; and a vocabulary projection (a Linear whose shape +mirrors an embedding in the same tree) is logits-family precision — a +separate, explicit binder decision, never a bulk-adoption side effect. +Scope is the module tree you hand in: pass the language stack, not the +whole multimodal shell, when towers outside it should keep their own +precision. + The calling convention matches the sibling door: the model is expected CPU-resident straight from its loader; each expert bank streams through the GPU in slabs as it packs, so peak footprint is the dense checkpoint @@ -31,7 +48,7 @@ __all__ = ["quantize_on_adopt"] -_FORMATS = ("moe_experts_nvfp4",) +_FORMATS = ("moe_experts_nvfp4", "linear_proj_nvfp4") def _is_moe_expert_bank(module: torch.nn.Module) -> bool: @@ -48,6 +65,68 @@ def _is_moe_expert_bank(module: torch.nn.Module) -> bool: and gu.shape[1] == 2 * dn.shape[2]) +def _is_gated_delta_layer(module: torch.nn.Module) -> bool: + """The state-layer profile (``conv1d`` + ``A_log``): its projections + are packed by the ``gated_delta_core`` scheme together with the + fused layer, never adopted piecemeal here.""" + return hasattr(module, "conv1d") and hasattr(module, "A_log") + + +def _vocab_signatures(model: torch.nn.Module) -> set: + """Shapes of every embedding in the tree: a Linear mirroring one is + a vocabulary projection (logits family), refused by this door.""" + return {(m.num_embeddings, m.embedding_dim) + for m in model.modules() + if isinstance(m, torch.nn.Embedding)} + + +def _adopt_linear_projections(model: torch.nn.Module, + report, *, verbose: bool) -> None: + from .impls.linear_proj import nvfp4_dynamic + + vocab_sigs = _vocab_signatures(model) + for name, module in list(model.named_modules()): + if _is_gated_delta_layer(module): + continue + for child_name, child in list(module.named_children()): + if not isinstance(child, torch.nn.Linear): + continue + w = getattr(child, "weight", None) + if (w is None or w.dim() != 2 + or not w.is_floating_point() + or hasattr(child, "weight_packed")): + continue + path = f"{name}.{child_name}" if name else child_name + if tuple(w.shape) in vocab_sigs: + report.retained[path] = "vocabulary projection" + continue + bias = getattr(child, "bias", None) + try: + bound, rel = nvfp4_dynamic.bind_proj_seam( + {"w": w.detach(), + "b": None if bias is None else bias.detach()}) + except ValueError as exc: + report.retained[path] = str(exc) + if verbose: + print(f"[quantize_on_adopt] {path}: retained " + f"({str(exc)[:80]})", flush=True) + continue + # release the dense projection before moving on: the + # streaming bind is only slab-peak if retired weights go + child.weight = None + setattr(module, child_name, bound) + report.replaced.append(path) + report.conversion_rel_l2[path] = rel + if verbose: + print(f"[quantize_on_adopt] {path}: relL2={rel:.4f}", + flush=True) + if not report.replaced: + raise ValueError( + "no dense projections found: this tree carries nothing " + "linear_proj_nvfp4 adopts (all out of profile, or already " + "packed)") + + @torch.no_grad() def quantize_on_adopt(model: torch.nn.Module, fmt: str = "moe_experts_nvfp4", *, @@ -59,9 +138,22 @@ def quantize_on_adopt(model: torch.nn.Module, f"unknown quantize-on-adopt format {fmt!r}; supported: " f"{', '.join(_FORMATS)}") + report = AdoptionReport(fmt=fmt) + if fmt == "linear_proj_nvfp4": + _adopt_linear_projections(model, report, verbose=verbose) + # sibling projections that read the same normed hidden share one + # activation quantization (identity-keyed, bit-identical by + # construction) — the adopted checkpoint's default execution form + from .impls.linear_proj import nvfp4_dynamic as _nv + n_shared = _nv.link_shared_producers(model) + torch.cuda.empty_cache() + if verbose: + print(f"[quantize_on_adopt] {report.summary()}; " + f"{n_shared} shared-producer groups", flush=True) + return report + from .impls.moe_experts import nvfp4_dynamic - report = AdoptionReport(fmt=fmt) for name, module in list(model.named_modules()): for child_name, child in list(module.named_children()): if not _is_moe_expert_bank(child): diff --git a/flash_rt/structures/schemes.py b/flash_rt/structures/schemes.py index e684b39e0..cf3635642 100644 --- a/flash_rt/structures/schemes.py +++ b/flash_rt/structures/schemes.py @@ -481,8 +481,32 @@ def resolve_auto() -> str: register("fp8_static_keep_outliers", Fp8Static(keep_outliers=20.0)) register("w8a16_decode", W8A16Decode()) register("w4a16_decode", W4A16Decode()) +class W4A4BalanceDecode(W4A4Decode): + """The W4A4 decode band with the balanced fold on the gated-delta + projections. + + Same band, same seams — the difference is the projection format: + ``nvfp4_balance`` folds a per-input-channel balance (fitted on + calibrated activation amax) into the packed weight, moving + quantization error out of the hot channels. Calibration is the + caller's precondition: run + ``gated_delta_core.fused_layer.calibrate_gdn_channel_amax`` over a + real forward first — an uncalibrated layer keeps the BF16 band and + is counted, never raised. + """ + + def __init__(self, release_host_weights: bool = False) -> None: + super().__init__(release_host_weights) + self.gdn_projection_format = "nvfp4_balance" + self.name = ("w4a4_balance_decode_release" + if release_host_weights else "w4a4_balance_decode") + + register("w4a4_decode", W4A4Decode()) register("w4a4_decode_release", W4A4Decode(release_host_weights=True)) +register("w4a4_balance_decode", W4A4BalanceDecode()) +register("w4a4_balance_decode_release", + W4A4BalanceDecode(release_host_weights=True)) register("none", NoQuant()) register("bf16_structural", Bf16Structural()) register("nvfp4_awq", Nvfp4Awq()) diff --git a/tests/test_structures_quantize_on_adopt.py b/tests/test_structures_quantize_on_adopt.py index 5046e1d1c..0b5869a12 100644 --- a/tests/test_structures_quantize_on_adopt.py +++ b/tests/test_structures_quantize_on_adopt.py @@ -95,3 +95,57 @@ def rec(name, host, chain, verdict="PASS", digest=True): assert "gated_delta_core: 2 host(s) — meets" in text assert "decode_loop: 1 host(s) — single-host" in text assert "HostC" not in text and "HostD" not in text + + +class _GdnLayer(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv1d = torch.nn.Conv1d(8, 8, 4) + self.A_log = torch.nn.Parameter(torch.zeros(4)) + self.in_proj_qkv = torch.nn.Linear(64, 128, bias=False) + + +class _DenseHost(torch.nn.Module): + def __init__(self): + super().__init__() + self.embed = torch.nn.Embedding(100, 64) + self.q_proj = torch.nn.Linear(64, 128, bias=False) + self.gdn = _GdnLayer() + self.head = torch.nn.Linear(64, 100, bias=False) + + +def test_gated_delta_profile_is_recognised_structurally(): + from flash_rt.structures.quantize_on_adopt import ( + _is_gated_delta_layer) + + assert _is_gated_delta_layer(_GdnLayer()) + assert not _is_gated_delta_layer(torch.nn.Linear(8, 8)) + + +def test_linear_adoption_skips_owned_families_and_releases(monkeypatch): + from flash_rt.structures.impls.linear_proj import nvfp4_dynamic + + bound_paths = [] + + def stub(weights): + bound_paths.append(tuple(weights["w"].shape)) + return torch.nn.Identity(), 0.1 + + monkeypatch.setattr(nvfp4_dynamic, "bind_proj_seam", stub) + host = _DenseHost() + report = quantize_on_adopt(host, "linear_proj_nvfp4") + # only the plain projection binds; the state layer's projection is + # the gated_delta_core scheme's to pack, the vocabulary projection + # is a logits-family decision, and the retired dense weight is gone + assert report.replaced == ["q_proj"] + assert bound_paths == [(128, 64)] + assert report.retained == {"head": "vocabulary projection"} + assert isinstance(host.q_proj, torch.nn.Identity) + assert isinstance(host.gdn.in_proj_qkv, torch.nn.Linear) + assert host.gdn.in_proj_qkv.weight is not None + + +def test_projectionless_tree_is_a_named_refusal(): + tree = torch.nn.Sequential(torch.nn.ReLU()) + with pytest.raises(ValueError, match="no dense projections"): + quantize_on_adopt(tree, "linear_proj_nvfp4")