sparse-delta-memory 2.0.0__py3-none-any.whl

sparse_delta_memory/cuda/px_qrec_cuda.py

@@ -0,0 +1,167 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+"""px_qrec_cuda: compact read records (qrec8) for the deterministic backward's read-side mover (GDS), built on a
+side stream, one chunk step at a time.
+
+qrec8[c, j] = {absolute slot, f32 bits of __fmul_rn(f32(q_val), expf(log_q))} for sorted read position j of chunk c
+(the order of Phase1.spay_q[c]); padding and invalid entries are {0, 0}. The records are bitwise equal to words 1-2 of
+Phase 1's qrec (DET_QREC) and to Phase 1's own qrec8 (DET_QREC8): the same inputs (the saved q_idx, q_val, log_q) and
+the same expression.
+
+build_qrec8_steps(...) returns (qrec8 int32 [Nt*P, C*R, 2], events). events[ti] is recorded once rows
+[ti*P, (ti+1)*P) are written (every step is queued up front, in backward order Nt-1..0), and the consumer
+(gx_gds_cuda) waits on events[ti] before step ti. events is None when one launch on the current stream builds every
+step (CUDA graph capture, or Nt == 1).
+
+Why the side-stream path is safe:
+- Fork: side.wait_stream(cur) orders every side launch after the work already queued on the current stream (the inputs
+  are complete; qrec8, allocated on the current stream just before, is not used by queued work).
+- Join: the consumer waits on events[ti] on the current stream before step ti, the only reader of those rows. Every
+  step is waited on, so all side work is complete, as seen from the current stream, once the backward loop has run.
+- Allocator: qrec8, every input and every contiguous copy made here are marked record_stream(side), so the caching
+  allocator keeps their blocks from reuse until the side work queued so far has completed. No host synchronization."""
+
+from __future__ import annotations
+
+import functools
+
+import torch
+
+from ._build import load_cuda_extension
+
+__all__ = ["build_qrec8", "build_qrec8_steps"]
+
+_CPP_SOURCE = r"""
+void px_qrec8(at::Tensor spay_q, at::Tensor q_idx, at::Tensor q_val, at::Tensor log_q, at::Tensor out,
+              int64_t c0, int64_t c1);
+"""
+
+_CUDA_SOURCE = r"""
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <cuda_bf16.h>
+
+namespace {
+constexpr uint32_t kInvalid = 1u << 17;
+__device__ __forceinline__ float f_of(float v) { return v; }
+__device__ __forceinline__ float f_of(__nv_bfloat16 v) { return __bfloat162float(v); }
+
+// One thread per sorted read position of chunks [c0, c1): the records of qrec words 1-2.
+template <typename IdxT, typename ValT>
+__global__ void __launch_bounds__(256) px_qrec8_records_kernel(
+    const int32_t* __restrict__ spay_q, const IdxT* __restrict__ q_idx, const ValT* __restrict__ q_val,
+    const float* __restrict__ log_q, int2* __restrict__ out, int64_t first, int64_t total, int n, int R) {
+    const int64_t i = first + static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+    if (i >= total) return;
+    const int64_t chunk = i / n;
+    const uint32_t e = static_cast<uint32_t>(spay_q[i]);
+    int slot = 0;
+    float qs = 0.0f;
+    if (!(e & kInvalid)) {
+        const int64_t o = chunk * n + static_cast<int64_t>((e >> 9) & 0xFFu) * R + (e & 0xFFu);
+        slot = static_cast<int>(q_idx[o]);
+        qs = __fmul_rn(f_of(q_val[o]), expf(log_q[o]));
+    }
+    __stcs(out + i, make_int2(slot, __float_as_int(qs)));
+}
+}  // namespace
+
+void px_qrec8(at::Tensor spay_q, at::Tensor q_idx, at::Tensor q_val, at::Tensor log_q, at::Tensor out,
+              int64_t c0, int64_t c1) {
+    const c10::cuda::CUDAGuard guard(spay_q.device());
+    TORCH_CHECK(spay_q.dim() == 2 && spay_q.scalar_type() == at::kInt && spay_q.is_contiguous(), "px_qrec8: spay_q");
+    const int64_t N = spay_q.size(0), n = spay_q.size(1);
+    TORCH_CHECK(q_idx.dim() == 3 && q_idx.size(0) == N && q_idx.size(1) * q_idx.size(2) == n &&
+                (q_idx.scalar_type() == at::kInt || q_idx.scalar_type() == at::kLong) && q_idx.is_contiguous(),
+                "px_qrec8: q_idx [N, C, R] int32/int64");
+    TORCH_CHECK(q_val.sizes() == q_idx.sizes() && q_val.is_contiguous() &&
+                (q_val.scalar_type() == at::kFloat || q_val.scalar_type() == at::kBFloat16), "px_qrec8: q_val");
+    TORCH_CHECK(log_q.sizes() == q_idx.sizes() && log_q.is_contiguous() && log_q.scalar_type() == at::kFloat,
+                "px_qrec8: log_q");
+    TORCH_CHECK(out.numel() == N * n * 2 && out.scalar_type() == at::kInt && out.is_contiguous() &&
+                (reinterpret_cast<uintptr_t>(out.data_ptr()) & 7) == 0, "px_qrec8: out [N, n, 2] int32");
+    for (const at::Tensor* t : {&q_idx, &q_val, &log_q, &out})
+        TORCH_CHECK(t->device() == spay_q.device(), "px_qrec8: one device");
+    TORCH_CHECK(0 <= c0 && c0 <= c1 && c1 <= N, "px_qrec8: chunk range");
+    const int64_t first = c0 * n, total = c1 * n;
+    if (total == first) return;
+    const int R = static_cast<int>(q_idx.size(2));
+    const int blocks = static_cast<int>((total - first + 255) / 256);
+    auto stream = at::cuda::getCurrentCUDAStream();
+    auto* o2 = reinterpret_cast<int2*>(out.data_ptr<int32_t>());
+#define PX_Q8(IDX, VAL)                                                                                   \
+    px_qrec8_records_kernel<IDX, VAL><<<blocks, 256, 0, stream>>>(                                        \
+        spay_q.data_ptr<int32_t>(), q_idx.data_ptr<IDX>(), reinterpret_cast<const VAL*>(q_val.data_ptr()),  \
+        log_q.data_ptr<float>(), o2, first, total, static_cast<int>(n), R)
+    const bool i64 = q_idx.scalar_type() == at::kLong, bf = q_val.scalar_type() == at::kBFloat16;
+    if (i64 && bf) { PX_Q8(int64_t, __nv_bfloat16); }
+    else if (i64) { PX_Q8(int64_t, float); }
+    else if (bf) { PX_Q8(int32_t, __nv_bfloat16); }
+    else { PX_Q8(int32_t, float); }
+#undef PX_Q8
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+"""
+
+
+@functools.cache
+def _module():
+    return load_cuda_extension(
+        "sdm_px_qrec",
+        cpp_sources=_CPP_SOURCE,
+        cuda_sources=_CUDA_SOURCE,
+        functions=("px_qrec8",),
+        extra_cuda_cflags=("-O3",),
+    )
+
+
+def build_qrec8(spay_q, q_idx, q_val, log_q):
+    """qrec8 int32 [N, C*R, 2] in one launch on the current stream (see the module docstring)."""
+    n = spay_q.shape[-1]
+    sq = spay_q.reshape(-1, n).contiguous()
+    N = sq.shape[0]
+    out = torch.empty(N, n, 2, device=sq.device, dtype=torch.int32)
+    _module().px_qrec8(sq, q_idx.reshape(N, -1, q_idx.shape[-1]).contiguous(),
+                       q_val.reshape(N, -1, q_val.shape[-1]).contiguous(),
+                       log_q.reshape(N, -1, log_q.shape[-1]).contiguous(), out, 0, N)
+    return out
+
+
+_SIDE = {}
+
+
+def build_qrec8_steps(spay_q, q_idx, q_val, log_q, *, P, Nt):
+    """(qrec8 int32 [Nt*P, C*R, 2], events or None): records built per chunk step on a side stream, steps queued in
+    backward order; events[ti] marks rows [ti*P, (ti+1)*P) as written. Under CUDA graph capture on the inputs' device,
+    or for Nt == 1, one launch on the current stream builds every step (events None). Same bits either way. Slots are
+    narrowed to int32 as in Phase 1's qrec: the forward that produced spay_q / log_q has checked the absolute-slot bound."""
+    n = spay_q.shape[-1]
+    N = spay_q.numel() // n
+    if N != Nt * P:
+        raise ValueError(f"build_qrec8_steps: {N} chunks != Nt * P = {Nt} * {P}")
+    dev = spay_q.device
+    with torch.cuda.device(dev):
+        if Nt <= 1 or torch.cuda.is_current_stream_capturing():
+            return build_qrec8(spay_q, q_idx, q_val, log_q), None
+        sq = spay_q.reshape(-1, n).contiguous()
+        qi = q_idx.reshape(N, -1, q_idx.shape[-1]).contiguous()
+        qv = q_val.reshape(N, -1, q_val.shape[-1]).contiguous()
+        lq = log_q.reshape(N, -1, log_q.shape[-1]).contiguous()
+        cur = torch.cuda.current_stream(dev)
+        side = _SIDE.get(dev)
+        if side is None:
+            side = _SIDE.setdefault(dev, torch.cuda.Stream(device=dev))
+        out = torch.empty(N, n, 2, device=dev, dtype=torch.int32)
+        mod = _module()
+        side.wait_stream(cur)
+        events = [None] * Nt
+        with torch.cuda.stream(side):
+            for ti in reversed(range(Nt)):
+                mod.px_qrec8(sq, qi, qv, lq, out, ti * P, (ti + 1) * P)
+                ev = torch.cuda.Event()
+                ev.record(side)
+                events[ti] = ev
+        for t in (out, sq, qi, qv, lq):
+            t.record_stream(side)
+    return out, events

