From 56eeef56be1752c08b63128ad7333dc476578bf0 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 14 Aug 2026 18:42:06 -0400 Subject: [PATCH 01/44] adopt: dense projections quantize on adoption - the full-precision sibling of the expert-bank door A checkpoint that ships full-precision linear projections too big for the card packs each one straight to the Hub NVFP4 layout through the same seam binder the pre-quantized door uses - no dequant-requant detour through someone else's grid, and both doors produce the same executable form. Two families are refused with their reasons on the receipt: gated-delta state layers belong to their own scheme's packing and release semantics, and a vocabulary projection is logits-family precision, an explicit binder decision rather than a bulk side effect. --- flash_rt/structures/prequantized.py | 3 + flash_rt/structures/quantize_on_adopt.py | 90 +++++++++++++++++++++- tests/test_structures_quantize_on_adopt.py | 54 +++++++++++++ 3 files changed, 145 insertions(+), 2 deletions(-) diff --git a/flash_rt/structures/prequantized.py b/flash_rt/structures/prequantized.py index 85771cdc..92a98c32 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 bcea0b62..2438989d 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,16 @@ 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) + torch.cuda.empty_cache() + if verbose: + print(f"[quantize_on_adopt] {report.summary()}", 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/tests/test_structures_quantize_on_adopt.py b/tests/test_structures_quantize_on_adopt.py index 5046e1d1..0b5869a1 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") From 8b601d14ae233d1a979bb619b314263693e7256a Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 15 Aug 2026 08:18:34 -0400 Subject: [PATCH 02/44] kernels: sm_120 decode/prefill tiers with in-tree verification harnesses Native mirrors of the three additive kernel tiers delivered to the hub packages, plus the harnesses that judge them, so packaging can verify against this tree directly: - fp4_w4a4_mma_warpsplit_ilv_sm120: interleaved-B warp-split GEMV and its bind-time repack; every block's global B reads become fully sequential, 89-91% of the achievable DRAM read roof on the wide decode shapes, bit-exact against the base kernel. - gated_delta_wy_kkt_mma: the WY K*KT chunk as one block per (chunk, k-head) with wmma accumulation and group-shared Gram tiles; 32.8x over the scalar kernel at S=2048, numerics in the bf16 reduction-order band, same entry surface as the scalar version. - cutlass_nvfp4_gemm_m256_sm120: 256x128x128 cooperative large-M tier; wins every prefill family over the 128-tile baseline. CUTLASS >= 4.5. csrc/kernels/checks/ carries the three judges: bit-exactness for the GEMV (including the device repack against a host reference), the scalar-vs-MMA numeric band and speed for the KKT, and the TFLOPS receipt for the GEMM tier. Build lines are in each file header. --- .../gemm/fp4/cutlass_nvfp4_gemm_m256_sm120.cu | 111 +++++++++++ .../fp4/cutlass_nvfp4_gemm_m256_sm120.cuh | 22 +++ csrc/kernels/checks/bench_gemm_m256.cu | 37 ++++ csrc/kernels/checks/check_warpsplit_ilv.cu | 68 +++++++ csrc/kernels/checks/check_wy_kkt_mma.cu | 82 ++++++++ .../fp4_w4a4_mma_warpsplit_ilv_sm120.cu | 187 ++++++++++++++++++ .../fp4_w4a4_mma_warpsplit_ilv_sm120.cuh | 24 +++ csrc/kernels/gated_delta_wy_kkt_mma.cu | 146 ++++++++++++++ csrc/kernels/gated_delta_wy_kkt_mma.cuh | 15 ++ 9 files changed, 692 insertions(+) create mode 100644 csrc/gemm/fp4/cutlass_nvfp4_gemm_m256_sm120.cu create mode 100644 csrc/gemm/fp4/cutlass_nvfp4_gemm_m256_sm120.cuh create mode 100644 csrc/kernels/checks/bench_gemm_m256.cu create mode 100644 csrc/kernels/checks/check_warpsplit_ilv.cu create mode 100644 csrc/kernels/checks/check_wy_kkt_mma.cu create mode 100644 csrc/kernels/fp4_w4a4_mma_warpsplit_ilv_sm120.cu create mode 100644 csrc/kernels/fp4_w4a4_mma_warpsplit_ilv_sm120.cuh create mode 100644 csrc/kernels/gated_delta_wy_kkt_mma.cu create mode 100644 csrc/kernels/gated_delta_wy_kkt_mma.cuh 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 00000000..b3f76830 --- /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 00000000..1ade2ae6 --- /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/checks/bench_gemm_m256.cu b/csrc/kernels/checks/bench_gemm_m256.cu new file mode 100644 index 00000000..bee77ca7 --- /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 00000000..11af9d07 --- /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_wy_kkt_mma.cu b/csrc/kernels/checks/check_wy_kkt_mma.cu new file mode 100644 index 00000000..a2f7e7ef --- /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_sm120.cu b/csrc/kernels/fp4_w4a4_mma_warpsplit_ilv_sm120.cu new file mode 100644 index 00000000..2afb4d40 --- /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 00000000..a33649f9 --- /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 00000000..db0ad0f1 --- /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 00000000..3f2194a0 --- /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); From e2f4b16cea32bd520c93119b142758693793dfeb Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 15 Aug 2026 08:40:43 -0400 Subject: [PATCH 03/44] decode_loop: block-draft speculative member (DFlash/DSpark family), first assembly One round = one block-draft forward plus one multi-token verify, replacing the MTP family's per-token draft chain. The draft is a small full-attention infill network conditioned on auxiliary features tapped from five target layers (forward pre-hooks over the loop's own passes); verify and re-advance ride _fwd_full with the GDN snapshot/rollback idiom, and the static-cache row mask makes cropping unnecessary. The geometry that survived measurement: the draft emits one logit row per noise slot including the seed slot - dropping the seed row, as the draft checkpoint's own reference implementation does, halves the acceptance length. The Markov bigram head applies sequentially, anchored at the seed. First-assembly receipts: deterministic repeat, coherent output, acceptance length 3.13 at 48 new tokens. Open case on the record: acceptance thins with generation depth (2.26 at 96) while the text stays coherent and easy - not stream hardness, not monotonic state corruption (rounds recover sporadically); the per-round diagnostic ladder is written up in the working notes. Graph capture of the round is the planned second stage, as with the MTP member. --- .../impls/decode_loop/dspark_block.py | 276 ++++++++++++++++++ 1 file changed, 276 insertions(+) create mode 100644 flash_rt/structures/impls/decode_loop/dspark_block.py 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 00000000..d4137fbf --- /dev/null +++ b/flash_rt/structures/impls/decode_loop/dspark_block.py @@ -0,0 +1,276 @@ +"""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. + +Aux features are the *inputs* of the configured target layers (the +serving-side capture point for this host family), collected by forward +pre-hooks that see every loop forward: the prompt pass and each +accepted-prefix re-advance feed the draft's context cache exactly once +per position. + +The target-side round rides the loop's own machinery: ``_fwd_full`` for +the verify and the re-advance (offset row mask, in-place cache slots, +gated-delta continuation via the cache's ``frt_continue`` contract), and +the GDN state snapshot/restore idiom the MTP member established. Static +KV slots beyond the accepted position need no cropping — the row mask +never attends past the current row. + +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 pathlib + +import torch +import torch.nn.functional as F + +__all__ = ["DSparkBlockDraft", "DSparkRunner"] + + +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.""" + + 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 + + t = {k: v.to(device, torch.bfloat16) + for k, v in load_file(str(d / "model.safetensors")).items()} + self._t = t + + # 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(max_len, 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 caches (roped keys), grown append-only + self._ck = [torch.empty(max_len, self.n_kv, self.hd, + device=device, dtype=torch.bfloat16) + for _ in range(self.n_layers)] + self._cv = [torch.empty(max_len, self.n_kv, self.hd, + device=device, dtype=torch.bfloat16) + for _ in range(self.n_layers)] + self._len = 0 + + def reset(self): + self._len = 0 + + 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 = (tgt @ t[p + "k_proj.weight"].T).view( + n, self.n_kv, self.hd) + k = self._rope(_rms(k, t[p + "k_norm.weight"]), pos) + v = (tgt @ t[p + "v_proj.weight"].T).view( + n, self.n_kv, self.hd) + self._ck[li][self._len:self._len + n] = k + self._cv[li][self._len:self._len + n] = v + self._len += n + + @torch.no_grad() + def propose(self, seed_id, start): + """One block forward -> gamma draft tokens (serial markov).""" + t = self._t + g = self.gamma + ids = torch.full((g,), self.mask_id, dtype=torch.long, + device=self._dev) + ids[0] = seed_id + pos = torch.arange(start, start + g, device=self._dev) + h = self._embed(ids.unsqueeze(0))[0] # [g, hidden] + L = self._len + rep = self.n_q // self.n_kv + 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 = (x @ t[a + "q_proj.weight"].T).view(g, self.n_q, self.hd) + q = self._rope(_rms(q, t[a + "q_norm.weight"]), pos) + kn = (x @ t[a + "k_proj.weight"].T).view(g, self.n_kv, self.hd) + kn = self._rope(_rms(kn, t[a + "k_norm.weight"]), pos) + vn = (x @ t[a + "v_proj.weight"].T).view(g, self.n_kv, self.hd) + K = torch.cat([self._ck[li][:L], kn], dim=0) + V = torch.cat([self._cv[li][:L], vn], dim=0) + K = K.repeat_interleave(rep, dim=1) + V = V.repeat_interleave(rep, dim=1) + # dual-source full attention: every noise slot sees all + # context and every noise sibling (bidirectional block) + att = torch.einsum("qhd,khd->hqk", q.float(), K.float()) + att = torch.softmax(att / (self.hd ** 0.5), dim=-1) + o = torch.einsum("hqk,khd->qhd", att, + V.float()).to(h.dtype).reshape(g, -1) + h = h + o @ t[a + "o_proj.weight"].T + x = _rms(h, t[p + "post_attention_layernorm.weight"]) + m = f"{p}mlp." + h = h + (F.silu(x @ t[m + "gate_proj.weight"].T) + * (x @ t[m + "up_proj.weight"].T)) \ + @ t[m + "down_proj.weight"].T + logits = self._head(_rms(h, t["norm.weight"]).unsqueeze(0))[0] + # serial markov: each slot biased by the previously sampled token + w1, w2 = t["markov_head.markov_w1.weight"], \ + t["markov_head.markov_w2.weight"] + drafts = torch.empty(g, dtype=torch.long, device=self._dev) + prev = int(seed_id) + for i in range(g): + bias = (w1[prev].float() @ w2.float().T) + tok = int((logits[i].float() + bias).argmax()) + drafts[i] = tok + prev = tok + return drafts + + +class DSparkRunner: + """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))) + self.last_acceptance = 0.0 + + def _make_hook(self, slot): + def hook(module, args): + self._tap_in[slot] = args[0] + return hook + + def _taps_cat(self): + d = self._draft + return torch.cat([self._tap_in[s][0] for s in + range(len(d.taps))], dim=-1) + + @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() + logits, _ = loop._fwd_full(input_ids, + torch.arange(L, device=dev)) + draft.append_ctx(self._taps_cat(), + torch.arange(L, device=dev)) + loop.cache.frt_continue = True + seed = int(logits[0, -1].float().argmax()) + seq = input_ids[0].tolist() + [seed] + start = L + rounds = 0 + accepted_total = 0 + while len(seq) - L < max_new_tokens: + drafts = draft.propose(seq[start], start) + blk = torch.tensor([[seq[start]] + drafts.tolist()], + device=dev) + pos = torch.arange(start, start + g + 1, device=dev) + conv = {i: loop.cache.conv_states[i].clone() + for i in loop._gdn_slots()} + rec = {i: loop.cache.recurrent_states[i].clone() + for i in loop._gdn_slots()} + vlog, _ = loop._fwd_full(blk, pos) + # 轮特征取自 verify 前向 (快照态下的正确前缀隐层, 与 + # serving 的 aux 捕获点同语义); re-advance 只养状态 + vfeats = self._taps_cat().clone() + post = vlog[0].float().argmax(-1) + match = (blk[0, 1:] == post[:-1]).long() + a = int(match.cumprod(0).sum()) + bonus = int(post[a]) + for i, t in conv.items(): + loop.cache.conv_states[i].copy_(t) + for i, t in rec.items(): + loop.cache.recurrent_states[i].copy_(t) + rpos = torch.arange(start, start + a + 1, device=dev) + loop._fwd_full(blk[:, :a + 1], rpos) + draft.append_ctx(vfeats[:a + 1], rpos) + seq.extend(blk[0, 1:a + 1].tolist() + [bonus]) + start += a + 1 + accepted_total += a + 1 + rounds += 1 + if not hasattr(self, "_round_log"): + self._round_log = [] + if rounds <= 40: + self._round_log.append(a + 1) + self.last_acceptance = accepted_total / max(rounds, 1) + self.last_rounds = getattr(self, "_round_log", None) + return torch.tensor([seq[:L + max_new_tokens]], device=dev) + + def detach(self): + for h in self._hooks: + h.remove() + self._hooks = [] From 96485cd420a2ba5e5e90d82bc61ef99f69eab4cd Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 15 Aug 2026 13:31:55 -0400 Subject: [PATCH 04/44] decode_loop: the block-draft round rides the graph families The draft forward is one captured graph: static context K/V take the block's own keys in place, one masked SDPA covers context and block siblings, and the serial Markov head unrolls on device. The verify captures with the gated-delta snapshot at its head; re-advance drops to the state sublayers alone - the verify already committed the accepted region's KV and carries the bonus logit at the cut - with rollback captured at each graph's head. One host sync per round (the arbiter), verify-input assembly rides inside the draft graph. Real-source coding stream, decode-only, same card and window: 38 -> 86 tok/s at AL 2.9 against the plain loop's 71; the rewrite-heavy stream 95.6 at AL 3.2. Repeat gate green; round profile identical to the whole-model re-advance form. --- .../impls/decode_loop/dspark_block.py | 354 +++++++++++++----- 1 file changed, 268 insertions(+), 86 deletions(-) diff --git a/flash_rt/structures/impls/decode_loop/dspark_block.py b/flash_rt/structures/impls/decode_loop/dspark_block.py index d4137fbf..7d21be08 100644 --- a/flash_rt/structures/impls/decode_loop/dspark_block.py +++ b/flash_rt/structures/impls/decode_loop/dspark_block.py @@ -15,18 +15,37 @@ *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 that see every loop forward: the prompt pass and each -accepted-prefix re-advance feed the draft's context cache exactly once -per position. - -The target-side round rides the loop's own machinery: ``_fwd_full`` for -the verify and the re-advance (offset row mask, in-place cache slots, -gated-delta continuation via the cache's ``frt_continue`` contract), and -the GDN state snapshot/restore idiom the MTP member established. Static -KV slots beyond the accepted position need no cropping — the row mask -never attends past the current row. +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 @@ -52,7 +71,12 @@ def _rms(x, w, eps=1e-6): class DSparkBlockDraft: - """The draft network, functional form over checkpoint tensors.""" + """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 @@ -71,6 +95,7 @@ def __init__(self, draft_dir, embed, lm_head, max_len, device="cuda"): 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()} @@ -89,20 +114,31 @@ def __init__(self, draft_dir, embed, lm_head, max_len, device="cuda"): 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(max_len, device=device)[None] + 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 caches (roped keys), grown append-only - self._ck = [torch.empty(max_len, self.n_kv, self.hd, + # 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.empty(max_len, self.n_kv, self.hd, + 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) + def reset(self): self._len = 0 @@ -128,22 +164,30 @@ def append_ctx(self, feats, pos): k = self._rope(_rms(k, t[p + "k_norm.weight"]), pos) v = (tgt @ t[p + "v_proj.weight"].T).view( n, self.n_kv, self.hd) - self._ck[li][self._len:self._len + n] = k - self._cv[li][self._len:self._len + n] = v + 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, seed_id, start): - """One block forward -> gamma draft tokens (serial markov).""" + 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 - ids = torch.full((g,), self.mask_id, dtype=torch.long, - device=self._dev) - ids[0] = seed_id - pos = torch.arange(start, start + g, device=self._dev) - h = self._embed(ids.unsqueeze(0))[0] # [g, hidden] - L = self._len - rep = self.n_q // self.n_kv + pos = self._start + self._arg + h = self._embed(self._ids.unsqueeze(0))[0] # [g, hidden] + lim = self._start + g + # boolean mask keeps SDPA on its fused backend: an additive + # float mask demotes the GQA path to materialising math (the + # profiler named the expand-clone of the whole context window) + m4 = (self._armax < lim).view(1, 1, 1, -1) + gqa = self.n_q != self.n_kv for li in range(self.n_layers): p = f"layers.{li}." a = f"{p}self_attn." @@ -153,16 +197,14 @@ def propose(self, seed_id, start): kn = (x @ t[a + "k_proj.weight"].T).view(g, self.n_kv, self.hd) kn = self._rope(_rms(kn, t[a + "k_norm.weight"]), pos) vn = (x @ t[a + "v_proj.weight"].T).view(g, self.n_kv, self.hd) - K = torch.cat([self._ck[li][:L], kn], dim=0) - V = torch.cat([self._cv[li][:L], vn], dim=0) - K = K.repeat_interleave(rep, dim=1) - V = V.repeat_interleave(rep, dim=1) - # dual-source full attention: every noise slot sees all - # context and every noise sibling (bidirectional block) - att = torch.einsum("qhd,khd->hqk", q.float(), K.float()) - att = torch.softmax(att / (self.hd ** 0.5), dim=-1) - o = torch.einsum("hqk,khd->qhd", att, - V.float()).to(h.dtype).reshape(g, -1) + self._ck[li].index_copy_(0, pos, kn) + self._cv[li].index_copy_(0, pos, vn) + o = F.scaled_dot_product_attention( + q.transpose(0, 1).unsqueeze(0), + self._ck[li].transpose(0, 1).unsqueeze(0), + self._cv[li].transpose(0, 1).unsqueeze(0), + attn_mask=m4, enable_gqa=gqa) + o = o.squeeze(0).transpose(0, 1).reshape(g, -1) h = h + o @ t[a + "o_proj.weight"].T x = _rms(h, t[p + "post_attention_layernorm.weight"]) m = f"{p}mlp." @@ -170,21 +212,23 @@ def propose(self, seed_id, start): * (x @ t[m + "up_proj.weight"].T)) \ @ t[m + "down_proj.weight"].T logits = self._head(_rms(h, t["norm.weight"]).unsqueeze(0))[0] - # serial markov: each slot biased by the previously sampled token - w1, w2 = t["markov_head.markov_w1.weight"], \ - t["markov_head.markov_w2.weight"] - drafts = torch.empty(g, dtype=torch.long, device=self._dev) - prev = int(seed_id) + # 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[prev].float() @ w2.float().T) - tok = int((logits[i].float() + bias).argmax()) - drafts[i] = tok + 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 - return drafts class DSparkRunner: - """Round driver over a built whole-step loop.""" + """Captured-round driver over a built whole-step loop.""" def __init__(self, loop, draft_dir, model=None): self._loop = loop @@ -198,6 +242,29 @@ def __init__(self, loop, draft_dir, model=None): 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): @@ -205,10 +272,111 @@ def hook(module, args): self._tap_in[slot] = args[0] return hook - def _taps_cat(self): - d = self._draft - return torch.cat([self._tap_in[s][0] for s in - range(len(d.taps))], dim=-1) + 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): + 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): @@ -224,51 +392,65 @@ def generate(self, input_ids, max_new_tokens): 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_() logits, _ = loop._fwd_full(input_ids, torch.arange(L, device=dev)) - draft.append_ctx(self._taps_cat(), - 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 = int(logits[0, -1].float().argmax()) - seq = input_ids[0].tolist() + [seed] + seed = logits[0, -1].float().argmax().view(1) + self._outbuf[0].copy_(seed[0]) + produced = 1 start = L rounds = 0 accepted_total = 0 - while len(seq) - L < max_new_tokens: - drafts = draft.propose(seq[start], start) - blk = torch.tensor([[seq[start]] + drafts.tolist()], - device=dev) - pos = torch.arange(start, start + g + 1, device=dev) - conv = {i: loop.cache.conv_states[i].clone() - for i in loop._gdn_slots()} - rec = {i: loop.cache.recurrent_states[i].clone() - for i in loop._gdn_slots()} - vlog, _ = loop._fwd_full(blk, pos) - # 轮特征取自 verify 前向 (快照态下的正确前缀隐层, 与 - # serving 的 aux 捕获点同语义); re-advance 只养状态 - vfeats = self._taps_cat().clone() - post = vlog[0].float().argmax(-1) - match = (blk[0, 1:] == post[:-1]).long() - a = int(match.cumprod(0).sum()) - bonus = int(post[a]) - for i, t in conv.items(): - loop.cache.conv_states[i].copy_(t) - for i, t in rec.items(): - loop.cache.recurrent_states[i].copy_(t) - rpos = torch.arange(start, start + a + 1, device=dev) - loop._fwd_full(blk[:, :a + 1], rpos) - draft.append_ctx(vfeats[:a + 1], rpos) - seq.extend(blk[0, 1:a + 1].tolist() + [bonus]) - start += a + 1 - accepted_total += a + 1 + if self._gp is None: + 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 not hasattr(self, "_round_log"): - self._round_log = [] if rounds <= 40: - self._round_log.append(a + 1) + self._round_log.append(rows) self.last_acceptance = accepted_total / max(rounds, 1) - self.last_rounds = getattr(self, "_round_log", None) - return torch.tensor([seq[:L + max_new_tokens]], device=dev) + 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: From 4b964082e455ee50a0696ecba977106df7a5c743 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 15 Aug 2026 14:02:42 -0400 Subject: [PATCH 05/44] kernels: multi-row tier for the interleaved warp-split NVFP4 GEMM 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 tier feeds up to 16 live rows through the same atom, same shared-memory plan, same accumulator budget - B stays the only DRAM-bound stream, so a multi-row call lands near single-GEMV cost. The spec-verify shapes (M = block + 1, and the shorter re-advance prefixes) are the customers. Measured on the wide decode shapes at M=8: 74-80% of the 1690 GB/s read roof, 5.3-6.0x over M sequential GEMV calls, 1.16-1.37x over the plain-B multi-row variant (the interleave advantage grows with M), and ~1.3x over the production tile GEMM's 59% at these shapes. Bit-exact per row against the M=1 interleaved kernel across all shapes, M in {1,2,4,7,8,16}, and warp configs (in-tree check included). SFA reads the production 512B-block layout (row r at tile*512 + r*16); B shares the M=1 entry's bind-time interleave. --- .../checks/check_warpsplit_ilv_mrows.cu | 83 +++++++++ .../fp4_w4a4_mma_warpsplit_ilv_mrows_sm120.cu | 159 ++++++++++++++++++ ...fp4_w4a4_mma_warpsplit_ilv_mrows_sm120.cuh | 27 +++ 3 files changed, 269 insertions(+) create mode 100644 csrc/kernels/checks/check_warpsplit_ilv_mrows.cu create mode 100644 csrc/kernels/fp4_w4a4_mma_warpsplit_ilv_mrows_sm120.cu create mode 100644 csrc/kernels/fp4_w4a4_mma_warpsplit_ilv_mrows_sm120.cuh 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 00000000..d156ba44 --- /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/fp4_w4a4_mma_warpsplit_ilv_mrows_sm120.cu b/csrc/kernels/fp4_w4a4_mma_warpsplit_ilv_mrows_sm120.cu new file mode 100644 index 00000000..9d483aa2 --- /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 00000000..57255964 --- /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 From 122b6944b09aec5b7d5f9c703c2b2603bb5f98ea Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 15 Aug 2026 14:46:04 -0400 Subject: [PATCH 06/44] adopt: small-M rows route to the native multi-row warp-split tier 2<=M<=16 activations - the spec-verify block and its re-advance prefixes - dispatch to the native build's multi-row warp-split GEMM where the local tree carries it: the 16x8x64 block-scaled atom computes a full 16-row tile, so all rows ride one weight stream. The pointer-style native entry registers as a torch custom op with a fake shim, so the compiled multi-row passes trace through it and the call lands on the capture stream. Per-shape launch config: deeper stages hide the strided-B latency the extra A-row loads expose. Layout compatibility is proven at the quantizer, not assumed: the production activation quantizer's multi-row SFA block feeds the tier bit-exactly against per-row GEMV calls across every projection family and M in {2,4,7,8}. Absence of the native build changes nothing - the tiled GEMM keeps serving every shape. Same card, same real-source coding stream: verify pass 21.4 -> 18.3ms, re-advance tiers 4.0-6.0ms, decode-only 86 -> 96.6 tok/s at AL 2.9 against the plain loop's 71.5; repeat gate green, round profile identical to the tiled-GEMM form. --- .../impls/linear_proj/nvfp4_dynamic.py | 72 ++++++++++++++++++- 1 file changed, 71 insertions(+), 1 deletion(-) diff --git a/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py b/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py index 98802ef6..92bf7461 100644 --- a/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py +++ b/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py @@ -55,6 +55,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: @@ -110,6 +160,21 @@ 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) + # 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. + mrows = _native_mrows() + 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) self._frt_arm(dtypes=CAST_OK, device=w_packed.device, k=int(k)) def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -119,8 +184,13 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: 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: + m = flat.shape[0] + if m == 1 and self._gemv is not None: y = self._gemv(a_packed, self._w_packed, a_sfa, self._w_sfb) + elif 2 <= m <= 16 and self._mrows is not None: + w_, s_ = self._mr_cfg + y = self._mrows(a_packed, self._w_packed, a_sfa, + self._w_sfb, self._n, self._k, w_, s_) else: y = self._gemm(a_packed, self._w_packed, a_sfa, self._w_sfb, variant=2) From 46467f6b8715f40b91333e53c35b774d6c9b4dd9 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 15 Aug 2026 15:13:44 -0400 Subject: [PATCH 07/44] decode_loop: the draft forward rides the vendored fa2 and keeps bf16 The draft's dual-source attention moves off SDPA onto the vendored fa2 seqused entry: the valid-K length (context + block) lives in a device buffer the kernel reads at run time, so the captured launch replays each round's window with no mask and no materialised attention math - the compiled SDPA path was demoting the masked GQA form to an expand-clone of the whole context window. Registered as a custom op with a fake shim so the compiled draft graph traces through it on the capture stream. Draft projection precision is decided by measurement, not appetite: W4A4 through the adoption door costs acceptance (2.93 -> 2.79 on the real-source coding stream) and W8A8 per-tensor costs the same (2.80) - a 1.4B draft's weight stream is not the binding constraint, its prediction quality is. bf16 stays the default; both arms remain behind FRT_DSPARK_DRAFT_QUANT for streams that judge differently. Same card, same stream: propose graph 6.95 -> 4.12 ms, decode-only 96.6 -> 100.9 tok/s at AL 2.93 against the plain loop's 71.5; draft tokens identical to the SDPA form round for round, repeat gate green. --- .../impls/decode_loop/dspark_block.py | 123 ++++++++++++++---- 1 file changed, 100 insertions(+), 23 deletions(-) diff --git a/flash_rt/structures/impls/decode_loop/dspark_block.py b/flash_rt/structures/impls/decode_loop/dspark_block.py index 7d21be08..3f1d9b38 100644 --- a/flash_rt/structures/impls/decode_loop/dspark_block.py +++ b/flash_rt/structures/impls/decode_loop/dspark_block.py @@ -56,6 +56,7 @@ from __future__ import annotations import json +import os import pathlib import torch @@ -64,6 +65,40 @@ __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) @@ -100,6 +135,41 @@ def __init__(self, draft_dir, embed, lm_head, max_len, device="cuda"): 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 @@ -138,10 +208,26 @@ def __init__(self, draft_dir, embed, lm_head, max_len, device="cuda"): 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) @@ -159,11 +245,9 @@ def append_ctx(self, feats, pos): n = feats.shape[0] for li in range(self.n_layers): p = f"layers.{li}.self_attn." - k = (tgt @ t[p + "k_proj.weight"].T).view( - n, self.n_kv, self.hd) + 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 = (tgt @ t[p + "v_proj.weight"].T).view( - n, self.n_kv, self.hd) + 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 @@ -182,35 +266,28 @@ def propose_(self): g = self.gamma pos = self._start + self._arg h = self._embed(self._ids.unsqueeze(0))[0] # [g, hidden] - lim = self._start + g - # boolean mask keeps SDPA on its fused backend: an additive - # float mask demotes the GQA path to materialising math (the - # profiler named the expand-clone of the whole context window) - m4 = (self._armax < lim).view(1, 1, 1, -1) - gqa = self.n_q != self.n_kv + # 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 = (x @ t[a + "q_proj.weight"].T).view(g, self.n_q, self.hd) + 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 = (x @ t[a + "k_proj.weight"].T).view(g, self.n_kv, self.hd) + 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 = (x @ t[a + "v_proj.weight"].T).view(g, self.n_kv, self.hd) + 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 = F.scaled_dot_product_attention( - q.transpose(0, 1).unsqueeze(0), - self._ck[li].transpose(0, 1).unsqueeze(0), - self._cv[li].transpose(0, 1).unsqueeze(0), - attn_mask=m4, enable_gqa=gqa) - o = o.squeeze(0).transpose(0, 1).reshape(g, -1) - h = h + o @ t[a + "o_proj.weight"].T + 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 + (F.silu(x @ t[m + "gate_proj.weight"].T) - * (x @ t[m + "up_proj.weight"].T)) \ - @ t[m + "down_proj.weight"].T + 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 From 6484484260315181acc263fe7a1250c13c684e2a Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 15 Aug 2026 15:40:45 -0400 Subject: [PATCH 08/44] gated_delta: per-row state stash - speculative rollback by selection The from-conv chunk core gains a stash arm: identical recurrence, gating math, and per-row bf16 state requantisation, with the carried state after each row additionally written to a stash slab. The stash write IS the carried requantisation, so row s is bit-equal to the final state a re-advance over rows 0..s would store - proven at the kernel level against the plain chunk entry across prefix lengths. A rejected speculative round then rolls back by selection: pick the stash row at the accepted length, rebuild the conv window from the snapshot tail plus the stashed raw rows (the epilogue's own semantics), and touch no projection at all. The re-advance graphs drop from forty-eight sublayer re-drives to a fixed set of copies: 4.0-6.0ms -> 0.27-0.39ms per shape, while the verify pass pays 0.4ms for the stash writes. The fused layer arms the stash eagerly (before any compiled pass traces) and only where the native build carries the kernel; absent that, rejected rounds keep the re-drive form. Same card, same real-source coding stream: decode-only 100.9 -> 116.4 tok/s at AL 2.97 against the plain loop's 71.5; repeat gate green. --- CMakeLists.txt | 5 + csrc/bindings.cpp | 29 +++ .../kernels/gdn_chunk_from_conv_smem_stash.cu | 208 ++++++++++++++++++ .../gdn_chunk_from_conv_smem_stash.cuh | 26 +++ .../impls/decode_loop/dspark_block.py | 37 +++- .../impls/gated_delta_core/fused_layer.py | 95 ++++++++ 6 files changed, 395 insertions(+), 5 deletions(-) create mode 100644 csrc/kernels/gdn_chunk_from_conv_smem_stash.cu create mode 100644 csrc/kernels/gdn_chunk_from_conv_smem_stash.cuh diff --git a/CMakeLists.txt b/CMakeLists.txt index 5c5a6fb0..12037410 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1347,6 +1347,11 @@ else() message(STATUS "SM120 NVFP4 W4A16 fp4_w4a4 kernels: SKIPPED (slim build)") endif() +# 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 ade0577e..3c8aeb5c 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -186,6 +186,7 @@ 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/gdn_chunk_from_conv_smem_stash.cuh" #include "quantize/fp8_block128_dequant.cuh" #ifdef FLASHRT_HAVE_NVFP4_SWIZZLE #include "quantize/fp8_block128_to_nvfp4_swizzled.cuh" @@ -6634,6 +6635,34 @@ graph-replay safe) to fill the SMs on long K. M in 1..16; N%8==0; K%64==0; #endif + 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/kernels/gdn_chunk_from_conv_smem_stash.cu b/csrc/kernels/gdn_chunk_from_conv_smem_stash.cu new file mode 100644 index 00000000..0b9203a7 --- /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 00000000..cb634e42 --- /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/flash_rt/structures/impls/decode_loop/dspark_block.py b/flash_rt/structures/impls/decode_loop/dspark_block.py index 3f1d9b38..9f82fa1e 100644 --- a/flash_rt/structures/impls/decode_loop/dspark_block.py +++ b/flash_rt/structures/impls/decode_loop/dspark_block.py @@ -448,11 +448,29 @@ def _capture(self): for m in range(1, g + 1): gr = torch.cuda.CUDAGraph() with torch.cuda.graph(gr): - 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) + 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() @@ -487,6 +505,15 @@ def generate(self, input_ids, max_new_tokens): 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) 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 4ae2ba22..a4af06b2 100644 --- a/flash_rt/structures/impls/gated_delta_core/fused_layer.py +++ b/flash_rt/structures/impls/gated_delta_core/fused_layer.py @@ -44,6 +44,53 @@ @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 @@ -221,6 +268,35 @@ 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. + """ + if _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. @@ -289,6 +365,25 @@ 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) + _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): From 7791fa121352909550a18274d34277432579bd9a Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 15 Aug 2026 16:06:52 -0400 Subject: [PATCH 09/44] structures: hub-first tiers, the wmma Gram swap, and the dspark member Capability probes now prefer the installed artifact and fall back to the local native build, tiled/scalar entries always the floor: - linear_proj: the multi-row warp-split tier resolves hub-first (fp4_w4a4_gemm_warpsplit_mrows_bf16, a compile-safe torch op) and keeps the native custom-op fallback for local builds. - gated_delta: the per-row stash arm resolves hub-first (gdn_chunk_from_conv_smem_stash_bf16); the whole-prompt WY chain takes the wmma Gram tier (gdn_wy_kkt_b64_mma_bf16) where the artifact carries it - same signature, same A layout, and the measured long-prompt Gram term drops off the profile's top set (42.1ms -> 1.3ms class at 2K). - decode_loop: enable_dspark(draft_dir) attaches the block-draft speculative runner as a loop member, mirroring enable_mtp. Real-source coding stream after the Gram swap: prefill at 5.2K drops ~100ms, decode round unchanged at ~25.5ms (AL moves within band with the prompt-numerics shift), repeat gate green. --- .../impls/decode_loop/whole_step.py | 15 ++++++++++ .../impls/gated_delta_core/fused_layer.py | 29 +++++++++++++++---- .../impls/linear_proj/nvfp4_dynamic.py | 15 ++++++++-- 3 files changed, 50 insertions(+), 9 deletions(-) diff --git a/flash_rt/structures/impls/decode_loop/whole_step.py b/flash_rt/structures/impls/decode_loop/whole_step.py index 2f9b378e..efcc2b3f 100644 --- a/flash_rt/structures/impls/decode_loop/whole_step.py +++ b/flash_rt/structures/impls/decode_loop/whole_step.py @@ -456,6 +456,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)) 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 a4af06b2..3ac1281b 100644 --- a/flash_rt/structures/impls/gated_delta_core/fused_layer.py +++ b/flash_rt/structures/impls/gated_delta_core/fused_layer.py @@ -277,7 +277,10 @@ def frt_enable_stash(self, rows, cache_params): branch is straight-line. Returns False (and arms nothing) when the native build does not carry the stash kernel. """ - if _native_stash_op() is None: + 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", @@ -377,10 +380,19 @@ def _prefill_chain(self, hidden_states, cache_params): core_out = torch.empty(S, self._hv, self._d, device=mixed.device, dtype=torch.bfloat16) - _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) + 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) @@ -430,7 +442,12 @@ def _wy_core(self, mixed, a_all, b_all, conv_state, 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 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) # 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 diff --git a/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py b/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py index 92bf7461..63fddd4a 100644 --- a/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py +++ b/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py @@ -166,7 +166,12 @@ def __init__(self, w_packed, w_sfb, bias, n, k): # 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. - mrows = _native_mrows() + # 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: @@ -189,8 +194,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: y = self._gemv(a_packed, self._w_packed, a_sfa, self._w_sfb) elif 2 <= m <= 16 and self._mrows is not None: w_, s_ = self._mr_cfg - y = self._mrows(a_packed, self._w_packed, a_sfa, - self._w_sfb, self._n, self._k, w_, s_) + if self._mrows_hub: + y = self._mrows(a_packed, self._w_packed, a_sfa, + self._w_sfb, warps=w_, stages=s_) + else: + y = self._mrows(a_packed, self._w_packed, a_sfa, + self._w_sfb, self._n, self._k, w_, s_) else: y = self._gemm(a_packed, self._w_packed, a_sfa, self._w_sfb, variant=2) From 917f202b51060930c7100aa48352f7935fab82aa Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 15 Aug 2026 16:12:21 -0400 Subject: [PATCH 10/44] adopt: prefill slabs take the cooperative 256-tile tier M>=512 activations dispatch to nvfp4_gemm_m256_bf16 where the artifact carries it - the tier wins every measured prefill family over the base tile, and the wrapper owns its workspace. Absence changes nothing. With the Gram swap and this tier together, the 2K prefill wall drops 214.5 -> 164.8 ms on the 27B host; the decode round and the token stream are unchanged, repeat gate green. --- flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py b/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py index 63fddd4a..5350ef67 100644 --- a/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py +++ b/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py @@ -180,6 +180,10 @@ def __init__(self, w_packed, w_sfb, bias, n, k): 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: @@ -200,6 +204,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: else: y = self._mrows(a_packed, self._w_packed, a_sfa, self._w_sfb, self._n, self._k, w_, s_) + elif m >= 512 and self._m256 is not None: + y = self._m256(a_packed, self._w_packed, a_sfa, + self._w_sfb) else: y = self._gemm(a_packed, self._w_packed, a_sfa, self._w_sfb, variant=2) From 97eb49974b2299447a7f29ec04d39a1fa818ef53 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 15 Aug 2026 18:44:57 -0400 Subject: [PATCH 11/44] structures: share activation quantization across sibling projections Sibling projections that read the same normed hidden (the attention q/k/v trio, the MLP gate/up pair) each quantized their input separately; inductor does not CSE the custom quantize op, so the same rows were packed up to three times per site. link_shared_producers hands each sibling group one identity-keyed cell: a seam reuses the stored quantization only when its input is the very tensor that produced it, so a fresh activation can never be served stale data. quantize_on_adopt links the groups after adoption - the checkpoint's default execution form. Bit-identical by construction; the sharing contract (functional activations) is documented on the cell. --- .../impls/linear_proj/nvfp4_dynamic.py | 75 ++++++++++++++++++- flash_rt/structures/quantize_on_adopt.py | 8 +- 2 files changed, 79 insertions(+), 4 deletions(-) diff --git a/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py b/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py index 5350ef67..22fa730b 100644 --- a/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py +++ b/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py @@ -137,6 +137,69 @@ def _quantize_activation(kern, flat: torch.Tensor): flat.to(torch.float16).contiguous()) +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.""" @@ -191,9 +254,15 @@ 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) - m = flat.shape[0] + 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 + m = a_packed.shape[0] if m == 1 and self._gemv is not None: y = self._gemv(a_packed, self._w_packed, a_sfa, self._w_sfb) elif 2 <= m <= 16 and self._mrows is not None: diff --git a/flash_rt/structures/quantize_on_adopt.py b/flash_rt/structures/quantize_on_adopt.py index 2438989d..276bf23f 100644 --- a/flash_rt/structures/quantize_on_adopt.py +++ b/flash_rt/structures/quantize_on_adopt.py @@ -141,9 +141,15 @@ def quantize_on_adopt(model: torch.nn.Module, 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()}", flush=True) + print(f"[quantize_on_adopt] {report.summary()}; " + f"{n_shared} shared-producer groups", flush=True) return report from .impls.moe_experts import nvfp4_dynamic From 8fe5ec9dbfcd8cf1cc048809191405c00e722dbb Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 15 Aug 2026 18:44:57 -0400 Subject: [PATCH 12/44] decode_loop: sliced prompt pass for deep windows A one-shot prompt pass over a deep window overruns the card: the whole-prompt activations, the GQA repeat-expand of the KV heads, and the full-rows lm_head each materialize window-sized transients. The loop (and the block-draft runner) now accept prefill_chunk: the prompt drives the same offset-mask forward in slices, gated-delta state carrying across slices through the continuation branch the verify batches already exercise. Slice attention routes through a compiled flex_attention arm over an offset-causal block mask (eager flex is the math fallback and materializes the score matrix; masked GQA SDPA does the same) - the mask depends only on slice length and start, so a prompt's handful of masks replay across layers, chunks and runs. The runner's slices take the last-row head: nothing in a prompt slice consumes the full-rows logits. Both arms are flag-gated to chunked loops; established short-window forms keep their exact paths. The sliced pass is not bitwise-identical to the one-shot form: two runs of the FP4 pipeline decorrelate at low mantissa bits (measured slice-vs-oneshot logits cosine equals the arbiter band squared), so the equivalence contract is the arbiter gate plus continuation identity, never bit equality. --- flash_rt/structures/__init__.py | 5 +- .../impls/decode_loop/dspark_block.py | 37 ++++++++++++-- .../structures/impls/decode_loop/fp8_kv.py | 49 +++++++++++++++++++ .../impls/decode_loop/whole_step.py | 40 +++++++++++++-- 4 files changed, 119 insertions(+), 12 deletions(-) diff --git a/flash_rt/structures/__init__.py b/flash_rt/structures/__init__.py index fce219b3..a9a66d70 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/impls/decode_loop/dspark_block.py b/flash_rt/structures/impls/decode_loop/dspark_block.py index 9f82fa1e..786e6434 100644 --- a/flash_rt/structures/impls/decode_loop/dspark_block.py +++ b/flash_rt/structures/impls/decode_loop/dspark_block.py @@ -492,11 +492,38 @@ def generate(self, input_ids, max_new_tokens): for i in loop._full: loop.cache.key_cache[i][:, :, L:].zero_() loop.cache.value_cache[i][:, :, L:].zero_() - 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)) + 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]) diff --git a/flash_rt/structures/impls/decode_loop/fp8_kv.py b/flash_rt/structures/impls/decode_loop/fp8_kv.py index ba2e6fd9..0f31de27 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 efcc2b3f..0670d754 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()] @@ -859,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) From 8f23cec08bf919df6f1cc1fdb3a710de048d541a Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 15 Aug 2026 18:44:57 -0400 Subject: [PATCH 13/44] gated_delta_core: calibrated balance arm for the projection band nvfp4_balance projection format: the same FP4 band with a per-input-channel balance fitted on calibrated activation amax, plus calibrate_gdn_channel_amax - the house Collector observing the two projection inputs over a caller-supplied real forward. A host layer without attached calibration keeps the BF16 band (counted, never raised). Schemes w4a4_balance_decode(_release) select the format; calibration is the caller's precondition. Measured on this host family: the out-projection seat cuts its quantization error 2.5x (the in-projection saturates at 1.3x), but the end-to-end teacher-forced gate is unchanged - the near-tie token flips are the aggregate of the whole pipeline's quantization noise, not the projection grid alone. The arm therefore ships as an explicit opt-in with this receipt, not as a default. --- .../impls/gated_delta_core/fused_layer.py | 77 +++++++++++++++++++ flash_rt/structures/schemes.py | 24 ++++++ 2 files changed, 101 insertions(+) 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 3ac1281b..1a63dab4 100644 --- a/flash_rt/structures/impls/gated_delta_core/fused_layer.py +++ b/flash_rt/structures/impls/gated_delta_core/fused_layer.py @@ -208,6 +208,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 " @@ -655,3 +674,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/schemes.py b/flash_rt/structures/schemes.py index e684b39e..cf363564 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()) From 09f0dac3dace6e3be4ea5b208e19c7957fe1f05a Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sun, 16 Aug 2026 18:00:47 -0400 Subject: [PATCH 14/44] kernels: fused producer and WY-chain tiers for the sm_120 decode/prefill path Four additive kernels with native bindings: a SwiGLU activation + NVFP4 quantize producer (quantize stage transcribed from the production bf16 quantize kernel, bit-exact against the split chain), a (1+w)-form RMSNorm + NVFP4 quantize producer, a batched 64x64 unit-lower-triangular inverse for the WY chain (column-independent forward substitution, identity/tril preparation folded in), and a v2 launch plan for the WY norm/pack + gate-cumsum pair (chunk-parallel cumsum, bit-exact per chunk). Also exposes the interleaved warp-split GEMV tiers through the native module. --- CMakeLists.txt | 27 +++ csrc/bindings.cpp | 127 ++++++++++++++ csrc/kernels/batched_unit_ltri_inv64.cu | 58 ++++++ csrc/kernels/batched_unit_ltri_inv64.cuh | 22 +++ csrc/kernels/gdn_wy_norm_cumsum_pack_qk_v2.cu | 138 +++++++++++++++ .../kernels/gdn_wy_norm_cumsum_pack_qk_v2.cuh | 28 +++ .../kernels/rms_norm_quantize_fp4_sfa_bf16.cu | 166 ++++++++++++++++++ .../rms_norm_quantize_fp4_sfa_bf16.cuh | 24 +++ .../kernels/silu_mul_quantize_fp4_sfa_bf16.cu | 133 ++++++++++++++ .../silu_mul_quantize_fp4_sfa_bf16.cuh | 23 +++ 10 files changed, 746 insertions(+) create mode 100644 csrc/kernels/batched_unit_ltri_inv64.cu create mode 100644 csrc/kernels/batched_unit_ltri_inv64.cuh create mode 100644 csrc/kernels/gdn_wy_norm_cumsum_pack_qk_v2.cu create mode 100644 csrc/kernels/gdn_wy_norm_cumsum_pack_qk_v2.cuh create mode 100644 csrc/kernels/rms_norm_quantize_fp4_sfa_bf16.cu create mode 100644 csrc/kernels/rms_norm_quantize_fp4_sfa_bf16.cuh create mode 100644 csrc/kernels/silu_mul_quantize_fp4_sfa_bf16.cu create mode 100644 csrc/kernels/silu_mul_quantize_fp4_sfa_bf16.cuh diff --git a/CMakeLists.txt b/CMakeLists.txt index 12037410..44745694 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1347,6 +1347,33 @@ 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) + # 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 diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index 3c8aeb5c..f5249baf 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -186,7 +186,13 @@ 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 "quantize/fp8_block128_dequant.cuh" #ifdef FLASHRT_HAVE_NVFP4_SWIZZLE #include "quantize/fp8_block128_to_nvfp4_swizzled.cuh" @@ -6631,10 +6637,131 @@ 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("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, diff --git a/csrc/kernels/batched_unit_ltri_inv64.cu b/csrc/kernels/batched_unit_ltri_inv64.cu new file mode 100644 index 00000000..e876b2bc --- /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 00000000..d18ff4fd --- /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/gdn_wy_norm_cumsum_pack_qk_v2.cu b/csrc/kernels/gdn_wy_norm_cumsum_pack_qk_v2.cu new file mode 100644 index 00000000..8da3ea37 --- /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 00000000..d79da5a8 --- /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_quantize_fp4_sfa_bf16.cu b/csrc/kernels/rms_norm_quantize_fp4_sfa_bf16.cu new file mode 100644 index 00000000..8bef58b0 --- /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 00000000..e74c70e8 --- /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 00000000..7241d492 --- /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 00000000..006d2b84 --- /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 From 8ffc58e7b948100bf23c6b6c2ed3593f3ce819f3 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sun, 16 Aug 2026 18:00:57 -0400 Subject: [PATCH 15/44] structures: fused-GLU decoder_ffn over merged gate/up projections The gate and up seams concatenate into one merged seam (one weight stream, one launch per tier) and the activation between them collapses into the fused silu-mul quantize producer, feeding the down projection its packed FP4 input directly. The projection seam grows a tier-dispatched packed-input entry and a measured per-shape launch config for the decode GEMV. --- .../impls/decoder_ffn/nvfp4_fused.py | 164 ++++++++++++++++++ .../impls/linear_proj/nvfp4_dynamic.py | 36 +++- 2 files changed, 198 insertions(+), 2 deletions(-) create mode 100644 flash_rt/structures/impls/decoder_ffn/nvfp4_fused.py 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 00000000..942d8bb1 --- /dev/null +++ b/flash_rt/structures/impls/decoder_ffn/nvfp4_fused.py @@ -0,0 +1,164 @@ +"""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) + +@lru_cache(maxsize=1) +def _native_silu_mul(): + """The local build's 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. Registered as a torch custom op with a fake so the + compiled prefill and the captured decode step trace through it. + """ + 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) + sfa = torch.empty( + ((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/linear_proj/nvfp4_dynamic.py b/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py index 22fa730b..b634842e 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 @@ -223,6 +224,22 @@ 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 @@ -262,9 +279,24 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: 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 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 m == 1 and self._gemv is not None: - y = self._gemv(a_packed, self._w_packed, a_sfa, self._w_sfb) + gw, gs = self._gemv_cfg + y = self._gemv(a_packed, self._w_packed, a_sfa, self._w_sfb, + warps=gw, stages=gs) elif 2 <= m <= 16 and self._mrows is not None: w_, s_ = self._mr_cfg if self._mrows_hub: @@ -281,7 +313,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: variant=2) 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() From a3ce038948ec99c0ef2e01665739777fdd13bbcd Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sun, 16 Aug 2026 18:34:35 -0400 Subject: [PATCH 16/44] kernels: step-batched causal conv1d update with GQA split outputs One thread computes 8 consecutive tokens with the taps rolling through registers, cutting the packaged kernel's K-fold DRAM re-read of every input element to (STEPS+K-1)/STEPS; the q/k/v channel split folds into the store. Tap order and fma chain match the packaged kernel. --- CMakeLists.txt | 4 + csrc/bindings.cpp | 22 ++++ .../causal_conv1d_update_steps_gqa_bf16.cu | 105 ++++++++++++++++++ .../causal_conv1d_update_steps_gqa_bf16.cuh | 26 +++++ 4 files changed, 157 insertions(+) create mode 100644 csrc/kernels/causal_conv1d_update_steps_gqa_bf16.cu create mode 100644 csrc/kernels/causal_conv1d_update_steps_gqa_bf16.cuh diff --git a/CMakeLists.txt b/CMakeLists.txt index 44745694..2b5870e2 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1374,6 +1374,10 @@ target_sources(flash_rt_kernels PRIVATE 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 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 diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index f5249baf..b1ffa7c3 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -193,6 +193,7 @@ extern "C" int cutlass_int8_rowwise_bf16out_t64x128( #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 "quantize/fp8_block128_dequant.cuh" #ifdef FLASHRT_HAVE_NVFP4_SWIZZLE #include "quantize/fp8_block128_to_nvfp4_swizzled.cuh" @@ -6760,6 +6761,27 @@ WY-chain q/k L2-norm + pack + per-chunk gate cumsum for the fixed 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("gdn_chunk_from_conv_smem_h_stash_bf16", 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 00000000..7e00885a --- /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 00000000..ed12a596 --- /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 From 2684a9182446898e66b77f67e1953959964b24d3 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sun, 16 Aug 2026 18:34:47 -0400 Subject: [PATCH 17/44] gated_delta_core: native WY-chain tiers in the prefill pass The triangular solve routes to the native batched unit-lower inverse (same forward-substitution recurrence as the batched cuBLAS solve, identity/tril preparation folded in), the norm/pack + gate-cumsum pair takes the chunk-parallel v2 launch, and the conv update goes through the GQA-split variants - the step-batched native arm where the build carries it, the packaged one otherwise. Projection seams now receive the caller's tensor instead of a view so an upstream producer's identity-keyed quantization handoff survives the call. --- .../impls/gated_delta_core/fused_layer.py | 254 ++++++++++++++++-- 1 file changed, 229 insertions(+), 25 deletions(-) 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 1a63dab4..0c401833 100644 --- a/flash_rt/structures/impls/gated_delta_core/fused_layer.py +++ b/flash_rt/structures/impls/gated_delta_core/fused_layer.py @@ -43,6 +43,176 @@ "version": ">=1"} +@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_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. @@ -186,6 +356,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", @@ -333,9 +508,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() @@ -449,32 +628,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) + 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) - # 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() + # 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) @@ -500,10 +705,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( @@ -579,9 +781,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] From 7382db24c92eeaf9acbf5b9bd2781449b2fae910 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sun, 16 Aug 2026 18:34:47 -0400 Subject: [PATCH 18/44] norm_fused: RMSNorm producer that emits its consumers' NVFP4 input The (1+w)-form norm computes exactly the host norm and quantizes the same values in one kernel, publishing packed FP4 + SFA through the identity-keyed share cell its consumer group already honors; the normed BF16 tensor still flows through the host graph, so a cell miss just means the consumer quantizes for itself. Decode-width calls keep the host norm - the group quantize is cheaper there than a standalone producer launch. Scale-atom tails are zero-filled: unwritten bytes in a partial 128-row atom were a measured nondeterminism channel. --- .../impls/decoder_ffn/nvfp4_fused.py | 5 +- .../impls/norm_fused/nvfp4_producer.py | 168 ++++++++++++++++++ 2 files changed, 172 insertions(+), 1 deletion(-) create mode 100644 flash_rt/structures/impls/norm_fused/nvfp4_producer.py diff --git a/flash_rt/structures/impls/decoder_ffn/nvfp4_fused.py b/flash_rt/structures/impls/decoder_ffn/nvfp4_fused.py index 942d8bb1..3886228e 100644 --- a/flash_rt/structures/impls/decoder_ffn/nvfp4_fused.py +++ b/flash_rt/structures/impls/decoder_ffn/nvfp4_fused.py @@ -51,7 +51,10 @@ def _op(merged: torch.Tensor) -> list[torch.Tensor]: h = two_h // 2 packed = torch.empty(m, h // 2, device=merged.device, dtype=torch.uint8) - sfa = torch.empty( + # 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(), 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 00000000..1ed9062b --- /dev/null +++ b/flash_rt/structures/impls/norm_fused/nvfp4_producer.py @@ -0,0 +1,168 @@ +"""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 + +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))) + # 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: at decode widths the host norm folds + # into the surrounding elementwise graph and the shared group + # quantize is cheaper than this producer's standalone launch; + # the fused pass pays from prompt-slab widths up + if x.numel() // x.shape[-1] < 64: + 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 From 641c56dbbed83c6ee93a89e5b61e510a6f979183 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Mon, 17 Aug 2026 02:00:56 -0400 Subject: [PATCH 19/44] vllm adapter: source-checkpoint binding, probe gate, measured tier parking Runtime weights in a mixed-precision engine are packed in formats the seam cannot read as rows, so binding takes its material from the source checkpoint the engine loaded (fused modules concatenated in the engine's own order) and gates each seat on a cosine probe against the host module it replaces. Optional tiers are parked by tracing them under a fake mode rather than by name, so an artifact that ships the missing fake impls lights them up with no code change; the gated-delta dependency pins the version whose build matrix covers the engine's runtime. --- flash_rt/structures/adapters/vllm_engine.py | 196 +++++++++++++++++- .../impls/gated_delta_core/fused_layer.py | 2 +- 2 files changed, 193 insertions(+), 5 deletions(-) diff --git a/flash_rt/structures/adapters/vllm_engine.py b/flash_rt/structures/adapters/vllm_engine.py index a529b5a3..ee01eda5 100644 --- a/flash_rt/structures/adapters/vllm_engine.py +++ b/flash_rt/structures/adapters/vllm_engine.py @@ -75,6 +75,117 @@ 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 + + +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 + 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() + + class _ProjSeat(nn.Module): """Preserves the engine's ``(out, bias)`` projection contract.""" @@ -167,6 +278,52 @@ def _is_projection(module) -> bool: and hasattr(module, "quant_method")) +def _park_untraceable_tiers(seam) -> list[str]: + """Drop the seam's optional tiers that cannot trace on fake tensors. + + The engine compiles seam forwards under a fake mode, so an entry + without a fake impl raises there and takes the whole graph with it. + Rather than name the offenders (they move with the installed + artifact — a hub release that adds the missing fakes should light + the tiers up with no code change), this traces each optional tier + once under a fake mode and parks only what actually fails. + """ + from torch._subclasses.fake_tensor import FakeTensorMode + + parked = [] + k = int(seam._k) + + def _traces(call) -> bool: + try: + with FakeTensorMode(allow_non_fake_inputs=True): + call() + return True + except Exception: + return False + + if getattr(seam, "_m256", None) is not None: + m = 512 + a = torch.empty(m, k // 2, device="cuda", dtype=torch.uint8) + sfa = torch.empty(((m + 127) // 128) * ((k + 63) // 64) * 512, + device="cuda", dtype=torch.uint8) + if not _traces(lambda: seam._m256(a, seam._w_packed, sfa, + seam._w_sfb)): + seam._m256 = None + parked.append("m256") + if getattr(seam, "_mrows", None) is not None and seam._mrows_hub: + m = 8 + a = torch.empty(m, k // 2, device="cuda", dtype=torch.uint8) + sfa = torch.empty(((m + 127) // 128) * ((k + 63) // 64) * 512, + device="cuda", dtype=torch.uint8) + w_, s_ = seam._mr_cfg + if not _traces(lambda: seam._mrows(a, seam._w_packed, sfa, + seam._w_sfb, warps=w_, + stages=s_)): + seam._mrows = None + parked.append("mrows") + return parked + + def _expert_holder(module): for _, child in module.named_modules(): if torch.is_tensor(getattr(child, "w13_weight", None)): @@ -228,6 +385,7 @@ def _init(self, *a, **kw): swaps: dict[str, nn.Module] = {} reverts: list = [] refused: list = [] + parked_tiers: dict[str, int] = {} modules = dict(model.named_modules()) # dense projections, smallest first: on tight cards early frees @@ -237,7 +395,33 @@ 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}) + mod._frt_seat_name = name + w_bind = _projection_weight(mod) + seam, _ = _linear.bind_proj_seam({"w": w_bind}) + # 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, w_bind.shape[1], 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 engine traces seam forwards with Meta tensors, so a + # tier whose entry carries no fake impl dies at trace time. + # Which tiers those are is a property of the installed + # artifact, not a constant: park by measuring it — trace + # each optional tier under a fake mode and keep the ones + # that survive. A capability parked, never a refusal. + for tier in _park_untraceable_tiers(seam): + parked_tiers[tier] = parked_tiers.get(tier, 0) + 1 swaps[name] = _ProjSeat(seam) except Exception as e: refused.append((name, repr(e)[:120])) @@ -308,10 +492,14 @@ def _init(self, *a, **kw): return handle handle = _swap.attach(model, swaps, revert=reverts) if verbose: + parked = (", parked " + ", ".join( + f"{t}x{n}" for t, n in sorted(parked_tiers.items())) + if parked_tiers else "") 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} + f"({head_slabs} head slabs), {len(refused)} refused" + f"{parked}", flush=True) + handle.notes = {"refused": refused, "head_slabs": head_slabs, + "parked_tiers": parked_tiers} return handle 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 0c401833..655f8909 100644 --- a/flash_rt/structures/impls/gated_delta_core/fused_layer.py +++ b/flash_rt/structures/impls/gated_delta_core/fused_layer.py @@ -36,7 +36,7 @@ from ...guard import CAST_OK, PROCEED, GuardedSeam GDA_DEP = {"provider": "hf", "repo": "flashrt/gated-delta-attention", - "version": ">=3"} + "version": ">=5"} CONV_DEP = {"provider": "hf", "repo": "flashrt/causal-conv1d-state", "version": ">=1"} FUSED_DEP = {"provider": "hf", "repo": "flashrt/transformer-fused-ops", From b5f07acb93be1cb99ee0f31771ffd204776ee8ee Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Mon, 17 Aug 2026 03:16:08 -0400 Subject: [PATCH 20/44] linear_proj: range-compiled callers take the tile tier; vllm fused-MLP seat A host that compiles a seam forward for a range of row counts hands the dispatch a symbolic M, so a tier branch would bake in whichever side the tracing sample took and run it for every replayed shape. Such callers now take the tiled GEMM, which serves every M; concrete row counts keep their tiers. The vLLM adapter gains a fused SwiGLU seat: the host's MLP is already gate_up -> activation -> down, so the three steps collapse into one seat where the fused activation-quantize producer is available, and it parks the two row-count-thresholded tiers this engine cannot replay safely. --- flash_rt/structures/adapters/vllm_engine.py | 128 +++++++++++------- .../impls/linear_proj/nvfp4_dynamic.py | 12 ++ 2 files changed, 91 insertions(+), 49 deletions(-) diff --git a/flash_rt/structures/adapters/vllm_engine.py b/flash_rt/structures/adapters/vllm_engine.py index ee01eda5..05812e65 100644 --- a/flash_rt/structures/adapters/vllm_engine.py +++ b/flash_rt/structures/adapters/vllm_engine.py @@ -197,6 +197,42 @@ def forward(self, x, *args, **kwargs): return self.seam(x), None +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): """Stands where the fused-MoE module stood: routing here, bank in the seam, the host's own shared-expert module added back (it owned @@ -278,49 +314,27 @@ def _is_projection(module) -> bool: and hasattr(module, "quant_method")) -def _park_untraceable_tiers(seam) -> list[str]: - """Drop the seam's optional tiers that cannot trace on fake tensors. +def _park_m_threshold_tiers(seam) -> list[str]: + """Park the tiers whose selection depends on the row count. - The engine compiles seam forwards under a fake mode, so an entry - without a fake impl raises there and takes the whole graph with it. - Rather than name the offenders (they move with the installed - artifact — a hub release that adds the missing fakes should light - the tiers up with no code change), this traces each optional tier - once under a fake mode and parks only what actually fails. + 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`` therefore bakes in whichever side the tracing sample + took and then runs it for every M — measured as a hard refusal from + the M256 tier at engine start. The two tiers that carry an M + threshold (the large-M cooperative tile, the small-M multi-row arm) + are parked here; the M=1 GEMV stays because this host's decode + graphs are captured at fixed batch sizes, and the tiled GEMM serves + every other M correctly. A capability parked, never a refusal. """ - from torch._subclasses.fake_tensor import FakeTensorMode - parked = [] - k = int(seam._k) - - def _traces(call) -> bool: - try: - with FakeTensorMode(allow_non_fake_inputs=True): - call() - return True - except Exception: - return False - if getattr(seam, "_m256", None) is not None: - m = 512 - a = torch.empty(m, k // 2, device="cuda", dtype=torch.uint8) - sfa = torch.empty(((m + 127) // 128) * ((k + 63) // 64) * 512, - device="cuda", dtype=torch.uint8) - if not _traces(lambda: seam._m256(a, seam._w_packed, sfa, - seam._w_sfb)): - seam._m256 = None - parked.append("m256") - if getattr(seam, "_mrows", None) is not None and seam._mrows_hub: - m = 8 - a = torch.empty(m, k // 2, device="cuda", dtype=torch.uint8) - sfa = torch.empty(((m + 127) // 128) * ((k + 63) // 64) * 512, - device="cuda", dtype=torch.uint8) - w_, s_ = seam._mr_cfg - if not _traces(lambda: seam._mrows(a, seam._w_packed, sfa, - seam._w_sfb, warps=w_, - stages=s_)): - seam._mrows = None - parked.append("mrows") + seam._m256 = None + parked.append("m256") + if getattr(seam, "_mrows", None) is not None: + seam._mrows = None + parked.append("mrows") return parked @@ -363,7 +377,8 @@ 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): """Seat a vLLM model: dense projections, expert banks, LM head. Call between weight load and the engine's first trace (see @@ -386,6 +401,7 @@ def _init(self, *a, **kw): reverts: list = [] refused: list = [] parked_tiers: dict[str, int] = {} + fused_mlps = 0 modules = dict(model.named_modules()) # dense projections, smallest first: on tight cards early frees @@ -414,18 +430,30 @@ def _init(self, *a, **kw): if float(cos) < 0.98: raise ValueError( f"bind probe cos {float(cos):.4f} < 0.98") - # the engine traces seam forwards with Meta tensors, so a - # tier whose entry carries no fake impl dies at trace time. - # Which tiers those are is a property of the installed - # artifact, not a constant: park by measuring it — trace - # each optional tier under a fake mode and keep the ones - # that survive. A capability parked, never a refusal. - for tier in _park_untraceable_tiers(seam): + for tier in _park_m_threshold_tiers(seam): parked_tiers[tier] = parked_tiers.get(tier, 0) + 1 swaps[name] = _ProjSeat(seam) except Exception as e: 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 + 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) for name, mod in modules.items(): @@ -495,11 +523,13 @@ def _init(self, *a, **kw): parked = (", parked " + ", ".join( f"{t}x{n}" for t, n in sorted(parked_tiers.items())) if parked_tiers else "") + fused = f", {fused_mlps} fused MLPs" if fused_mlps else "" print(f"[structures.vllm] {len(swaps)} seats " f"({head_slabs} head slabs), {len(refused)} refused" - f"{parked}", flush=True) + f"{parked}{fused}", flush=True) handle.notes = {"refused": refused, "head_slabs": head_slabs, - "parked_tiers": parked_tiers} + "parked_tiers": parked_tiers, + "fused_mlps": fused_mlps} return handle diff --git a/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py b/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py index b634842e..7f8ebc0a 100644 --- a/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py +++ b/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py @@ -293,6 +293,18 @@ def _mm_packed(self, a_packed: torch.Tensor, Returns the (m, n) BF16 product with bias applied. """ m = a_packed.shape[0] + if not isinstance(m, int): + # a symbolic row count: the host is tracing this call for a + # *range* of M, so a Python-level tier branch would bake in + # whichever side the tracing sample happened to take and + # then run it for every replayed M. The tiled GEMM serves + # every M correctly, so the range-compiled path takes it and + # the specialized ones (concrete M at capture) keep theirs. + y = self._gemm(a_packed, self._w_packed, a_sfa, self._w_sfb, + 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, From 637c76b103d16af39af23646833beea732f4a2ea Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Mon, 17 Aug 2026 04:12:19 -0400 Subject: [PATCH 21/44] kernels: decode-step gated-delta recurrence and gated-norm producer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The recurrence gives every value column its own thread, so the columns now spread across blocks instead of tying block width to the 128 columns, and the column state streams in two passes rather than sitting in a per-thread array the hardware spills to local memory — 2x on the step's own shape, output bit-identical. The gated norm gains a variant that also emits its consumer's NVFP4 rows: the output projection is its only consumer and quantizes what it receives, and the quantizer's blocks tile a head's lanes exactly, so the whole step fits in the block that produced the row (bit-identical normed rows, packed bytes and scales alike). --- CMakeLists.txt | 8 + csrc/bindings.cpp | 47 ++++++ .../gdn_recurrent_inout_vsplit_bf16.cu | 144 ++++++++++++++++++ .../gdn_recurrent_inout_vsplit_bf16.cuh | 36 +++++ .../rms_norm_gated_silu_quant_fp4_bf16.cu | 142 +++++++++++++++++ .../rms_norm_gated_silu_quant_fp4_bf16.cuh | 27 ++++ 6 files changed, 404 insertions(+) create mode 100644 csrc/kernels/gdn_recurrent_inout_vsplit_bf16.cu create mode 100644 csrc/kernels/gdn_recurrent_inout_vsplit_bf16.cuh create mode 100644 csrc/kernels/rms_norm_gated_silu_quant_fp4_bf16.cu create mode 100644 csrc/kernels/rms_norm_gated_silu_quant_fp4_bf16.cuh diff --git a/CMakeLists.txt b/CMakeLists.txt index 2b5870e2..b1d509f0 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1378,6 +1378,14 @@ target_sources(flash_rt_kernels PRIVATE target_sources(flash_rt_kernels PRIVATE csrc/kernels/causal_conv1d_update_steps_gqa_bf16.cu) +# Gated-delta recurrent decode step, V-split launch plan. +target_sources(flash_rt_kernels PRIVATE + csrc/kernels/gdn_recurrent_inout_vsplit_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 diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index b1ffa7c3..591c2705 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -194,6 +194,8 @@ extern "C" int cutlass_int8_rowwise_bf16out_t64x128( #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_vsplit_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" @@ -6782,6 +6784,51 @@ 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_vsplit_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_vsplit_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 V-split launch plan: 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", diff --git a/csrc/kernels/gdn_recurrent_inout_vsplit_bf16.cu b/csrc/kernels/gdn_recurrent_inout_vsplit_bf16.cu new file mode 100644 index 00000000..9fc694a1 --- /dev/null +++ b/csrc/kernels/gdn_recurrent_inout_vsplit_bf16.cu @@ -0,0 +1,144 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Gated-delta recurrent decode step, V-split launch plan. See header. +#include "kernels/gdn_recurrent_inout_vsplit_bf16.cuh" + +#include + +namespace flash_rt { +namespace kernels { +namespace { + +constexpr int kHD = 128; +constexpr int kCols = 32; // value columns per block (one warp) +constexpr float kEps = 1e-6f; + +__device__ __forceinline__ float warp_sum(float v) { + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + v += __shfl_xor_sync(0xffffffffu, v, off); + return v; +} + +__global__ void recurrent_vsplit_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 vb = blockIdx.y; // which 32-column slice + const int lane = threadIdx.x; // 0..31 + const int t = vb * kCols + lane; // this thread's value column + + const size_t hv_off = ((size_t)blockIdx.z * num_v_heads + h) * kHD; + + // q/k stage in registers: each lane owns 4 of the 128 entries, and + // the L2 norms reduce across the warp + float qs[4], ks[4]; + #pragma unroll + for (int j = 0; j < 4; ++j) { + const int i = j * kCols + lane; + qs[j] = static_cast(q_in[hv_off + i]); + ks[j] = static_cast(k_in[hv_off + i]); + } + if (use_qk_l2norm) { + float q_sq = 0.f, k_sq = 0.f; + #pragma unroll + for (int j = 0; j < 4; ++j) { + q_sq += qs[j] * qs[j]; + k_sq += ks[j] * ks[j]; + } + const float q_inv = rsqrtf(warp_sum(q_sq) + kEps); + const float k_inv = rsqrtf(warp_sum(k_sq) + kEps); + #pragma unroll + for (int j = 0; j < 4; ++j) { + qs[j] *= q_inv; + ks[j] *= k_inv; + } + } + const float qscale = rsqrtf(static_cast(kHD)); + #pragma unroll + for (int j = 0; j < 4; ++j) qs[j] *= qscale; + + // broadcast the staged vectors so every lane sees all 128 entries + // in the packaged kernel's index order + __shared__ float sq[kHD], sk[kHD]; + #pragma unroll + for (int j = 0; j < 4; ++j) { + sq[j * kCols + lane] = qs[j]; + sk[j * kCols + lane] = ks[j]; + } + __syncwarp(); + + const float g_t = + __expf(static_cast(g_in[blockIdx.z * num_v_heads + h])); + const float beta_t = + static_cast(beta_in[blockIdx.z * 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_vsplit_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, kHD / kCols, B); + recurrent_vsplit_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_vsplit_bf16.cuh b/csrc/kernels/gdn_recurrent_inout_vsplit_bf16.cuh new file mode 100644 index 00000000..e3152f35 --- /dev/null +++ b/csrc/kernels/gdn_recurrent_inout_vsplit_bf16.cuh @@ -0,0 +1,36 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Gated-delta recurrent decode step, V-split launch plan. +// +// The recurrence gives every value column its own thread: that thread +// holds state[:, v] in registers and reduces over K entirely on its +// own, so the only cross-thread work in the step is the q/k L2 norm. +// The packaged kernel still launches one block per head — 48 blocks on +// a 170-SM part, a quarter of the machine — because it ties the block +// width to the 128 value columns. This entry splits the columns across +// blocks instead (one warp per 32-column slice), keeping the total +// thread count identical while spreading it over 4x the blocks. +// +// The per-column arithmetic is transcribed unchanged, so the state and +// output a column receives are bit-identical to the packaged kernel's. +// The L2 norm reduces over a warp rather than a 128-thread block, so +// its summation order differs — a fp32 rounding difference on the q/k +// scale, judged by the model's own arbiter band, not claimed bitwise. +// 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_vsplit_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/rms_norm_gated_silu_quant_fp4_bf16.cu b/csrc/kernels/rms_norm_gated_silu_quant_fp4_bf16.cu new file mode 100644 index 00000000..cb5bce9b --- /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 00000000..8740cfac --- /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 From e250f3462fad24ea9eaca216eafa10d8e6e2f24c Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Mon, 17 Aug 2026 04:12:33 -0400 Subject: [PATCH 22/44] gated_delta_core: decode step takes the native recurrence and norm producer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The recurrent step routes to the V-split entry where the build carries it, and the gated norm hands the output projection packed rows directly instead of writing a row its only consumer re-reads to quantize. Buffers for the producer arm are allocated before any capture so the graph records stable addresses, and the norm weight is copied detached to keep the op off an autograd edge. Kernel resolution also stops trusting a single release: a version range now walks the published tags down to its floor and takes the newest release whose build matrix covers the caller's torch/CUDA pair. A publisher adding variants release by release can drop a pair an older release carried, and a resolver too old to read versions lands on the repo default regardless — either way the range, not one pinned tag, is what the dependency actually asked for. --- flash_rt/structures/impls/__init__.py | 79 ++++++--- .../impls/gated_delta_core/fused_layer.py | 158 +++++++++++++++++- 2 files changed, 207 insertions(+), 30 deletions(-) diff --git a/flash_rt/structures/impls/__init__.py b/flash_rt/structures/impls/__init__.py index e85ab215..0a3993bf 100644 --- a/flash_rt/structures/impls/__init__.py +++ b/flash_rt/structures/impls/__init__.py @@ -144,6 +144,36 @@ def _check_arch(repo: str, module) -> None: #: refusal path must not manufacture that error on retry) _LOADED: dict[tuple[str, str], object] = {} +#: 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): + """The newest release at or above ``version``'s floor that loads. + + ``first`` is an optional resolution to try before the walk (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 + last = None + 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): @@ -166,31 +196,32 @@ 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): + # 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) except (OSError, RuntimeError, ValueError) as unavailable: _record_unavailable(repo, version, unavailable) raise KernelUnavailable( 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 655f8909..a9b12e92 100644 --- a/flash_rt/structures/impls/gated_delta_core/fused_layer.py +++ b/flash_rt/structures/impls/gated_delta_core/fused_layer.py @@ -35,8 +35,12 @@ 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": ">=5"} + "version": ">=3"} CONV_DEP = {"provider": "hf", "repo": "flashrt/causal-conv1d-state", "version": ">=1"} FUSED_DEP = {"provider": "hf", "repo": "flashrt/transformer-fused-ops", @@ -88,6 +92,100 @@ def _(big_a): _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 + 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_vsplit(): + """The V-split launch of the gated-delta recurrent decode step. + + Per-column arithmetic is the packaged kernel's; the columns just + spread over four times the blocks (and the column state streams + instead of sitting in a per-thread array the hardware would spill). + Measured 2x on the step's own shape, output bit-identical; the q/k + L2 norm reduces over a warp, so its fp32 rounding can differ. + """ + if _os.environ.get("FRT_GDN_VSPLIT", "1") == "0": + return None + try: + from flash_rt import flash_rt_kernels as _fk + except ImportError: + return None + fn = getattr(_fk, "gdn_recurrent_inout_vsplit_bf16", None) + if fn is None: + return None + + from torch.library import custom_op + + @custom_op("flashrt_native::gdn_recurrent_vsplit", + 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_vsplit 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. @@ -410,6 +508,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 = ( @@ -814,15 +914,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) + vsplit = (_native_recurrent_vsplit() if self._d == 128 + else None) + if vsplit is not None: + core_out, new_state = vsplit( + 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) @@ -832,6 +959,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, From 53bd3ac97688243faa278037b1ddac68239f1658 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Mon, 17 Aug 2026 04:19:57 -0400 Subject: [PATCH 23/44] gated_delta_core: keep the recurrent step bitwise while unspilling it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The measured win in the streaming rewrite was never the column split — it was not spilling the state column to local memory. Splitting the columns across narrower blocks also moved the q/k norms onto a warp reduction, which changes their fp32 summation order for a further 0.03 ms/step; the block width goes back to 128 so the whole kernel is bit-identical to the packaged one, state and output alike, and the step still stops paying DRAM for state it believes is resident. --- CMakeLists.txt | 4 +- csrc/bindings.cpp | 8 +- ....cu => gdn_recurrent_inout_stream_bf16.cu} | 100 +++++++++--------- .../gdn_recurrent_inout_stream_bf16.cuh | 37 +++++++ .../gdn_recurrent_inout_vsplit_bf16.cuh | 36 ------- .../impls/gated_delta_core/fused_layer.py | 32 +++--- 6 files changed, 108 insertions(+), 109 deletions(-) rename csrc/kernels/{gdn_recurrent_inout_vsplit_bf16.cu => gdn_recurrent_inout_stream_bf16.cu} (60%) create mode 100644 csrc/kernels/gdn_recurrent_inout_stream_bf16.cuh delete mode 100644 csrc/kernels/gdn_recurrent_inout_vsplit_bf16.cuh diff --git a/CMakeLists.txt b/CMakeLists.txt index b1d509f0..3fdb85f2 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1378,9 +1378,9 @@ target_sources(flash_rt_kernels PRIVATE target_sources(flash_rt_kernels PRIVATE csrc/kernels/causal_conv1d_update_steps_gqa_bf16.cu) -# Gated-delta recurrent decode step, V-split launch plan. +# Gated-delta recurrent decode step, streaming-column form. target_sources(flash_rt_kernels PRIVATE - csrc/kernels/gdn_recurrent_inout_vsplit_bf16.cu) + csrc/kernels/gdn_recurrent_inout_stream_bf16.cu) # Gated norm that also emits its consumer's NVFP4 input. target_sources(flash_rt_kernels PRIVATE diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index 591c2705..e7952a17 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -194,7 +194,7 @@ extern "C" int cutlass_int8_rowwise_bf16out_t64x128( #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_vsplit_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 @@ -6805,12 +6805,12 @@ 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_vsplit_bf16", + 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_vsplit_bf16( + 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, @@ -6822,7 +6822,7 @@ kernel; quantize stage is the production path verbatim. dim must be 128. py::arg("head_dim"), py::arg("use_qk_l2norm") = true, py::arg("stream") = 0, R"pbdoc( -Gated-delta recurrent decode step over a V-split launch plan: one warp +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 diff --git a/csrc/kernels/gdn_recurrent_inout_vsplit_bf16.cu b/csrc/kernels/gdn_recurrent_inout_stream_bf16.cu similarity index 60% rename from csrc/kernels/gdn_recurrent_inout_vsplit_bf16.cu rename to csrc/kernels/gdn_recurrent_inout_stream_bf16.cu index 9fc694a1..be510416 100644 --- a/csrc/kernels/gdn_recurrent_inout_vsplit_bf16.cu +++ b/csrc/kernels/gdn_recurrent_inout_stream_bf16.cu @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // -// Gated-delta recurrent decode step, V-split launch plan. See header. -#include "kernels/gdn_recurrent_inout_vsplit_bf16.cuh" +// Gated-delta recurrent decode step, streaming-column form. See header. +#include "kernels/gdn_recurrent_inout_stream_bf16.cuh" #include @@ -10,17 +10,34 @@ namespace kernels { namespace { constexpr int kHD = 128; -constexpr int kCols = 32; // value columns per block (one warp) constexpr float kEps = 1e-6f; -__device__ __forceinline__ float warp_sum(float v) { +// 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) - v += __shfl_xor_sync(0xffffffffu, v, off); - return v; + 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_vsplit_kernel( +__global__ void recurrent_stream_kernel( const __nv_bfloat16* __restrict__ q_in, const __nv_bfloat16* __restrict__ k_in, const __nv_bfloat16* __restrict__ v_in, @@ -31,54 +48,35 @@ __global__ void recurrent_vsplit_kernel( __nv_bfloat16* __restrict__ out_, int num_v_heads, bool use_qk_l2norm) { const int h = blockIdx.x; - const int vb = blockIdx.y; // which 32-column slice - const int lane = threadIdx.x; // 0..31 - const int t = vb * kCols + lane; // this thread's value column + 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)blockIdx.z * num_v_heads + h) * kHD; + const size_t hv_off = ((size_t)b * num_v_heads + h) * kHD; - // q/k stage in registers: each lane owns 4 of the 128 entries, and - // the L2 norms reduce across the warp - float qs[4], ks[4]; - #pragma unroll - for (int j = 0; j < 4; ++j) { - const int i = j * kCols + lane; - qs[j] = static_cast(q_in[hv_off + i]); - ks[j] = static_cast(k_in[hv_off + i]); - } - if (use_qk_l2norm) { - float q_sq = 0.f, k_sq = 0.f; - #pragma unroll - for (int j = 0; j < 4; ++j) { - q_sq += qs[j] * qs[j]; - k_sq += ks[j] * ks[j]; - } - const float q_inv = rsqrtf(warp_sum(q_sq) + kEps); - const float k_inv = rsqrtf(warp_sum(k_sq) + kEps); - #pragma unroll - for (int j = 0; j < 4; ++j) { - qs[j] *= q_inv; - ks[j] *= k_inv; - } - } - const float qscale = rsqrtf(static_cast(kHD)); - #pragma unroll - for (int j = 0; j < 4; ++j) qs[j] *= qscale; + __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(); - // broadcast the staged vectors so every lane sees all 128 entries - // in the packaged kernel's index order - __shared__ float sq[kHD], sk[kHD]; - #pragma unroll - for (int j = 0; j < 4; ++j) { - sq[j * kCols + lane] = qs[j]; - sk[j * kCols + lane] = ks[j]; + 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(); } - __syncwarp(); + sq[t] *= rsqrtf(static_cast(kHD)); + __syncthreads(); const float g_t = - __expf(static_cast(g_in[blockIdx.z * num_v_heads + h])); + __expf(static_cast(g_in[b * num_v_heads + h])); const float beta_t = - static_cast(beta_in[blockIdx.z * num_v_heads + h]); + 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 @@ -117,7 +115,7 @@ __global__ void recurrent_vsplit_kernel( } // namespace -int gdn_recurrent_inout_vsplit_bf16( +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, @@ -125,8 +123,8 @@ int gdn_recurrent_inout_vsplit_bf16( 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, kHD / kCols, B); - recurrent_vsplit_kernel<<>>( + dim3 grid(num_v_heads, B); + recurrent_stream_kernel<<>>( reinterpret_cast(q), reinterpret_cast(k), reinterpret_cast(v), 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 00000000..713db0f7 --- /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_recurrent_inout_vsplit_bf16.cuh b/csrc/kernels/gdn_recurrent_inout_vsplit_bf16.cuh deleted file mode 100644 index e3152f35..00000000 --- a/csrc/kernels/gdn_recurrent_inout_vsplit_bf16.cuh +++ /dev/null @@ -1,36 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// -// Gated-delta recurrent decode step, V-split launch plan. -// -// The recurrence gives every value column its own thread: that thread -// holds state[:, v] in registers and reduces over K entirely on its -// own, so the only cross-thread work in the step is the q/k L2 norm. -// The packaged kernel still launches one block per head — 48 blocks on -// a 170-SM part, a quarter of the machine — because it ties the block -// width to the 128 value columns. This entry splits the columns across -// blocks instead (one warp per 32-column slice), keeping the total -// thread count identical while spreading it over 4x the blocks. -// -// The per-column arithmetic is transcribed unchanged, so the state and -// output a column receives are bit-identical to the packaged kernel's. -// The L2 norm reduces over a warp rather than a 128-thread block, so -// its summation order differs — a fp32 rounding difference on the q/k -// scale, judged by the model's own arbiter band, not claimed bitwise. -// 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_vsplit_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/flash_rt/structures/impls/gated_delta_core/fused_layer.py b/flash_rt/structures/impls/gated_delta_core/fused_layer.py index a9b12e92..61a2a799 100644 --- a/flash_rt/structures/impls/gated_delta_core/fused_layer.py +++ b/flash_rt/structures/impls/gated_delta_core/fused_layer.py @@ -137,28 +137,28 @@ def _(x, gate, weight, out, packed, sfa, eps): @lru_cache(maxsize=1) -def _native_recurrent_vsplit(): - """The V-split launch of the gated-delta recurrent decode step. - - Per-column arithmetic is the packaged kernel's; the columns just - spread over four times the blocks (and the column state streams - instead of sitting in a per-thread array the hardware would spill). - Measured 2x on the step's own shape, output bit-identical; the q/k - L2 norm reduces over a warp, so its fp32 rounding can differ. +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_VSPLIT", "1") == "0": + if _os.environ.get("FRT_GDN_STREAM", "1") == "0": return None try: from flash_rt import flash_rt_kernels as _fk except ImportError: return None - fn = getattr(_fk, "gdn_recurrent_inout_vsplit_bf16", 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_vsplit", + @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, @@ -171,7 +171,7 @@ def _op(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, torch.cuda.current_stream().cuda_stream) if rc != 0: raise RuntimeError( - f"gdn_recurrent_vsplit refused rc={rc} H={h} D={d}") + 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): @@ -914,10 +914,10 @@ 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 - vsplit = (_native_recurrent_vsplit() if self._d == 128 - else None) - if vsplit is not None: - core_out, new_state = vsplit( + 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, From 02a52230db2637dbff19eac0e2791702bbb16ec6 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Mon, 17 Aug 2026 04:29:03 -0400 Subject: [PATCH 24/44] norm_fused: make the producer's width band a measured knob The band where the fused producer beats the host norm is a property of what the compiler folded around that norm, not a constant: at decode width the host path carries the residual add into the same kernel, and replacing it there costs more than the quantize it saves (measured +0.13 ms/step). The threshold stays where the measurement put it and is now nameable, so the next host can be measured rather than assumed. --- .../structures/impls/norm_fused/nvfp4_producer.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/flash_rt/structures/impls/norm_fused/nvfp4_producer.py b/flash_rt/structures/impls/norm_fused/nvfp4_producer.py index 1ed9062b..9c9ee11f 100644 --- a/flash_rt/structures/impls/norm_fused/nvfp4_producer.py +++ b/flash_rt/structures/impls/norm_fused/nvfp4_producer.py @@ -13,6 +13,7 @@ from __future__ import annotations +import os from functools import lru_cache import torch @@ -79,6 +80,7 @@ def __init__(self, host_norm: torch.nn.Module, cell: _ShareCell): 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 @@ -89,11 +91,11 @@ def __init__(self, host_norm: torch.nn.Module, cell: _ShareCell): 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: at decode widths the host norm folds - # into the surrounding elementwise graph and the shared group - # quantize is cheaper than this producer's standalone launch; - # the fused pass pays from prompt-slab widths up - if x.numel() // x.shape[-1] < 64: + # 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 From e2d3edfc5768c298b0b8dd95b9d5b2d197d4d573 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Mon, 17 Aug 2026 06:23:45 -0400 Subject: [PATCH 25/44] impls: resolve a version range from the resolver's own report A resolver that accepts only an exact major still names the majors it has when it rejects a range, so take the newest one at or above the floor from that report instead of probing tags downward. The blind walk stays as the fallback for a resolver that says nothing. --- flash_rt/structures/impls/__init__.py | 26 ++++++++++++++++++++++---- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/flash_rt/structures/impls/__init__.py b/flash_rt/structures/impls/__init__.py index 0a3993bf..ea5391e1 100644 --- a/flash_rt/structures/impls/__init__.py +++ b/flash_rt/structures/impls/__init__.py @@ -150,10 +150,16 @@ def _check_arch(repo: str, module) -> None: _TAG_SEARCH_TOP = 32 -def _newest_loadable(get_kernel, repo, version, kw, first=None): +def _newest_loadable(get_kernel, repo, version, kw, first=None, + reported=None): """The newest release at or above ``version``'s floor that loads. - ``first`` is an optional resolution to try before the walk (the + 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). """ @@ -165,7 +171,15 @@ def _newest_loadable(get_kernel, repo, version, kw, first=None): 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) @@ -208,7 +222,7 @@ def hub_kernel(repo: str, version: str): _LOADED[key] = _newest_loadable( get_kernel, repo, version, _kw, first=lambda: get_kernel(repo)) - except (ValueError, FileNotFoundError): + 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 @@ -221,7 +235,11 @@ def hub_kernel(repo: str, version: str): # 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) + 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) raise KernelUnavailable( From ea740a73865872a5efdcd4a8c934b9f7ac299c76 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Mon, 17 Aug 2026 07:54:07 -0400 Subject: [PATCH 26/44] structures: take the fused producers from the Hub artifact first MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Each of these three kernels now ships in a published package, and the packaged entries are torch ops carrying fakes — a host that compiles the call traces them unaided, and a serving process that cannot load this build's native extension (different torch/CUDA pair) still gets the fused form instead of falling back to the split chain. The local build stays as the second source, wrapped as before. --- .../impls/decoder_ffn/nvfp4_fused.py | 24 ++++++++++++++--- .../impls/gated_delta_core/fused_layer.py | 26 +++++++++++++++++++ 2 files changed, 47 insertions(+), 3 deletions(-) diff --git a/flash_rt/structures/impls/decoder_ffn/nvfp4_fused.py b/flash_rt/structures/impls/decoder_ffn/nvfp4_fused.py index 3886228e..c4deb77b 100644 --- a/flash_rt/structures/impls/decoder_ffn/nvfp4_fused.py +++ b/flash_rt/structures/impls/decoder_ffn/nvfp4_fused.py @@ -26,15 +26,33 @@ 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 local build's fused SwiGLU + NVFP4 quantize producer. + """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. Registered as a torch custom op with a fake so the - compiled prefill and the captured decode step trace through it. + 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: + return lambda merged: hub_fn(merged) try: from flash_rt import flash_rt_kernels as _fk except ImportError: 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 61a2a799..9f1de001 100644 --- a/flash_rt/structures/impls/gated_delta_core/fused_layer.py +++ b/flash_rt/structures/impls/gated_delta_core/fused_layer.py @@ -105,6 +105,20 @@ def _native_gated_norm_quant(): """ 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: @@ -148,6 +162,18 @@ def _native_recurrent_stream(): """ 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: From 031a53f7b9f9c5406e289799a444655be021a0fe Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Mon, 17 Aug 2026 12:33:27 -0400 Subject: [PATCH 27/44] impls: bound hub resolution and cache the unresolvable Version resolution walks releases over the network, and every step of that walk is an unbounded round trip. A host binding a whole model asks for the same missing package once per seat, so one unpublished build variant used to cost that walk hundreds of times over, with the engine silent throughout. Default the hub timeouts (callers that set their own keep them) and remember resolution failures for the life of the process: absence is as cacheable as presence and far more expensive to re-establish. --- flash_rt/structures/impls/__init__.py | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/flash_rt/structures/impls/__init__.py b/flash_rt/structures/impls/__init__.py index ea5391e1..b3c09990 100644 --- a/flash_rt/structures/impls/__init__.py +++ b/flash_rt/structures/impls/__init__.py @@ -144,6 +144,10 @@ 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. @@ -193,7 +197,22 @@ def _newest_loadable(get_kernel, repo, version, kw, first=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) + 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 @@ -242,6 +261,13 @@ def hub_kernel(repo: str, version: str): 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__}: " From d83ae6313d7a4ea6984ec68b9d32bb58d4497261 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Mon, 17 Aug 2026 12:33:27 -0400 Subject: [PATCH 28/44] linear_proj: adopt already-packed weights without regridding A host holding a packed checkpoint does not need the regridding entry 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. bind_proj_seam_packed adopts the e2m1 bytes by reference (the block-scale relayout is a permutation, it loses nothing) and carries a checkpoint's per-tensor factor on each tier's alpha, where the epilogue is already writing the result; the separate pass it replaces cost a full-size read-modify-write per projection. Also let the fp8 no-bias form bind the quantize's return value: on the traced path the persistent out-buffer is None and the returned tensor is the only handle to the result. The fused silu-mul hub entry sizes its outputs at the call site for the same tracing reason: the wrapper's size helper is a custom op without a fake, and a host tracing on Meta tensors dies inside it. --- .../impls/decoder_ffn/nvfp4_fused.py | 17 +- .../impls/linear_proj/fp8_static.py | 9 +- .../impls/linear_proj/nvfp4_dynamic.py | 204 +++++++++++++++++- 3 files changed, 216 insertions(+), 14 deletions(-) diff --git a/flash_rt/structures/impls/decoder_ffn/nvfp4_fused.py b/flash_rt/structures/impls/decoder_ffn/nvfp4_fused.py index c4deb77b..5b5f538f 100644 --- a/flash_rt/structures/impls/decoder_ffn/nvfp4_fused.py +++ b/flash_rt/structures/impls/decoder_ffn/nvfp4_fused.py @@ -52,7 +52,22 @@ def _native_silu_mul(): hub = None hub_fn = getattr(hub, "silu_mul_quantize_fp4_sfa_bf16", None) if hub_fn is not None: - return lambda merged: hub_fn(merged) + 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: diff --git a/flash_rt/structures/impls/linear_proj/fp8_static.py b/flash_rt/structures/impls/linear_proj/fp8_static.py index dfed873d..a6f8c050 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 7f8ebc0a..3cd2e53b 100644 --- a/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py +++ b/flash_rt/structures/impls/linear_proj/nvfp4_dynamic.py @@ -138,6 +138,101 @@ 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. @@ -206,13 +301,22 @@ class LinearProjNvfp4Dynamic(GuardedSeam, torch.nn.Module): _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 @@ -284,6 +388,18 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: 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 @@ -293,14 +409,24 @@ def _mm_packed(self, a_packed: torch.Tensor, 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: the host is tracing this call for a - # *range* of M, so a Python-level tier branch would bake in - # whichever side the tracing sample happened to take and - # then run it for every replayed M. The tiled GEMM serves - # every M correctly, so the range-compiled path takes it and - # the specialized ones (concrete M at capture) keep theirs. + # 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 @@ -308,26 +434,82 @@ def _mm_packed(self, a_packed: torch.Tensor, 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, - warps=gw, stages=gs) + 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, warps=w_, stages=s_) + 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) + self._w_sfb, alpha=al) else: y = self._gemm(a_packed, self._w_packed, a_sfa, self._w_sfb, - variant=2) + 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 +@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() def bind_proj_seam( weights: Mapping[str, torch.Tensor], From 3bbf0aaef952030ee726a863a310377902050859 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Mon, 17 Aug 2026 12:33:27 -0400 Subject: [PATCH 29/44] vllm adapter: precision tiers with pack adoption A mixed-precision checkpoint is a per-projection decision made with calibration data, and seating every position at the adapter's default width overrides that decision silently. Read the choice off what the host actually holds and honour it: positions already in the seam's width adopt the host's own tensors by reference (no regrid, no second copy on the card), FP8 positions get a seat on their own storage (the engine keeps its weight transposed for its cutlass entry, so the orientation comes from the layer's declaration, and one transposed swap leaves total memory unchanged), and unquantized positions are left alone -- a draft model a checkpoint excluded from quantization stays excluded, because a speculative host pays for numerics in acceptance length, not in any number a throughput probe reads. The auto tier adds one opinion on top: a W8 position is carried to W4, sourced from the host's own FP8 rows (every FP8 code is exact in BF16, so no external checkpoint is involved), and the attach report says so in plain text -- the precision change is the caller's to validate. Staged release of replaced host weights now actually releases them: the revert closure used to pin the very tensor it was freeing through a default argument, which is why a consume pass reported the same bytes freed whether the staging ran or not. --- flash_rt/structures/adapters/vllm_engine.py | 495 +++++++++++++++++--- 1 file changed, 438 insertions(+), 57 deletions(-) diff --git a/flash_rt/structures/adapters/vllm_engine.py b/flash_rt/structures/adapters/vllm_engine.py index 05812e65..3d8f1e41 100644 --- a/flash_rt/structures/adapters/vllm_engine.py +++ b/flash_rt/structures/adapters/vllm_engine.py @@ -146,6 +146,36 @@ def read(key): 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. @@ -159,6 +189,21 @@ def _projection_weight(mod) -> torch.Tensor: 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 @@ -186,6 +231,122 @@ def _projection_weight(mod) -> torch.Tensor: 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: + gs = float(v.data.reshape(-1)[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.""" @@ -291,7 +452,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 @@ -299,7 +467,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) @@ -314,28 +483,21 @@ def _is_projection(module) -> bool: and hasattr(module, "quant_method")) -def _park_m_threshold_tiers(seam) -> list[str]: - """Park the tiers whose selection depends on the row count. +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`` therefore bakes in whichever side the tracing sample - took and then runs it for every M — measured as a hard refusal from - the M256 tier at engine start. The two tiers that carry an M - threshold (the large-M cooperative tile, the small-M multi-row arm) - are parked here; the M=1 GEMV stays because this host's decode - graphs are captured at fixed batch sizes, and the tiled GEMM serves - every other M correctly. A capability parked, never a refusal. + 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. """ - parked = [] - if getattr(seam, "_m256", None) is not None: - seam._m256 = None - parked.append("m256") - if getattr(seam, "_mrows", None) is not None: - seam._mrows = None - parked.append("mrows") - return parked + _linear.register_runtime_dispatch(seam) def _expert_holder(module): @@ -378,7 +540,9 @@ def summary(self): def attach_engine(model, *, seats=DENSE_SEAT_SUFFIXES, experts=True, head=True, use_gemv=None, verbose=True, strict=False, - fused_mlp=True): + 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 @@ -402,8 +566,64 @@ def _init(self, *a, **kw): refused: list = [] 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) + 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() @@ -411,14 +631,68 @@ def _init(self, *a, **kw): targets.sort(key=lambda t: t[1].weight.numel()) for name, mod in targets: try: - mod._frt_seat_name = name - w_bind = _projection_weight(mod) - seam, _ = _linear.bind_proj_seam({"w": w_bind}) + mod._frt_seat_name = (ckpt_rename(name) if ckpt_rename + else name) + kind = (_host_precision(mod) + if precision in ("mirror", "auto") 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) + 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, w_bind.shape[1], device="cuda", + probe = torch.randn(4, k_in, device="cuda", dtype=torch.bfloat16) with torch.no_grad(): host_out = mod(probe) @@ -430,11 +704,54 @@ def _init(self, *a, **kw): if float(cos) < 0.98: raise ValueError( f"bind probe cos {float(cos):.4f} < 0.98") - for tier in _park_m_threshold_tiers(seam): - parked_tiers[tier] = parked_tiers.get(tier, 0) + 1 + # 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) + 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 @@ -448,6 +765,12 @@ def _init(self, *a, **kw): 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") @@ -479,29 +802,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: @@ -519,21 +819,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: - parked = (", parked " + ", ".join( - f"{t}x{n}" for t, n in sorted(parked_tiers.items())) - if parked_tiers else "") + parked = "" fused = f", {fused_mlps} fused MLPs" if fused_mlps else "" - print(f"[structures.vllm] {len(swaps)} seats " + 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}", flush=True) + 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, + **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 @@ -559,6 +917,29 @@ def install_load_hook(*, on_attached=None, **attach_kwargs): def load_model(self, *a, __orig=orig, **kw): __orig(self, *a, **kw) 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" + 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 From e8f8e9387e4b4852cbbab57856b44d6216d4a2f6 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Mon, 17 Aug 2026 14:22:05 -0400 Subject: [PATCH 30/44] impls: load a build variant straight from a directory An air-gapped serving container cannot walk the hub at all: the xet transport ignores the process proxy and hangs, offline resolution insists on a complete snapshot when the cache holds only the variant that was ever downloaded, and pinning the resolver library's version trades one incompatibility for another. FRT_KERNEL_DIR_ names the build variant directory directly and imports it as the package it already is -- exactly what the resolver would do after its network walk, minus the walk. The arch check still runs; a wrong variant refuses the same way a resolved one would. --- flash_rt/structures/impls/__init__.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/flash_rt/structures/impls/__init__.py b/flash_rt/structures/impls/__init__.py index b3c09990..853cff01 100644 --- a/flash_rt/structures/impls/__init__.py +++ b/flash_rt/structures/impls/__init__.py @@ -208,6 +208,30 @@ def hub_kernel(repo: str, version: str): 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( From 20d03e2c9cde00ec90a6b513a457360239141b28 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Mon, 17 Aug 2026 14:22:05 -0400 Subject: [PATCH 31/44] adapters: carry the precision tiers across the process boundary The sglang hook now routes to the same tier assembly the vllm adapter runs: the two engines share module lineage down to the quant-method attribute names, so per-position adoption and the auto tier's W8->W4 carry assemble unchanged, and the engine-specific part of that adapter stays what it was -- the sitecustomize carrier. Two host behaviours surfaced by the second engine, fixed where they belong: a projection seat now falls through attribute reads to the module it replaced (an engine that planned a fusion around the projection reads scale attributes off the module object, and a seat answering only forward turned that into a startup crash far from the seam); and a merged projection carrying one global factor per constituent refuses adoption under a single alpha unless the factors agree -- pretending otherwise would scale one half by the other's factor. --- flash_rt/structures/adapters/sglang_engine.py | 20 +++++++++-- flash_rt/structures/adapters/vllm_engine.py | 35 ++++++++++++++++--- 2 files changed, 49 insertions(+), 6 deletions(-) diff --git a/flash_rt/structures/adapters/sglang_engine.py b/flash_rt/structures/adapters/sglang_engine.py index 9dcc9b8e..904697c0 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 3d8f1e41..88ca041f 100644 --- a/flash_rt/structures/adapters/vllm_engine.py +++ b/flash_rt/structures/adapters/vllm_engine.py @@ -319,7 +319,17 @@ def _adopt_nvfp4_pack(mod): for attr in ("weight_global_scale", "weight_scale_2"): v = getattr(mod, attr, None) if v is not None: - gs = float(v.data.reshape(-1)[0]) + 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( @@ -348,11 +358,28 @@ def _bind_fp8_seam(mod, w, rows_hint): 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 + 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 @@ -715,7 +742,7 @@ def _init(self, *a, **kw): if kind == "nvfp4": _arm_runtime_dispatch(seam) kinds[kind] = kinds.get(kind, 0) + 1 - swaps[name] = _ProjSeat(seam) + 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 From 153d2d797f3d8bb5aea6045231d4753780e5db8b Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Mon, 17 Aug 2026 17:09:22 -0400 Subject: [PATCH 32/44] vllm adapter: verify-attention seat and opt-in head/draft tiers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The verify-attention seat routes speculative verify steps through the paged BF16-Q/FP8-KV kernel behind a full guard ladder: metadata must be exactly a single-request causal verify batch on a supported layer geometry, anything else falls through to the host. Page counts come from a hook on the prefill wrapper's plan (the device buffers hold the same values, but reading them there is a host sync per layer, which stalls the launch pipeline for more than the kernel saves), K/V claims are as_strided views over the raw cache storage (flattening the permuted cache view silently copies the whole pool), and capture-time calls are declined so replayed graphs never bake a stale shape. The head and draft seats are explicit opt-ins rather than defaults: the seam's row tiers quantize activations to FP4, and on the logits family and the proposal path that measurably moves speculative acceptance — the mechanical step-rate win is paid back with interest. A checkpoint-packed NVFP4 head can be adopted zero-copy behind FRT_HEAD_W4A4; a BF16 draft can be re-gridded via draft_precision. Projection seats now honor the host's return_bias so a bare-tensor call site is answered in kind. --- flash_rt/structures/adapters/vllm_engine.py | 245 +++++++++++++++++++- 1 file changed, 243 insertions(+), 2 deletions(-) diff --git a/flash_rt/structures/adapters/vllm_engine.py b/flash_rt/structures/adapters/vllm_engine.py index 88ca041f..507381a1 100644 --- a/flash_rt/structures/adapters/vllm_engine.py +++ b/flash_rt/structures/adapters/vllm_engine.py @@ -369,6 +369,13 @@ class _ProjSeat(nn.Module): 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) @@ -382,7 +389,8 @@ def __getattr__(self, name): 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): @@ -534,6 +542,194 @@ 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 + 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) + 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. @@ -591,6 +787,12 @@ 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 @@ -619,6 +821,35 @@ def _init(self, *a, **kw): 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) if lm is not None: try: from ..impls.linear_proj import w8a16_static as _w8 @@ -918,7 +1149,7 @@ def _draft_ckpt_rename(name: str) -> str: def install_load_hook(*, on_attached=None, seat_draft=True, - **attach_kwargs): + 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 @@ -958,6 +1189,16 @@ def load_model(self, *a, __orig=orig, **kw): 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 From ca8d0be0fd8e1b092bf19945f51c5dad1684dbc4 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Mon, 17 Aug 2026 17:41:47 -0400 Subject: [PATCH 33/44] vllm adapter: seat path must also write the step's K/V MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The host forward the verify-attention seat replaces is not just the attention call — it is also where the step's new K/V rows enter the paged cache. Taking the call without that write leaves the tokens' K/V unwritten: the seat's own kernel reads stale slots for them, and every later step is missing them as prefix, so the damage compounds for the rest of the sequence. The seat now performs the host's cache update before running the kernel, and refuses calls that carry fused output-quant arguments it does not implement. --- flash_rt/structures/adapters/vllm_engine.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/flash_rt/structures/adapters/vllm_engine.py b/flash_rt/structures/adapters/vllm_engine.py index 507381a1..d6b166d0 100644 --- a/flash_rt/structures/adapters/vllm_engine.py +++ b/flash_rt/structures/adapters/vllm_engine.py @@ -634,6 +634,11 @@ def forward(self, layer, query, key, value, kv_cache, attn_metadata, 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) @@ -672,6 +677,13 @@ def forward(self, layer, query, key, value, kv_cache, attn_metadata, 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") From 1fa7d00325b42f7720b32aecc918d523af3389f5 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Tue, 18 Aug 2026 08:20:56 -0400 Subject: [PATCH 34/44] vllm adapter: relay the checkpoint's W4A16 head for small M MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The head's verify-step calls are M<=16 logits projections, served by the host on a tiled kernel sized for prefill. The relay adopts the checkpoint's own NVFP4 pack into the small-M layout at bind time — same grid, same dequantized values — and routes calls inside the band there, host method for everything else. Judged across three boots the step rate moves +7-8% with acceptance unchanged. The bind receipt compares against an E2M1 table-decode of the checkpoint's rows rather than the host's apply: the host quantizes activations on its own path, so an exact relay still reads ~0.92 against it while sitting at 0.999999 against the rows themselves. A 'w8' precision tier joins the projection binder for the draft band: per-channel weight-only INT8 from the module's dense rows, the middle ground between BF16's weight stream and W4's rounding. --- flash_rt/structures/adapters/vllm_engine.py | 115 +++++++++++++++++++- 1 file changed, 114 insertions(+), 1 deletion(-) diff --git a/flash_rt/structures/adapters/vllm_engine.py b/flash_rt/structures/adapters/vllm_engine.py index d6b166d0..ae42e13a 100644 --- a/flash_rt/structures/adapters/vllm_engine.py +++ b/flash_rt/structures/adapters/vllm_engine.py @@ -512,6 +512,40 @@ 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 @@ -862,6 +896,78 @@ def _init(self, *a, **kw): 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 @@ -904,7 +1010,8 @@ def _init(self, *a, **kw): mod._frt_seat_name = (ckpt_rename(name) if ckpt_rename else name) kind = (_host_precision(mod) - if precision in ("mirror", "auto") else "nvfp4") + 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, @@ -945,6 +1052,12 @@ def _init(self, *a, **kw): 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}) From cb5dd95c15c0f7f62e1a4d4ece024b7799f571ff Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Tue, 18 Aug 2026 12:51:11 -0400 Subject: [PATCH 35/44] vllm adapter: head relay steps aside under speculative decode The engine's MTP draft shares the target's lm_head, and acceptance is an agreement test between the two. Relaying only the target's head to the small-M layout changes one side of that agreement: judged paired over ten prompts the acceptance length lands lower on all ten (about -0.5) while the step rate gain is smaller than the loss. With a speculative config present the relay now defaults off; without one it remains the head's default, where it is pure step-rate. The paired-by-prompt protocol is the judging change that surfaced this: single-prompt greedy runs drift to different continuations per arm and bury systematic acceptance effects in content variance. --- flash_rt/structures/adapters/vllm_engine.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/flash_rt/structures/adapters/vllm_engine.py b/flash_rt/structures/adapters/vllm_engine.py index ae42e13a..b7c303c7 100644 --- a/flash_rt/structures/adapters/vllm_engine.py +++ b/flash_rt/structures/adapters/vllm_engine.py @@ -1299,6 +1299,17 @@ def install_load_hook(*, on_attached=None, seat_draft=True, 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 From 40abea058b904cc4a1896986bd30af0ec56e9493 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Tue, 18 Aug 2026 12:55:12 -0400 Subject: [PATCH 36/44] docs: serving-engine attach guide (vLLM / SGLang) Usage for the engine adapters with every configuration in its measured form: the load-hook window and precision tiers, the spec-decode head-relay rule, the explicit-KV long-context recipe, the SGLang sitecustomize/container bridge with the DSpark memory floor, and the judging protocol the receipts require (paired-by- prompt for speculative arms, multi-boot medians at long context, real-text prompts with difference-based decode measurement). --- docs/serving_engine_attach.md | 165 ++++++++++++++++++++++++++++++++++ 1 file changed, 165 insertions(+) create mode 100644 docs/serving_engine_attach.md diff --git a/docs/serving_engine_attach.md b/docs/serving_engine_attach.md new file mode 100644 index 00000000..8dff8dd0 --- /dev/null +++ b/docs/serving_engine_attach.md @@ -0,0 +1,165 @@ +# 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. + +Hardware/software baseline for the numbers quoted here: one RTX 5090 +(32 GB), vLLM 0.27.x, SGLang 0.5.x, a mixed-precision NVFP4/FP8 +checkpoint of a hybrid-attention 27B model, real code/text prompts. + +## 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. + +### 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. + +### 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. From ae731f903f5dd6a68c6a794d22091d0667fc1e2e Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Tue, 18 Aug 2026 13:10:54 -0400 Subject: [PATCH 37/44] docs: requirements section for the engine attach guide MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The bind-or-refuse determinants in one place: SM120 for the sm120a kernel tiers (elsewhere the seats refuse and the host runs unmodified), the vLLM/SGLang series the adapters read, the torch/CUDA build variant the published kernels ship for, and the ModelOpt mixed-precision checkpoint family the adopt tiers expect — with the re-grid tiers as the road for dense checkpoints. Records the measured speculative configuration and the single-stream scope of the receipts. --- docs/serving_engine_attach.md | 46 ++++++++++++++++++++++++++++++++--- 1 file changed, 43 insertions(+), 3 deletions(-) diff --git a/docs/serving_engine_attach.md b/docs/serving_engine_attach.md index 8dff8dd0..56b2420f 100644 --- a/docs/serving_engine_attach.md +++ b/docs/serving_engine_attach.md @@ -5,9 +5,49 @@ 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. -Hardware/software baseline for the numbers quoted here: one RTX 5090 -(32 GB), vLLM 0.27.x, SGLang 0.5.x, a mixed-precision NVFP4/FP8 -checkpoint of a hybrid-attention 27B model, real code/text prompts. +## 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: **Qwen3.8-27B** (Qwen3.5 backbone: 48 gated-delta + + 16 full-attention layers), served from a **ModelOpt mixed-precision + checkpoint** — 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). +- 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 From f500edb8f742b8daabd3657d6ec264df55d6b40f Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Tue, 18 Aug 2026 13:12:42 -0400 Subject: [PATCH 38/44] docs: SGLang long-context receipt joins the guide MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Same multiplier, second engine: the stock DSpark server's pool is 34,659 tokens and refuses a 60K request; the attached server's pool is 87,128 (2.51x) and decodes an 80K real prompt at 210.8 tok/s. Absolute ceilings differ from vLLM's because the hybrid-state cache and draft KV price each token higher — the released weight memory multiplies whichever budget the engine's accounting allows. --- docs/serving_engine_attach.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/docs/serving_engine_attach.md b/docs/serving_engine_attach.md index 56b2420f..9bf2fc47 100644 --- a/docs/serving_engine_attach.md +++ b/docs/serving_engine_attach.md @@ -120,6 +120,14 @@ 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, identical memory +fraction): the stock server's pool is **34,659** tokens — a 60K +request is refused outright — while the attached server's pool is +**87,128** (2.51x) and an 80K real-prompt request decodes at +**210.8 tok/s** (AL 4.49). The multiplier is the released weight +memory; the absolute ceilings differ because SGLang's hybrid-state +cache and draft KV cost more per token than vLLM's pools. + ### Judging protocol (what the receipts require) - Speculative arms drift to different greedy continuations per arm, and From eddf8b13c237445fd70562852d8f51fd7c868811 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Tue, 18 Aug 2026 13:20:01 -0400 Subject: [PATCH 39/44] docs: scope disclaimer, checkpoint provenance, SGLang context table MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The guide now opens with what it is and is not: measured hot-plug configurations on each engine's documented recipe, not an exhaustive engine tuning and not a community comparison — the ceiling figures reflect the engines' default memory accounting under these settings. Names the measured checkpoints (the ModelOpt mixed-precision quantization of the 27B host, and the DSpark speculator) so the receipts are reproducible from the hub. The SGLang long-context receipt becomes a table with the full curve: pool 34,659 vs 87,128 tokens, the stock refusal at 60K, and 147.7 / 210.8 tok/s decode at 60K / 80K on the attached server. --- docs/serving_engine_attach.md | 52 +++++++++++++++++++++++++++-------- 1 file changed, 41 insertions(+), 11 deletions(-) diff --git a/docs/serving_engine_attach.md b/docs/serving_engine_attach.md index 9bf2fc47..69dca1f9 100644 --- a/docs/serving_engine_attach.md +++ b/docs/serving_engine_attach.md @@ -5,6 +5,20 @@ 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. +> **Scope disclaimer — read first.** The vLLM/SGLang configurations +> measured here are **not an exhaustive tuning of either engine**, and +> 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 our protocol needed; +> both engines have many knobs we did not 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. + ## Requirements These determine whether the seats bind at all — check them before @@ -32,12 +46,19 @@ anything else. `FRT_KERNEL_DIR_*`. **Model / checkpoint** -- Measured host: **Qwen3.8-27B** (Qwen3.5 backbone: 48 gated-delta + - 16 full-attention layers), served from a **ModelOpt mixed-precision - checkpoint** — 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 +- 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. @@ -121,12 +142,21 @@ 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, identical memory -fraction): the stock server's pool is **34,659** tokens — a 60K -request is refused outright — while the attached server's pool is -**87,128** (2.51x) and an 80K real-prompt request decodes at -**210.8 tok/s** (AL 4.49). The multiplier is the released weight -memory; the absolute ceilings differ because SGLang's hybrid-state -cache and draft KV cost more per token than vLLM's pools. +fraction, real code prompts): + +| context | stock server | attached server | +|---|---|---| +| 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 multiplier is the released weight memory; the absolute ceilings +differ from vLLM's because SGLang's hybrid-state cache and draft KV +price each token higher. 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) From 2c3d575ed3efb66ffacb7a33adb1257b5de71f49 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Tue, 18 Aug 2026 14:15:30 -0400 Subject: [PATCH 40/44] docs: record the stock-arm tuning effort behind the SGLang table MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The comparison now states what was tried before it was written: the draft model card's own serving flags on both arms, a memory-fraction sweep (a higher fraction grows the paper pool to ~56K tokens but 48K requests then fail on runtime headroom, and neither arm boots above it), CUDA-graph trims, and the hybrid-cache knobs — none of which move the pool, whose 2.25 GB intermediate state cache is insensitive to all three. Notes that higher published context figures for this family come from the card's multi-GPU recipe, keeping the single-card table apples-to-apples. --- docs/serving_engine_attach.md | 31 ++++++++++++++++++++++--------- 1 file changed, 22 insertions(+), 9 deletions(-) diff --git a/docs/serving_engine_attach.md b/docs/serving_engine_attach.md index 69dca1f9..5555797b 100644 --- a/docs/serving_engine_attach.md +++ b/docs/serving_engine_attach.md @@ -141,22 +141,35 @@ 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, identical memory -fraction, real code prompts): +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 | attached server | +| 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 multiplier is the released weight memory; the absolute ceilings -differ from vLLM's because SGLang's hybrid-state cache and draft KV -price each token higher. 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. +We did try to tune the stock arm higher before writing this table: +raising the memory fraction to 0.95 grows the paper pool to 55,822 +tokens, but a 48K request then fails server-side at both prefill +chunk sizes we tried — the extra fraction is exactly the runtime +headroom the request needed (the attached arm cannot boot at 0.95 +either; 0.92 is the stable envelope for both). Trimming CUDA-graph +allocations, shrinking `context-length`, and widening +`mamba-track-interval` did not move the pool; the hybrid line's +2.25 GB intermediate state cache is insensitive to all three. 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) From 5fafa0221ed9d0005ed9ffb01a6ea68a3820f335 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Tue, 18 Aug 2026 14:17:44 -0400 Subject: [PATCH 41/44] docs: acknowledgment note and the stock-arm sweep as a table MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The disclaimer now says plainly what is owed: these adapters exist because both engines are excellent hosts, and better configurations than ours are welcome — the numbers will be updated for them. The SGLang stock-arm tuning narrative becomes a checkable table: every configuration attempted, its pool, and what a long request actually did, with the paper-pool footnote and the stable-envelope conclusion the comparison table rests on. --- docs/serving_engine_attach.md | 36 ++++++++++++++++++++++++++--------- 1 file changed, 27 insertions(+), 9 deletions(-) diff --git a/docs/serving_engine_attach.md b/docs/serving_engine_attach.md index 5555797b..e5092e8c 100644 --- a/docs/serving_engine_attach.md +++ b/docs/serving_engine_attach.md @@ -18,6 +18,14 @@ each command is the exact form its receipts were produced with. > possible when the structures layer attaches to an engine *as > configured*, with everything reverting to the untouched host on > detach or refusal. +> +> We are grateful to the vLLM and SGLang teams — these adapters exist +> because both engines are excellent hosts, and everything here runs +> *through* their serving stacks, not around them. If a configuration +> we did not cover serves these workloads better (several knobs on +> both engines were outside our sweep), we would genuinely like to +> hear about it and will update the numbers; the tuning attempts we +> did make are recorded below so they can be checked and improved on. ## Requirements @@ -152,15 +160,25 @@ draft `unquant`, `mamba-radix-cache-strategy extra_buffer`): | 60K request | **refused** (exceeds pool) | **147.7 tok/s** (AL 3.05) | | 80K request | refused | **210.8 tok/s** (AL 4.49) | -We did try to tune the stock arm higher before writing this table: -raising the memory fraction to 0.95 grows the paper pool to 55,822 -tokens, but a 48K request then fails server-side at both prefill -chunk sizes we tried — the extra fraction is exactly the runtime -headroom the request needed (the attached arm cannot boot at 0.95 -either; 0.92 is the stable envelope for both). Trimming CUDA-graph -allocations, shrinking `context-length`, and widening -`mamba-track-interval` did not move the pool; the hybrid line's -2.25 GB intermediate state cache is insensitive to all three. Higher +We tried to tune the stock arm higher before writing this table, and +record the sweep 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 From 1ee0126824ae4f340434e9d27008f055ed0314b8 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Tue, 18 Aug 2026 14:24:19 -0400 Subject: [PATCH 42/44] docs: the note is a fairness explanation, in one author's voice MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Retitled from a disclaimer to what it actually is — a note on how to read the numbers fairly: not a community benchmark, an exploration of what hot-plugging adds to each engine as configured. The voice now matches the project (single author), and the invitation stands: better configurations are welcome and the numbers will be updated for them. --- docs/serving_engine_attach.md | 43 +++++++++++++++++------------------ 1 file changed, 21 insertions(+), 22 deletions(-) diff --git a/docs/serving_engine_attach.md b/docs/serving_engine_attach.md index e5092e8c..2a7648af 100644 --- a/docs/serving_engine_attach.md +++ b/docs/serving_engine_attach.md @@ -5,27 +5,26 @@ 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. -> **Scope disclaimer — read first.** The vLLM/SGLang configurations -> measured here are **not an exhaustive tuning of either engine**, and -> 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 our protocol needed; -> both engines have many knobs we did not 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. +> **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. > -> We are grateful to the vLLM and SGLang teams — these adapters exist -> because both engines are excellent hosts, and everything here runs -> *through* their serving stacks, not around them. If a configuration -> we did not cover serves these workloads better (several knobs on -> both engines were outside our sweep), we would genuinely like to -> hear about it and will update the numbers; the tuning attempts we -> did make are recorded below so they can be checked and improved on. +> 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 @@ -160,8 +159,8 @@ draft `unquant`, `mamba-radix-cache-strategy extra_buffer`): | 60K request | **refused** (exceeds pool) | **147.7 tok/s** (AL 3.05) | | 80K request | refused | **210.8 tok/s** (AL 4.49) | -We tried to tune the stock arm higher before writing this table, and -record the sweep so it can be checked and improved on: +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 | |---|---|---| From e91c48797afa0a64de003af06a372a8cce217e7a Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Wed, 19 Aug 2026 06:35:32 -0400 Subject: [PATCH 43/44] docs+examples: stock-CLI serving attach, and the demo that shows it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The guide gains the zero-code-change path for a stock vllm serve: a sitecustomize on PYTHONPATH installs the load hook behind an env gate, with the measured server-side receipts (code +7.5%, text +9%, TTFT -13% against back-to-back stock boots; the spec-decode head rule carries over unchanged). The example is the demo built on it: the same serve command twice, where the stock boot refuses the model's native 262144 context on a 32 GB card and the attached boot serves it — released weight memory becoming the KV pool — plus a streaming client with a live decode meter for asking questions of a 200K-token codebase. --- docs/serving_engine_attach.md | 41 +++++++++ examples/serving_attach_demo/README.md | 48 +++++++++++ examples/serving_attach_demo/ask.py | 84 +++++++++++++++++++ .../serving_attach_demo/serve_attached.sh | 15 ++++ examples/serving_attach_demo/serve_stock.sh | 10 +++ examples/serving_attach_demo/sitecustomize.py | 26 ++++++ 6 files changed, 224 insertions(+) create mode 100644 examples/serving_attach_demo/README.md create mode 100644 examples/serving_attach_demo/ask.py create mode 100755 examples/serving_attach_demo/serve_attached.sh create mode 100755 examples/serving_attach_demo/serve_stock.sh create mode 100644 examples/serving_attach_demo/sitecustomize.py diff --git a/docs/serving_engine_attach.md b/docs/serving_engine_attach.md index 2a7648af..876cc7bc 100644 --- a/docs/serving_engine_attach.md +++ b/docs/serving_engine_attach.md @@ -99,6 +99,47 @@ 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 | diff --git a/examples/serving_attach_demo/README.md b/examples/serving_attach_demo/README.md new file mode 100644 index 00000000..cb0b9fdf --- /dev/null +++ b/examples/serving_attach_demo/README.md @@ -0,0 +1,48 @@ +# 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. diff --git a/examples/serving_attach_demo/ask.py b/examples/serving_attach_demo/ask.py new file mode 100644 index 00000000..58f4c344 --- /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/serve_attached.sh b/examples/serving_attach_demo/serve_attached.sh new file mode 100755 index 00000000..27badc23 --- /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_stock.sh b/examples/serving_attach_demo/serve_stock.sh new file mode 100755 index 00000000..6c94c6ad --- /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/sitecustomize.py b/examples/serving_attach_demo/sitecustomize.py new file mode 100644 index 00000000..3b62c05a --- /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) From 12d902946b9f8f707fed02b69638bc0a387c8866 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Wed, 19 Aug 2026 07:08:16 -0400 Subject: [PATCH 44/44] =?UTF-8?q?examples:=20the=20context=20race=20?= =?UTF-8?q?=E2=80=94=20a=20growing=20conversation=20to=20the=20ceiling?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit One conversation that keeps growing beats one giant prompt on camera: with prefix caching on, each turn prefills only its increment, so the demo keeps a conversational pace while the context meter climbs. The stock arm dies at its ceiling with the server's own 400 printed as the crash row; the attached arm walks the same turns to the model's native maximum. Arms are recorded separately (one card cannot hold two copies of the weights) and the client prints the same table for both, so the recordings align turn for turn. The race form runs without speculative decode, which this vLLM series requires for prefix caching; the one-shot scenes keep the speculative speed numbers. --- examples/serving_attach_demo/README.md | 30 ++++++ examples/serving_attach_demo/race.py | 96 +++++++++++++++++++ .../serve_attached_race.sh | 13 +++ .../serving_attach_demo/serve_stock_race.sh | 12 +++ 4 files changed, 151 insertions(+) create mode 100644 examples/serving_attach_demo/race.py create mode 100755 examples/serving_attach_demo/serve_attached_race.sh create mode 100755 examples/serving_attach_demo/serve_stock_race.sh diff --git a/examples/serving_attach_demo/README.md b/examples/serving_attach_demo/README.md index cb0b9fdf..f48819db 100644 --- a/examples/serving_attach_demo/README.md +++ b/examples/serving_attach_demo/README.md @@ -46,3 +46,33 @@ 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/race.py b/examples/serving_attach_demo/race.py new file mode 100644 index 00000000..9be1f6bf --- /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_race.sh b/examples/serving_attach_demo/serve_attached_race.sh new file mode 100755 index 00000000..b478e0c4 --- /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_race.sh b/examples/serving_attach_demo/serve_stock_race.sh new file mode 100755 index 00000000..23f55bf5 --- /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