sparse_delta_memory/cuda/sdm_copy.cuh

@@ -0,0 +1,90 @@
+// sdm_copy.cuh: the routing predicate's device mirror and the portable copy layer of the fast kernels.
+//
+// SDM_PORTABLE_ARCH is defined in every device pass except the modern ones (sm_90 / sm_100 / sm_103: H100, B200,
+// GB300), mirroring cuda/_device.py backend(): the modern passes keep their preprocessed text, every other device pass
+// (sm_80, 86, 89, 120, ...) compiles the portable side of the kernels' arch guards, so an sm_120 build runs the same
+// device branches as sm_86 / 89. Guards: `#ifdef SDM_PORTABLE_ARCH` (a portable device pass) and
+// `#if defined(__CUDA_ARCH__) && !defined(SDM_PORTABLE_ARCH)` (a modern device pass); the host pass defines neither.
+//
+// sdmcp: cooperative global -> shared copies with cp.async for the portable side. Every copying thread issues its pieces
+// and commits its own groups; a thread's cp.async wait covers its own copies only, so the bytes reach the other threads
+// through a barrier passed AFTER every copier waited (publish_cta / publish_named), or through an mbarrier the copies
+// arrive on (arrive_on). No bulk copies, no transaction counts, no PDL, no clusters. Every call is made by all threads of
+// the copying group at the same program point, never from a thread-0 branch (the barriers would deadlock). Only portable
+// device passes see this code (no modern pass gets a token of it); below sm_80 (no cp.async) every call traps, so
+// multi-architecture builds keep compiling.
+#pragma once
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ != 900 && __CUDA_ARCH__ != 1000 && __CUDA_ARCH__ != 1030
+#define SDM_PORTABLE_ARCH 1
+#endif
+
+#ifdef SDM_PORTABLE_ARCH
+#include <cstdint>
+#if __CUDA_ARCH__ >= 800
+#define SDMCP_ASM(...) asm volatile(__VA_ARGS__)
+#else
+#define SDMCP_ASM(...) __trap()
+#endif
+namespace sdmcp {
+__device__ __forceinline__ uint32_t saddr(const void* p) { return static_cast<uint32_t>(__cvta_generic_to_shared(p)); }
+// 16 bytes through L2 only (coherent with what the previous kernel wrote); the bytes past src_bytes (0..16) are
+// zero-filled (src_bytes 0 reads nothing, but `g` must still be a valid global address).
+__device__ __forceinline__ void cp16(void* s, const void* g, int src_bytes = 16) {
+    SDMCP_ASM("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"r"(saddr(s)), "l"(g), "r"(src_bytes) : "memory");
+}
+// 8 / 4 bytes (rows that are not whole 16-byte vectors), through L1 and L2; zero fill past src_bytes.
+__device__ __forceinline__ void cp8(void* s, const void* g, int src_bytes = 8) {
+    SDMCP_ASM("cp.async.ca.shared.global [%0], [%1], 8, %2;\n" ::"r"(saddr(s)), "l"(g), "r"(src_bytes) : "memory");
+}
+__device__ __forceinline__ void cp4(void* s, const void* g, int src_bytes = 4) {
+    SDMCP_ASM("cp.async.ca.shared.global [%0], [%1], 4, %2;\n" ::"r"(saddr(s)), "l"(g), "r"(src_bytes) : "memory");
+}
+__device__ __forceinline__ void commit() { SDMCP_ASM("cp.async.commit_group;\n" ::: "memory"); }
+// This thread's groups: at most N of the newest still pending (wait<N>), or none (wait_all).
+template <int N>
+__device__ __forceinline__ void wait() { SDMCP_ASM("cp.async.wait_group %0;\n" ::"n"(N) : "memory"); }
+__device__ __forceinline__ void wait_all() { SDMCP_ASM("cp.async.wait_all;\n" ::: "memory"); }
+// Cooperative copy of `rows` rows of `row_bytes` bytes (a multiple of 16; every row start 16-byte aligned on both
+// sides): row r from g + r * g_stride to s + r * s_stride (bytes), the 16-byte pieces dealt round-robin over the threads
+// tid = 0 .. nthr - 1 of the copying group. Issues only: the caller commits.
+__device__ __forceinline__ void copy_rows16(void* s, int s_stride, const void* g, int64_t g_stride, int rows,
+                                            int row_bytes, int tid, int nthr) {
+    const int per = row_bytes >> 4;
+    for (int i = tid; i < rows * per; i += nthr) {
+        const int r = i / per;
+        const int c = (i - r * per) << 4;
+        cp16(static_cast<char*>(s) + r * s_stride + c, static_cast<const char*>(g) + r * g_stride + c);
+    }
+}
+// Every thread of the CTA: its own groups have landed, then the CTA barrier publishes every copier's bytes.
+__device__ __forceinline__ void publish_cta() {
+    wait_all();
+    __syncthreads();
+}
+// The same for the nthr threads (a multiple of 32) that use named barrier `id` (1..15; 0 is __syncthreads').
+__device__ __forceinline__ void publish_named(int id, int nthr) {
+    wait_all();
+    asm volatile("bar.sync %0, %1;\n" ::"r"(id), "r"(nthr) : "memory");
+}
+// mbarrier forms (sm_80): init with the number of copying threads; arrive_on = one arrival of this thread that the
+// system makes when ALL of its prior cp.async copies have landed (.noinc: it counts against the init count); consumers
+// spin on test_parity(bar, phase & 1).
+__device__ __forceinline__ void mbar_init(uint64_t* bar, uint32_t count) {
+    SDMCP_ASM("mbarrier.init.shared.b64 [%0], %1;\n" ::"r"(saddr(bar)), "r"(count) : "memory");
+}
+__device__ __forceinline__ void arrive_on(uint64_t* bar) {
+    SDMCP_ASM("cp.async.mbarrier.arrive.noinc.shared.b64 [%0];\n" ::"r"(saddr(bar)) : "memory");
+}
+__device__ __forceinline__ bool test_parity(uint64_t* bar, uint32_t parity) {
+    uint32_t done = 0u;
+#if __CUDA_ARCH__ >= 800
+    asm volatile("{\n .reg .pred p;\n mbarrier.test_wait.parity.shared.b64 p, [%1], %2;\n selp.u32 %0, 1, 0, p;\n}\n"
+                 : "=r"(done) : "r"(saddr(bar)), "r"(parity) : "memory");
+#else
+    __trap();
+#endif
+    return done != 0u;
+}
+}  // namespace sdmcp
+#undef SDMCP_ASM
+#endif  // SDM_PORTABLE_ARCH

sparse_delta_memory/cuda/sdm_plancap.cuh

@@ -0,0 +1,150 @@
+// sdm_plancap.cuh: launch capture for the backward's per-step plans (host code only: no device code).
+//
+// A plan builder runs a per-step launcher's own host function ONCE, at the step loop's largest step, with this thread's
+// capture slot set: every check of the host function runs as for a real call, and its single kernel launch is RECORDED
+// instead of issued -- the kernel (its host stub address), grid, block, dynamic shared memory and every parameter by
+// value, converted to the kernel's own parameter type as <<<>>> and cudaLaunchKernelEx(&cfg, kernel, args...) convert the
+// arguments. A plan's step launch patches only the parameters that vary per step (the step index, row offsets, valid
+// tokens) and issues the launch with cudaLaunchKernel (what <<<>>> issues) or, with PDL, with cudaLaunchKernelExC and
+// programmatic stream serialization (what cudaLaunchKernelEx(&cfg, kernel, args...) issues with that attribute): the
+// same kernel, configuration, parameters and stream as the host function's own launch.
+#pragma once
+#include <ATen/ATen.h>
+#include <cuda_runtime.h>
+#include <c10/util/Exception.h>
+#include <cstdint>
+#include <cstring>
+#include <type_traits>
+#include <utility>
+
+namespace sdm_plancap {
+
+constexpr int kMaxParams = 32;
+constexpr int kBufBytes = 768;
+
+struct Launch {
+    int32_t set;                      // 1 once the planned call recorded its launch (0: it launched nothing)
+    int32_t nparams;
+    const void* fn;                   // the kernel's host stub address (what <<<>>> passes to cudaLaunchKernel)
+    uint32_t grid[3];
+    uint32_t block[3];
+    uint64_t smem;
+    int32_t off[kMaxParams];          // byte offset of parameter i in buf
+    int32_t size[kMaxParams];         // sizeof(parameter i's type)
+    alignas(16) unsigned char buf[kBufBytes];
+};
+
+// This thread's capture slot: non-null only while a plan builder runs its host function.
+inline Launch*& slot() {
+    static thread_local Launch* p = nullptr;
+    return p;
+}
+
+template <typename T>
+inline void put(Launch& c, int& pos, const T& v) {
+    static_assert(std::is_trivially_copyable<T>::value, "plan capture: kernel parameters must be trivially copyable");
+    constexpr int al = alignof(T) < 16 ? static_cast<int>(alignof(T)) : 16;
+    pos = (pos + al - 1) / al * al;
+    TORCH_CHECK(c.nparams < kMaxParams && pos + static_cast<int>(sizeof(T)) <= kBufBytes,
+                "plan capture: too many kernel parameter bytes");
+    std::memcpy(c.buf + pos, &v, sizeof(T));
+    c.off[c.nparams] = pos;
+    c.size[c.nparams] = static_cast<int32_t>(sizeof(T));
+    ++c.nparams;
+    pos += static_cast<int>(sizeof(T));
+}
+
+// Inside a plan builder: record the launch (each argument converted to the kernel's parameter type) and return true; the
+// host function then returns without launching. Outside one: false (the host function launches as it always did).
+template <typename... P, typename... A>
+inline bool take(void (*kernel)(P...), dim3 grid, dim3 block, size_t smem, A&&... args) {
+    static_assert(sizeof...(P) == sizeof...(A), "plan capture: argument count != kernel parameter count");
+    Launch* c = slot();
+    if (c == nullptr) return false;
+    TORCH_CHECK(c->set == 0, "plan capture: the planned call launched more than one kernel");
+    c->set = 1;
+    c->nparams = 0;
+    c->fn = reinterpret_cast<const void*>(kernel);
+    c->grid[0] = grid.x; c->grid[1] = grid.y; c->grid[2] = grid.z;
+    c->block[0] = block.x; c->block[1] = block.y; c->block[2] = block.z;
+    c->smem = static_cast<uint64_t>(smem);
+    int pos = 0;
+    (put<typename std::decay<P>::type>(*c, pos, static_cast<typename std::decay<P>::type>(std::forward<A>(args))), ...);
+    return true;
+}
+
+// Parameter i of a captured launch as type T (nullptr when the recorded size differs: the caller refuses the plan).
+template <typename T>
+inline T* param(Launch& c, int i) {
+    if (i < 0 || i >= c.nparams || c.size[i] != static_cast<int32_t>(sizeof(T))) return nullptr;
+    return reinterpret_cast<T*>(c.buf + c.off[i]);
+}
+
+// Issue a captured launch on `stream`: cudaLaunchKernel (= <<<>>>), or with pdl cudaLaunchKernelExC + programmatic stream
+// serialization (= cudaLaunchKernelEx(&cfg, kernel, args...) with that attribute). Returns 0 or a cudaError_t (the
+// launch's, else cudaGetLastError()'s right after it, as C10_CUDA_KERNEL_LAUNCH_CHECK probes); an error state is consumed.
+inline int64_t launch(const Launch& c, bool pdl, cudaStream_t stream) {
+    void* params[kMaxParams];
+    for (int i = 0; i < c.nparams; ++i) params[i] = const_cast<unsigned char*>(c.buf) + c.off[i];
+    const dim3 grid(c.grid[0], c.grid[1], c.grid[2]);
+    const dim3 block(c.block[0], c.block[1], c.block[2]);
+    cudaError_t e;
+    if (pdl) {
+        cudaLaunchAttribute attr[1];
+        attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+        attr[0].val.programmaticStreamSerializationAllowed = 1;
+        cudaLaunchConfig_t cfg = {};
+        cfg.gridDim = grid;
+        cfg.blockDim = block;
+        cfg.dynamicSmemBytes = static_cast<size_t>(c.smem);
+        cfg.stream = stream;
+        cfg.attrs = attr;
+        cfg.numAttrs = 1;
+        e = cudaLaunchKernelExC(&cfg, c.fn, params);
+    } else {
+        e = cudaLaunchKernel(c.fn, grid, block, params, static_cast<size_t>(c.smem), stream);
+    }
+    if (e == cudaSuccess) e = cudaGetLastError();
+    else (void)cudaGetLastError();
+    return static_cast<int64_t>(e);
+}
+
+// RAII: set this thread's slot to `c` for one planned call (refuses nesting).
+struct Scope {
+    explicit Scope(Launch* c) {
+        TORCH_CHECK(slot() == nullptr, "plan capture: nested plan builders");
+        std::memset(c, 0, sizeof(Launch));
+        slot() = c;
+    }
+    ~Scope() { slot() = nullptr; }
+    Scope(const Scope&) = delete;
+    Scope& operator=(const Scope&) = delete;
+};
+
+// The current device, or -1 (error state consumed).
+inline int current_device() {
+    int dev = -1;
+    if (cudaGetDevice(&dev) != cudaSuccess) {
+        (void)cudaGetLastError();
+        return -1;
+    }
+    return dev;
+}
+
+// A plan as a CPU uint8 tensor (the bytes of the trivially copyable plan struct), and back.
+template <typename T>
+inline at::Tensor to_tensor(const T& pl) {
+    static_assert(std::is_trivially_copyable<T>::value, "plan: trivially copyable");
+    at::Tensor t = at::empty({static_cast<int64_t>(sizeof(T))}, at::TensorOptions().dtype(at::kByte));
+    std::memcpy(t.data_ptr<uint8_t>(), &pl, sizeof(T));
+    return t;
+}
+
+template <typename T>
+inline bool from_bytes(T& pl, const void* bytes, int64_t n, uint64_t magic) {
+    if (n != static_cast<int64_t>(sizeof(T))) return false;
+    std::memcpy(&pl, bytes, sizeof(T));
+    return pl.magic == magic;
+}
+
+}  // namespace sdm_plancap