diff --git a/docs/request-row-gather.md b/docs/request-row-gather.md new file mode 100644 index 000000000..7cf28562a --- /dev/null +++ b/docs/request-row-gather.md @@ -0,0 +1,179 @@ +# Per-HookPoint request-row gather and in-graph heads + +In packed (vLLM-style) execution one activation tensor per HookPoint holds the +rows of every scheduled request. Request-row gather lets each HookPoint copy +only the rows it needs: the rows of some requests, or only some token +positions of them, chosen independently per HookPoint and per step. The +selection is data in a per-HookPoint row table, so a captured CUDA graph +replays it without recapture. + +In-graph projection heads go one step further: instead of copying the +selected rows, a fused node computes `K` fp32 affine scores per selected row +inside the graph, and only the scores leave the GPU. + +## The row-table contract + +A row table is a stable int64 device tensor read by the producer at every +launch or replay: + +| word | meaning | +|---|---| +| `t[0]` | number of segments `n` (clamped to `GATHER_MAX_SEGS` = 256 and to the segments the table's storage holds); `-1` disables the producer for the step | +| `t[1 + 2k]` | source row of segment `k` | +| `t[2 + 2k]` | rows in segment `k` | + +The output is always packed: destination offsets are the running prefix of +the segment lengths. `n = -1` publishes nothing (no reservation, no task, no +payload) and needs no metadata; `n = 0` publishes an empty task. A producer +clamps every segment to the captured tensor, so a host-side bug cannot read +or write out of bounds. The torch ops pass the table's element count to the +kernels (never a host read of `t[0]`, which would sync and be baked into a +captured graph), and the device loader reads at most +`min(n, GATHER_MAX_SEGS, (numel - 1) / 2)` segments +(`native/csrc/ring/row_table.h`), so a count word larger than the table +never reads past it. + +The table is loaded by all threads in parallel and the destination prefix is +one warp's `__shfl_up_sync` scan; a single segment (a full-row table) takes a +flat whole-grid copy, the cost of the prefix producer. + +`configure_hook_gather(hook_point, table, row_bytes)` (in `dmi.hook_rows`) +arms gather mode on an installed HookPoint: it takes precedence over the +strip modes, and `row_bytes` is the packed row stride. Under the eager safety net +a gather HookPoint reserves its full tensor, so on an eager step the +integration writes a full-row table. + +Gather tables and heads belong to one attachment: `install_ring_hooks` and +`uninstall_ring_hooks` clear them (`clear_hook_row_features`), releasing the +table, arena and accumulator references and restoring the strip stride that +`configure_hook_gather` replaced. Integrations configure gather and heads +after installing the HookPoints. + +`RingEngine.set_gather_host_mirror(dev_base, host_base, row_bytes, rows)` +registers the pinned host copies of the tables. Outside stream capture a +launch whose mirrored row is `-1` is skipped (`gather_host_skips()`); a +capture always launches, because the graph re-reads the device row at every +replay. + +## Planning rows: `dmi.hook_rows` + +A `HookRows` plan holds the producer's `segments` and the consumer's +attribution: `req_ids`, `token_ranges`, destination `offsets`, +`source_offsets` and `orig_positions`. + +* `plan_hook_rows` plans one HookPoint from a request selection and optional + absolute token windows (`canonical_token_windows` validates and unions + them; `None` as an end is unbounded). +* `plan_hook_groups` / `HookRowPlanner` plan many HookPoints at once and + share one plan object between HookPoints whose selections agree. + `HookRowPlanner` caches request membership for a repeated request order, + never positions or rows. +* `PackedDemandPlanner` plans declared per-request demand with NumPy array + operations: a request's rule (HookPoint index and window pieces) is + compiled once per demand change and reused every step; the step derives + every HookPoint's rows for the whole batch with one clip, compaction and + merge. Its plans compare equal to the reference planner's and expose + `layout_arrays()` so the native side ingests int64 arrays instead of one + Python tuple per row. +* `full_row_plan` is the static full-row baseline: the whole physical + tensor `(0, total_rows)` in one run with the step's physical offsets, so a + gather producer fed it copies exactly what an ungathered producer copies + (padding rows included) and attributes the same rows to each request. A + step without scheduled requests gives the empty plan. +* `encode_row_table` writes a plan (or `None` for disabled) into the table + layout above; `segment_array` gives a plan's segments as an `[n, 2]` array + and `row_signature` its value identity, so an integration can skip + rewriting a table whose plan did not change. + +A plan with more segments than the table holds falls back to the +un-compacted full payload while keeping the consumer-side selection +(`capacity_fallback`). + +## Publishing per-hook layouts + +`StepContext.hook_rows` (and `RingTransport.set_step_context(hook_rows=...)`) +maps `(hook_type, layer_no)` to a plan. A planned hook reserves and publishes +its plan's `row_count` rows and carries its own request layout in its +metadata; other hooks use the step layout. Equal plans share one native +layout per step. `orig_positions` keeps each request's index in the +unfiltered step, which position-indexed hooks (`final_logits`) use. + +## The toggle and the table gate + +The gather, head and fused producers record their kernel nodes inside the +node-toggle capture window (see [node-toggle.md](node-toggle.md)); a fused +node carries both of its keys. A gather producer captured without its node +would keep firing after its key was disabled, while metadata and reservation +excluded it. + +`RingTransport.set_prefill_table_gate(True)` is an integration's guarantee that +every producer reads a per-HookPoint table whose `-1` marker also gates eager +and partial-model execution. The per-step spec set then follows the toggle on +every step, and excluded graphs may replay on any step. + +In device-toggle mode, a capture outside the toggle window (graphs the host +does not toggle) captures gather producers as device-updatable nodes behind +one planner node per captured forward, each gated by its own table row. + +## In-graph projection heads + +`ring::producer_fused` runs a HookPoint's row producer and its heads in one +captured node. Each part has its own table and toggle key: the rows under +the HookPoint's key, the heads under `HOOK_TYPE_*_HEAD` (`resid_pre_head`, +`final_ln_head`) with shape `[rows, head_slots]` fp32. The node arena (int32 +words, written only between steps) holds per-slot modes, accumulator +columns, fp32 biases and `K x hidden` weights (bf16 or fp32): + +| mode | output column | +|---|---| +| `FREE` | NaN | +| `SINGLE` | `w . h + b` | +| `FIRST` | NaN; writes `acc[row][a] = w . h` | +| `MID` | NaN; adds `w . h` to `acc[row][a]` | +| `LAST` | `acc[row][a] + w . h + b` | + +A linear head over several layers holds one slot per participating node; +layers run in order inside a forward, so the sum has a fixed order. +Registering or retiring a head rewrites the arena; nothing is recaptured. +Each `(row, slot)` score uses a fixed reduction order (no atomics, no +split-K), so it is bit-identical across batch compositions and grids. + +`ring::producer_head` is the standalone head node (fp32 arena: +`K x hidden` weights, `K` biases, split scratch). Selections of at most 16 +rows split each row's dot products over 8 hidden slices so a one-row decode +step is not bound by one SM. + +`configure_hook_heads(hook_point, head_type=..., layer_no=..., table=..., +arena=..., acc=..., slots=..., weight_bf16=...)` arms the heads. The +heads' spec names the same HookPoint and follows its row spec +(`BackendAdapter.replace_active_specs`), because the heads publish after the +rows. Dispatch keeps the planned tasks and their order on the fast path and +under the eager safety net alike: + +| HookPoint | launches | +|---|---| +| row spec + gather table + heads | one fused node (rows, then heads) | +| row spec (static, prefix or chunked) + heads | the row producer, then a head-only fused node | +| heads only (no row spec) | a head-only fused node | + +A head-only node carries no row table and is recorded under the head key +alone, so the row and head toggle states stay independent. A head spec +passed to `install_ring_hooks` binds the payload but never replaces the +HookPoint's row key. + +## Tests + +* `tests/native/ring/test_producer_equivalence.cu`: static, prefix, chunked + and gather producers (and planner-gated device nodes) against a host model + of the whole ring state over thousands of eager launches and graph replays. +* `tests/native/ring/test_inline_head.cu`, `test_fused_head.cu`: heads + against double-precision references, multi-layer chains, batch invariance, + graph replays with tables and arenas rewritten between replays. Each takes + `NVCC_EXTRA=-DDMI_TEST_MUTATE_*` mutants that must fail. +* `tests/test_hook_rows.py`: planners against independent enumerations. +* `tests/test_hook_row_features.py`: gather/head attachment lifecycle and + plan / fast path / eager task agreement for every row mode with heads. +* `tests/test_row_table_clamp.py`: the row-table segment-count clamp + against table storage, on the CPU. +* `tests/test_per_hook_gather_gpu.py`: end-to-end delivery through + HookPoints, the transport and the native consumer. diff --git a/native/csrc/bindings.cpp b/native/csrc/bindings.cpp index 7f5db58ab..6e448a63d 100644 --- a/native/csrc/bindings.cpp +++ b/native/csrc/bindings.cpp @@ -6,8 +6,12 @@ #include #endif #include +#include #include #include +#include +#include +#include #include #include #include @@ -259,6 +263,83 @@ std::vector CopyRecordDescriptors( } #endif +#ifndef DMI_HOST_ONLY +using HookRowsPtr = std::shared_ptr>; + +// One hook's request layout for push_all_metas. Two encodings: +// * a list of (req_id, start, end, dim0_offset, kv_offset, orig_position) +// tuples; +// * a tuple of five int64 arrays (index into the step's request list, +// start, end, dim0_offset, orig_position), which avoids building one +// Python tuple per row. +HookRowsPtr parse_hook_rows(py::handle rows_obj, + const std::vector& requests) { + auto parsed = std::make_shared>(); + if (py::isinstance(rows_obj)) { + auto fields = py::reinterpret_borrow(rows_obj); + if (py::len(fields) != 5) + throw py::value_error("hook request arrays need five fields"); + using I64 = py::array_t; + I64 idx = fields[0].cast(), lo = fields[1].cast(); + I64 hi = fields[2].cast(), off = fields[3].cast(); + I64 pos = fields[4].cast(); + if (idx.ndim() != 1 || lo.ndim() != 1 || hi.ndim() != 1 || + off.ndim() != 1 || pos.ndim() != 1) + throw py::value_error("hook request arrays must be one-dimensional"); + const py::ssize_t n = idx.shape(0); + if (lo.shape(0) != n || hi.shape(0) != n || off.shape(0) != n || + pos.shape(0) != n) + throw py::value_error("hook request arrays differ in length"); + auto pi = idx.unchecked<1>(); + auto pl = lo.unchecked<1>(); + auto ph = hi.unchecked<1>(); + auto po = off.unchecked<1>(); + auto pp = pos.unchecked<1>(); + const int64_t n_req = static_cast(requests.size()); + auto int32_field = [](int64_t v) { + if (v < INT32_MIN || v > INT32_MAX) + throw py::value_error("hook request field exceeds int32"); + return static_cast(v); + }; + parsed->reserve(static_cast(n)); + for (py::ssize_t i = 0; i < n; ++i) { + const int64_t r = pi(i); + if (r < 0 || r >= n_req) + throw py::value_error("hook request index outside the step's requests"); + ring_py::RequestMeta rm; + rm.req_id = requests[static_cast(r)].req_id; + rm.start_token = int32_field(pl(i)); + rm.end_token = int32_field(ph(i)); + rm.dim0_offset = po(i); + rm.kv_offset = 0; + rm.orig_position = int32_field(pp(i)); + if (rm.end_token < rm.start_token || rm.dim0_offset < 0) + throw py::value_error("invalid hook request range or offset"); + parsed->push_back(std::move(rm)); + } + return parsed; + } + auto rows = rows_obj.cast(); + parsed->reserve(py::len(rows)); + for (auto row : rows) { + auto fields = row.cast(); + if (py::len(fields) != 6) + throw py::value_error("hook request row needs six fields"); + ring_py::RequestMeta rm; + rm.req_id = fields[0].cast(); + rm.start_token = fields[1].cast(); + rm.end_token = fields[2].cast(); + rm.dim0_offset = fields[3].cast(); + rm.kv_offset = fields[4].cast(); + rm.orig_position = fields[5].cast(); + if (rm.end_token < rm.start_token || rm.dim0_offset < 0) + throw py::value_error("invalid hook request range or offset"); + parsed->push_back(std::move(rm)); + } + return parsed; +} +#endif + } // namespace PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { @@ -895,9 +976,16 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::list req_ids_py, py::list token_ranges_py, py::list dim0_offsets_py, - py::list kv_offsets_py) { + py::list kv_offsets_py, + py::list orig_positions_py, + py::list hook_requests_py) { + const size_t n_hooks = static_cast(py::len(hook_types_py)); + if (py::len(hook_requests_py) && + static_cast(py::len(hook_requests_py)) != n_hooks) + throw py::value_error("hook_requests must have one entry per hook"); // Build step context (heap-allocated, ownership to FIFO/p2p) - auto* ctx = new ring_py::StepContext(); + auto owned_ctx = std::make_unique(); + auto* ctx = owned_ctx.get(); ctx->model_id = model_id; ctx->tp_rank = tp_rank; ctx->dp_rank = dp_rank; @@ -914,12 +1002,20 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { rm.dim0_offset = py::cast(dim0_offsets_py[i]); if (i < static_cast(py::len(kv_offsets_py))) rm.kv_offset = py::cast(kv_offsets_py[i]); + // Only when the integration filtered the request list + // (empty list = unfiltered; -1 = use the list index). + if (i < static_cast(py::len(orig_positions_py))) + rm.orig_position = py::cast(orig_positions_py[i]); ctx->requests.push_back(std::move(rm)); } // Build per-hook metas size_t n = static_cast(py::len(hook_types_py)); std::vector metas; metas.reserve(n); + // Equal per-hook layouts are passed as the same Python object: + // parse each once while the arguments keep it alive, and never + // retain Python objects in the FIFO. + std::unordered_map layouts; for (size_t i = 0; i < n; ++i) { ring_py::TensorMeta meta; meta.hook_type = py::cast(hook_types_py[i]); @@ -930,11 +1026,22 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::list shape = shapes_py[i].cast(); for (auto d : shape) meta.shape.push_back(py::cast(d)); + if (py::len(hook_requests_py) && !hook_requests_py[i].is_none()) { + py::handle rows = hook_requests_py[i]; + auto found = layouts.find(rows.ptr()); + if (found != layouts.end()) { + meta.requests = found->second; + } else { + auto parsed = parse_hook_rows(rows, ctx->requests); + meta.requests = parsed; + layouts.emplace(rows.ptr(), std::move(parsed)); + } + } metas.push_back(std::move(meta)); } // Release GIL, push context + metas in single lock py::gil_scoped_release release; - self.push_step(ctx, metas); + self.push_step(owned_ctx.release(), metas); }, py::arg("hook_types"), py::arg("layer_nos"), py::arg("shapes"), py::arg("dtypes"), py::arg("flags"), @@ -944,7 +1051,16 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("flattened"), py::arg("req_ids"), py::arg("token_ranges"), py::arg("dim0_offsets"), - py::arg("kv_offsets") = py::list()) + py::arg("kv_offsets") = py::list(), + py::arg("orig_positions") = py::list(), + py::arg("hook_requests") = py::list()) + .def_static("supports_hook_request_arrays", []() { return true; }) + .def("set_gather_host_mirror", + &ring_py::RingEnginePy::set_gather_host_mirror, + py::arg("dev_base"), py::arg("host_base"), py::arg("row_bytes"), + py::arg("rows")) + .def("gather_host_skips", &ring_py::RingEnginePy::gather_host_skips) + .def("head_launches", &ring_py::RingEnginePy::head_launches) .def("set_null_mode", &ring_py::RingEnginePy::set_null_mode, py::arg("enabled"), @@ -1021,5 +1137,27 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("ring_clear_active_engine", []() { ring_set_active_engine(nullptr); }); + + // CPU-only view of push_all_metas' per-hook layout parsing: returns the + // parsed (req_id, start, end, dim0_offset, kv_offset, orig_position) rows. + m.def("parse_hook_requests", + [](py::list req_ids_py, py::handle rows) { + std::vector requests; + requests.reserve(py::len(req_ids_py)); + for (auto rid : req_ids_py) { + ring_py::RequestMeta rm; + rm.req_id = rid.cast(); + requests.push_back(std::move(rm)); + } + py::list out; + const HookRowsPtr parsed = parse_hook_rows(rows, requests); + for (const auto& rm : *parsed) { + out.append(py::make_tuple(rm.req_id, rm.start_token, rm.end_token, + rm.dim0_offset, rm.kv_offset, + rm.orig_position)); + } + return out; + }, + py::arg("req_ids"), py::arg("rows")); #endif } diff --git a/native/csrc/ring/p2p_thread.cpp b/native/csrc/ring/p2p_thread.cpp index fb73ca8d1..4278d6645 100644 --- a/native/csrc/ring/p2p_thread.cpp +++ b/native/csrc/ring/p2p_thread.cpp @@ -251,7 +251,9 @@ void P2PThread::do_post_processing(at::Tensor& tensor, const DrainTask& first_ta int32_t shard_rank = resolve_shard_rank(meta.hook_type, *current_ctx_); bool is_attn = ring_py::is_attn_weight_matrix(meta.hook_type); - const auto& requests = current_ctx_->requests; + // An independently gathered hook carries its own request layout; every + // other hook uses the step's. + const auto& requests = meta.requests ? *meta.requests : current_ctx_->requests; bool should_clone = cfg_.clone_slices && requests.size() > 1; for (size_t j = 0; j < requests.size(); ++j) { @@ -267,10 +269,15 @@ void P2PThread::do_post_processing(at::Tensor& tensor, const DrainTask& first_ta if (meta.hook_type == ring_py::HOOK_TYPE_FINAL_LOGITS) { // Contract: vLLM compute_logits returns exactly one logit // per request, flattened to [num_reqs, vocab]. The meta - // shape matches this (logits_to_keep=num_reqs passed from - // execute_model). We index by request position j (not - // the token-based dim0_offset) and extract one row. - slice = slice_flattened(tensor, /*dim0_offset=*/j, + // shape matches this (logits_to_keep counts every request + // of the step, even when the request list was filtered). + // Index by the request's position in the unfiltered step + // (orig_position; the list index j when unfiltered), never + // by the token-based dim0_offset, and extract one row. + const int64_t row = (req.orig_position >= 0) + ? static_cast(req.orig_position) + : static_cast(j); + slice = slice_flattened(tensor, /*dim0_offset=*/row, /*start=*/0, /*end=*/1); } else { slice = slice_flattened(tensor, req.dim0_offset, diff --git a/native/csrc/ring/producer.cu b/native/csrc/ring/producer.cu index ebe68ba72..a374050a2 100644 --- a/native/csrc/ring/producer.cu +++ b/native/csrc/ring/producer.cu @@ -33,6 +33,9 @@ #include "payload_ring.cuh" #include "task_ring.cuh" #include "ring_config.h" +#include "toggle_device.h" + +#include #include #include @@ -652,6 +655,640 @@ __global__ void record_producer_segmented_pack_kernel( ring, task_head, payload_head, alloc_bytes, actual_bytes); } + +// --------------------------------------------------------------------------- +// Row tables: shared by the gather, head and fused producers. +// --------------------------------------------------------------------------- + +// Load a gather-contract table into shared memory and turn its per-segment +// row counts into a destination prefix, clamped at upper_rows. Every thread +// of the block must call this; it ends with a barrier. +// +// Defensive clamps: a host-side bug must not let a kernel read or write past +// the captured tensor (upper_rows). The serial form of the clamp, +// n'_k = min(n_k, upper_rows - prefix_k), +// telescopes to prefix_{k+1} = min(sum_{j<=k} n_j, upper_rows), so the table +// is loaded by all threads in parallel and the prefix is one warp's +// __shfl_up_sync scan. A single serial thread-0 loop costs ~0.05 us per +// segment in every block (31 us for 150 one-row segments). +// +// The segment count is bounded by GATHER_MAX_SEGS and by the table's +// physical storage (table_numel words, row_table_segment_count), so a count +// word larger than the table never reads past it. +__device__ inline uint32_t load_row_table(const int64_t* table, + int64_t table_numel, + int64_t upper_rows, + int64_t* s_src, + int64_t* s_prefix) +{ + const uint32_t n_segs = row_table_segment_count(table[0], table_numel); + for (uint32_t k = threadIdx.x; k < n_segs; k += blockDim.x) { + int64_t rs = table[1 + 2 * k]; + int64_t nr = table[2 + 2 * k]; + if (rs < 0) rs = 0; + if (rs > upper_rows) rs = upper_rows; + if (nr < 0) nr = 0; + if (nr > upper_rows - rs) nr = upper_rows - rs; + s_src[k] = rs; + s_prefix[k + 1] = nr; // per-segment count; scanned below + } + if (threadIdx.x == 0) s_prefix[0] = 0; + __syncthreads(); + // A single segment's count is already <= upper_rows, so its scan is the + // identity. n_segs is block-uniform, so the barrier below is safe. + if (n_segs > 1) { + if (threadIdx.x < 32) { + // Warp 0: each lane owns a contiguous run of PER_LANE counts. + constexpr uint32_t PER_LANE = (GATHER_MAX_SEGS + 31) / 32; + const uint32_t lane = threadIdx.x; + const uint32_t begin = lane * PER_LANE; + int64_t local = 0; + for (uint32_t j = 0; j < PER_LANE; ++j) + if (begin + j < n_segs) local += s_prefix[begin + j + 1]; + int64_t incl = local; + for (uint32_t off = 1; off < 32; off <<= 1) { + const int64_t v = __shfl_up_sync(0xffffffffu, incl, off); + if (lane >= off) incl += v; + } + int64_t run = incl - local; // exclusive prefix of this lane's run + for (uint32_t j = 0; j < PER_LANE; ++j) { + if (begin + j < n_segs) { + run += s_prefix[begin + j + 1]; + s_prefix[begin + j + 1] = (run < upper_rows) ? run : upper_rows; + } + } + } + __syncthreads(); + } + return n_segs; +} + +// Source row of packed output row r: the largest k with s_prefix[k] <= r. +__device__ inline uint64_t table_src_row(const int64_t* s_src, + const int64_t* s_prefix, + uint32_t n_segs, uint64_t r) +{ + uint32_t lo = 0, hi = n_segs; + while (hi - lo > 1) { + const uint32_t mid = (lo + hi) >> 1; + if (static_cast(s_prefix[mid]) <= r) lo = mid; else hi = mid; + } + return static_cast(s_src[lo]) + + (r - static_cast(s_prefix[lo])); +} + +__device__ inline void store_f32_with_wrap(const RingState& ring, + const TwoSpan& spans, + uint64_t offset, float value) +{ + uint8_t* dst = (offset < spans.len1) + ? ring.payload_buf + spans.off1 + offset + : ring.payload_buf + spans.off2 + (offset - spans.len1); + *reinterpret_cast(dst) = value; +} + +// Block-wide last-block-arrives join for kernels whose last block has work +// left before publishing (a combine step, or several tasks). Every thread of +// a participating block calls it after its payload writes; it returns true in +// every thread of the last of `participants` blocks to arrive. The same +// release/acquire fence pairing as publish_last_block_arrives. +__device__ inline bool join_last_block(uint32_t participants) +{ + __shared__ bool s_last; + __syncthreads(); + if (threadIdx.x == 0) { + cuda::atomic_thread_fence(cuda::memory_order_release, + cuda::thread_scope_device); + const uint32_t finished = atomicAdd(&g_block_done_counter, 1); + s_last = finished == participants - 1; + if (s_last) { + cuda::atomic_thread_fence(cuda::memory_order_acquire, + cuda::thread_scope_device); + } + } + __syncthreads(); + return s_last; +} + +// Publish `count` consecutive tasks from thread 0 of the last block, after +// join_last_block() and a block barrier covering the last block's own writes. +__device__ inline void publish_tasks(RingState& ring, + uint64_t task_head, + uint64_t payload_head, + const uint64_t* totals, + const uint64_t* allocs, + uint32_t count) +{ + if (threadIdx.x != 0) return; + cuda::atomic_thread_fence(cuda::memory_order_release, + cuda::thread_scope_device); + uint64_t th = task_head; + uint64_t ph = payload_head; + uint64_t total = 0; + for (uint32_t i = 0; i < count; ++i) { + task_publish(ring.publication_slots, ring.task_cap, th, totals[i]); + ++th; + payload_advance_head(ph, allocs[i]); + total += totals[i]; + } + *ring.task_head = th; + *ring.payload_head = ph; + atomicAdd(reinterpret_cast(ring.actual_bytes_counter), + static_cast(total)); + atomicExch(&g_block_done_counter, 0); +} + +// --------------------------------------------------------------------------- +// producer_gather_kernel -- copy the table's row runs, packed. +// +// Output ROWS are distributed to BLOCKS grid-stride (a row's source is a +// binary search over the shared prefix) and the block's threads copy that +// row vectorized. Row-to-block (not segment-to-block) mapping keeps one +// large segment (a whole prefill chunk) spread across the grid, and keeps the +// grid independent of the runtime segment count -- required under CUDA-graph +// capture, where gridDim is baked but the table is re-read each replay. +// --------------------------------------------------------------------------- +__global__ void producer_gather_kernel( + RingState ring, + const uint8_t* src, + uint64_t nbytes_upper, + const int64_t* seg_table_dev_ptr, + int64_t table_numel, + uint64_t row_bytes, + uint32_t hook_type) +{ + if (g_ring_null_mode) return; + // -1 disables this producer for the step, including task publication. + // Unlike 0 (an active, empty observation) it has no reservation or + // metadata entry. The table is stable for the launch. + if (seg_table_dev_ptr[0] == -1) return; + + __shared__ int64_t s_src[GATHER_MAX_SEGS]; + __shared__ int64_t s_prefix[GATHER_MAX_SEGS + 1]; + + const int64_t upper_rows = + (row_bytes > 0) ? static_cast(nbytes_upper / row_bytes) : 0; + const uint32_t n_segs = load_row_table(seg_table_dev_ptr, table_numel, + upper_rows, s_src, s_prefix); + + const uint64_t out_rows = static_cast(s_prefix[n_segs]); + const uint64_t actual_total = out_rows * row_bytes; + const uint64_t alloc_bytes = align_up(actual_total, PAYLOAD_ALIGN); + const uint64_t task_head = *ring.task_head; + const uint64_t payload_head = *ring.payload_head; + const TwoSpan spans = payload_compute_spans(payload_head, + ring.payload_cap, + alloc_bytes); + + if (n_segs == 1) { + // One contiguous run (a full-row table or one merged window): a flat + // whole-grid copy, the same cost as producer_prefix. The row-mapped + // loop is paid only when the selection is actually fragmented. + const uint64_t gtid = uint64_t(blockIdx.x) * blockDim.x + threadIdx.x; + const uint64_t stride = uint64_t(gridDim.x) * blockDim.x; + copy_chunk_with_wrap(ring.payload_buf, spans, 0, actual_total, + src + static_cast(s_src[0]) * row_bytes, + gtid, stride); + } else { + for (uint64_t r = blockIdx.x; r < out_rows; r += gridDim.x) { + const uint64_t src_row = table_src_row(s_src, s_prefix, n_segs, r); + copy_chunk_with_wrap(ring.payload_buf, spans, r * row_bytes, + row_bytes, src + src_row * row_bytes, + threadIdx.x, blockDim.x); + } + } + + __threadfence(); + __syncthreads(); + publish_last_block_arrives(ring, task_head, payload_head, alloc_bytes, + actual_total); +} + +// --------------------------------------------------------------------------- +// producer_head_kernel -- in-graph linear heads (see producer.cuh). +// +// The table setup is the gather kernel's, so a head node and a gather node +// fed the same table select the same rows in the same packed order. One +// output row (or, for decode-sized selections, one (row, hidden slice)) per +// block iteration; the block's threads split the hidden dimension, each +// accumulating K partial dot products in fp32 registers, then a warp-shuffle +// and shared-memory reduction yields the K scores. +// --------------------------------------------------------------------------- +__device__ inline void head_load8(const __nv_bfloat16* p, float* x) { + const uint4 v = *reinterpret_cast(p); + const __nv_bfloat162* h2 = reinterpret_cast(&v); +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float2 f = __bfloat1622float2(h2[j]); + x[2 * j] = f.x; + x[2 * j + 1] = f.y; + } +} +__device__ inline void head_load8(const float* p, float* x) { + const float4 a = *reinterpret_cast(p); + const float4 b = *reinterpret_cast(p + 4); + x[0] = a.x; x[1] = a.y; x[2] = a.z; x[3] = a.w; + x[4] = b.x; x[5] = b.y; x[6] = b.z; x[7] = b.w; +} +__device__ inline float head_elem(const __nv_bfloat16* p) { return __bfloat162float(*p); } +__device__ inline float head_elem(const float* p) { return *p; } + +template +__global__ void producer_head_kernel( + RingState ring, + const T* src, + uint64_t upper_rows_u, + const int64_t* seg_table_dev_ptr, + int64_t table_numel, + uint32_t hidden, + const float* arena, + uint32_t K, + uint32_t hook_type) +{ + if (g_ring_null_mode) return; + if (seg_table_dev_ptr[0] == -1) return; // disabled: no reservation/task + + __shared__ int64_t s_src[GATHER_MAX_SEGS]; + __shared__ int64_t s_prefix[GATHER_MAX_SEGS + 1]; + __shared__ float s_part[PRODUCER_BLOCK_DIM / 32][HEAD_MAX_SLOTS]; + + const int64_t upper_rows = static_cast(upper_rows_u); + if (K > HEAD_MAX_SLOTS) K = HEAD_MAX_SLOTS; // the launcher refuses this + const uint32_t n_segs = load_row_table(seg_table_dev_ptr, table_numel, + upper_rows, s_src, s_prefix); + + const uint64_t out_rows = static_cast(s_prefix[n_segs]); + const uint64_t row_out = static_cast(K) * sizeof(float); + const uint64_t actual_total = out_rows * row_out; + const uint64_t alloc_bytes = align_up(actual_total, PAYLOAD_ALIGN); + const uint64_t task_head = *ring.task_head; + const uint64_t payload_head = *ring.payload_head; + const TwoSpan spans = payload_compute_spans(payload_head, + ring.payload_cap, + alloc_bytes); + const float* bias = arena + static_cast(K) * hidden; + float* scratch = const_cast(bias + K); + // 8-wide loads need 16-byte aligned bases; with hidden % 8 == 0 every row + // and slot offset is then a multiple of 16 bytes as well. A tensor view + // at an odd storage offset takes the scalar path. + const bool vec8 = (hidden % 8u) == 0u + && (reinterpret_cast(src) & 15u) == 0u + && (reinterpret_cast(arena) & 15u) == 0u; + const uint32_t warp = threadIdx.x >> 5, lane = threadIdx.x & 31u; + const uint32_t n_warps = blockDim.x >> 5; + // Decode-sized selections split each row over HEAD_SPLITS hidden slices; + // the partials go to the arena's scratch and the last block sums them in + // a fixed order (deterministic) and adds the bias. + const uint32_t S = (out_rows <= HEAD_SPLIT_ROWS) ? HEAD_SPLITS : 1u; + uint32_t chunk = (hidden + S - 1) / S; + if (vec8) chunk = (chunk + 7u) & ~7u; + const uint64_t items = out_rows * S; + + for (uint64_t it = blockIdx.x; it < items; it += gridDim.x) { + const uint64_t r = it / S; + const uint32_t sl = static_cast(it - r * S); + const uint32_t i_lo = sl * chunk; + const uint32_t i_hi = (i_lo + chunk < hidden) ? i_lo + chunk : hidden; + const T* h = src + table_src_row(s_src, s_prefix, n_segs, r) * hidden; + + float acc[HEAD_MAX_SLOTS]; +#pragma unroll + for (uint32_t k = 0; k < HEAD_MAX_SLOTS; ++k) acc[k] = 0.f; + if (vec8) { + for (uint32_t i = i_lo + threadIdx.x * 8u; i < i_hi; i += blockDim.x * 8u) { + float x[8]; + head_load8(h + i, x); +#pragma unroll + for (uint32_t k = 0; k < HEAD_MAX_SLOTS; ++k) { + if (k < K) { + const float* w = arena + static_cast(k) * hidden + i; + const float4 w0 = *reinterpret_cast(w); + const float4 w1 = *reinterpret_cast(w + 4); + float a = acc[k]; + a = fmaf(w0.x, x[0], a); a = fmaf(w0.y, x[1], a); + a = fmaf(w0.z, x[2], a); a = fmaf(w0.w, x[3], a); + a = fmaf(w1.x, x[4], a); a = fmaf(w1.y, x[5], a); + a = fmaf(w1.z, x[6], a); a = fmaf(w1.w, x[7], a); + acc[k] = a; + } + } + } + } else { + for (uint32_t i = i_lo + threadIdx.x; i < i_hi; i += blockDim.x) { + const float xi = head_elem(h + i); +#pragma unroll + for (uint32_t k = 0; k < HEAD_MAX_SLOTS; ++k) + if (k < K) + acc[k] = fmaf(arena[static_cast(k) * hidden + i], xi, acc[k]); + } + } +#pragma unroll + for (uint32_t k = 0; k < HEAD_MAX_SLOTS; ++k) { + if (k < K) { + float v = acc[k]; + for (int off = 16; off > 0; off >>= 1) + v += __shfl_down_sync(0xffffffffu, v, off); + if (lane == 0) s_part[warp][k] = v; + } + } + __syncthreads(); + if (threadIdx.x < K) { + float v = 0.f; + for (uint32_t w = 0; w < n_warps; ++w) v += s_part[w][threadIdx.x]; + if (S == 1) { + store_f32_with_wrap(ring, spans, + r * row_out + threadIdx.x * sizeof(float), + v + bias[threadIdx.x]); + } else { + scratch[(r * S + sl) * K + threadIdx.x] = v; + } + } + __syncthreads(); // s_part is reused by the next item + } + + if (!join_last_block(gridDim.x)) return; + if (S > 1) { + // Fixed-order split-K combine. __ldcg bypasses L1, which is not + // coherent with the other blocks' scratch writes. + for (uint64_t j = threadIdx.x; j < out_rows * K; j += blockDim.x) { + const uint64_t r = j / K; + const uint32_t k = static_cast(j - r * K); + float v = 0.f; + for (uint32_t sl = 0; sl < S; ++sl) + v += __ldcg(&scratch[(r * S + sl) * K + k]); + store_f32_with_wrap(ring, spans, j * sizeof(float), v + bias[k]); + } + __syncthreads(); + } + publish_tasks(ring, task_head, payload_head, &actual_total, &alloc_bytes, 1); +} + +// --------------------------------------------------------------------------- +// producer_fused_kernel -- a HookPoint's row producer AND its heads in ONE +// captured node. +// +// Two independent parts, each driven by its own gather-contract table +// (-1 = part disabled: no reservation, no task, no payload): +// row part: copy the selected rows (exactly producer_gather's bytes); +// head part: K fp32 columns per selected row; only the node's active +// slots are computed, free slots' columns are NaN. +// Publish order: the row task (if on), then the head task (if on), payloads +// back to back, each aligned. +// +// Head math: an item is a tile of HEAD_RT consecutive output rows x a group +// of HEAD_ST active slots. The block's warps split the hidden dimension into +// fixed chunks; each lane accumulates fp32 FMAs in increasing order, each +// warp reduces its lanes with a shfl_down tree, and one thread sums the warp +// partials in warp order. A (row, slot) result therefore depends only on +// that row's values and that slot's weights -- bit-identical across batch +// compositions and grids (fixed reduction order, no atomics, no split-K). +// +// Participating blocks: P = max(blocks the row part needs, blocks the head +// part needs) <= gridDim.x, computed identically by every block from the +// stable tables; the others return before reading the ring heads and never +// join the arrival count. +// --------------------------------------------------------------------------- +template +__global__ void producer_fused_kernel( + RingState ring, + const T* src, + uint64_t upper_rows_u, + uint32_t hidden, + const int64_t* gtab, + int64_t gtab_numel, + const int64_t* htab, + int64_t htab_numel, + const int32_t* node, + uint32_t K, + float* acc, + uint64_t acc_rows, + uint32_t hook_type) +{ + if (g_ring_null_mode) return; + const bool g_on = gtab != nullptr && gtab[0] != -1; + const bool h_on = htab != nullptr && htab[0] != -1; + if (!g_on && !h_on) return; + + __shared__ int64_t g_src[GATHER_MAX_SEGS]; + __shared__ int64_t g_pre[GATHER_MAX_SEGS + 1]; + __shared__ int64_t h_src[GATHER_MAX_SEGS]; + __shared__ int64_t h_pre[GATHER_MAX_SEGS + 1]; + __shared__ int32_t s_list[HEAD_MAX_SLOTS]; // compact active slot ids + __shared__ int32_t s_mode[HEAD_MAX_SLOTS]; // by slot id + __shared__ int32_t s_acc[HEAD_MAX_SLOTS]; // by slot id + __shared__ float s_bias[HEAD_MAX_SLOTS]; // by slot id + __shared__ uint32_t s_nact; + + const int64_t upper_rows = static_cast(upper_rows_u); + const uint64_t row_bytes = static_cast(hidden) * sizeof(T); + if (K > HEAD_MAX_SLOTS) K = HEAD_MAX_SLOTS; // the launcher refuses this + + // Head header: the last warp loads modes, accumulator columns and biases + // while the tables load. + const uint32_t hdr_t0 = blockDim.x - 32; + if (h_on && threadIdx.x >= hdr_t0 && threadIdx.x < hdr_t0 + HEAD_MAX_SLOTS) { + const HeadNodeHeader* hdr = reinterpret_cast(node); + const float* bias = reinterpret_cast(node + HEAD_BIAS_WORD); + const uint32_t k = threadIdx.x - hdr_t0; + int32_t m = (k < K) ? hdr->mode[k] : static_cast(HEAD_FREE); + const int32_t ac = hdr->acc[k]; + if (m < HEAD_FREE || m > HEAD_LAST) m = HEAD_FREE; + if ((m == HEAD_FIRST || m == HEAD_MID || m == HEAD_LAST) + && (ac < 0 || ac >= static_cast(HEAD_MAX_ACC))) + m = HEAD_FREE; + s_mode[k] = m; + s_acc[k] = ac; + s_bias[k] = bias[k]; + } + const uint32_t gn = g_on ? load_row_table(gtab, gtab_numel, upper_rows, g_src, g_pre) : 0; + const uint32_t hn = h_on ? load_row_table(htab, htab_numel, upper_rows, h_src, h_pre) : 0; + if (h_on && threadIdx.x == 0) { + // Compact list of computed slots in ascending slot id, derived from + // the modes so no column is left unwritten. + uint32_t c = 0; + for (uint32_t k = 0; k < K; ++k) + if (s_mode[k] != HEAD_FREE) s_list[c++] = static_cast(k); + s_nact = c; + } + __syncthreads(); + + const uint64_t g_rows = g_on ? static_cast(g_pre[gn]) : 0; + const uint64_t h_rows = h_on ? static_cast(h_pre[hn]) : 0; + const uint64_t g_total = g_rows * row_bytes; + const uint64_t h_total = h_rows * K * sizeof(float); + const uint64_t g_alloc = align_up(g_total, PAYLOAD_ALIGN); + const uint64_t h_alloc = align_up(h_total, PAYLOAD_ALIGN); + const uint32_t nact = h_on ? s_nact : 0; + + // Participating blocks (block-uniform). + const uint32_t NW = blockDim.x >> 5; + uint64_t gb = 0; + if (g_on) + gb = (gn == 1) + ? (g_total + PRODUCER_BYTES_PER_BLOCK - 1) / PRODUCER_BYTES_PER_BLOCK + : g_rows; + const uint64_t n_sg = (nact + HEAD_ST - 1) / HEAD_ST; + const uint64_t h_items = ((h_rows + HEAD_RT - 1) / HEAD_RT) * n_sg; + uint64_t hb = h_items; // one block per head item + const uint64_t fill = h_rows * K; // free-column NaN fill + const uint64_t fb = (fill + blockDim.x * 64 - 1) / (blockDim.x * 64); + if (fb > hb) hb = fb; + uint64_t P64 = gb > hb ? gb : hb; + if (P64 < 1) P64 = 1; + if (P64 > gridDim.x) P64 = gridDim.x; + const uint32_t P = static_cast(P64); + if (blockIdx.x >= P) return; + + const uint64_t task_head = *ring.task_head; + const uint64_t payload_head = *ring.payload_head; + const TwoSpan gsp = payload_compute_spans(payload_head, ring.payload_cap, g_alloc); + const TwoSpan hsp = payload_compute_spans(payload_head + (g_on ? g_alloc : 0), + ring.payload_cap, h_alloc); + + // ---- row part (byte-identical to producer_gather) ---- + if (g_rows) { + const uint8_t* src8 = reinterpret_cast(src); + if (gn == 1) { + const uint64_t gtid = uint64_t(blockIdx.x) * blockDim.x + threadIdx.x; + const uint64_t stride = uint64_t(P) * blockDim.x; + copy_chunk_with_wrap(ring.payload_buf, gsp, 0, g_total, + src8 + static_cast(g_src[0]) * row_bytes, + gtid, stride); + } else { + for (uint64_t r = blockIdx.x; r < g_rows; r += P) { + const uint64_t sr = table_src_row(g_src, g_pre, gn, r); + copy_chunk_with_wrap(ring.payload_buf, gsp, r * row_bytes, row_bytes, + src8 + sr * row_bytes, threadIdx.x, blockDim.x); + } + } + } + + // ---- head part ---- + if (h_rows) { + const uint32_t warp = threadIdx.x >> 5, lane = threadIdx.x & 31u; + const W* wbase = reinterpret_cast(node + HEAD_HDR_WORDS); + // As in producer_head_kernel: 8-wide loads only from 16-byte aligned + // bases (row and slot offsets are multiples of 16 bytes then). + const bool vec8 = (hidden % 8u) == 0u + && (reinterpret_cast(src) & 15u) == 0u + && (reinterpret_cast(wbase) & 15u) == 0u; + uint32_t chunk = (hidden + NW - 1) / NW; + if (vec8) chunk = (chunk + 7u) & ~7u; + const uint32_t w_lo = warp * chunk; + const uint32_t w_hi = (w_lo + chunk < hidden) ? w_lo + chunk : hidden; + __shared__ uint64_t s_hrow[HEAD_RT]; + __shared__ float s_red[PRODUCER_BLOCK_DIM / 32][HEAD_RT * HEAD_ST]; + for (uint64_t it = blockIdx.x; it < h_items; it += P) { + const uint64_t rt = it / n_sg; + const uint32_t sg = static_cast(it - rt * n_sg); + const uint64_t r0 = rt * HEAD_RT; + const uint32_t nr = static_cast( + (h_rows - r0) < HEAD_RT ? (h_rows - r0) : HEAD_RT); + const uint32_t s0 = sg * HEAD_ST; + const uint32_t ns = (nact - s0) < HEAD_ST ? (nact - s0) : HEAD_ST; + if (threadIdx.x < HEAD_RT) + s_hrow[threadIdx.x] = (threadIdx.x < nr) + ? table_src_row(h_src, h_pre, hn, r0 + threadIdx.x) : 0; + __syncthreads(); + const T* hp[HEAD_RT]; + const W* wp[HEAD_ST]; +#pragma unroll + for (uint32_t rr = 0; rr < HEAD_RT; ++rr) + hp[rr] = src + s_hrow[rr < nr ? rr : 0] * hidden; +#pragma unroll + for (uint32_t ss = 0; ss < HEAD_ST; ++ss) + wp[ss] = wbase + static_cast(s_list[s0 + (ss < ns ? ss : 0)]) * hidden; + float a[HEAD_RT][HEAD_ST]; +#pragma unroll + for (uint32_t rr = 0; rr < HEAD_RT; ++rr) +#pragma unroll + for (uint32_t ss = 0; ss < HEAD_ST; ++ss) a[rr][ss] = 0.f; + if (vec8) { +#pragma unroll 2 + for (uint32_t i = w_lo + lane * 8u; i < w_hi; i += 256u) { + float x[HEAD_RT][8], w[HEAD_ST][8]; +#pragma unroll + for (uint32_t rr = 0; rr < HEAD_RT; ++rr) + if (rr < nr) head_load8(hp[rr] + i, x[rr]); +#pragma unroll + for (uint32_t ss = 0; ss < HEAD_ST; ++ss) + if (ss < ns) head_load8(wp[ss] + i, w[ss]); +#pragma unroll + for (uint32_t rr = 0; rr < HEAD_RT; ++rr) +#pragma unroll + for (uint32_t ss = 0; ss < HEAD_ST; ++ss) + if (rr < nr && ss < ns) { + float v = a[rr][ss]; +#pragma unroll + for (int e = 0; e < 8; ++e) v = fmaf(w[ss][e], x[rr][e], v); + a[rr][ss] = v; + } + } + } else { + for (uint32_t i = w_lo + lane; i < w_hi; i += 32u) { + float x[HEAD_RT], w[HEAD_ST]; +#pragma unroll + for (uint32_t rr = 0; rr < HEAD_RT; ++rr) + x[rr] = (rr < nr) ? head_elem(hp[rr] + i) : 0.f; +#pragma unroll + for (uint32_t ss = 0; ss < HEAD_ST; ++ss) + w[ss] = (ss < ns) ? head_elem(wp[ss] + i) : 0.f; +#pragma unroll + for (uint32_t rr = 0; rr < HEAD_RT; ++rr) +#pragma unroll + for (uint32_t ss = 0; ss < HEAD_ST; ++ss) + if (rr < nr && ss < ns) a[rr][ss] = fmaf(w[ss], x[rr], a[rr][ss]); + } + } +#pragma unroll + for (uint32_t rr = 0; rr < HEAD_RT; ++rr) +#pragma unroll + for (uint32_t ss = 0; ss < HEAD_ST; ++ss) { + float v = a[rr][ss]; +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + v += __shfl_down_sync(0xffffffffu, v, off); + if (lane == 0) s_red[warp][rr * HEAD_ST + ss] = v; + } + __syncthreads(); + const uint32_t pidx = threadIdx.x, rr = pidx / HEAD_ST, ss = pidx % HEAD_ST; + if (pidx < HEAD_RT * HEAD_ST && rr < nr && ss < ns) { + float v = 0.f; + for (uint32_t w = 0; w < NW; ++w) v += s_red[w][pidx]; + const uint32_t k = static_cast(s_list[s0 + ss]); + const int32_t m = s_mode[k]; + const uint64_t sr = s_hrow[rr]; + float outv = __int_as_float(0x7fc00000); // NaN: no score at this node + if (m == HEAD_SINGLE) { + outv = v + s_bias[k]; + } else if (sr < acc_rows) { + float* cell = acc + sr * HEAD_MAX_ACC + static_cast(s_acc[k]); + if (m == HEAD_FIRST) *cell = v; + else if (m == HEAD_MID) *cell = *cell + v; + else outv = (*cell + v) + s_bias[k]; // LAST + } + store_f32_with_wrap(ring, hsp, ((r0 + rr) * K + k) * sizeof(float), outv); + } + __syncthreads(); // s_hrow / s_red are reused by the next item + } + // Free slots' columns: NaN (never delivered; the host routes placed slots). + for (uint64_t j = uint64_t(blockIdx.x) * blockDim.x + threadIdx.x; j < fill; + j += uint64_t(P) * blockDim.x) { + const uint32_t k = static_cast(j % K); + if (s_mode[k] == HEAD_FREE) + store_f32_with_wrap(ring, hsp, j * sizeof(float), __int_as_float(0x7fc00000)); + } + } + + if (!join_last_block(P)) return; + uint64_t totals[2]; + uint64_t allocs[2]; + uint32_t count = 0; + if (g_on) { totals[count] = g_total; allocs[count] = g_alloc; ++count; } + if (h_on) { totals[count] = h_total; allocs[count] = h_alloc; ++count; } + publish_tasks(ring, task_head, payload_head, totals, allocs, count); +} + // --------------------------------------------------------------------------- // Host launchers // --------------------------------------------------------------------------- @@ -697,6 +1334,172 @@ void launch_producer_chunked( ring, d_src, nbytes_upper, chunk_bytes_dev_ptr, K, hook_type); } +void launch_producer_gather( + const RingState& ring, + const uint8_t* d_src, + uint64_t nbytes_upper, + const int64_t* seg_table_dev_ptr, + int64_t table_numel, + uint64_t row_bytes, + uint32_t hook_type, + cudaStream_t stream) +{ + if (seg_table_dev_ptr == nullptr || table_numel < 1) { + throw std::invalid_argument( + "gather producer requires a row table holding its segment count"); + } + // Grid sized by the upper bound (a capture-time constant); the kernel + // strides fewer rows when the table selects fewer. + const uint32_t n_blocks = producer_grid_blocks(nbytes_upper); + producer_gather_kernel<<>>( + ring, d_src, nbytes_upper, seg_table_dev_ptr, table_numel, row_bytes, + hook_type); +} + +cudaError_t launch_producer_gather_device_updatable( + const RingState& ring, + const uint8_t* d_src, + uint64_t nbytes_upper, + const int64_t* seg_table_dev_ptr, + int64_t table_numel, + uint64_t row_bytes, + uint32_t hook_type, + cudaStream_t stream) +{ +#if !DMX_HAS_DEVICE_GRAPH_UPDATE + (void)ring; (void)d_src; (void)nbytes_upper; (void)seg_table_dev_ptr; + (void)table_numel; (void)row_bytes; (void)hook_type; (void)stream; + return cudaErrorNotSupported; +#else + if (seg_table_dev_ptr == nullptr || table_numel < 1) + return cudaErrorInvalidValue; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(producer_grid_blocks(nbytes_upper)); + cfg.blockDim = dim3(PRODUCER_BLOCK_DIM); + cfg.dynamicSmemBytes = 0; + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeDeviceUpdatableKernelNode; + attr[0].val.deviceUpdatableKernelNode.deviceUpdatable = 1; + attr[0].val.deviceUpdatableKernelNode.devNode = nullptr; + cfg.attrs = attr; + cfg.numAttrs = 1; + return cudaLaunchKernelEx(&cfg, producer_gather_kernel, + ring, d_src, nbytes_upper, seg_table_dev_ptr, + table_numel, row_bytes, hook_type); +#endif +} + +void launch_producer_head( + const RingState& ring, + const uint8_t* d_src, + uint64_t nbytes_upper, + const int64_t* seg_table_dev_ptr, + int64_t table_numel, + uint32_t hidden, + uint32_t src_elem_bytes, + const float* arena, + uint32_t K, + uint32_t hook_type, + cudaStream_t stream) +{ + if (K == 0 || K > HEAD_MAX_SLOTS || hidden == 0 || + (src_elem_bytes != 2 && src_elem_bytes != 4) || + seg_table_dev_ptr == nullptr || table_numel < 1) { + throw std::invalid_argument( + "head producer requires 1 <= K <= HEAD_MAX_SLOTS, a positive " + "hidden size, a bf16 or fp32 source and a row table holding its " + "segment count"); + } + const uint64_t row_bytes = static_cast(hidden) * src_elem_bytes; + const uint64_t upper_rows = nbytes_upper / row_bytes; + // One block per possible row (the capture-time bound) or, for + // decode-sized bounds, per (row, hidden slice); capped at one wave. + uint64_t blocks = upper_rows; + const uint64_t split_rows = upper_rows < HEAD_SPLIT_ROWS ? upper_rows : HEAD_SPLIT_ROWS; + if (blocks < split_rows * HEAD_SPLITS) blocks = split_rows * HEAD_SPLITS; + if (blocks < 1) blocks = 1; + if (blocks > PRODUCER_MAX_BLOCKS) blocks = PRODUCER_MAX_BLOCKS; + const uint32_t grid = static_cast(blocks); + if (src_elem_bytes == 2) { + producer_head_kernel<__nv_bfloat16><<>>( + ring, reinterpret_cast(d_src), upper_rows, + seg_table_dev_ptr, table_numel, hidden, arena, K, hook_type); + } else { + producer_head_kernel<<>>( + ring, reinterpret_cast(d_src), upper_rows, + seg_table_dev_ptr, table_numel, hidden, arena, K, hook_type); + } +} + +uint32_t producer_fused_grid(uint64_t nbytes_upper, uint64_t upper_rows) +{ + uint64_t blocks = producer_grid_blocks(nbytes_upper); + // Head capacity at the capture-time bound: one block per (row tile, slot + // group); at most HEAD_MAX_SLOTS / HEAD_ST groups. + const uint64_t hb = + ((upper_rows + HEAD_RT - 1) / HEAD_RT) * (HEAD_MAX_SLOTS / HEAD_ST); + if (hb > blocks) blocks = hb; + if (blocks < 1) blocks = 1; + if (blocks > PRODUCER_MAX_BLOCKS) blocks = PRODUCER_MAX_BLOCKS; + return static_cast(blocks); +} + +void launch_producer_fused( + const RingState& ring, + const uint8_t* d_src, + uint64_t nbytes_upper, + uint32_t hidden, + uint32_t src_elem_bytes, + const int64_t* gather_table, + int64_t gather_table_numel, + const int64_t* head_table, + int64_t head_table_numel, + const int32_t* node_arena, + uint32_t K, + uint32_t weight_bf16, + float* acc, + uint64_t acc_rows, + uint32_t hook_type, + cudaStream_t stream) +{ + if (hidden == 0 || (src_elem_bytes != 2 && src_elem_bytes != 4) || + (head_table != nullptr && (K == 0 || K > HEAD_MAX_SLOTS)) || + (gather_table != nullptr && gather_table_numel < 1) || + (head_table != nullptr && head_table_numel < 1)) { + throw std::invalid_argument( + "fused producer requires a positive hidden size, a bf16 or fp32 " + "source, row tables holding their segment counts and, with " + "heads, 1 <= K <= HEAD_MAX_SLOTS"); + } + const uint64_t row_bytes = static_cast(hidden) * src_elem_bytes; + const uint64_t upper_rows = nbytes_upper / row_bytes; + const uint32_t blocks = producer_fused_grid(nbytes_upper, upper_rows); + if (src_elem_bytes == 2) { + if (weight_bf16) + producer_fused_kernel<__nv_bfloat16, __nv_bfloat16><<>>( + ring, reinterpret_cast(d_src), upper_rows, hidden, + gather_table, gather_table_numel, head_table, head_table_numel, + node_arena, K, acc, acc_rows, hook_type); + else + producer_fused_kernel<__nv_bfloat16, float><<>>( + ring, reinterpret_cast(d_src), upper_rows, hidden, + gather_table, gather_table_numel, head_table, head_table_numel, + node_arena, K, acc, acc_rows, hook_type); + } else { + if (weight_bf16) + producer_fused_kernel<<>>( + ring, reinterpret_cast(d_src), upper_rows, hidden, + gather_table, gather_table_numel, head_table, head_table_numel, + node_arena, K, acc, acc_rows, hook_type); + else + producer_fused_kernel<<>>( + ring, reinterpret_cast(d_src), upper_rows, hidden, + gather_table, gather_table_numel, head_table, head_table_numel, + node_arena, K, acc, acc_rows, hook_type); + } +} + void launch_record_producer_static( const RingState& ring, const uint8_t* d_src, diff --git a/native/csrc/ring/producer.cuh b/native/csrc/ring/producer.cuh index 1a518a8bb..0330f0c51 100644 --- a/native/csrc/ring/producer.cuh +++ b/native/csrc/ring/producer.cuh @@ -26,6 +26,8 @@ #include +#include "row_table.h" + namespace ring { // Grid geometry for the D2D producers: one block per @@ -88,6 +90,59 @@ RING_PRODUCER_HD inline uint32_t producer_chunked_grid_blocks( // chunked kernel and host-side argument validation. static constexpr uint32_t PRODUCER_MAX_K = 64; +// GATHER_MAX_SEGS and the row-table segment-count clamp live in +// row_table.h (CUDA-free, shared with host code and CPU tests). + +// In-graph linear projection heads (producer_head / producer_fused). +// +// A head node reads the same per-step row table as a gather producer and, +// instead of copying each selected row, writes K fp32 affine scores per +// selected row: out[r, k] = sum_i w[k, i] * h[src_row(r), i] + b[k], with +// fp32 accumulation. Weights and biases are data in a stable device arena: +// registering or retiring a slot writes the arena between steps (ordered on +// the model stream) and no graph is recaptured. +static constexpr uint32_t HEAD_MAX_SLOTS = 16; + +// producer_head arena (fp32): K x hidden weights, K biases, then a scratch of +// HEAD_SPLIT_ROWS x HEAD_SPLITS x K partials. Selections of at most +// HEAD_SPLIT_ROWS rows split each row's dot products over HEAD_SPLITS hidden +// slices so a one-row decode step is not bound by one SM reading K x hidden +// weights; the last block combines the partials in a fixed order. A free +// slot holds w = 0, b = NaN, so its column is NaN. +static constexpr uint32_t HEAD_SPLIT_ROWS = 16; +static constexpr uint32_t HEAD_SPLITS = 8; +inline uint64_t head_arena_floats(uint32_t K, uint32_t hidden) { + return static_cast(K) * hidden + K + + static_cast(HEAD_SPLIT_ROWS) * HEAD_SPLITS * K; +} + +// producer_fused node arena (int32 words, stable address, written only +// between steps): +// words [0, 16) mode[k] HEAD_FREE / SINGLE / FIRST / MID / LAST +// words [16, 32) acc[k] accumulator column of a multi-layer slot +// words [32, 48) bias[k] fp32 +// words [64, ...) weights K x hidden, bf16 or fp32 (a launch constant) +// Only non-free slots are computed. A linear head over several layers holds +// one slot at every participating node: FIRST writes acc[row][a] = w.h, MID +// adds, LAST emits acc + w.h + b in fixed layer order; FIRST, MID and FREE +// columns are NaN. acc is an fp32 [acc_rows, HEAD_MAX_ACC] buffer indexed by +// the row's position in the forward's token batch (the same at every layer). +enum HeadSlotMode : int32_t { + HEAD_FREE = 0, HEAD_SINGLE = 1, HEAD_FIRST = 2, HEAD_MID = 3, HEAD_LAST = 4 +}; +struct HeadNodeHeader { int32_t mode[16]; int32_t acc[16]; }; +static constexpr uint32_t HEAD_BIAS_WORD = 32; +static constexpr uint32_t HEAD_HDR_WORDS = 64; +static constexpr uint32_t HEAD_MAX_ACC = 16; +static constexpr uint32_t HEAD_RT = 4; // rows per head work item +static constexpr uint32_t HEAD_ST = 4; // slots per head work item +inline uint64_t head_node_words(uint32_t K, uint32_t hidden, bool weight_bf16) { + const uint64_t w = weight_bf16 + ? (static_cast(K) * hidden + 1) / 2 + : static_cast(K) * hidden; + return (HEAD_HDR_WORDS + w + 3) & ~uint64_t(3); +} + } // namespace ring #ifdef __CUDACC__ @@ -155,6 +210,94 @@ void launch_producer_chunked( uint32_t hook_type, cudaStream_t stream = 0); +// Gather: copy the row segments listed in a per-step row table, packed. +// Table layout (int64, stable device address, values re-read each replay -- +// the same contract as row_count_dev_ptr): +// table[0] = n_segs (clamped to GATHER_MAX_SEGS and to the table's +// storage, row_table_segment_count); -1 disables the +// producer for the step: no reservation, no task, no +// payload. 0 publishes an empty task. +// table[1 + 2k] = source row of segment k +// table[2 + 2k] = rows in segment k +// table_numel is the table's int64 element count (>= 1), passed at launch +// so a count larger than the storage never reads past it. +// Destination offsets are the running prefix sum: the output is always +// packed, so the table holds (src, len) pairs, not (src, dst, len). Each +// HookPoint passes its own table and row_bytes. +__global__ void producer_gather_kernel( + RingState ring, + const uint8_t* src, + uint64_t nbytes_upper, + const int64_t* seg_table_dev_ptr, + int64_t table_numel, + uint64_t row_bytes, + uint32_t hook_type); + +void launch_producer_gather( + const RingState& ring, + const uint8_t* d_src, + uint64_t nbytes_upper, + const int64_t* seg_table_dev_ptr, + int64_t table_numel, + uint64_t row_bytes, + uint32_t hook_type, + cudaStream_t stream = 0); + +// The same kernel and geometry, launched as a device-updatable kernel node +// (cudaLaunchAttributeDeviceUpdatableKernelNode) for device-side toggling +// (toggle_device.h). Valid only under stream capture. Returns the launch +// error; cudaErrorNotSupported when built without device graph updates. +cudaError_t launch_producer_gather_device_updatable( + const RingState& ring, + const uint8_t* d_src, + uint64_t nbytes_upper, + const int64_t* seg_table_dev_ptr, + int64_t table_numel, + uint64_t row_bytes, + uint32_t hook_type, + cudaStream_t stream); + +// In-graph linear heads at one HookPoint (see HEAD_* above): K fp32 scores +// per row selected by the gather-contract table. Payload bytes are +// selected_rows x K x 4. src_elem_bytes: 2 = bf16, 4 = fp32. +void launch_producer_head( + const RingState& ring, + const uint8_t* d_src, + uint64_t nbytes_upper, + const int64_t* seg_table_dev_ptr, + int64_t table_numel, + uint32_t hidden, + uint32_t src_elem_bytes, + const float* arena, + uint32_t K, + uint32_t hook_type, + cudaStream_t stream = 0); + +// The HookPoint's row producer AND its heads in ONE node. Each part has its +// own table (gather contract; -1 = that part publishes nothing); either table +// pointer may be null. Publish order: the row task, then the head task +// (the FIFO metadata order of the two specs). Head results are a fixed +// reduction order per (row, slot): bit-identical across batch compositions. +uint32_t producer_fused_grid(uint64_t nbytes_upper, uint64_t upper_rows); + +void launch_producer_fused( + const RingState& ring, + const uint8_t* d_src, + uint64_t nbytes_upper, + uint32_t hidden, + uint32_t src_elem_bytes, + const int64_t* gather_table, + int64_t gather_table_numel, + const int64_t* head_table, + int64_t head_table_numel, + const int32_t* node_arena, + uint32_t K, + uint32_t weight_bf16, + float* acc, + uint64_t acc_rows, + uint32_t hook_type, + cudaStream_t stream = 0); + // --------------------------------------------------------------------------- // Additive generic-record producers. // diff --git a/native/csrc/ring/ring_engine_py.cu b/native/csrc/ring/ring_engine_py.cu index a60b8fa04..b9c29f139 100644 --- a/native/csrc/ring/ring_engine_py.cu +++ b/native/csrc/ring/ring_engine_py.cu @@ -12,6 +12,7 @@ #include "ring/ring_debug.h" #include "ring/record_stall_budget.h" #include "ring/toggle_registry.h" +#include "ring/toggle_device.h" #include // at::cuda::getCurrentCUDAStream #include #include @@ -177,6 +178,27 @@ struct RingEnginePy::Impl { // in the registry; the engine methods delegate. ToggleRegistry toggle; + // Host mirror of the per-HookPoint row tables (set_gather_host_mirror). + // gather_dev_base is the arming word: written last, cleared first. + std::atomic gather_dev_base{0}; + uint64_t gather_host_base{0}; + uint64_t gather_row_bytes{0}; + uint64_t gather_rows{0}; + std::atomic gather_host_skips{0}; + std::atomic head_launches{0}; + + // True when the mirrored row at `table` holds the disabled sentinel. + // A pointer outside the mirror is never treated as disabled. + bool mirrored_row_disabled(uint64_t table) const { + const uint64_t base = gather_dev_base.load(std::memory_order_acquire); + if (base == 0 || table < base) return false; + const uint64_t off = table - base; + return off % gather_row_bytes == 0 + && off / gather_row_bytes < gather_rows + && *reinterpret_cast( + gather_host_base + off) == -1; + } + // Step-gated drain: ring records reserved since the last allow_flush_step() // (the current step's gap). Written by the reserving thread only. uint64_t gap_records{0}; @@ -982,6 +1004,185 @@ at::Tensor RingEnginePy::payload_tensor() const { // as much room as we computed. // --------------------------------------------------------------------------- +namespace { + +bool stream_is_capturing(cudaStream_t stream) { + cudaStreamCaptureStatus status = cudaStreamCaptureStatusActive; + return !(cudaStreamIsCapturing(stream, &status) == cudaSuccess && + status == cudaStreamCaptureStatusNone); +} + +} // namespace + +void RingEnginePy::hook_no_notify_gather(uint64_t d_ptr, uint64_t nbytes_upper, + uint64_t seg_table_dev_ptr, + int64_t seg_table_numel, + uint64_t row_bytes, + uint32_t hook_type, + uint64_t stream_handle) +{ + refuse_on_record_ring(impl_->record_mode, "legacy gather producer"); + cudaStream_t stream = reinterpret_cast(stream_handle); + impl_->current_hook_idx++; + // A disabled row on an uncaptured launch: the kernel's first read is the + // row's n, uploaded on this stream from the pinned row read here, and -1 + // makes it return with no reservation, task or payload, so not launching + // is the same observation. A capture always launches: the graph bakes + // the node and re-reads the device row at every replay. + if (impl_->mirrored_row_disabled(seg_table_dev_ptr) && + !stream_is_capturing(stream)) { + impl_->gather_host_skips.fetch_add(1, std::memory_order_relaxed); + return; + } +#if DMX_HAS_DEVICE_GRAPH_UPDATE + // Device-side toggling: a capture outside the toggle capture window (the + // graphs the host does not toggle, e.g. partial-model segment graphs) + // launches the producer as a device-updatable node, preceded once per + // captured forward by the planner node. Graphs inside the window keep + // host toggling, where SetEnabled calls are rare (set changes only). + if (impl_->toggle.device_armed() && !impl_->toggle.capture_enabled()) { + cudaStreamCaptureStatus st = cudaStreamCaptureStatusNone; + unsigned long long cid = 0; + cudaGraph_t g = nullptr; + const cudaGraphNode_t* deps = nullptr; + size_t nd = 0; + if (dmx_capture_info(stream, &st, &cid, &g, &deps, &nd) == cudaSuccess + && st == cudaStreamCaptureStatusActive && g != nullptr) { + bool fresh = false; + const int slot = impl_->toggle.device_slot( + reinterpret_cast(g), + reinterpret_cast(seg_table_dev_ptr), &fresh); + if (slot >= 0) { + if (fresh) + ring::launch_toggle_planner(impl_->toggle.device_pool(), + slot, stream); + const cudaError_t le = ring::launch_producer_gather_device_updatable( + impl_->engine.ring_state(), + reinterpret_cast(d_ptr), nbytes_upper, + reinterpret_cast(seg_table_dev_ptr), + seg_table_numel, row_bytes, hook_type, stream); + if (le != cudaSuccess) { + throw std::runtime_error( + std::string("device-updatable producer launch failed: ") + + cudaGetErrorString(le)); + } + if (dmx_capture_info(stream, &st, &cid, &g, &deps, &nd) != cudaSuccess + || nd < 1) { + impl_->toggle.note_anomaly(); + return; + } + cudaGraphNode_t node = deps[nd - 1]; + cudaKernelNodeAttrValue value = {}; + cudaGraphDeviceNode_t dev = nullptr; + if (cudaGraphKernelNodeGetAttribute( + node, cudaKernelNodeAttributeDeviceUpdatableKernelNode, + &value) == cudaSuccess && + value.deviceUpdatableKernelNode.deviceUpdatable) + dev = value.deviceUpdatableKernelNode.devNode; + impl_->toggle.register_device_entry( + slot, reinterpret_cast(g), + reinterpret_cast(node), + reinterpret_cast(seg_table_dev_ptr), dev); + return; + } + // Pool exhausted: launch normally; the node stays enabled and + // early-exits on a -1 row, but activation reports the anomaly. + impl_->toggle.note_anomaly(); + } + } +#endif + ring::launch_producer_gather( + impl_->engine.ring_state(), + reinterpret_cast(d_ptr), nbytes_upper, + reinterpret_cast(seg_table_dev_ptr), + seg_table_numel, row_bytes, hook_type, stream); +} + +void RingEnginePy::hook_no_notify_head(uint64_t d_ptr, uint64_t nbytes_upper, + uint64_t seg_table_dev_ptr, + int64_t seg_table_numel, + uint32_t hidden, uint32_t src_elem_bytes, + uint64_t arena_dev_ptr, uint32_t K, + uint32_t hook_type, + uint64_t stream_handle) +{ + refuse_on_record_ring(impl_->record_mode, "legacy head producer"); + cudaStream_t stream = reinterpret_cast(stream_handle); + impl_->current_hook_idx++; + if (impl_->mirrored_row_disabled(seg_table_dev_ptr) && + !stream_is_capturing(stream)) { + impl_->gather_host_skips.fetch_add(1, std::memory_order_relaxed); + return; + } + impl_->head_launches.fetch_add(1, std::memory_order_relaxed); + ring::launch_producer_head( + impl_->engine.ring_state(), + reinterpret_cast(d_ptr), nbytes_upper, + reinterpret_cast(seg_table_dev_ptr), seg_table_numel, + hidden, src_elem_bytes, + reinterpret_cast(arena_dev_ptr), K, hook_type, stream); +} + +void RingEnginePy::hook_no_notify_fused(uint64_t d_ptr, uint64_t nbytes_upper, + uint32_t hidden, uint32_t src_elem_bytes, + uint64_t gather_table_dev_ptr, + int64_t gather_table_numel, + uint64_t head_table_dev_ptr, + int64_t head_table_numel, + uint64_t arena_dev_ptr, uint32_t K, + uint32_t weight_bf16, + uint64_t acc_dev_ptr, uint64_t acc_rows, + uint32_t hook_type, + uint64_t stream_handle) +{ + refuse_on_record_ring(impl_->record_mode, "legacy fused producer"); + cudaStream_t stream = reinterpret_cast(stream_handle); + impl_->current_hook_idx++; + // Skip an uncaptured launch only when every present part is disabled. + auto disabled = [&](uint64_t table) { + return table == 0 || impl_->mirrored_row_disabled(table); + }; + if (disabled(gather_table_dev_ptr) && disabled(head_table_dev_ptr) && + !stream_is_capturing(stream)) { + impl_->gather_host_skips.fetch_add(1, std::memory_order_relaxed); + return; + } + impl_->head_launches.fetch_add(1, std::memory_order_relaxed); + ring::launch_producer_fused( + impl_->engine.ring_state(), + reinterpret_cast(d_ptr), nbytes_upper, hidden, + src_elem_bytes, + reinterpret_cast(gather_table_dev_ptr), gather_table_numel, + reinterpret_cast(head_table_dev_ptr), head_table_numel, + reinterpret_cast(arena_dev_ptr), K, weight_bf16, + reinterpret_cast(acc_dev_ptr), acc_rows, hook_type, stream); +} + +uint64_t RingEnginePy::head_launches() const { + return impl_->head_launches.load(std::memory_order_relaxed); +} + +void RingEnginePy::set_gather_host_mirror(uint64_t dev_base, uint64_t host_base, + uint64_t row_bytes, uint64_t rows) +{ + if (rows != 0 && (dev_base == 0 || host_base == 0 || + row_bytes < sizeof(int64_t) || + row_bytes % sizeof(int64_t) != 0)) { + throw std::invalid_argument( + "gather host mirror needs base pointers and whole int64 rows"); + } + // Disarm first so a reader never pairs old and new fields. + impl_->gather_dev_base.store(0, std::memory_order_release); + impl_->gather_rows = rows; + impl_->gather_row_bytes = row_bytes; + impl_->gather_host_base = host_base; + impl_->gather_dev_base.store(rows ? dev_base : 0, std::memory_order_release); +} + +uint64_t RingEnginePy::gather_host_skips() const { + return impl_->gather_host_skips.load(std::memory_order_relaxed); +} + // --------------------------------------------------------------------------- // Runtime CUDA-graph node toggling -- thin delegates to ToggleRegistry. The // engine adds only the process-global capture flag the producer ops read diff --git a/native/csrc/ring/ring_engine_py.h b/native/csrc/ring/ring_engine_py.h index 915f54810..2baccf519 100644 --- a/native/csrc/ring/ring_engine_py.h +++ b/native/csrc/ring/ring_engine_py.h @@ -161,6 +161,61 @@ class RingEnginePy { uint32_t hook_type, uint64_t stream_handle); + // Gather: copy the row runs listed in a per-step row table, packed (see + // producer.cuh for the table contract). Used by ring::producer_gather. + // Outside stream capture, a launch whose table row holds the disabled + // sentinel (-1) in the registered host mirror is skipped (the kernel + // would return on the same value). Under capture in device-toggle mode + // and outside the toggle capture window, the producer is captured as a + // device-updatable node gated by the in-graph planner. Every + // *_table_numel is the table's int64 element count (table.numel(), a + // launch constant); the kernels bound their segment reads by it. + void hook_no_notify_gather(uint64_t d_ptr, uint64_t nbytes_upper, + uint64_t seg_table_dev_ptr, + int64_t seg_table_numel, + uint64_t row_bytes, + uint32_t hook_type, + uint64_t stream_handle); + + // In-graph linear heads: K fp32 scores per selected row (producer.cuh, + // launch_producer_head). The same table contract and eager skip as + // hook_no_notify_gather. + void hook_no_notify_head(uint64_t d_ptr, uint64_t nbytes_upper, + uint64_t seg_table_dev_ptr, + int64_t seg_table_numel, + uint32_t hidden, uint32_t src_elem_bytes, + uint64_t arena_dev_ptr, uint32_t K, + uint32_t hook_type, + uint64_t stream_handle); + + // A HookPoint's row producer and its heads as ONE node + // (launch_producer_fused). Either table may be 0 (part absent). + // Outside capture the launch is skipped when every present part's + // mirrored row is the disabled sentinel. + void hook_no_notify_fused(uint64_t d_ptr, uint64_t nbytes_upper, + uint32_t hidden, uint32_t src_elem_bytes, + uint64_t gather_table_dev_ptr, + int64_t gather_table_numel, + uint64_t head_table_dev_ptr, + int64_t head_table_numel, + uint64_t arena_dev_ptr, uint32_t K, + uint32_t weight_bf16, + uint64_t acc_dev_ptr, uint64_t acc_rows, + uint32_t hook_type, + uint64_t stream_handle); + // Head and fused launches issued (diagnostic). + uint64_t head_launches() const; + + // Host mirror of the per-HookPoint row tables: `rows` device rows of + // `row_bytes` bytes at dev_base, each uploaded from the pinned row at the + // same offset from host_base. The caller rewrites a pinned row only + // after its previous upload completed, and uploads every rewrite on the + // producer stream before the forward. rows = 0 clears the mirror. + void set_gather_host_mirror(uint64_t dev_base, uint64_t host_base, + uint64_t row_bytes, uint64_t rows); + // Uncaptured launches skipped because their mirrored row was disabled. + uint64_t gather_host_skips() const; + // Additive schema-driven producers. These methods are reachable only // through the separately registered ring::record_producer* operations. void record_no_notify(uint64_t d_ptr, uint64_t nbytes, diff --git a/native/csrc/ring/ring_torch_op.cpp b/native/csrc/ring/ring_torch_op.cpp index 762c499d3..7a4aaae46 100644 --- a/native/csrc/ring/ring_torch_op.cpp +++ b/native/csrc/ring/ring_torch_op.cpp @@ -21,8 +21,17 @@ static std::atomic g_host_calls[HOST_HOOK_MAX] = {}; // Producer kernels launched by the ring producer ops, eager or captured. The // CUDA-graph replay guard compares it across a capture to tell a graph that // captured producers from one that did not (dmi.node_toggle). +// Every producer op (static, prefix, chunked, gather, head, fused) counts +// each kernel it launches through count_producer_launch(), once per launch: a +// fused rows+heads node is one launch. An op that launches without counting +// would make a capture holding only that op look producer-free, so its graph +// would skip the device-toggle upload and the replay guard. static std::atomic g_producer_launches{0}; +static inline void count_producer_launch() { + g_producer_launches.fetch_add(1, std::memory_order_relaxed); +} + uint64_t ring_producer_launch_count() { return g_producer_launches.load(std::memory_order_relaxed); } @@ -157,7 +166,7 @@ void ring_producer_impl( static_cast(tensor.nbytes()), static_cast(hook_type), reinterpret_cast(stream.stream())); - g_producer_launches.fetch_add(1, std::memory_order_relaxed); + count_producer_launch(); dmx_record_capture_node(stream.stream(), static_cast(hook_type), static_cast(hook_id), /*supported=*/true); } @@ -184,7 +193,7 @@ void ring_producer_prefix_impl( static_cast(row_bytes), static_cast(hook_type), reinterpret_cast(stream.stream())); - g_producer_launches.fetch_add(1, std::memory_order_relaxed); + count_producer_launch(); dmx_record_capture_node(stream.stream(), static_cast(hook_type), static_cast(hook_id), /*supported=*/true); } @@ -212,13 +221,198 @@ void ring_producer_chunked_impl( K, static_cast(hook_type), reinterpret_cast(stream.stream())); - g_producer_launches.fetch_add(1, std::memory_order_relaxed); + count_producer_launch(); // Not toggle-managed: fail closed if captured inside the window. dmx_record_capture_node(stream.stream(), static_cast(hook_type), static_cast(hook_id), /*supported=*/false); } } +// A producer row table ([n, src0, len0, ...] int64) as the kernels read it: +// on the HookPoint tensor's device, contiguous, and holding at least its +// segment count. A wrong table is a caller bug: refuse it rather than skip +// the producer, which would leave its reserved task unpublished. The count +// word itself is never read here (that would be a host sync per launch and +// would bake a capture-time value into graphs); the kernels receive +// table.numel() and bound the segments they read by the table's storage +// (row_table.h, row_table_segment_count). +static void check_row_table(const at::Tensor& table, const at::Tensor& tensor, + const char* op) { + TORCH_CHECK(table.is_cuda() && table.device() == tensor.device(), + op, ": row table must be on the HookPoint tensor's device"); + TORCH_CHECK(table.scalar_type() == at::kLong && table.is_contiguous(), + op, ": row table must be contiguous int64"); + TORCH_CHECK(table.numel() >= 1, op, ": row table must hold a segment count"); +} + +// ring::producer_gather(ring_payload, x, seg_table, row_bytes, hook_type, +// hook_id) +// Copies the row runs listed in seg_table (producer.cuh table contract), +// packed. A single kernel launch, so it records its node like prefix. +void ring_producer_gather_impl( + const at::Tensor& /*ring_payload*/, + const at::Tensor& tensor, + const at::Tensor& seg_table, + int64_t row_bytes, + int64_t hook_type, int64_t hook_id) +{ + if (!g_active_engine) { return; } + if (hook_type >= 0 && hook_type < HOST_HOOK_MAX) + g_host_calls[hook_type].fetch_add(1); + + if (tensor.is_cuda() && tensor.is_contiguous() && seg_table.defined()) { + check_row_table(seg_table, tensor, "ring::producer_gather"); + TORCH_CHECK(row_bytes > 0, + "ring::producer_gather: row_bytes must be positive"); + auto stream = at::cuda::getCurrentCUDAStream(tensor.device().index()); + g_active_engine->hook_no_notify_gather( + reinterpret_cast(tensor.data_ptr()), + static_cast(tensor.nbytes()), + reinterpret_cast(seg_table.data_ptr()), + static_cast(seg_table.numel()), + static_cast(row_bytes), + static_cast(hook_type), + reinterpret_cast(stream.stream())); + count_producer_launch(); + // Without this, a gather producer captured under the toggle would + // fire unregistered: disabling its key could not disable its node + // while metadata and reservation exclude it. + dmx_record_capture_node(stream.stream(), static_cast(hook_type), + static_cast(hook_id), /*supported=*/true); + } +} + +// ring::producer_head(ring_payload, x, seg_table, arena, K, hook_type, hook_id) +// In-graph linear heads at one HookPoint: K fp32 scores w_k . x[row] + b_k +// for each row seg_table selects, from the stable fp32 arena (K x hidden +// weights, K biases, split scratch). Recorded as its own toggle node under +// (hook_type, hook_id), independent of the HookPoint's row producer. +void ring_producer_head_impl( + const at::Tensor& /*ring_payload*/, + const at::Tensor& tensor, + const at::Tensor& seg_table, + const at::Tensor& arena, + int64_t K, + int64_t hook_type, int64_t hook_id) +{ + if (!g_active_engine) { return; } + if (hook_type >= 0 && hook_type < HOST_HOOK_MAX) + g_host_calls[hook_type].fetch_add(1); + if (!(tensor.is_cuda() && tensor.is_contiguous() && tensor.dim() >= 1)) return; + TORCH_CHECK(seg_table.defined(), "ring::producer_head: a row table is required"); + check_row_table(seg_table, tensor, "ring::producer_head"); + TORCH_CHECK(arena.defined() && arena.is_cuda() && arena.device() == tensor.device() + && arena.is_contiguous() && arena.scalar_type() == at::kFloat, + "ring::producer_head: arena must be contiguous fp32 on the " + "HookPoint tensor's device"); + const auto st = tensor.scalar_type(); + TORCH_CHECK(st == at::kBFloat16 || st == at::kFloat, + "ring::producer_head supports bf16 and fp32 HookPoints only"); + TORCH_CHECK(K >= 1 && K <= static_cast(ring::HEAD_MAX_SLOTS), + "ring::producer_head: K must be in [1, HEAD_MAX_SLOTS]"); + const int64_t hidden = tensor.size(tensor.dim() - 1); + TORCH_CHECK(arena.numel() == static_cast(ring::head_arena_floats( + static_cast(K), static_cast(hidden))), + "ring::producer_head: arena must hold K*hidden weights, K " + "biases and the split scratch (head_arena_floats)"); + auto stream = at::cuda::getCurrentCUDAStream(tensor.device().index()); + g_active_engine->hook_no_notify_head( + reinterpret_cast(tensor.data_ptr()), + static_cast(tensor.nbytes()), + reinterpret_cast(seg_table.data_ptr()), + static_cast(seg_table.numel()), + static_cast(hidden), + static_cast(tensor.element_size()), + reinterpret_cast(arena.data_ptr()), + static_cast(K), + static_cast(hook_type), + reinterpret_cast(stream.stream())); + count_producer_launch(); + dmx_record_capture_node(stream.stream(), static_cast(hook_type), + static_cast(hook_id), /*supported=*/true); +} + +// ring::producer_fused(ring_payload, x, gather_table?, row_bytes, head_table?, +// arena, acc, K, weight_bf16, hook_type, hook_id, +// head_hook_type) +// The HookPoint's row producer AND its heads in ONE node. Each present +// part has its own table and its own toggle key: (hook_type, hook_id) for +// the rows, (head_hook_type, hook_id) for the heads. Both keys are +// registered to the same captured node, which is enabled iff either key +// is; each part early-exits on its own -1 row, so "metadata pushed <=> +// task published" still holds per key. +// arena: int32 [head_node_words(K, hidden, weight_bf16)]; +// acc: fp32 [acc_rows, HEAD_MAX_ACC], shared by all nodes. +void ring_producer_fused_impl( + const at::Tensor& /*ring_payload*/, + const at::Tensor& tensor, + const c10::optional& gather_table, + int64_t row_bytes, + const c10::optional& head_table, + const at::Tensor& arena, + const at::Tensor& acc, + int64_t K, int64_t weight_bf16, + int64_t hook_type, int64_t hook_id, int64_t head_hook_type) +{ + if (!g_active_engine) { return; } + if (hook_type >= 0 && hook_type < HOST_HOOK_MAX) + g_host_calls[hook_type].fetch_add(1); + const bool g = gather_table.has_value() && gather_table->defined(); + const bool h = head_table.has_value() && head_table->defined(); + if (!(tensor.is_cuda() && tensor.is_contiguous() && tensor.dim() >= 1)) return; + if (!g && !h) return; + const auto st = tensor.scalar_type(); + TORCH_CHECK(st == at::kBFloat16 || st == at::kFloat, + "ring::producer_fused supports bf16 and fp32 HookPoints only"); + const int64_t hidden = tensor.size(tensor.dim() - 1); + if (g) { + check_row_table(*gather_table, tensor, "ring::producer_fused (rows)"); + TORCH_CHECK(row_bytes == hidden * static_cast(tensor.element_size()), + "ring::producer_fused: row_bytes must equal hidden * element size"); + TORCH_CHECK(row_bytes % 16 == 0, + "ring::producer_fused: rows must be 16-byte multiples"); + } + if (h) { + check_row_table(*head_table, tensor, "ring::producer_fused (heads)"); + TORCH_CHECK(K >= 1 && K <= static_cast(ring::HEAD_MAX_SLOTS), + "ring::producer_fused: K must be in [1, HEAD_MAX_SLOTS]"); + TORCH_CHECK(arena.defined() && arena.is_cuda() + && arena.device() == tensor.device() && arena.is_contiguous() + && arena.scalar_type() == at::kInt + && arena.numel() == static_cast(ring::head_node_words( + static_cast(K), static_cast(hidden), + weight_bf16 != 0)), + "ring::producer_fused: arena must be int32 " + "[head_node_words(K, hidden, weight_bf16)]"); + TORCH_CHECK(acc.defined() && acc.is_cuda() && acc.device() == tensor.device() + && acc.is_contiguous() + && acc.scalar_type() == at::kFloat && acc.dim() == 2 + && acc.size(1) == static_cast(ring::HEAD_MAX_ACC), + "ring::producer_fused: acc must be fp32 [rows, HEAD_MAX_ACC]"); + } + auto stream = at::cuda::getCurrentCUDAStream(tensor.device().index()); + g_active_engine->hook_no_notify_fused( + reinterpret_cast(tensor.data_ptr()), + static_cast(tensor.nbytes()), + static_cast(hidden), + static_cast(tensor.element_size()), + g ? reinterpret_cast(gather_table->data_ptr()) : 0, + g ? static_cast(gather_table->numel()) : 0, + h ? reinterpret_cast(head_table->data_ptr()) : 0, + h ? static_cast(head_table->numel()) : 0, + h ? reinterpret_cast(arena.data_ptr()) : 0, + static_cast(K), static_cast(weight_bf16 != 0), + h ? reinterpret_cast(acc.data_ptr()) : 0, + h ? static_cast(acc.size(0)) : 0, + static_cast(hook_type), + reinterpret_cast(stream.stream())); + count_producer_launch(); // one node for both parts + if (g) dmx_record_capture_node(stream.stream(), static_cast(hook_type), + static_cast(hook_id), /*supported=*/true); + if (h) dmx_record_capture_node(stream.stream(), static_cast(head_hook_type), + static_cast(hook_id), /*supported=*/true); +} + namespace { const int32_t* checked_record_gate( @@ -419,6 +613,16 @@ TORCH_LIBRARY(ring, m) { "int hook_type, int hook_id) -> ()"); m.def("producer_chunked(Tensor(a!) ring_payload, Tensor x, " "Tensor chunk_bytes, int hook_type, int hook_id) -> ()"); + m.def("producer_gather(Tensor(a!) ring_payload, Tensor x, " + "Tensor seg_table, int row_bytes, " + "int hook_type, int hook_id) -> ()"); + m.def("producer_head(Tensor(a!) ring_payload, Tensor x, " + "Tensor seg_table, Tensor arena, int K, " + "int hook_type, int hook_id) -> ()"); + m.def("producer_fused(Tensor(a!) ring_payload, Tensor x, " + "Tensor? gather_table, int row_bytes, Tensor? head_table, " + "Tensor arena, Tensor acc, int K, int weight_bf16, " + "int hook_type, int hook_id, int head_hook_type) -> ()"); m.def("record_producer(Tensor(a!) ring_payload, Tensor x, " "Tensor? emit_gate=None, int emit_value=0) -> ()"); m.def("record_producer_prefix(Tensor(a!) ring_payload, Tensor x, " @@ -438,6 +642,9 @@ TORCH_LIBRARY_IMPL(ring, CUDA, m) { m.impl("producer", ring_producer_impl); m.impl("producer_prefix", ring_producer_prefix_impl); m.impl("producer_chunked", ring_producer_chunked_impl); + m.impl("producer_gather", ring_producer_gather_impl); + m.impl("producer_head", ring_producer_head_impl); + m.impl("producer_fused", ring_producer_fused_impl); m.impl("record_producer", ring_record_producer_impl); m.impl("record_producer_prefix", ring_record_producer_prefix_impl); m.impl("record_producer_chunked", ring_record_producer_chunked_impl); diff --git a/native/csrc/ring/row_table.h b/native/csrc/ring/row_table.h new file mode 100644 index 000000000..2a02c97f8 --- /dev/null +++ b/native/csrc/ring/row_table.h @@ -0,0 +1,47 @@ +// Producer row tables: the segment-count clamp shared by the gather, head and +// fused producers. Plain C++ (no CUDA dependency) so host code and CPU unit +// tests use the exact device arithmetic. +#pragma once + +#include + +#if defined(__CUDACC__) +#define DMX_ROW_TABLE_HD __host__ __device__ +#else +#define DMX_ROW_TABLE_HD +#endif + +namespace ring { + +// Upper bound on the number of row segments in a gather table +// (shared-memory sizing). Segments are merged runs of selected rows, so the +// count is bounded by the number of concurrently scheduled requests; a host +// with more must fall back to a full-row table (over-capture is safe). +static constexpr uint32_t GATHER_MAX_SEGS = 256; + +// Segments a table of `table_numel` int64 words can physically hold: word 0 +// is the count, segment k occupies words 1 + 2k and 2 + 2k. +DMX_ROW_TABLE_HD inline int64_t row_table_capacity_segs(int64_t table_numel) { + return table_numel > 0 ? (table_numel - 1) / 2 : 0; +} + +// The number of segments a producer reads from a table whose first word is +// `n_raw`: a negative count reads none, and the count is bounded by both the +// shared-memory limit and the table's own storage, so the segment reads +// table[1 + 2k] and table[2 + 2k] stay below table_numel for every k read. +// The count stays a device-side value: callers pass the table's element +// count (known on the host at launch/capture time), never a host read of +// table[0], so graph replays still re-read the count. +DMX_ROW_TABLE_HD inline uint32_t row_table_segment_count(int64_t n_raw, + int64_t table_numel) { + if (n_raw < 0) n_raw = 0; + if (n_raw > static_cast(GATHER_MAX_SEGS)) + n_raw = static_cast(GATHER_MAX_SEGS); + const int64_t capacity = row_table_capacity_segs(table_numel); + if (n_raw > capacity) n_raw = capacity; + return static_cast(n_raw); +} + +} // namespace ring + +#undef DMX_ROW_TABLE_HD diff --git a/native/csrc/ring/tensor_meta.h b/native/csrc/ring/tensor_meta.h index 1f054fab0..ba56afca9 100644 --- a/native/csrc/ring/tensor_meta.h +++ b/native/csrc/ring/tensor_meta.h @@ -10,6 +10,7 @@ #pragma once #include #include +#include #include #include #include @@ -53,7 +54,11 @@ enum HookType : int { HOOK_TYPE_ROUTER_LOGITS= 21, HOOK_TYPE_TOPK_IDS = 22, HOOK_TYPE_TOPK_WEIGHTS = 23, - HOOK_TYPE_COUNT = 24, + // In-graph linear projection heads: K fp32 scores per selected row of + // the HookPoint named by the act_name prefix (producer_head/_fused). + HOOK_TYPE_RESID_PRE_HEAD = 24, + HOOK_TYPE_FINAL_LN_HEAD = 25, + HOOK_TYPE_COUNT = 26, }; // Hook group — which sub-block produces this tensor. @@ -71,10 +76,12 @@ enum HookGroup : int { GROUP_ATTN = 0, GROUP_MLP = 1, GROUP_OTHER = 2 }; // SHAPE_ROUTER_LOGITS : [batch, q_len, num_experts] // SHAPE_TOPK_IDS : [batch, q_len, top_k] // SHAPE_TOPK_WEIGHTS : [batch, q_len, top_k] +// SHAPE_HEAD_SCORES : [batch, q_len, head_slots] (fp32 projection heads) enum ShapeClass : int { SHAPE_HIDDEN = 0, SHAPE_QKV_Q = 1, SHAPE_QKV_KV = 2, SHAPE_QKV_Z = 3, SHAPE_ATTN_WT = 4, SHAPE_MLP_POST = 5, SHAPE_TOKEN_IDS = 6, SHAPE_LOGITS = 7, SHAPE_ROUTER_LOGITS = 8, SHAPE_TOPK_IDS = 9, SHAPE_TOPK_WEIGHTS = 10, + SHAPE_HEAD_SCORES = 11, }; // Pipeline-parallel stage placement. @@ -116,6 +123,8 @@ static constexpr HookDef HOOK_DEFS[] = { {HOOK_TYPE_ROUTER_LOGITS,"mlp.hook_router_logits", "router_logits",true, GROUP_OTHER, false, SHAPE_ROUTER_LOGITS, PP_ANY }, {HOOK_TYPE_TOPK_IDS, "mlp.hook_topk_ids", "topk_ids", true, GROUP_OTHER, false, SHAPE_TOPK_IDS, PP_ANY }, {HOOK_TYPE_TOPK_WEIGHTS,"mlp.hook_topk_weights", "topk_weights", true, GROUP_OTHER, false, SHAPE_TOPK_WEIGHTS, PP_ANY }, + {HOOK_TYPE_RESID_PRE_HEAD,"hook_resid_pre_head", "resid_pre_head", true, GROUP_OTHER, false, SHAPE_HEAD_SCORES, PP_ANY }, + {HOOK_TYPE_FINAL_LN_HEAD, "hook_final_ln_head", "final_ln_head", false, GROUP_OTHER, false, SHAPE_HEAD_SCORES, PP_LAST}, }; static constexpr int HOOK_DEFS_COUNT = sizeof(HOOK_DEFS) / sizeof(HOOK_DEFS[0]); @@ -180,6 +189,11 @@ struct RequestMeta { int32_t end_token = 0; int64_t dim0_offset = 0; // HF: batch_idx, vLLM: token offset int32_t kv_offset = 0; // attn kv-dim start (HF dynamic cache: pad_len, else 0) + // Position of this request in the step's UNFILTERED request list. + // Position-indexed hooks (final_logits: row j of [num_all_reqs, vocab]) + // index the tensor by it, not by the request's index in the (possibly + // filtered) request list. -1 = the list is unfiltered: use the index. + int32_t orig_position = -1; }; // Per-step shared context. Heap-allocated by push_step caller. @@ -211,6 +225,12 @@ struct TensorMeta { int dtype = 0; // at::ScalarType integer value bool last_in_step = false; uint8_t flags = 0; // META_FLAG_* bitmask + // Optional request layout of an independently gathered hook. The step + // keeps one StepContext for the whole forward, zero-row hooks included. + // nullptr inherits the step's request list; a non-null empty vector + // selects no request. Identical layouts share one immutable vector + // within a submitted step. + std::shared_ptr> requests; }; // Thread-safe dual FIFO for TensorMeta + StepContext. diff --git a/native/csrc/ring/toggle_registry.cu b/native/csrc/ring/toggle_registry.cu index d4c6b15c5..0c459b807 100644 --- a/native/csrc/ring/toggle_registry.cu +++ b/native/csrc/ring/toggle_registry.cu @@ -41,7 +41,16 @@ void ToggleRegistry::register_node(uint64_t graph, int hook_type, int layer_no, std::lock_guard lk(mu_); RegEntry re{hook_type, layer_no, reinterpret_cast(node)}; re.device = dev_nodes_.count(re.node) > 0; - reg_nodes_[reinterpret_cast(graph)].push_back(re); + auto& entries = reg_nodes_[reinterpret_cast(graph)]; + if (!entries.empty() && entries.back().node == re.node && + entries.back().alt_hook_type < 0 && + !(entries.back().hook_type == hook_type && + entries.back().layer_no == layer_no)) { + entries.back().alt_hook_type = hook_type; // fused producer's second key + entries.back().alt_layer_no = layer_no; + } else { + entries.push_back(re); + } registered_hooks_.insert({hook_type, layer_no}); ++version_; } @@ -188,7 +197,10 @@ bool ToggleRegistry::uniform() const { std::set> ref; for (const auto& kv : reg_nodes_) { std::set> s; - for (const auto& e : kv.second) s.insert({e.hook_type, e.layer_no}); + for (const auto& e : kv.second) { + s.insert({e.hook_type, e.layer_no}); + if (e.alt_hook_type >= 0) s.insert({e.alt_hook_type, e.alt_layer_no}); + } if (!have_ref) { ref = std::move(s); have_ref = true; } else if (s != ref) return false; } diff --git a/native/csrc/ring/toggle_registry.h b/native/csrc/ring/toggle_registry.h index 08c3b5c3b..4156477ec 100644 --- a/native/csrc/ring/toggle_registry.h +++ b/native/csrc/ring/toggle_registry.h @@ -35,6 +35,8 @@ class ToggleRegistry { void set_capture(bool enabled); bool capture_enabled() const; // Called from the producer op during capture (single kernel just enqueued). + // A second, different key recorded for the node just registered (a fused + // producer) is attached to that node instead of adding an entry. void register_node(uint64_t graph, int hook_type, int layer_no, uint64_t node); // Called when a capture-time recording could not be validated (non-kernel // tail, unreadable capture state, unsupported producer) -> fail closed. @@ -115,6 +117,10 @@ class ToggleRegistry { // the next apply sets the node unconditionally. bool state_known{true}; bool device{false}; // toggled by the in-graph planner, never by the host + // A fused producer node carries two keys (rows + heads); it is + // enabled iff either key is enabled. -1 = no second key. + int alt_hook_type{-1}; + int alt_layer_no{-1}; }; cudaError_t upload_device_graph_locked(cudaGraph_t g, cudaGraphExec_t exec); // The apply loops skip a node only when its exec state is known to match. @@ -122,7 +128,9 @@ class ToggleRegistry { return e.state_known && on == e.last_enabled; } bool entry_on(const RegEntry& e) const { - return enabled_hooks_.count({e.hook_type, e.layer_no}) > 0; + return enabled_hooks_.count({e.hook_type, e.layer_no}) > 0 + || (e.alt_hook_type >= 0 && + enabled_hooks_.count({e.alt_hook_type, e.alt_layer_no}) > 0); } std::unordered_map> reg_nodes_; // per captured graph std::unordered_map reg_exec_; // graph -> exec diff --git a/src/dmi/adapters/base.py b/src/dmi/adapters/base.py index 619c0b071..6f3130c3e 100644 --- a/src/dmi/adapters/base.py +++ b/src/dmi/adapters/base.py @@ -487,6 +487,22 @@ def commit_step( def close(self) -> None: self.engine.close() + def replace_active_specs(self, specs: List[HookSpec]) -> None: + """Replace the attached spec list, keeping the transport in sync. + + For integrations that add specs over HookPoints already installed by + ``attach_model`` -- the in-graph head spec of a HookPoint + (``dmi.hook_rows.configure_hook_heads``) is inserted right after the + HookPoint's own spec, because the fused node launches the heads after + the rows and metadata order must equal launch order. Every spec must + name a HookPoint of the current list; call before graph capture.""" + installed = {id(spec.module) for spec in self.active_specs} + specs = list(specs) + if any(id(spec.module) not in installed for spec in specs): + raise ValueError("replacement specs must name installed HookPoints") + self.active_specs = specs + self.transport._active_specs = specs + def plan_step(self, ctx: StepContext) -> StepPlan: """Build a side-effect-free reservation plan for one step. @@ -517,6 +533,7 @@ def plan_step(self, ctx: StepContext) -> StepPlan: transport = getattr(self, "transport", None) if getattr(transport, "toggle_gate_active", False) is True: specs = transport.specs_for_step() + hook_rows = getattr(ctx, "hook_rows", None) or {} for spec in specs: # When the adapter has populated actual_q_len AND the spec is # prefix-strip-eligible, size the byte budget for the unpadded @@ -525,6 +542,10 @@ def plan_step(self, ctx: StepContext) -> StepPlan: spec_q_len = (ctx.actual_q_len if ctx.actual_q_len is not None and spec.dim0_is_actual_tokens else ctx.q_len) + # A hook with a request-row plan writes exactly its plan's rows. + rows = hook_rows.get((spec.hook_type, spec.layer_no)) + if rows is not None: + spec_q_len = rows.row_count shape = compute_hook_shape( spec.hook_type, self.model_cfg, ctx.batch, spec_q_len, ctx.kv_dim, diff --git a/src/dmi/adapters/types.py b/src/dmi/adapters/types.py index 5c111dd5f..ae85155b1 100644 --- a/src/dmi/adapters/types.py +++ b/src/dmi/adapters/types.py @@ -56,10 +56,18 @@ class StepContext: # adapter driver: a schedule with capture_prefill=False drops prefill # steps, capture_decode=False drops decode steps. None captures as before. phase: Optional[str] = None + # Each request's index in the step's unfiltered request list, when the + # adapter filtered req_ids/token_ranges (see RingTransport); None = + # unfiltered. + orig_positions: Optional[List[int]] = None + # Per-HookPoint request-row plans (dmi.hook_rows) keyed by + # (hook_type, layer_no). A planned hook reserves and publishes its + # plan's row_count rows; None = every hook uses the step layout. + hook_rows: Optional[dict] = None def transport_kwargs(self) -> dict: """Return kwargs accepted by ``RingTransport.set_step_context``.""" - return dict( + kwargs = dict( model_id=self.model_id, req_ids=self.req_ids, token_ranges=self.token_ranges, @@ -71,6 +79,11 @@ def transport_kwargs(self) -> dict: pp_rank=self.pp_rank, flattened=self.flattened, ) + if self.orig_positions is not None: + kwargs["orig_positions"] = self.orig_positions + if self.hook_rows is not None: + kwargs["hook_rows"] = self.hook_rows + return kwargs __all__ = ["StepContext"] diff --git a/src/dmi/configuration/architecture.py b/src/dmi/configuration/architecture.py index 3b8f1fc72..6d87a5184 100644 --- a/src/dmi/configuration/architecture.py +++ b/src/dmi/configuration/architecture.py @@ -44,7 +44,7 @@ class ArchitectureNode: id="residual", label="Residual stream", scope=SCOPE_LAYER, - hooks=("resid_pre", "ln1", "resid_mid", "ln2"), + hooks=("resid_pre", "ln1", "resid_mid", "ln2", "resid_pre_head"), ), ArchitectureNode( id="attention", @@ -68,7 +68,7 @@ class ArchitectureNode: id="output", label="Final norm & logits", scope=SCOPE_GLOBAL, - hooks=("resid_final", "final_ln", "final_logits"), + hooks=("resid_final", "final_ln", "final_logits", "final_ln_head"), ), ) diff --git a/src/dmi/hook_rows.py b/src/dmi/hook_rows.py new file mode 100644 index 000000000..87d8ed249 --- /dev/null +++ b/src/dmi/hook_rows.py @@ -0,0 +1,849 @@ +"""Per-HookPoint request-row plans for the gather producer. + +A plan says which rows of a step's packed activation tensor one HookPoint +copies (``segments``: merged ``(source_row, rows)`` runs, the producer's row +table) and how the consumer attributes the packed result to requests +(``req_ids``, ``token_ranges``, destination ``offsets``, ``orig_positions``). +Plans are immutable and carry no GPU or framework dependency. + +* :func:`plan_hook_rows` plans one HookPoint from an explicit request + selection and optional absolute token windows. +* :func:`plan_hook_groups` and :class:`HookRowPlanner` plan many HookPoints + at once, sharing one plan between HookPoints whose selections agree. +* :class:`PackedDemandPlanner` plans declared per-request demand with array + operations over rules compiled once per demand change. +* :func:`full_row_plan` is the static full-row baseline: the whole physical + payload as one run, physical offsets. + +A plan whose segments exceed the producer's table capacity falls back to the +un-compacted full payload while keeping the consumer-side selection +(``capacity_fallback``). +""" +from bisect import bisect_right +from dataclasses import dataclass, replace +from itertools import chain as _chain, islice, repeat as _repeat +from operator import itemgetter +from types import MappingProxyType + +import numpy as _np + + +# Row-table capacity of the gather producer (GATHER_MAX_SEGS in +# native/csrc/ring/producer.cuh): a table holds at most this many segments. +GATHER_MAX_SEGS = 256 + +_window_start = itemgetter(0) + + +def canonical_token_windows(windows): + """Validate and union half-open absolute token intervals at admission. + + ``None`` as an interval's end means unbounded. An empty list is an empty + selection, not the default (all positions). No sorting/validation is needed + in the per-step planner after this conversion. + """ + parsed = [] + if not isinstance(windows, (list, tuple)): + raise ValueError("token windows must be a list of (start, end) pairs") + for interval in windows: + if not isinstance(interval, (list, tuple)) or len(interval) != 2: + raise ValueError("token windows must be (start, end) pairs") + start, end = interval + if (type(start) is not int or start < 0 or + (end is not None and (type(end) is not int or end <= start))): + raise ValueError("token windows require 0 <= start < end (or end=None)") + parsed.append((start, end)) + merged = [] + for start, end in sorted(parsed, key=lambda pair: pair[0]): + if merged and (merged[-1][1] is None or start <= merged[-1][1]): + previous_start, previous_end = merged[-1] + merged[-1] = (previous_start, None if previous_end is None or end is None + else max(previous_end, end)) + else: + merged.append((start, end)) + return tuple(merged) + + +@dataclass(frozen=True) +class TokenWindowIndex: + """Hook-indexed immutable filters, shared when their request maps agree.""" + by_hook: object + long_hooks: frozenset = frozenset() + + +def compile_token_windows(windows): + """Transpose and intern on selection changes, not on every serving step. + + Input intervals must already be canonical. Values are owned immutable maps; + caller mutation cannot alter a snapshot while a forward is using it. + """ + by_hook, long_hooks = {}, set() + for rid, hooks in windows.items(): + for hook, spans in hooks.items(): + by_hook.setdefault(hook, {})[rid] = spans + if len(spans) > 8: + long_hooks.add(hook) + interned = {} + for hook, by_request in by_hook.items(): + signature = tuple(sorted(by_request.items())) + if signature not in interned: + interned[signature] = MappingProxyType(by_request) + by_hook[hook] = interned[signature] + return TokenWindowIndex(MappingProxyType(by_hook), frozenset(long_hooks)) + + +@dataclass(frozen=True) +class HookRows: + segments: tuple + row_count: int + req_ids: tuple + token_ranges: tuple + offsets: tuple + source_offsets: tuple + orig_positions: tuple + capacity_fallback: bool = False + + def full_payload(self, total_rows): + """Keep consumer selection, but describe an un-compacted fallback.""" + return replace(self, segments=((0, total_rows),), row_count=total_rows, + offsets=self.source_offsets) + + +def plan_hook_rows(req_ids, token_ranges, offsets, orig_positions, selected, + total_rows, max_segments, token_windows=None): + """Select scheduled request spans or their declared token windows. + + `selected=None` means all supplied requests, not graph-padding rows. + Metadata offsets address the compacted destination, never the source. + On table overflow, capture full actual rows but preserve consumer filtering. + ``token_windows`` maps request IDs to canonical absolute intervals. Adjacent + source ranges can merge for copying even when they belong to different + requests; metadata continues to identify each interval separately. + """ + scheduled = _checked_layout(req_ids, token_ranges, offsets, orig_positions, total_rows) + chosen = (row for row in scheduled if selected is None or row[0] in selected) + long_windows = bool(token_windows) and any(len(w) > 8 for w in token_windows.values()) + return _plan_selected_rows(chosen, total_rows, max_segments, token_windows, long_windows) + + +def _checked_layout(req_ids, token_ranges, offsets, orig_positions, total_rows): + """Validate all scheduled rows, including rows no monitor selects.""" + if not (len(req_ids) == len(token_ranges) == len(offsets) == len(orig_positions)): + raise ValueError("request row metadata lengths differ") + scheduled = tuple(zip(req_ids, token_ranges, offsets, orig_positions)) + previous_end = 0 + for _, tr, off, _ in scheduled: + length = tr[1] - tr[0] + if length <= 0 or off < previous_end or off + length > total_rows: + raise ValueError("invalid or overlapping scheduled request rows") + previous_end = off + length + return scheduled + + +def _plan_selected_rows(chosen, total_rows, max_segments, token_windows, long_windows=False): + """Build one profile from a checked layout; no repeated batch validation.""" + segs, ids, ranges, src, dst, positions = [], [], [], [], [], [] + count = 0 + for rid, tr, off, pos in chosen: + windows = token_windows.get(rid) if token_windows else None + if long_windows and windows and len(windows) > 8: + if (windows[0][0] >= tr[1] or + (windows[-1][1] is not None and windows[-1][1] <= tr[0])): + continue + # Seek to the current chunk instead of rescanning all earlier + # sampled positions on every prefill/decode step. No cursor is + # retained, so reordering or replaying an earlier chunk is safe. + first = max(0, bisect_right(windows, tr[0], key=_window_start) - 1) + windows = islice(windows, first, None) + for start, end in (tr,) if windows is None else windows: + # Canonical windows are sorted and disjoint. In particular, a + # prompt-only window must not leak into subsequent decode rows. + if windows is not None: + if start >= tr[1]: + break + if start < tr[0]: + start = tr[0] + if end is None or end > tr[1]: + end = tr[1] + if end <= start: + continue + source = off + start - tr[0] + length = end - start + if segs and segs[-1][0] + segs[-1][1] == source: + segs[-1] = (segs[-1][0], segs[-1][1] + length) + else: + segs.append((source, length)) + ids.append(rid) + ranges.append((start, end)) + src.append(source) + dst.append(count) + positions.append(pos) + count += length + plan = HookRows(tuple(segs), count, tuple(ids), tuple(ranges), tuple(dst), + tuple(src), tuple(positions)) + if len(segs) > max_segments: + plan = replace(plan.full_payload(total_rows), capacity_fallback=True) + return plan + + +def _sparse_rows(scheduled, bits): + """Visit selected request indices in packed order without a full scan.""" + while bits: + lowest = bits & -bits + yield scheduled[lowest.bit_length() - 1] + bits ^= lowest + + +class HookRowPlanner: + """Own a selection snapshot and cache only its request-membership groups. + + Graph row indices and absolute positions are deliberately NOT cached. + A cold request order uses one-step grouping; reusable indices are built + only if that order repeats. Changed selections/windows install a new + planner. Only one ordering is retained, bounding memory. + The serving thread owns plan(), just as it owns the current batch layout. + """ + + def __init__(self, keys, selection, default=None, token_windows=None): + self.keys = tuple(keys) + self.selection = {rid: None if hooks is None else frozenset(hooks) + for rid, hooks in selection.items()} + self.default = None if default is None else frozenset(default) + self.windows = (token_windows if isinstance(token_windows, TokenWindowIndex) + else compile_token_windows(token_windows) if token_windows else None) + self._request_order = None + self._groups = None + self.group_builds = 0 + self.group_reuses = 0 + + def _group(self, req_ids): + profiles = {} + for i, rid in enumerate(req_ids): + hooks = self.selection.get(rid, self.default) + profiles[hooks] = profiles.get(hooks, 0) | (1 << i) + masks = dict.fromkeys(self.keys, 0) + for hooks, bits in profiles.items(): + for key in self.keys if hooks is None else hooks: + if key in masks: + masks[key] |= bits + filters = self.windows.by_hook if self.windows else {} + long_hooks = self.windows.long_hooks if self.windows else frozenset() + groups = {} + all_rows = (1 << len(req_ids)) - 1 + for key, bits in masks.items(): + windows = filters.get(key) + identity = (bits, id(windows)) if windows else bits + if identity not in groups: + # Compute dense and sparse membership just once, then visit + # only selected rows on all subsequent steps in this batch. + indices = (None if bits == all_rows else + tuple(_sparse_rows(tuple(range(len(req_ids))), bits))) + groups[identity] = ([], indices, windows, key in long_hooks) + groups[identity][0].append(key) + return tuple((tuple(keys), indices, windows, long) + for keys, indices, windows, long in groups.values()) + + def plan(self, req_ids, token_ranges, offsets, orig_positions, total_rows, max_segments): + order = tuple(req_ids) + if order != self._request_order: + # A one-off layout needs no persistent membership-index cache. + # Validate and finish planning before replacing the previous cache. + plans = plan_hook_groups(self.keys, order, token_ranges, offsets, + orig_positions, self.selection, self.default, + total_rows, max_segments, self.windows) + self._groups = None + self._request_order = order + self.group_builds += 1 + return plans + scheduled = _checked_layout(order, token_ranges, offsets, orig_positions, total_rows) + if self._groups is None: + self._groups = self._group(order) + self.group_reuses += 1 + plans = {} + for keys, indices, windows, long_windows in self._groups: + chosen = scheduled if indices is None else map(scheduled.__getitem__, indices) + rows = _plan_selected_rows(chosen, total_rows, max_segments, windows, long_windows) + plans.update(dict.fromkeys(keys, rows)) + return plans + + +def plan_hook_groups(keys, req_ids, token_ranges, offsets, orig_positions, + selection, default, total_rows, max_segments, token_windows=None): + """Intern equal request selections within ONE batch layout. + + A bit identifies a scheduled request, not a token row. Requests sharing a + profile are grouped first, then transposed to HookPoint request bitsets. + Absolute token positions belong to this step and are never cached across + steps, even when the physical row intervals are unchanged. + """ + scheduled = _checked_layout(req_ids, token_ranges, offsets, orig_positions, total_rows) + profiles = {} + resolved = [selection.get(rid, default) for rid in req_ids] + for i, hooks in enumerate(resolved): + profiles[hooks] = profiles.get(hooks, 0) | (1 << i) + masks = dict.fromkeys(keys, 0) + for hooks, bits in profiles.items(): + for key in keys if hooks is None else hooks: + if key in masks: + masks[key] |= bits + if token_windows and not isinstance(token_windows, TokenWindowIndex): + token_windows = compile_token_windows(token_windows) + filters = token_windows.by_hook if token_windows else {} + long_hooks = token_windows.long_hooks if token_windows else frozenset() + interned, plans = {}, {} + all_rows = (1 << len(scheduled)) - 1 + for key, bits in masks.items(): + # Keep the whole-span fast path allocation-free. With position filters, + # two hooks share a plan only when both membership and windows agree. + windows = filters.get(key) + group = (bits, id(windows)) if windows else bits + if group not in interned: + # Membership is already resolved above. Do not construct a set of + # IDs, then scan and validate the same batch again for each group. + if not bits: + chosen = () + elif bits == all_rows: + chosen = scheduled + elif bits.bit_count() * 4 < len(scheduled): + chosen = _sparse_rows(scheduled, bits) + else: + chosen = (row for row, hooks in zip(scheduled, resolved) + if hooks is None or key in hooks) + interned[group] = _plan_selected_rows(chosen, total_rows, max_segments, windows, key in long_hooks) + plans[key] = interned[group] + return plans + + +# --------------------------------------------------------------------------- +# Packed (array) planner for declared monitor demand. +# +# A request's demand is fixed while its demand entry is unchanged, so its +# per-HookPoint row rule (HookPoint index + absolute window pieces) is +# compiled once and reused on every later step. Each step then derives the +# rows of every HookPoint for the whole batch with vectorised NumPy over the +# concatenated rules: clip to the scheduled token range, drop empty pieces, +# compact destination offsets and merge adjacent source runs. The plans +# expose exactly HookRows' fields, values and order; tuples are built only if +# a consumer reads them, and the native metadata rows of all HookPoints are +# built in one pass per step. HookRowPlanner/plan_hook_groups remain the +# reference planner for explicit selections. +# --------------------------------------------------------------------------- + +_UNBOUNDED = 1 << 62 +_EMPTY_I64 = _np.zeros(0, dtype=_np.int64) +EMPTY_HOOK_ROWS = HookRows((), 0, (), (), (), (), ()) + + +def row_signature(rows): + """Value identity of a plan's source segments. + + Two plans with equal signatures produce the same row table, so an + integration can skip rewriting (and re-uploading) a table whose + signature did not change.""" + sig = getattr(rows, "segment_signature", None) + return rows.segments if sig is None else sig + + +def segment_array(rows): + """A plan's ``(source row, rows)`` segments as an int64 ``[n, 2]`` array.""" + arr = getattr(rows, "segment_array", None) + if arr is not None: + return arr + return _np.asarray(rows.segments, dtype=_np.int64).reshape(-1, 2) + + +_hook_signature = row_signature +_segment_array = segment_array + + +class _PackedStep: + """One step's selected pieces for all HookPoints, sorted by HookPoint.""" + __slots__ = ("ids", "b", "lo", "hi", "dst", "src", "pos", "_rows") + + def __init__(self, ids, b, lo, hi, dst, src, pos): + self.ids, self.b, self.lo, self.hi = ids, b, lo, hi + self.dst, self.src, self.pos = dst, src, pos + self._rows = None + + def rows(self): + """Native metadata rows ``(rid, start, end, dst, 0, orig_position)``.""" + if self._rows is None: + ids = self.ids + self._rows = list(zip(map(ids.__getitem__, self.b.tolist()), self.lo.tolist(), + self.hi.tolist(), self.dst.tolist(), _repeat(0), self.pos.tolist())) + return self._rows + + +class PackedHookRows: + """One HookPoint's rows for one step: a slice of a shared _PackedStep. + + Field semantics are exactly HookRows'. ``full`` (total rows) marks the + un-compacted payload of a table-capacity or eager fallback: segments + cover all rows, offsets become source offsets. + """ + __slots__ = ("_step", "_r0", "_r1", "_seg", "row_count", "capacity_fallback", "_full", + "segment_signature", "_cache") + + def __init__(self, step, r0, r1, seg, row_count, capacity_fallback=False, full=None): + self._step, self._r0, self._r1, self._seg = step, r0, r1, seg + self.capacity_fallback = bool(capacity_fallback) + self._full = None if full is None else int(full) + self.row_count = int(row_count) if full is None else self._full + self.segment_signature = self.segment_array.tobytes() + self._cache = None + + @property + def segment_array(self): + if self._full is not None: + return _np.array([[0, self._full]], dtype=_np.int64) + return self._seg + + def _field(self, name): + cache = self._cache + if cache is None: + cache = self._cache = {} + value = cache.get(name) + if value is None: + s, r0, r1 = self._step, self._r0, self._r1 + if name == "segments": + value = tuple(map(tuple, self.segment_array.tolist())) + elif name == "req_ids": + value = tuple(map(s.ids.__getitem__, s.b[r0:r1].tolist())) + elif name == "token_ranges": + value = tuple(zip(s.lo[r0:r1].tolist(), s.hi[r0:r1].tolist())) + elif name == "offsets": + value = tuple((s.src if self._full is not None else s.dst)[r0:r1].tolist()) + elif name == "source_offsets": + value = tuple(s.src[r0:r1].tolist()) + elif name == "orig_positions": + value = tuple(s.pos[r0:r1].tolist()) + elif name == "layout": + if self._full is None: + value = s.rows()[r0:r1] + else: + value = list(zip(self.req_ids, s.lo[r0:r1].tolist(), s.hi[r0:r1].tolist(), + s.src[r0:r1].tolist(), _repeat(0), s.pos[r0:r1].tolist())) + cache[name] = value + return value + + segments = property(lambda self: self._field("segments")) + req_ids = property(lambda self: self._field("req_ids")) + token_ranges = property(lambda self: self._field("token_ranges")) + offsets = property(lambda self: self._field("offsets")) + source_offsets = property(lambda self: self._field("source_offsets")) + orig_positions = property(lambda self: self._field("orig_positions")) + + def layout(self): + """This HookPoint's native metadata rows (shared, do not mutate).""" + return self._field("layout") + + def layout_arrays(self): + """The same rows as int64 arrays for native ingestion: + (index into the step's request list, start, end, offset, orig_position).""" + s, r0, r1 = self._step, self._r0, self._r1 + return (s.b[r0:r1], s.lo[r0:r1], s.hi[r0:r1], + (s.src if self._full is not None else s.dst)[r0:r1], s.pos[r0:r1]) + + def full_payload(self, total_rows): + """Keep consumer selection, but describe an un-compacted fallback.""" + return PackedHookRows(self._step, self._r0, self._r1, self._seg, 0, + self.capacity_fallback, total_rows) + + def as_hook_rows(self): + return HookRows(self.segments, self.row_count, self.req_ids, self.token_ranges, + self.offsets, self.source_offsets, self.orig_positions, + self.capacity_fallback) + + def __eq__(self, other): + if isinstance(other, PackedHookRows): + other = other.as_hook_rows() + if not isinstance(other, HookRows): + return NotImplemented + return self.as_hook_rows() == other + + def __hash__(self): + return hash(self.as_hook_rows()) + + def __repr__(self): + return f"Packed{self.as_hook_rows()!r}" + + +def compile_request_rule(hooks, windows, key_index, validate_hook): + """Admission-time rule of one request: sorted (HookPoint index, start, end). + + ``hooks`` are the request's selected HookPoints (None = every gather key); + ``windows`` maps a HookPoint to raw intervals (absent = all positions, + [] = none). Validation and canonicalization match the reference install; + windows never enable an unselected HookPoint. + """ + selected = validate_hook(hooks) + if selected is None: + selected = frozenset(key_index) + parsed = {} + for hook, intervals in (windows or {}).items(): + (parsed_hook,) = validate_hook([hook]) + parsed[parsed_hook] = canonical_token_windows(intervals) + h, ws, we = [], [], [] + for hook in sorted(selected, key=key_index.__getitem__): + k = key_index[hook] + spans = parsed.get(hook) + for start, end in ((0, None),) if spans is None else spans: + h.append(k) + ws.append(start) + we.append(_UNBOUNDED if end is None else end) + return (_np.asarray(h, dtype=_np.int64), _np.asarray(ws, dtype=_np.int64), + _np.asarray(we, dtype=_np.int64)) + + +class PackedDemandPlanner: + """Per-step HookPoint row planning for declared (per-request) demand. + + Rules are cached per request by the identity of its demand entries, so a + request is compiled once per demand change (admission, terminal position), + not once per step. The concatenated, HookPoint-sorted batch rule is cached + by (request order, demand object); per step only the position clip, + compaction and segment merge run, as array operations. + The serving thread owns plan(), just as it owns the current batch layout. + """ + + def __init__(self, keys, validate_hook): + self.keys = tuple(keys) + self.key_index = {key: i for i, key in enumerate(self.keys)} + if len(self.keys) >= 1 << 15: + raise ValueError("too many HookPoints for the packed planner") + self._validate = validate_hook + # rid -> (hooks entry, windows entry, pool start, piece count). Rules + # live in append-only pools, compacted when mostly dead. + self._rules = {} + self._pool = _np.zeros((3, 1024), dtype=_np.int64) # rows: HookPoint index, start, end + self._pool_used = 0 + self._live = 0 + self._batch_key = None + self._batch = None + self.present_keys = [] + self.group_builds = 0 + self.group_reuses = 0 + self.rule_compiles = 0 + + def _compile(self, rid, hooks, windows): + h, ws, we = compile_request_rule(hooks, windows, self.key_index, self._validate) + self.rule_compiles += 1 + n = len(h) + if self._pool_used + n > self._pool.shape[1]: + # Growing keeps every existing position valid. + grown = _np.zeros((3, max(2 * self._pool.shape[1], self._pool_used + n)), dtype=_np.int64) + grown[:, :self._pool_used] = self._pool[:, :self._pool_used] + self._pool = grown + start = self._pool_used + self._pool[0, start:start + n] = h + self._pool[1, start:start + n] = ws + self._pool[2, start:start + n] = we + self._pool_used = start + n + old = self._rules.get(rid) + self._live += n - (old[3] if old is not None else 0) + entry = self._rules[rid] = (hooks, windows, start, n) + return entry + + def _maintain(self, req_ids): + """Between batches: drop rules of departed requests, compact the pool.""" + rules = self._rules + if len(rules) > 1024 and len(rules) > 4 * len(req_ids): + keep = {rid: rules[rid] for rid in req_ids if rid in rules} + rules.clear() + rules.update(keep) + self._live = sum(entry[3] for entry in keep.values()) + if self._pool_used > 4096 and self._pool_used > 2 * self._live: + old, used = self._pool, 0 + pool = _np.zeros((3, max(1024, 2 * self._live)), dtype=_np.int64) + for rid, (hooks, windows, start, n) in list(rules.items()): + pool[:, used:used + n] = old[:, start:start + n] + rules[rid] = (hooks, windows, used, n) + used += n + self._pool, self._pool_used = pool, used + + def _build_batch(self, req_ids, rows, windows, default): + self._maintain(req_ids) + rules, compile_ = self._rules, self._compile + get_rows, get_windows = rows.get, windows.get + starts, counts = [], [] + for rid in req_ids: + hooks, spans = get_rows(rid, default), get_windows(rid) + entry = rules.get(rid) + if entry is None or entry[0] is not hooks or entry[1] is not spans: + entry = compile_(rid, hooks, spans) + starts.append(entry[2]) + counts.append(entry[3]) + n = len(starts) + count = _np.fromiter(counts, dtype=_np.int64, count=n) + total = int(count.sum()) if n else 0 + if not total: + return _EMPTY_I64, _EMPTY_I64, _EMPTY_I64, _EMPTY_I64 + first = _np.fromiter(starts, dtype=_np.int64, count=n) + excl = _np.cumsum(count) - count + index = _np.repeat(first - excl, count) + _np.arange(total, dtype=_np.int64) + b = _np.repeat(_np.arange(n, dtype=_np.int64), count) + pool_h, pool_ws, pool_we = self._pool + # Stable (radix on int16): within a HookPoint, batch order then window order. + order = _np.argsort(pool_h[index].astype(_np.int16), kind="stable") + index = index[order] + return pool_h[index], b[order], pool_ws[index], pool_we[index] + + def plan(self, req_ids, token_ranges, offsets, orig_positions, total_rows, max_segments, + rows, windows, default=frozenset(), demand_identity=None): + n = len(req_ids) + if not (n == len(token_ranges) == len(offsets) == len(orig_positions)): + raise ValueError("request row metadata lengths differ") + ids = req_ids if isinstance(req_ids, tuple) else tuple(req_ids) + if n: + tr = _np.fromiter(_chain.from_iterable(token_ranges), dtype=_np.int64, + count=2 * n).reshape(n, 2) + tr_lo, tr_hi = tr[:, 0], tr[:, 1] + off = _np.fromiter(offsets, dtype=_np.int64, count=n) + length = tr_hi - tr_lo + prev_end = _np.empty(n, dtype=_np.int64) + prev_end[0] = 0 + _np.add(off[:-1], length[:-1], out=prev_end[1:]) + if ((length <= 0) | (off < prev_end) | (off + length > total_rows)).any(): + raise ValueError("invalid or overlapping scheduled request rows") + key = None if demand_identity is None else (ids, demand_identity) + if key is not None and key == self._batch_key: + self.group_reuses += 1 + h, b, ws, we = self._batch + else: + h, b, ws, we = self._build_batch(ids, rows, windows, default) + self._batch_key, self._batch = key, (h, b, ws, we) + self.group_builds += 1 + plans = dict.fromkeys(self.keys, EMPTY_HOOK_ROWS) + # Keys given a non-empty plan this step (every other key maps to + # EMPTY_HOOK_ROWS): callers derive the enabled set in O(present). + self.present_keys = present = [] + if not len(h): + return plans + lo = _np.maximum(ws, tr_lo[b]) + hi = _np.minimum(we, tr_hi[b]) + keep = _np.flatnonzero(lo < hi) + if not len(keep): + return plans + if len(keep) != len(h): + h, b, lo, hi = h[keep], b[keep], lo[keep], hi[keep] + size = hi - lo + src = off[b] + (lo - tr_lo[b]) + pos = _np.fromiter(orig_positions, dtype=_np.int64, count=n)[b] + m = len(h) + excl = _np.cumsum(size) + excl -= size + hook_change = _np.empty(m, dtype=bool) + hook_change[0] = True + _np.not_equal(h[1:], h[:-1], out=hook_change[1:]) + first = _np.flatnonzero(hook_change) # first piece of each present HookPoint + bounds = _np.append(first, m) + dst = excl - _np.repeat(excl[first], _np.diff(bounds)) + new_seg = hook_change.copy() + new_seg[1:] |= src[1:] != (src[:-1] + size[:-1]) + seg_idx = _np.flatnonzero(new_seg) + seg = _np.empty((len(seg_idx), 2), dtype=_np.int64) + seg[:, 0] = src[seg_idx] + seg[:, 1] = _np.add.reduceat(size, seg_idx) + seg_bounds = _np.searchsorted(seg_idx, bounds).tolist() + totals = _np.add.reduceat(size, first).tolist() + step = _PackedStep(ids, b, lo, hi, dst, src, pos) + bl = bounds.tolist() + keys = self.keys + coarse, interned, made = {}, {}, {} + for j, k in enumerate(h[first].tolist()): + r0, r1, s0, s1 = bl[j], bl[j + 1], seg_bounds[j], seg_bounds[j + 1] + # Equal (request, interval) pieces imply equal offsets and segments: + # such HookPoints share one plan. Compare contents only when the + # piece/row/segment counts already agree. + plan, shape = None, (r1 - r0, totals[j], s1 - s0) + candidates = coarse.get(shape) + if candidates is not None: + sig = (b[r0:r1].tobytes(), lo[r0:r1].tobytes(), hi[r0:r1].tobytes()) + for other in candidates: + if other not in interned: + o0, o1 = bl[other], bl[other + 1] + interned[other] = (b[o0:o1].tobytes(), lo[o0:o1].tobytes(), hi[o0:o1].tobytes()) + if interned[other] == sig: + plan = made[other] + break + if plan is None: + if s1 - s0 > max_segments: + plan = PackedHookRows(step, r0, r1, seg[s0:s1], 0, True, total_rows) + else: + plan = PackedHookRows(step, r0, r1, seg[s0:s1], totals[j]) + coarse.setdefault(shape, []).append(j) + made[j] = plan + plans[keys[k]] = plan + present.append(keys[k]) + return plans + + +def full_row_plan(req_ids, token_ranges, offsets, orig_positions, total_rows): + """Static full-row baseline: the ungathered producer's payload. + + The plan copies the whole physical tensor, ``((0, total_rows),)``, so a + gather producer fed this plan writes exactly what an ungathered producer + writes -- including leading, interior and trailing padding rows -- and + the destination offsets are the step's physical ``offsets``: each + request's rows are ``payload[off : off + (hi - lo)]``, the same rows the + ungathered consumer attributes. Used to compare request-row gather + against plain full-row capture without changing the consumer. + + A step with no scheduled request (or no rows) is the empty plan: no + segments, no rows and nothing to attribute. + """ + scheduled = _checked_layout(req_ids, token_ranges, offsets, orig_positions, total_rows) + plan = _plan_selected_rows(scheduled, total_rows, max_segments=1 << 30, + token_windows=None) + if not plan.segments: + return plan + return plan.full_payload(total_rows) + + +# HookPoint feature state armed by configure_hook_gather/configure_hook_heads. +# It belongs to one attachment: install_ring_hooks and uninstall_ring_hooks +# clear it (clear_hook_row_features), so a detached or re-attached HookPoint +# never runs a previous attachment's gather table or heads. +_GATHER_SAVED_STRIP_ROW_BYTES = "_gather_saved_strip_row_bytes" +_HEAD_DEFAULTS = ( + ("_head_hook_type", None), + ("_head_hook_id", None), + ("_head_table", None), + ("_head_arena", None), + ("_head_acc", None), + ("_head_slots", 0), + ("_head_wbf16", 1), +) + + +def clear_hook_gather(hook_point): + """Disarm a HookPoint's gather producer and restore its strip stride. + + ``configure_hook_gather`` overwrites ``_strip_row_bytes`` with the gather + row stride (a chunked strip keeps 0 there); the value it replaced is + restored so the HookPoint returns to its own static/prefix/chunked mode. + Releases the table reference. Idempotent.""" + saved = getattr(hook_point, _GATHER_SAVED_STRIP_ROW_BYTES, None) + if saved is not None: + hook_point._strip_row_bytes = saved + try: + delattr(hook_point, _GATHER_SAVED_STRIP_ROW_BYTES) + except AttributeError: # pragma: no cover - exotic attribute store + setattr(hook_point, _GATHER_SAVED_STRIP_ROW_BYTES, None) + hook_point._gather_table = None + + +def clear_hook_heads(hook_point): + """Disarm a HookPoint's in-graph heads, releasing the table, arena and + accumulator references. Idempotent.""" + for name, default in _HEAD_DEFAULTS: + setattr(hook_point, name, default) + + +def clear_hook_row_features(hook_point): + """Clear every gather/head feature of a HookPoint's attachment.""" + clear_hook_heads(hook_point) + clear_hook_gather(hook_point) + + +def configure_hook_gather(hook_point, table, row_bytes): + """Arm an installed HookPoint's request-row gather producer. + + ``table`` is the HookPoint's stable int64 row-table view (the layout + above, ``1 + 2 * max_segments`` words) and ``row_bytes`` the stride of + one packed row. Configure before graph capture: a captured producer + keeps the table's address and the stride. ``table=None`` disarms the + gather producer (``clear_hook_gather``). The configuration lasts for + the current attachment: ``install_ring_hooks`` and + ``uninstall_ring_hooks`` clear it.""" + import torch + if table is None: + clear_hook_gather(hook_point) + return + if table.dtype != torch.int64 or table.dim() != 1 or not table.is_contiguous(): + raise ValueError("a row table is a contiguous 1-D int64 tensor") + if table.numel() < 1: + raise ValueError("a row table holds at least its segment count") + if int(row_bytes) <= 0: + raise ValueError("gather needs a positive row stride") + if getattr(hook_point, _GATHER_SAVED_STRIP_ROW_BYTES, None) is None: + setattr(hook_point, _GATHER_SAVED_STRIP_ROW_BYTES, + int(getattr(hook_point, "_strip_row_bytes", 0))) + hook_point._gather_table = table + hook_point._strip_row_bytes = int(row_bytes) + + +def configure_hook_heads(hook_point, *, head_type, layer_no, table, arena, acc, + slots, weight_bf16): + """Arm in-graph projection heads on an installed HookPoint. + + The heads score the rows selected by ``table`` under their own toggle key + ``(head_type, layer_no)``. ``arena`` is the node's int32 weight arena + (``head_node_words`` words), ``acc`` the shared fp32 ``[rows, 16]`` + multi-layer accumulator, ``slots`` the head count ``K`` and + ``weight_bf16`` the arena's weight dtype. + + Dispatch: with a gather table and a row spec the rows and the heads run + as one fused node; otherwise the HookPoint's own row producer (static, + prefix or chunked, when it has a row spec) launches first and a + head-only node follows, so the published tasks always match the spec + order (row, then head). The configuration lasts for the current + attachment (see ``configure_hook_gather``).""" + import torch + if not 0 < int(slots) <= 16: + raise ValueError("in-graph heads take 1..16 slots") + if arena.dtype != torch.int32 or acc.dtype != torch.float32: + raise ValueError("head arenas are int32 words and accumulators fp32") + if (table is None or table.dtype != torch.int64 or table.dim() != 1 + or not table.is_contiguous() or table.numel() < 1): + raise ValueError("a head row table is a contiguous 1-D int64 tensor " + "holding at least its segment count") + hook_point._head_hook_type = int(head_type) + hook_point._head_hook_id = int(layer_no) + hook_point._head_table = table + hook_point._head_arena = arena + hook_point._head_acc = acc + hook_point._head_slots = int(slots) + hook_point._head_wbf16 = int(bool(weight_bf16)) + + +def encode_row_table(rows, max_segments): + """The producer row table of one plan: ``[n, src0, len0, src1, len1, + ...]`` padded to ``1 + 2 * max_segments`` int64 words. ``rows=None`` + disables the producer for the step (``n = -1``).""" + width = 1 + 2 * max_segments + table = _np.zeros(width, dtype=_np.int64) + if rows is None: + table[0] = -1 + return table + segments = _segment_array(rows) + if len(segments) > max_segments: + raise ValueError("plan has more segments than the row table holds") + table[0] = len(segments) + table[1:1 + 2 * len(segments)] = segments.reshape(-1) + return table + + +__all__ = [ + "EMPTY_HOOK_ROWS", + "GATHER_MAX_SEGS", + "HookRowPlanner", + "HookRows", + "PackedDemandPlanner", + "PackedHookRows", + "TokenWindowIndex", + "canonical_token_windows", + "clear_hook_gather", + "clear_hook_heads", + "clear_hook_row_features", + "compile_request_rule", + "compile_token_windows", + "configure_hook_gather", + "configure_hook_heads", + "encode_row_table", + "full_row_plan", + "plan_hook_groups", + "plan_hook_rows", + "row_signature", + "segment_array", +] diff --git a/src/dmi/hooks/catalog.py b/src/dmi/hooks/catalog.py index 4d2224686..cc13b59ee 100644 --- a/src/dmi/hooks/catalog.py +++ b/src/dmi/hooks/catalog.py @@ -15,6 +15,7 @@ SHAPE_HIDDEN, SHAPE_QKV_Q, SHAPE_QKV_KV, SHAPE_QKV_Z = 0, 1, 2, 3 SHAPE_ATTN_WT, SHAPE_MLP_POST, SHAPE_TOKEN_IDS, SHAPE_LOGITS = 4, 5, 6, 7 SHAPE_ROUTER_LOGITS, SHAPE_TOPK_IDS, SHAPE_TOPK_WEIGHTS = 8, 9, 10 +SHAPE_HEAD_SCORES = 11 # Pipeline-parallel placement. PP_ANY, PP_FIRST, PP_LAST = 0, 1, 2 @@ -45,6 +46,8 @@ (21, "mlp.hook_router_logits", "router_logits", True, GROUP_OTHER, False, SHAPE_ROUTER_LOGITS, PP_ANY), (22, "mlp.hook_topk_ids", "topk_ids", True, GROUP_OTHER, False, SHAPE_TOPK_IDS, PP_ANY), (23, "mlp.hook_topk_weights", "topk_weights", True, GROUP_OTHER, False, SHAPE_TOPK_WEIGHTS, PP_ANY), + (24, "hook_resid_pre_head", "resid_pre_head", True, GROUP_OTHER, False, SHAPE_HEAD_SCORES, PP_ANY), + (25, "hook_final_ln_head", "final_ln_head", False, GROUP_OTHER, False, SHAPE_HEAD_SCORES, PP_LAST), ) @@ -64,6 +67,7 @@ "SHAPE_ROUTER_LOGITS", "SHAPE_TOPK_IDS", "SHAPE_TOPK_WEIGHTS", + "SHAPE_HEAD_SCORES", "PP_ANY", "PP_FIRST", "PP_LAST", diff --git a/src/dmi/hooks/dispatch.py b/src/dmi/hooks/dispatch.py index 69f57c2ce..c62fbe5d2 100644 --- a/src/dmi/hooks/dispatch.py +++ b/src/dmi/hooks/dispatch.py @@ -7,7 +7,8 @@ import torch -from .specs import HookSpec +from ..hook_rows import clear_hook_row_features +from .specs import HEAD_TYPES, HookSpec def dispatch_producer( @@ -17,10 +18,20 @@ def dispatch_producer( strip_row_bytes: int, hook_type: int, hook_id: int, + gather_table: Optional[torch.Tensor] = None, ) -> None: - """Dispatch a tensor to the producer matching its configured strip mode.""" + """Dispatch a tensor to the producer matching its configured mode. - if strip_tensor is None: + A per-HookPoint row table (gather mode) takes precedence over strip + mode: the table subsumes the prefix strip (padding rows are never + listed), and ``strip_row_bytes`` doubles as the gather row stride. + """ + + if gather_table is not None: + torch.ops.ring.producer_gather( + ring_payload, tensor, gather_table, strip_row_bytes, hook_type, hook_id + ) + elif strip_tensor is None: torch.ops.ring.producer(ring_payload, tensor, hook_type, hook_id) elif strip_row_bytes > 0: torch.ops.ring.producer_prefix( @@ -36,17 +47,38 @@ def install_ring_hooks( specs: Sequence[HookSpec], ring_payload: Optional[torch.Tensor] = None, ) -> None: - """Bind hook specifications to their executable hook-point modules.""" + """Bind hook specifications to their executable hook-point modules. + + Installing starts a new attachment: any gather table or in-graph heads a + previous attachment configured on these HookPoints are cleared (and the + strip stride gather replaced is restored), so a plain re-attach runs the + HookPoint's ordinary producer. Integrations re-arm gather/heads after + installing (``dmi.hook_rows.configure_hook_gather`` / + ``configure_hook_heads``). + + An in-graph head spec (``HEAD_TYPES``) names the HookPoint whose rows it + scores but is not that HookPoint's row producer: it binds the payload + only and never replaces the row key, so a ``[row, head]`` spec pair keeps + the row producer and a lone head spec leaves a head-only HookPoint. + """ for spec in specs: - hook_point = spec.module - if hook_point is None: + if spec.module is None: raise RuntimeError( "install_ring_hooks received an unbound model-wide HookSpec" ) + for spec in specs: + # A new attachment: drop the previous row key and gather/head state, + # so a HookPoint that now receives only a head spec is head-only. + spec.module._ring_hook_type = None + clear_hook_row_features(spec.module) + for spec in specs: + hook_point = spec.module + hook_point._ring_payload = ring_payload + if spec.hook_type in HEAD_TYPES: + continue hook_point._ring_hook_type = spec.hook_type hook_point._ring_hook_id = spec.layer_no - hook_point._ring_payload = ring_payload def uninstall_ring_hooks(specs: Sequence[HookSpec]) -> None: @@ -70,6 +102,10 @@ def uninstall_ring_hooks(specs: Sequence[HookSpec]) -> None: a recompile. That is the price of not corrupting the ring on the next ordinary forward, and ``attach_model`` re-runs ``install_ring_hooks`` anyway, so the pair stays symmetric. + + Gather and in-graph head state (row tables, head arena and accumulator) + belongs to the attachment and is cleared here too, releasing those tensor + references and restoring the strip stride gather replaced. """ for spec in specs: hook_point = spec.module @@ -77,6 +113,7 @@ def uninstall_ring_hooks(specs: Sequence[HookSpec]) -> None: continue hook_point._ring_hook_type = None hook_point._ring_payload = None + clear_hook_row_features(hook_point) __all__ = ["dispatch_producer", "install_ring_hooks", "uninstall_ring_hooks"] diff --git a/src/dmi/hooks/point.py b/src/dmi/hooks/point.py index 3f202eb70..dfbe761b7 100644 --- a/src/dmi/hooks/point.py +++ b/src/dmi/hooks/point.py @@ -53,6 +53,18 @@ def _nvtx_range(message: str): yield +def _source_rows(x: Any) -> int: + """Physical source rows of a contiguous HookPoint tensor: the product of + every dimension before hidden (1 for a 1-D ``[hidden]`` tensor). The + native producers derive the same count as nbytes / (hidden x element + size).""" + shape = tuple(int(d) for d in x.shape) + rows = 1 + for d in shape[:-1]: + rows *= d + return rows + + @dataclass class LensHandle: """Dataclass that holds information about a PyTorch hook.""" @@ -148,6 +160,41 @@ def __init__(self): self._strip_tensor: Optional[torch.Tensor] = None self._strip_row_bytes: int = 0 + # Gather mode: when set, takes precedence over strip mode. The + # producer copies the row runs listed in this per-HookPoint row table + # (dmi.hook_rows.encode_row_table layout), packed, re-reading the + # VALUES each replay under the same stable-address contract as + # _strip_tensor; _strip_row_bytes is the row stride. A table whose + # first word is -1 disables the producer for the step. + # Armed by dmi.hook_rows.configure_hook_gather for one attachment: + # install_ring_hooks / uninstall_ring_hooks clear it. + self._gather_table: Optional[torch.Tensor] = None + + # In-graph projection heads. Set by the integration + # (dmi.hook_rows.configure_hook_heads, never by install_ring_hooks) + # for one attachment; install_ring_hooks / uninstall_ring_hooks + # clear them. K fp32 scores per row selected by _head_table, under + # their own spec and toggle key (_head_hook_type, _head_hook_id). + # Dispatch (forward): + # row spec + gather table + heads -> ring::producer_fused, the rows + # and the heads in one node; + # row spec (static / prefix / chunked) + heads + # -> the row producer, then a + # head-only producer_fused node; + # heads only (no row spec) -> the head-only node. + # Either way the row task is published before the head task: the + # spec (metadata) order. _head_arena is the node's int32 weight + # arena, _head_acc the shared fp32 [rows, 16] multi-layer + # accumulator, and _head_wbf16 the arena's weight dtype (a launch + # constant). + self._head_hook_type: Optional[int] = None + self._head_hook_id: Optional[int] = None + self._head_table: Optional[torch.Tensor] = None + self._head_arena: Optional[torch.Tensor] = None + self._head_acc: Optional[torch.Tensor] = None + self._head_slots: int = 0 + self._head_wbf16: int = 1 + def add_perma_hook(self, hook: HookFunction, dir: Literal["fwd", "bwd"] = "fwd") -> None: self.add_hook(hook, dir=dir, is_permanent=True) @@ -268,7 +315,9 @@ def forward(self, x: Tensor) -> Tensor: if transport is not None and not transport.capture_step: return x - if self._ring_hook_type is None or not x.is_cuda: + row_on = self._ring_hook_type is not None + head_on = self._head_hook_type is not None + if not (row_on or head_on) or not x.is_cuda: return x x_cont = x.contiguous() @@ -285,84 +334,164 @@ def forward(self, x: Tensor) -> Tensor: # never changes); torch.compile traces only the live branch. strip_t = self._strip_tensor strip_rb = self._strip_row_bytes + gather_t = self._gather_table if row_on else None ring_payload = self._ring_payload from ..transport import ring as _rt transport = _rt._active_transport - if transport is not None and transport.force_eager: - # Node toggle: the eager safety net fires producers outside any - # graph, where cudaGraphNodeSetEnabled cannot gate them. On a - # gated step that fell back to eager (a step larger than the - # ring), a disabled hook must be skipped here, or it fires while - # metadata and reservation counted only the enabled subset. - # This branch is never captured or compiled, so a host check is - # safe; toggle_gate_active keeps the stock eager path untouched. - if (getattr(transport, "toggle_gate_active", False) - and getattr(transport, "_gated_step", True) - and not transport._ring_engine.is_hook_enabled( - self._ring_hook_type, self._ring_hook_id)): - return x_cont + if (transport is not None and transport.force_eager + and transport._ring_engine is not None): + # The eager safety net publishes exactly what the fast path + # publishes, in the same order: the row task (when this + # HookPoint has a row spec), then the head task. engine = transport._ring_engine - if engine is not None: - nbytes = x_cont.nbytes - transport_bytes = align_up_py(nbytes, 16) - if strip_t is not None: - # CPU-direct bypasses ring/staging accounting. Flush older - # ring tasks first so the shared metadata FIFO stays in - # producer order. The legacy consumer applies its normal - # CPU-side request slicing to the full tensor. - engine.flush_and_wait() - transport.submit_cpu_direct( - x_cont.cpu(), - self._ring_hook_type, self._ring_hook_id) - return x_cont - # The ring's real per-step ceiling is min(payload, staging), - # not payload alone. The drain assembles each flush batch per - # WHOLE entry and breaks when the entry does not fit staging - # (drain_thread.cpp), so a tensor larger than staging that is - # admitted here is never drained: the capture is lost and its - # payload reservation is never released, permanently shrinking - # ring capacity, while flush_and_wait() still returns success. - # Same ceiling the native prepare_step/reserve_record use, and - # the one RingCapacities.effective_bytes publishes. Both caps - # are engine-lifetime constants queried over pybind, so the - # min is cached on the transport and shared by every hook -- - # querying the pair per hook would repeat it for the whole - # active set on the first eager forward. - # A free task entry is needed too: a step with more hooks - # than task entries lands here with the payload ring nearly - # empty, and reserving past task_cap makes a producer - # overwrite an unread slot. A flush frees both. - effective_cap = transport.effective_cap - if (transport_bytes <= min(engine.available_capacity(), - effective_cap) - and engine.available_task_slots() > 0): - engine.reserve_one(nbytes) - dispatch_producer(ring_payload, x_cont, strip_t, strip_rb, - self._ring_hook_type, self._ring_hook_id) - elif transport_bytes <= effective_cap: - engine.flush_and_wait() - engine.reserve_one(nbytes) - dispatch_producer(ring_payload, x_cont, strip_t, strip_rb, - self._ring_hook_type, self._ring_hook_id) - else: - # Larger than the ring's effective capacity (payload OR - # staging) -- bypass via cpu_direct. Flush first so submit_cpu_direct consumes - # the FIFO meta for THIS hook (prior ring entries finish - # first; this hook's pre-pushed meta becomes the head). - engine.flush_and_wait() - transport.submit_cpu_direct( - x_cont.cpu(), - self._ring_hook_type, self._ring_hook_id) - return x_cont + if row_on: + self._eager_row(transport, engine, ring_payload, x_cont, + strip_t, strip_rb, gather_t) + if head_on: + self._eager_head(transport, engine, ring_payload, x_cont) + return x_cont # Fast path: torch.compile-serializable, CUDA-graph captureable. # C++ ring_producer_* impls no-op when g_active_engine is null # (monitoring inactive); otherwise launch the producer kernel. - dispatch_producer(ring_payload, x_cont, strip_t, strip_rb, - self._ring_hook_type, self._ring_hook_id) + if head_on and gather_t is not None: + # One node for the rows and the heads; each part is gated by its + # own table row. + torch.ops.ring.producer_fused( + ring_payload, x_cont, gather_t, strip_rb, self._head_table, + self._head_arena, self._head_acc, self._head_slots, + self._head_wbf16, self._ring_hook_type, self._ring_hook_id, + self._head_hook_type) + return x_cont + if row_on: + dispatch_producer(ring_payload, x_cont, strip_t, strip_rb, + self._ring_hook_type, self._ring_hook_id, gather_t) + if head_on: + self._launch_head_only(ring_payload, x_cont) return x_cont + def _launch_head_only(self, ring_payload, x_cont): + """The head part alone: a producer_fused node without a row table, + recorded and toggled under the head key only, so the row producer + and the heads keep independent toggle states.""" + torch.ops.ring.producer_fused( + ring_payload, x_cont, None, 0, self._head_table, + self._head_arena, self._head_acc, self._head_slots, + self._head_wbf16, self._head_hook_type, self._head_hook_id, + self._head_hook_type) + + def _eager_row(self, transport, engine, ring_payload, x_cont, + strip_t, strip_rb, gather_t): + """Eager safety net for the row producer (never captured).""" + # Node toggle: the eager safety net fires producers outside any + # graph, where cudaGraphNodeSetEnabled cannot gate them. On a + # gated step that fell back to eager (a step larger than the + # ring), a disabled hook must be skipped here, or it fires while + # metadata and reservation counted only the enabled subset. + # This branch is never captured or compiled, so a host check is + # safe; toggle_gate_active keeps the stock eager path untouched. + if (getattr(transport, "toggle_gate_active", False) + and (getattr(transport, "_gated_step", True) + or getattr(transport, "_prefill_table_gate", False)) + and not engine.is_hook_enabled( + self._ring_hook_type, self._ring_hook_id)): + return + nbytes = x_cont.nbytes + transport_bytes = align_up_py(nbytes, 16) + # Gather mode reserves the full tensor here: on an eager step + # the integration writes the full-row table, so the producer + # copies exactly nbytes and the reservation matches. + if strip_t is not None and gather_t is None: + # CPU-direct bypasses ring/staging accounting. Flush older + # ring tasks first so the shared metadata FIFO stays in + # producer order. The legacy consumer applies its normal + # CPU-side request slicing to the full tensor. + engine.flush_and_wait() + transport.submit_cpu_direct( + x_cont.cpu(), + self._ring_hook_type, self._ring_hook_id) + return + # The ring's real per-step ceiling is min(payload, staging), + # not payload alone. The drain assembles each flush batch per + # WHOLE entry and breaks when the entry does not fit staging + # (drain_thread.cpp), so a tensor larger than staging that is + # admitted here is never drained: the capture is lost and its + # payload reservation is never released, permanently shrinking + # ring capacity, while flush_and_wait() still returns success. + # Same ceiling the native prepare_step/reserve_record use, and + # the one RingCapacities.effective_bytes publishes. Both caps + # are engine-lifetime constants queried over pybind, so the + # min is cached on the transport and shared by every hook -- + # querying the pair per hook would repeat it for the whole + # active set on the first eager forward. + # A free task entry is needed too: a step with more hooks + # than task entries lands here with the payload ring nearly + # empty, and reserving past task_cap makes a producer + # overwrite an unread slot. A flush frees both. + effective_cap = transport.effective_cap + if (transport_bytes <= min(engine.available_capacity(), + effective_cap) + and engine.available_task_slots() > 0): + engine.reserve_one(nbytes) + dispatch_producer(ring_payload, x_cont, strip_t, strip_rb, + self._ring_hook_type, self._ring_hook_id, + gather_t) + elif transport_bytes <= effective_cap: + engine.flush_and_wait() + engine.reserve_one(nbytes) + dispatch_producer(ring_payload, x_cont, strip_t, strip_rb, + self._ring_hook_type, self._ring_hook_id, + gather_t) + else: + # Larger than the ring's effective capacity (payload OR + # staging) -- bypass via cpu_direct. Flush first so submit_cpu_direct consumes + # the FIFO meta for THIS hook (prior ring entries finish + # first; this hook's pre-pushed meta becomes the head). + engine.flush_and_wait() + transport.submit_cpu_direct( + x_cont.cpu(), + self._ring_hook_type, self._ring_hook_id) + + def _eager_head(self, transport, engine, ring_payload, x_cont): + """Eager safety net for the head part (never captured). + + The same gate and reservation order as the row producer: a disabled + head reserves nothing; otherwise the integration wrote the full-row + head table on eager steps, so the heads write rows x K x 4 bytes. + + ``rows`` is every physical source row: the product of all dimensions + before hidden, as the native launcher derives it (nbytes / (hidden x + element size) of the contiguous source) and as the step planner + budgets the head spec's ``[..., K]`` shape. A batched ``[B, S, H]`` + HookPoint therefore has B x S rows, not B. + """ + if engine is None: + return + if (getattr(transport, "toggle_gate_active", False) + and (getattr(transport, "_gated_step", True) + or getattr(transport, "_prefill_table_gate", False)) + and not engine.is_hook_enabled(self._head_hook_type, + self._head_hook_id)): + return + rows = _source_rows(x_cont) + nbytes = rows * int(self._head_slots) * 4 + transport_bytes = align_up_py(nbytes, 16) + # As for the rows: an entry larger than min(payload, staging) is never + # drained and would wedge the ring. Heads have no CPU-direct path + # (the scores exist only on the GPU), so refuse before reserving. + if transport_bytes > transport.effective_cap: + raise RuntimeError( + f"in-graph head output ({transport_bytes} bytes for " + f"{rows} rows x {int(self._head_slots)} slots) " + f"exceeds the ring's effective capacity " + f"({transport.effective_cap} bytes)") + if (transport_bytes > engine.available_capacity() + or engine.available_task_slots() == 0): + engine.flush_and_wait() + engine.reserve_one(nbytes) + self._launch_head_only(ring_payload, x_cont) + def layer(self): # Returns the layer index if the name has the form 'blocks.{layer}.{...}' # Helper function that's mainly useful on HookedTransformer diff --git a/src/dmi/hooks/specs.py b/src/dmi/hooks/specs.py index f9efe3670..5debdabfb 100644 --- a/src/dmi/hooks/specs.py +++ b/src/dmi/hooks/specs.py @@ -22,6 +22,7 @@ PP_FIRST, PP_LAST, SHAPE_ATTN_WT, + SHAPE_HEAD_SCORES, SHAPE_HIDDEN, SHAPE_LOGITS, SHAPE_MLP_POST, @@ -109,6 +110,18 @@ class HookRowBasis(Enum): PP_LAST_ONLY: frozenset = frozenset( _id for _id, _act, _short, _pl, _grp, _tp, _sc, _pp in _HOOK_DEFS if _pp == PP_LAST ) +# In-graph projection head score rows (fp32 [rows, head_slots]). +HEAD_TYPES: frozenset = frozenset( + _id for _id, _act, _short, _pl, _grp, _tp, _sc, _pp in _HOOK_DEFS + if _sc == SHAPE_HEAD_SCORES +) +# Head hook type -> hook type of the HookPoint whose rows it scores. +HEAD_SOURCE_TYPE: Dict[int, int] = { + _id: _id_by_short[_short[:-len("_head")]] + for _id, _act, _short, _pl, _grp, _tp, _sc, _pp in _HOOK_DEFS + if _sc == SHAPE_HEAD_SCORES and _short.endswith("_head") + and _short[:-len("_head")] in _id_by_short +} def hook_row_basis(hook_type: int) -> HookRowBasis: """Return the canonical row basis derived from `_HOOK_DEFS.shape_class`. @@ -181,6 +194,7 @@ class ModelShapeConfig: top_k: int = 0 # topk_ids / topk_weights final dim tp_size: int = 1 # tensor parallel world size tp_rank: int = 0 # this rank's TP index + head_slots: int = 0 # K score slots per in-graph head node (0 = none) # --------------------------------------------------------------------------- @@ -237,6 +251,9 @@ def compute_hook_shape( if hook_type in _HIDDEN_DIM_TYPES: return b + [q_len, cfg.hidden_dim] + if hook_type in HEAD_TYPES: + k = getattr(cfg, "head_slots", 0) + return (b + [q_len, k]) if k > 0 else [] # The HOOK_TYPE_* names are bound by the globals() loop over # _HOOK_DEFS above, which static analysis cannot follow: hence the # F821 suppressions, one per line rather than one for the file, so the @@ -297,6 +314,8 @@ def compute_hook_shape( ) __all__ = [ + "HEAD_SOURCE_TYPE", + "HEAD_TYPES", "HookRowBasis", "HookSpec", "ModelShapeConfig", diff --git a/src/dmi/transport/native.py b/src/dmi/transport/native.py index 989cf7a23..0af261988 100644 --- a/src/dmi/transport/native.py +++ b/src/dmi/transport/native.py @@ -34,6 +34,7 @@ "RingEngine", "ring_set_active_engine", "ring_clear_active_engine", + "parse_hook_requests", } ) _NATIVE_EXPORTS = _HOST_EXPORTS | _RING_EXPORTS diff --git a/src/dmi/transport/ring.py b/src/dmi/transport/ring.py index 4838d68db..dd68ff34b 100644 --- a/src/dmi/transport/ring.py +++ b/src/dmi/transport/ring.py @@ -78,6 +78,32 @@ def _ring_producer_chunked_fake( ) -> None: return None + @torch.library.register_fake("ring::producer_gather") + def _ring_producer_gather_fake( + ring_payload: torch.Tensor, tensor: torch.Tensor, + seg_table: torch.Tensor, row_bytes: int, + hook_type: int, hook_id: int, + ) -> None: + return None + + @torch.library.register_fake("ring::producer_head") + def _ring_producer_head_fake( + ring_payload: torch.Tensor, tensor: torch.Tensor, + seg_table: torch.Tensor, arena: torch.Tensor, K: int, + hook_type: int, hook_id: int, + ) -> None: + return None + + @torch.library.register_fake("ring::producer_fused") + def _ring_producer_fused_fake( + ring_payload: torch.Tensor, tensor: torch.Tensor, + gather_table: Optional[torch.Tensor], row_bytes: int, + head_table: Optional[torch.Tensor], + arena: torch.Tensor, acc: torch.Tensor, K: int, weight_bf16: int, + hook_type: int, hook_id: int, head_hook_type: int, + ) -> None: + return None + @torch.library.register_fake("ring::record_producer") def _record_producer_fake( ring_payload: torch.Tensor, @@ -222,6 +248,11 @@ def __init__(self, ring_engine: Any) -> None: self._current_token_ranges: Optional[List[Tuple[int, int]]] = None self._current_dim0_offsets: Optional[List[int]] = None self._current_kv_offsets: Optional[List[int]] = None + self._current_orig_positions: Optional[List[int]] = None + # Per-HookPoint request-row plans (dmi.hook_rows) for this step, + # keyed by (hook_type, layer_no); None = every hook uses the step + # layout. + self._current_hook_rows: Optional[dict] = None # When True: meta pushes are skipped so the FIFO stays empty. # Producer kernel still fires (for CUDA graph capture) but as no-ops. @@ -265,6 +296,11 @@ def __init__(self, ring_engine: Any) -> None: # the full active set fires). Default True so a caller that only # replays toggled graphs keeps gating. self._gated_step: bool = True + # Opt-in integration guarantee: every producer reads a per-HookPoint + # row table whose disabled marker (-1) gates eager and partial-graph + # execution as well as toggled graphs. The per-step spec set then + # follows the toggle on every step. + self._prefill_table_gate: bool = False # Hook selection preset name (e.g. "full", "hidden-states", "logits"). # Set by the active adapter before hook installation. @@ -300,6 +336,8 @@ def set_step_context( ep_rank: int = 0, pp_rank: int = 0, flattened: bool = False, + orig_positions: Optional[List[int]] = None, + hook_rows: Optional[dict] = None, ) -> None: """Called before each forward pass to provide per-step batch metadata. @@ -315,6 +353,14 @@ def set_step_context( Static-cache batched / packed: 0 (real keys at the start). None = auto-generate zeros. flattened: False = batched [batch, q_len, ...], True = packed [total_tokens, ...]. + orig_positions: each request's index in the step's UNFILTERED + request list, when the integration filtered the lists. + Position-indexed hooks (final_logits) select rows by it. + None = unfiltered. + hook_rows: per-HookPoint request-row plans (``dmi.hook_rows``) + keyed by ``(hook_type, layer_no)``. A hook with a plan + publishes ``row_count`` rows and its own request layout; others + use the step layout. """ self._current_model_id = model_id self._current_tp_rank = tp_rank @@ -332,6 +378,8 @@ def set_step_context( kv_offsets if kv_offsets is not None else [0] * len(req_ids) ) + self._current_orig_positions = orig_positions + self._current_hook_rows = hook_rows def set_model_cfg(self, cfg: ModelShapeConfig) -> None: """Set the model shape config for analytical shape computation.""" @@ -355,7 +403,8 @@ def specs_for_step(self) -> "List[HookSpec]": publication THIS step: ``effective_specs`` on a gated step, the full active set on an ungated (eager or prefill) step. Identical when the toggle is inactive.""" - if getattr(self, "_gated_step", True): + if (getattr(self, "_gated_step", True) + or getattr(self, "_prefill_table_gate", False)): return self.effective_specs return self._active_specs @@ -368,10 +417,23 @@ def set_gated_step(self, gated: bool) -> None: active set fires).""" self._gated_step = bool(gated) + def set_prefill_table_gate(self, enabled: bool) -> None: + """Declare (or withdraw) the per-HookPoint table-gate guarantee. + + An integration enables this only after every installed producer + reads its own row table (``dmi.hook_rows.configure_hook_gather``): + a disabled table row (``-1``) then gates eager and partial-graph + producers too, so the per-step spec set follows the toggle on every + step.""" + self._prefill_table_gate = bool(enabled) + def gated_step_requires_registry(self) -> bool: """True when this step's producers must all be toggle-managed, i.e. - a graph excluded from the registry must not replay.""" - return bool(getattr(self, "_gated_step", True)) + a graph excluded from the registry must not replay. With the + per-HookPoint table gate armed, the -1 table rows gate such graphs' + producers, so they may replay on any step.""" + return (bool(getattr(self, "_gated_step", True)) + and not getattr(self, "_prefill_table_gate", False)) def _toggle_ctrl(self) -> "NodeToggleController": if self._toggle is None: @@ -447,10 +509,14 @@ def pre_push_all_metas(self, batch: int, q_len: int, kv_dim: int, flags = [] # The per-step spec set: the metadata published here, the device # node state and the capacity reservation stay in lockstep. + hook_rows = getattr(self, "_current_hook_rows", None) or {} for spec in self.specs_for_step(): spec_q_len = (actual_q_len if actual_q_len is not None and spec.dim0_is_actual_tokens else q_len) + rows = hook_rows.get((spec.hook_type, spec.layer_no)) + if rows is not None: + spec_q_len = rows.row_count shape = compute_hook_shape( spec.hook_type, self._model_cfg, batch, spec_q_len, kv_dim, logits_to_keep=logits_to_keep, @@ -470,6 +536,14 @@ def pre_push_all_metas(self, batch: int, q_len: int, kv_dim: int, flags.append(1 if spec.allow_token_cnt_mismatch else 0) if hook_types: + extra = {} + orig_positions = getattr(self, "_current_orig_positions", None) + if orig_positions is not None: + extra["orig_positions"] = list(orig_positions) + if hook_rows: + extra["hook_requests"] = self._hook_request_layouts( + hook_types, layer_nos, hook_rows) + # One submission keeps one step context for the whole forward. self._ring_engine.push_all_metas( hook_types, layer_nos, shapes, dtypes, flags, self._current_model_id, @@ -482,8 +556,37 @@ def pre_push_all_metas(self, batch: int, q_len: int, kv_dim: int, list(self._current_token_ranges), list(self._current_dim0_offsets), list(self._current_kv_offsets) if self._current_kv_offsets else [], + **extra, ) + def _hook_request_layouts(self, hook_types, layer_nos, hook_rows): + """Per-hook request layouts for push_all_metas, in metadata order. + + Equal plans (the same object) share one layout object, which the + native side parses once per step. Array-backed plans hand over + int64 arrays instead of one Python tuple per row. + """ + arrays = getattr(self, "_hook_request_arrays", None) + if arrays is None: + probe = getattr(self._ring_engine, "supports_hook_request_arrays", None) + arrays = self._hook_request_arrays = bool(probe is not None and probe()) + layouts = [] + shared = {} + for key in zip(hook_types, layer_nos): + rows = hook_rows.get(key) + if rows is None: + layouts.append(None) + continue + identity = id(rows) + if identity not in shared: + build = getattr(rows, "layout_arrays" if arrays else "layout", None) + shared[identity] = build() if build is not None else [ + (rid, tr[0], tr[1], off, 0, pos) + for rid, tr, off, pos in zip(rows.req_ids, rows.token_ranges, + rows.offsets, rows.orig_positions)] + layouts.append(shared[identity]) + return layouts + def submit_cpu_direct(self, cpu_tensor: torch.Tensor, hook_type: int, hook_id: int) -> None: """Submit a CPU-tensor to the drain -> p2p pipeline. diff --git a/tests/native/ring/Makefile b/tests/native/ring/Makefile index ee32064eb..3f7246b74 100644 --- a/tests/native/ring/Makefile +++ b/tests/native/ring/Makefile @@ -68,10 +68,14 @@ TARGETS := $(BUILD)/test_rings $(BUILD)/test_producer \ $(BUILD)/test_ring_engine $(BUILD)/test_null_mode \ $(BUILD)/test_record_consumer \ $(BUILD)/test_clickhouse_record_sink \ - $(BUILD)/test_toggle_registry + $(BUILD)/test_toggle_registry \ + $(BUILD)/test_producer_equivalence \ + $(BUILD)/test_inline_head \ + $(BUILD)/test_fused_head .PHONY: all test_rings test_producer test_ring_engine test_null_mode \ test_record_consumer test_clickhouse_record_sink test_toggle_registry \ + test_producer_equivalence test_inline_head test_fused_head \ run clean FORCE \ bench_producer_grid @@ -101,6 +105,17 @@ $(BUILD)/test_ring_engine: test_ring_engine.cu $(PRODUCER_CU) $(DRAIN_SRC) \ $(BUILD)/test_toggle_registry: test_toggle_registry.cu $(TOGGLE_REGISTRY_CU) $(CUDA_CONFIG_STAMP) | $(BUILD) $(CUDA_NVCC) $(NVCC_FLAGS) $(filter-out $(CUDA_CONFIG_STAMP),$^) -o $@ +# Row-table producers against independent host models. The head tests take +# NVCC_EXTRA=-DDMI_TEST_MUTATE_* for mutation checks (each must FAIL). +$(BUILD)/test_producer_equivalence: test_producer_equivalence.cu $(PRODUCER_CU) $(TOGGLE_REGISTRY_CU) $(CUDA_CONFIG_STAMP) | $(BUILD) + $(CUDA_NVCC) $(NVCC_FLAGS) $(NVCC_EXTRA) $(filter-out $(CUDA_CONFIG_STAMP),$^) -o $@ + +$(BUILD)/test_inline_head: test_inline_head.cu $(PRODUCER_CU) $(CUDA_CONFIG_STAMP) | $(BUILD) + $(CUDA_NVCC) $(NVCC_FLAGS) $(NVCC_EXTRA) $(filter-out $(CUDA_CONFIG_STAMP),$^) -o $@ + +$(BUILD)/test_fused_head: test_fused_head.cu $(PRODUCER_CU) $(CUDA_CONFIG_STAMP) | $(BUILD) + $(CUDA_NVCC) $(NVCC_FLAGS) $(NVCC_EXTRA) $(filter-out $(CUDA_CONFIG_STAMP),$^) -o $@ + $(BUILD)/test_null_mode: test_null_mode.cu $(PRODUCER_CU) $(CUDA_CONFIG_STAMP) | $(BUILD) $(CUDA_NVCC) $(NVCC_FLAGS) $(filter-out $(CUDA_CONFIG_STAMP),$^) -o $@ @@ -135,6 +150,15 @@ test_clickhouse_record_sink: $(BUILD)/test_clickhouse_record_sink test_toggle_registry: $(BUILD)/test_toggle_registry $(BUILD)/test_toggle_registry +test_producer_equivalence: $(BUILD)/test_producer_equivalence + $(BUILD)/test_producer_equivalence + +test_inline_head: $(BUILD)/test_inline_head + $(BUILD)/test_inline_head + +test_fused_head: $(BUILD)/test_fused_head + $(BUILD)/test_fused_head + # Manual microbenchmark; not part of `all` or `run`. $(BUILD)/bench_producer_grid: bench_producer_grid.cu $(PRODUCER_CU) $(CUDA_CONFIG_STAMP) | $(BUILD) $(CUDA_NVCC) $(NVCC_FLAGS) bench_producer_grid.cu -o $@ @@ -150,6 +174,9 @@ run: all $(BUILD)/test_record_consumer $(BUILD)/test_clickhouse_record_sink $(BUILD)/test_toggle_registry + $(BUILD)/test_producer_equivalence + $(BUILD)/test_inline_head + $(BUILD)/test_fused_head clean: rm -rf $(BUILD) diff --git a/tests/native/ring/cpu_cuda_stub/ATen/cuda/CUDAContext.h b/tests/native/ring/cpu_cuda_stub/ATen/cuda/CUDAContext.h new file mode 100644 index 000000000..c80c8e30c --- /dev/null +++ b/tests/native/ring/cpu_cuda_stub/ATen/cuda/CUDAContext.h @@ -0,0 +1,24 @@ +// CPU-only stand-in for , enough for +// ring/ring_torch_op.cpp: the current-stream query its producer ops make. +// Placed ahead of the Torch include directories, it shadows the real header, +// so the ops compile against CPU-only Torch without the CUDA toolkit. No +// CUDA library is linked or initialized. +#pragma once +#include +#include + +namespace at::cuda { + +// Device index the last getCurrentCUDAStream() asked for. +inline c10::DeviceIndex fake_last_stream_device = -1; + +struct FakeCUDAStream { + cudaStream_t stream() const { return reinterpret_cast(0x5); } +}; + +inline FakeCUDAStream getCurrentCUDAStream(c10::DeviceIndex device = -1) { + fake_last_stream_device = device; + return {}; +} + +} // namespace at::cuda diff --git a/tests/native/ring/cpu_cuda_stub/cuda_runtime.h b/tests/native/ring/cpu_cuda_stub/cuda_runtime.h index 08f6e8e91..ba02b4660 100644 --- a/tests/native/ring/cpu_cuda_stub/cuda_runtime.h +++ b/tests/native/ring/cpu_cuda_stub/cuda_runtime.h @@ -1,7 +1,8 @@ // CPU-only stand-in for the CUDA runtime API used by the host half of -// ring/toggle_registry.cu. No CUDA library is linked or initialized: each -// call records what a real exec would hold so tests can compare the -// registry's view with the exec's actual node state. +// ring/toggle_registry.cu and by ring/ring_torch_op.cpp's producer ops. No +// CUDA library is linked or initialized: each call records what a real exec +// would hold so tests can compare the registry's view with the exec's actual +// node state. #pragma once #include #include @@ -17,6 +18,12 @@ using cudaStream_t = void*; using cudaEvent_t = void*; using cudaStreamCaptureStatus = int; constexpr int cudaSuccess = 0; +constexpr int cudaErrorUnknown = 999; +constexpr cudaStreamCaptureStatus cudaStreamCaptureStatusNone = 0; +constexpr cudaStreamCaptureStatus cudaStreamCaptureStatusActive = 1; +using cudaGraphNodeType = int; +constexpr cudaGraphNodeType cudaGraphNodeTypeKernel = 0; +constexpr cudaGraphNodeType cudaGraphNodeTypeMemcpy = 1; constexpr unsigned cudaStreamNonBlocking = 1; constexpr unsigned cudaEventDisableTiming = 2; @@ -41,9 +48,30 @@ inline int cudaGraphNodeSetEnabled(cudaGraphExec_t ex, cudaGraphNode_t node, reinterpret_cast(node)}] = on != 0; return 0; } -inline int cudaStreamGetCaptureInfo(cudaStream_t, cudaStreamCaptureStatus*, - unsigned long long*, cudaGraph_t*, - const cudaGraphNode_t**, size_t*) { +// Stream capture as the producer ops observe it. fake_capture_graph != null +// means "the stream is capturing into that graph"; every query then reports +// one tail dependency, a fresh kernel node (fake_tail_type says otherwise). +inline cudaGraph_t fake_capture_graph = nullptr; +inline cudaGraphNodeType fake_tail_type = cudaGraphNodeTypeKernel; +inline uintptr_t fake_next_node = 0x1000; +inline cudaGraphNode_t fake_tail_node = nullptr; +inline int cudaStreamGetCaptureInfo(cudaStream_t, cudaStreamCaptureStatus* st, + unsigned long long* id, cudaGraph_t* g, + const cudaGraphNode_t** deps, size_t* nd) { + if (fake_capture_graph == nullptr) { + if (st) *st = cudaStreamCaptureStatusNone; + return 0; + } + fake_tail_node = reinterpret_cast(fake_next_node++); + if (st) *st = cudaStreamCaptureStatusActive; + if (id) *id = 1; + if (g) *g = fake_capture_graph; + if (deps) *deps = &fake_tail_node; + if (nd) *nd = 1; + return 0; +} +inline int cudaGraphNodeGetType(cudaGraphNode_t, cudaGraphNodeType* type) { + *type = fake_tail_type; return 0; } inline int cudaMallocManaged(void** p, size_t n) { diff --git a/tests/native/ring/test_fused_head.cu b/tests/native/ring/test_fused_head.cu new file mode 100644 index 000000000..812c5c861 --- /dev/null +++ b/tests/native/ring/test_fused_head.cu @@ -0,0 +1,459 @@ +// GPU equivalence test for the fused producer (producer_fused_kernel / +// launch_producer_fused): the HookPoint's +// row producer and its in-graph linear heads in one node, incl. multi-layer +// heads as per-layer partial sums. +// +// Phases: +// 1. eager: random hidden sizes (vector + scalar paths), bf16/fp32 sources, +// bf16/fp32 weights, K in [1,16], random slot modes, random gather and +// head tables (-1 / 0 / null part / garbage / > GATHER_MAX_SEGS) -- MID +// only with disjoint tables (planner contract) -- interleaved with plain +// producer_gather launches into one small wrapping ring; +// 2. multi-layer forwards (eager): L fused nodes per forward, slots holding +// FIRST/MID.../LAST chains across nodes, rows routed per node; +// 3. batch invariance: a row's scores are bit-identical when computed alone, +// inside a 37-row tile mix, and with a different capture bound (grid); +// 4. CUDA graph: a 6-node forward captured once, replayed 200 times with the +// tables, slot modes and weights rewritten between replays (register / +// retire as data, no recapture). +// Host model: exact ring bytes for rows, publication words, task/payload heads and +// actual_bytes_counter; scores in double from the same (bf16-rounded) weights, +// within 1e-5 * sum|w h| + 1e-6 (fp32 accumulation); NaN for free/FIRST/MID +// columns. The model mirrors the accumulator cell by cell (stale cells too). +// +// Mutants (NVCC_EXTRA=...; each must FAIL): -DDMI_TEST_MUTATE_SLOT +// (neighbour slot's weights), -DDMI_TEST_MUTATE_ACC (MID overwrites instead +// of adding), -DDMI_TEST_MUTATE_ORDER (head task published before the row +// task). +// +// Environment: FUSED_TEST_PHASES (default "1234"), FUSED_TEST_N1 (phase-1 +// launches, default 500), FUSED_TEST_KINDS (bit mask of launch kinds), +// FUSED_TEST_CHECK_EACH (check the ring after every phase-1 launch). +#include "ring/producer.cuh" +#include "ring/publication_word.h" +#include "ring/ring_alloc.h" + +#include +#include +#include +#include +#include +#include +#include + +using namespace ring; + +#define CK(x) do { cudaError_t e_ = (x); if (e_ != cudaSuccess) { \ + fprintf(stderr, "%s:%d %s: %s\n", __FILE__, __LINE__, #x, cudaGetErrorString(e_)); \ + exit(2); } } while (0) + +static uint64_t g_rng = 0x9E3779B97F4A7C15ull; +static uint64_t rnd() { g_rng ^= g_rng << 13; g_rng ^= g_rng >> 7; g_rng ^= g_rng << 17; return g_rng; } +static int64_t rint_(int64_t lo, int64_t hi) { return lo + int64_t(rnd() % uint64_t(hi - lo + 1)); } +static float rflt() { return float(int64_t(rnd() % 2000001) - 1000000) / 1000000.f; } + +static const uint64_t CAP = 3ull * 1024 * 1024 + 4096 + 16; +static const uint64_t TCAP = 64; +static const int TBL_N = 300; +static const int64_t TBL_W = 1 + 2 * TBL_N; // table words (table_numel) +static const int MAXROWS = 700; +static const int MAXH = 4096; +static const uint64_t ACC_ROWS = 600; // < MAXROWS: rows beyond it exercise the bound + +struct NodeW { + uint32_t K = 1, hidden = 8; bool wbf16 = true; + int mode[16]; int acc[16]; float bias[16]; + std::vector w; // K*hidden, exactly the device values + std::vector words() const { + std::vector v(head_node_words(K, hidden, wbf16), 0); + for (int k = 0; k < 16; ++k) { v[k] = mode[k]; v[16 + k] = acc[k]; std::memcpy(&v[HEAD_BIAS_WORD + k], &bias[k], 4); } + if (wbf16) { + __nv_bfloat16* b = reinterpret_cast<__nv_bfloat16*>(&v[HEAD_HDR_WORDS]); + for (size_t i = 0; i < w.size(); ++i) b[i] = __float2bfloat16(w[i]); + } else std::memcpy(&v[HEAD_HDR_WORDS], w.data(), 4 * w.size()); + return v; + } + void randomize_slot(uint32_t k, int m, int a) { + mode[k] = m; acc[k] = a; bias[k] = (m == HEAD_FREE) ? NAN : rflt(); + for (uint32_t i = 0; i < hidden; ++i) { + float x = (m == HEAD_FREE) ? 0.f : rflt() / 64.f; + if (wbf16) x = __bfloat162float(__float2bfloat16(x)); + w[size_t(k) * hidden + i] = x; + } + } + void init(uint32_t K_, uint32_t h_, bool bf) { + K = K_; hidden = h_; wbf16 = bf; w.assign(size_t(K) * hidden, 0.f); + for (int k = 0; k < 16; ++k) { mode[k] = HEAD_FREE; acc[k] = 0; bias[k] = NAN; } + } +}; + +struct Model { + std::vector payload; + std::vector words; // publication word per task slot + std::map> score; // payload pos -> (ref, scale) + std::vector acc, accs; // [ACC_ROWS * HEAD_MAX_ACC] + uint64_t th = 0, ph = 0, actual = 0; + + uint64_t pos(const TwoSpan& sp, uint64_t o) { return (o < sp.len1) ? sp.off1 + o : sp.off2 + (o - sp.len1); } + void publish_entry(uint64_t seq, const TwoSpan&, uint64_t total) { + words[seq % TCAP] = PUBLICATION_READY | total; + } + static std::vector rows_of(const std::vector& t, int64_t U) { + std::vector out; + int64_t n = t[0]; if (n < 0) n = 0; if (n > int64_t(GATHER_MAX_SEGS)) n = GATHER_MAX_SEGS; + int64_t remaining = U; + for (int64_t k = 0; k < n; ++k) { + int64_t rs = t[1 + 2 * k], nr = t[2 + 2 * k]; + if (rs < 0) rs = 0; if (rs > U) rs = U; + if (nr < 0) nr = 0; if (nr > U - rs) nr = U - rs; if (nr > remaining) nr = remaining; + for (int64_t j = 0; j < nr; ++j) out.push_back(rs + j); + remaining -= nr; + } + return out; + } + // plain producer_gather between fused nodes + void gather(const std::vector& src, uint64_t upper, const std::vector& t, uint64_t rb) { + if (t[0] == -1) return; + std::vector rows = rows_of(t, int64_t(upper / rb)); + const uint64_t total = rows.size() * rb, a = align_up(total, PAYLOAD_ALIGN); + TwoSpan sp = payload_compute_spans(ph, CAP, a); + write_rows(src, rows, rb, sp); + publish_entry(th, sp, total); th += 1; ph += a; actual += total; + } + void write_rows(const std::vector& src, const std::vector& rows, uint64_t rb, const TwoSpan& sp) { + for (size_t r = 0; r < rows.size(); ++r) + for (uint64_t i = 0; i < rb; ++i) { + const uint64_t p = pos(sp, r * rb + i); + payload[p] = src[uint64_t(rows[r]) * rb + i]; + score.erase(p - (p % 4)); + } + } + // fused node: gtab/htab may be null (part absent) + template + void fused(const std::vector& src_bytes, H h, uint32_t elem, int64_t U, + const std::vector* gt, const std::vector* ht, const NodeW& nd) { + const bool g_on = gt && (*gt)[0] != -1, h_on = ht && (*ht)[0] != -1; + if (!g_on && !h_on) return; + const uint64_t rb = uint64_t(nd.hidden) * elem; + std::vector grows = g_on ? rows_of(*gt, U) : std::vector(); + std::vector hrows = h_on ? rows_of(*ht, U) : std::vector(); + const uint64_t gtot = grows.size() * rb, ga = align_up(gtot, PAYLOAD_ALIGN); + const uint64_t htot = hrows.size() * nd.K * 4, ha = align_up(htot, PAYLOAD_ALIGN); + const TwoSpan gsp = payload_compute_spans(ph, CAP, ga); + const TwoSpan hsp = payload_compute_spans(ph + (g_on ? ga : 0), CAP, ha); + if (g_on) write_rows(src_bytes, grows, rb, gsp); + for (size_t r = 0; r < hrows.size(); ++r) { + const int64_t sr = hrows[r]; + for (uint32_t k = 0; k < nd.K; ++k) { +#ifdef DMI_TEST_MUTATE_SLOT + const uint32_t kw = (k + 1) % nd.K; +#else + const uint32_t kw = k; +#endif + const int m = nd.mode[k]; + double dot = 0, sc = 0; + if (m != HEAD_FREE) + for (uint32_t i = 0; i < nd.hidden; ++i) { + const double v = double(nd.w[size_t(kw) * nd.hidden + i]) * h(sr, i); + dot += v; sc += std::fabs(v); + } + double ref = NAN, scale = 0; + const bool in_acc = uint64_t(sr) < ACC_ROWS; + const size_t c = size_t(sr) * HEAD_MAX_ACC + size_t(nd.acc[k]); + if (m == HEAD_SINGLE) { ref = dot + nd.bias[k]; scale = sc + std::fabs(nd.bias[k]); } + else if (m == HEAD_FIRST && in_acc) { acc[c] = dot; accs[c] = sc; } +#ifdef DMI_TEST_MUTATE_ACC + else if (m == HEAD_MID && in_acc) { acc[c] = dot; accs[c] = sc; } +#else + else if (m == HEAD_MID && in_acc) { acc[c] += dot; accs[c] += sc; } +#endif + else if (m == HEAD_LAST && in_acc) { ref = acc[c] + dot + nd.bias[k]; scale = accs[c] + sc + std::fabs(nd.bias[k]); } + const uint64_t p = pos(hsp, (r * nd.K + k) * 4); + score[p] = {ref, scale}; + const float f = float(ref); std::memcpy(&payload[p], &f, 4); + } + } + uint64_t t0 = th; +#ifdef DMI_TEST_MUTATE_ORDER + if (h_on) { publish_entry(t0++, hsp, htot); } + if (g_on) { publish_entry(t0++, gsp, gtot); } +#else + if (g_on) { publish_entry(t0++, gsp, gtot); } + if (h_on) { publish_entry(t0++, hsp, htot); } +#endif + th = t0; + ph += (g_on ? ga : 0) + (h_on ? ha : 0); + actual += (g_on ? gtot : 0) + (h_on ? htot : 0); + } +}; + +static std::vector random_table(int64_t U, bool disjoint) { + std::vector t(1 + 2 * TBL_N, 0); + const int mode = int(rint_(0, 9)); + if (mode == 0) { t[0] = -1; return t; } + if (mode == 1) { t[0] = 0; return t; } + if (mode == 2) { t[0] = 1; t[1] = 0; t[2] = U; return t; } + if (!disjoint && mode == 3) { t[0] = 1; t[1] = rint_(0, U); t[2] = rint_(0, U); return t; } + if (!disjoint && mode == 4) { t[0] = TBL_N; for (int k = 0; k < TBL_N; ++k) { t[1+2*k] = rint_(0, U); t[2+2*k] = rint_(0, 2); } return t; } + if (!disjoint && mode == 5) { + int64_t n = rint_(2, 40); t[0] = n; + for (int k = 0; k < n; ++k) { t[1+2*k] = rint_(-5, U + 5); t[2+2*k] = rint_(-3, U); } + return t; + } + int64_t n = rint_(2, 200), r = 0, k = 0; + while (k < n && r < U) { + r += rint_(0, 3); if (r >= U) break; + int64_t len = rint_(1, 4); if (len > U - r) len = U - r; + t[1+2*k] = r; t[2+2*k] = len; r += len; ++k; + } + t[0] = k; return t; +} + +// Disjoint table containing every row of `must` plus random extra rows. +static std::vector table_with(int64_t U, const std::vector& must, int extra_pct) { + std::vector sel(U, 0); + for (int64_t r = 0; r < U; ++r) sel[r] = must[r] || int(rnd() % 100) < extra_pct; + std::vector t(1 + 2 * TBL_N, 0); + int64_t k = 0; + for (int64_t r = 0; r < U && k < int64_t(GATHER_MAX_SEGS); ) { + if (!sel[r]) { ++r; continue; } + int64_t s = r; while (r < U && sel[r]) ++r; + t[1 + 2 * k] = s; t[2 + 2 * k] = r - s; ++k; + } + t[0] = k; return t; +} + +int main() { + CK(cudaSetDevice(0)); + RingConfig cfg; cfg.task_ring_entries = TCAP; cfg.payload_ring_bytes = CAP; + AllocatedRing ar(cfg); ar.init(); CK(cudaDeviceSynchronize()); + RingState rs = ar.state(); + CK(cudaMemset(rs.payload_buf, 0xCD, CAP)); + Model m; m.payload.assign(CAP, 0xCD); m.words.assign(TCAP, 0); + m.acc.assign(ACC_ROWS * HEAD_MAX_ACC, 0.0); m.accs.assign(ACC_ROWS * HEAD_MAX_ACC, 0.0); + + const uint64_t NSRC = uint64_t(MAXROWS) * MAXH; + std::vector<__nv_bfloat16> hb(NSRC); std::vector hf(NSRC); + for (uint64_t i = 0; i < NSRC; ++i) { hf[i] = rflt() * 4.f; hb[i] = __float2bfloat16(hf[i]); } + std::vector bbytes(NSRC * 2), fbytes(NSRC * 4); + std::memcpy(bbytes.data(), hb.data(), NSRC * 2); std::memcpy(fbytes.data(), hf.data(), NSRC * 4); + __nv_bfloat16* db; CK(cudaMalloc(&db, NSRC * 2)); CK(cudaMemcpy(db, hb.data(), NSRC * 2, cudaMemcpyHostToDevice)); + float* df; CK(cudaMalloc(&df, NSRC * 4)); CK(cudaMemcpy(df, hf.data(), NSRC * 4, cudaMemcpyHostToDevice)); + float* dacc; CK(cudaMalloc(&dacc, 4 * ACC_ROWS * HEAD_MAX_ACC)); CK(cudaMemset(dacc, 0, 4 * ACC_ROWS * HEAD_MAX_ACC)); + const int NT = 16; + int64_t* dt[NT]; for (int i = 0; i < NT; ++i) CK(cudaMalloc(&dt[i], 8 * (1 + 2 * TBL_N))); + const int NN = 8; + int32_t* dn[NN]; for (int i = 0; i < NN; ++i) CK(cudaMalloc(&dn[i], 4 * head_node_words(16, MAXH, 0))); + cudaStream_t s; CK(cudaStreamCreateWithFlags(&s, cudaStreamNonBlocking)); + auto hb_at = [&](uint32_t hidden) { return [&, hidden](int64_t r, uint32_t i) { return double(__bfloat162float(hb[uint64_t(r) * hidden + i])); }; }; + auto hf_at = [&](uint32_t hidden) { return [&, hidden](int64_t r, uint32_t i) { return double(hf[uint64_t(r) * hidden + i]); }; }; + auto up_table = [&](int i, const std::vector& t) { CK(cudaMemcpyAsync(dt[i], t.data(), 8 * t.size(), cudaMemcpyHostToDevice, s)); }; + auto up_node = [&](int i, const NodeW& nd) { std::vector w = nd.words(); CK(cudaMemcpyAsync(dn[i], w.data(), 4 * w.size(), cudaMemcpyHostToDevice, s)); CK(cudaStreamSynchronize(s)); }; + + int n_fused = 0, n_ops = 0; + const char* phases = getenv("FUSED_TEST_PHASES") ? getenv("FUSED_TEST_PHASES") : "1234"; + auto on = [&](char c) { return std::strchr(phases, c) != nullptr; }; + const int n1 = getenv("FUSED_TEST_N1") ? atoi(getenv("FUSED_TEST_N1")) : 500; + const int kinds = getenv("FUSED_TEST_KINDS") ? atoi(getenv("FUSED_TEST_KINDS")) : 15; // bit0 gather, bit1 bf16 head, bit2/3 fp32 head + // ---- 1. eager random single nodes ---- + for (int it = 0; it < (on('1') ? n1 : 0); ++it) { + const int kind = int(rnd() % 4); + if (!((kinds >> kind) & 1)) continue; + if (kind == 0) { + const uint64_t rb = 8192; const int64_t U = rint_(1, 128); const uint64_t upper = U * rb; + std::vector t = random_table(U, false); up_table(0, t); + launch_producer_gather(rs, reinterpret_cast(db), upper, dt[0], TBL_W, rb, 0, s); + m.gather(bbytes, upper, t, rb); + } else { + const bool fp32 = kind == 3, wbf = rnd() % 2; + const uint32_t hidden = fp32 ? uint32_t(rint_(1, 1100)) : (rnd() % 2 ? 4096u : (rnd() % 4 ? uint32_t(rint_(1, 512)) * 8u : uint32_t(rint_(1, 1100)))); + const uint32_t K = (rnd() % 3 == 0) ? 16u : uint32_t(rint_(1, 16)); + const int64_t maxU = int64_t(NSRC / hidden) < MAXROWS ? int64_t(NSRC / hidden) : MAXROWS; + int64_t U = (rnd() % 8 == 0) ? rint_(256, maxU) : rint_(1, 48); + if (uint64_t(U) * hidden * (fp32 ? 4 : 2) > CAP / 2) U = int64_t(CAP / 2 / (uint64_t(hidden) * (fp32 ? 4 : 2))); + const uint64_t upper = uint64_t(U) * hidden * (fp32 ? 4 : 2); + const bool garbage = rnd() % 2; + NodeW nd; nd.init(K, hidden, wbf); + for (uint32_t k = 0; k < K; ++k) { + int md = int(rnd() % 5); + if (garbage && md == HEAD_MID) md = HEAD_SINGLE; // MID needs disjoint rows + nd.randomize_slot(k, md, int(k)); // production: accumulator column = slot + } + std::vector gt = random_table(U, false), ht = random_table(U, !garbage); + const int parts = int(rnd() % 4); // 0: both, 1: gather null, 2: head null, 3: both + // The row copy keeps the gather contract: uint4 rows (row_bytes % 16 == 0). + const bool gnull = parts == 1 || (uint64_t(hidden) * (fp32 ? 4 : 2)) % 16 != 0, hnull = parts == 2; + up_table(0, gt); up_table(1, ht); up_node(0, nd); + launch_producer_fused(rs, fp32 ? reinterpret_cast(df) : reinterpret_cast(db), + upper, hidden, fp32 ? 4 : 2, gnull ? nullptr : dt[0], TBL_W, hnull ? nullptr : dt[1], TBL_W, + dn[0], K, wbf, dacc, ACC_ROWS, 0, s); + if (fp32) m.fused(fbytes, hf_at(hidden), 4, U, gnull ? nullptr : >, hnull ? nullptr : &ht, nd); + else m.fused(bbytes, hb_at(hidden), 2, U, gnull ? nullptr : >, hnull ? nullptr : &ht, nd); + ++n_fused; + } + CK(cudaStreamSynchronize(s)); + ++n_ops; + } + // ---- 2. multi-layer forwards (eager): chains FIRST -> MID* -> LAST ---- + int chains = 0; + for (int f = 0; f < (on('2') ? 60 : 0); ++f) { + const uint32_t hidden = (f % 3 == 0) ? 4096u : uint32_t(rint_(1, 64)) * 8u; + const uint32_t K = uint32_t(rint_(2, 16)); const bool wbf = f % 2; + const int L = int(rint_(2, 6)); + // A launch never exceeds the ring (the host reservation guarantees it): big row + // counts only with narrow rows. + const int64_t U = (f % 5 == 0 && hidden <= 1024) ? int64_t(rint_(256, 650)) : rint_(1, 48); + const uint64_t upper = uint64_t(U) * hidden * 2; + std::vector nodes(L); + for (int l = 0; l < L; ++l) { nodes[l].init(K, hidden, wbf); } + // a few chains over random layer subsets, each with its own acc column + std::vector need(U, 0); + const int nch = int(rint_(1, K < 4 ? K : 4)); + for (int c = 0; c < nch; ++c) { + const uint32_t k = uint32_t(c); const int a = int(k); + int first = int(rint_(0, L - 2)), last = int(rint_(first + 1, L - 1)); + for (int l = first; l <= last; ++l) { + const int md = (l == first) ? HEAD_FIRST : (l == last) ? HEAD_LAST : (rnd() % 3 ? HEAD_MID : HEAD_FREE); + nodes[l].randomize_slot(k, md, a); + } + ++chains; + } + for (int l = 0; l < L; ++l) + for (uint32_t k = uint32_t(nch); k < K; ++k) nodes[l].randomize_slot(k, rnd() % 2 ? HEAD_SINGLE : HEAD_FREE, 0); + for (int64_t r = 0; r < U; ++r) need[r] = int(rnd() % 100) < 30; + for (int l = 0; l < L; ++l) { + std::vector ht = table_with(U, need, 10), gt = random_table(U, false); + up_table(2 + l, ht); up_table(8 + l, gt); up_node(1 + l, nodes[l]); + launch_producer_fused(rs, reinterpret_cast(db), upper, hidden, 2, dt[8 + l], TBL_W, dt[2 + l], TBL_W, + dn[1 + l], K, wbf, dacc, ACC_ROWS, 0, s); + m.fused(bbytes, hb_at(hidden), 2, U, >, &ht, nodes[l]); + CK(cudaStreamSynchronize(s)); + ++n_fused; ++n_ops; + if (getenv("FUSED_TEST_CHECK_EACH")) { + std::vector dv(CAP); CK(cudaMemcpy(dv.data(), rs.payload_buf, CAP, cudaMemcpyDeviceToHost)); + uint64_t badb = 0, badsc = 0; + std::vector isc(CAP, 0); for (auto& kv : m.score) for (int j = 0; j < 4; ++j) isc[kv.first + j] = 1; + for (uint64_t i = 0; i < CAP; ++i) if (!isc[i] && dv[i] != m.payload[i]) ++badb; + for (auto& kv : m.score) { float d; std::memcpy(&d, &dv[kv.first], 4); const double ref = kv.second.first; + if (std::isnan(ref) ? !std::isnan(d) : !(std::fabs(d - ref) <= 1e-5 * kv.second.second + 1e-6)) ++badsc; } + if (badb || badsc) { printf("first bad: forward %d layer %d hidden %u K %u U %ld gt0 %ld ht0 %ld badbytes %lu badscores %lu th %lu/%lu\n", + f, l, hidden, K, (long)U, (long)gt[0], (long)ht[0], (unsigned long)badb, (unsigned long)badsc, + (unsigned long)*rs.task_head, (unsigned long)m.th); exit(3); } + } + } + } + // ---- 3. batch invariance (bit-identical scores for one row) ---- + int inv_fail = 0; + if (on('3')) { + const uint32_t hidden = 4096, K = 16; NodeW nd; nd.init(K, hidden, true); + for (uint32_t k = 0; k < K; ++k) nd.randomize_slot(k, HEAD_SINGLE, 0); + up_node(0, nd); + const int64_t row = 77; + std::vector> got; + for (int cfgi = 0; cfgi < 3; ++cfgi) { + const int64_t U = cfgi == 2 ? 600 : 128; + std::vector t(1 + 2 * TBL_N, 0); + uint64_t idx = 0; // packed index of `row` + if (cfgi == 0) { t[0] = 1; t[1] = row; t[2] = 1; } + else { t[0] = 3; t[1] = 3; t[2] = 30; t[3] = row - 2; t[4] = 5; t[5] = 100; t[6] = 2; idx = 32; } + up_table(1, t); + const uint64_t ph0 = *rs.payload_head; + launch_producer_fused(rs, reinterpret_cast(db), uint64_t(U) * hidden * 2, hidden, 2, nullptr, 0, dt[1], TBL_W, + dn[0], K, 1, dacc, ACC_ROWS, 0, s); + CK(cudaStreamSynchronize(s)); + m.fused(bbytes, hb_at(hidden), 2, U, nullptr, &t, nd); + std::vector v(K); + const uint64_t total = Model::rows_of(t, U).size() * K * 4; + const TwoSpan sp = payload_compute_spans(ph0, CAP, align_up(total, PAYLOAD_ALIGN)); + for (uint32_t k = 0; k < K; ++k) { + const uint64_t o = (idx * K + k) * 4, p = (o < sp.len1) ? sp.off1 + o : sp.off2 + (o - sp.len1); + CK(cudaMemcpy(&v[k], rs.payload_buf + p, 4, cudaMemcpyDeviceToHost)); + } + got.push_back(v); + n_fused++; n_ops++; + } + for (int c = 1; c < 3; ++c) if (std::memcmp(got[0].data(), got[c].data(), 4 * K) != 0) ++inv_fail; + } + // ---- 4. graph: 6-node forward, tables/modes/weights rewritten per replay ---- + int replays = 0; + if (on('4')) { + const uint32_t hidden = 4096, K = 8; const int L = 6; const int64_t U = 48; + const uint64_t upper = uint64_t(U) * hidden * 2; + std::vector nodes(L); + for (int l = 0; l < L; ++l) { nodes[l].init(K, hidden, true); for (uint32_t k = 2; k < K; ++k) nodes[l].randomize_slot(k, rnd() % 2 ? HEAD_SINGLE : HEAD_FREE, 0); } + auto set_chain = [&](uint32_t k, int a) { + int first = int(rint_(0, L - 2)), last = int(rint_(first + 1, L - 1)); + for (int l = 0; l < L; ++l) + nodes[l].randomize_slot(k, l == first ? HEAD_FIRST : l == last ? HEAD_LAST : (l > first && l < last) ? HEAD_MID : HEAD_FREE, a); + }; + set_chain(0, 0); set_chain(1, 1); + for (int l = 0; l < L; ++l) up_node(1 + l, nodes[l]); + cudaGraph_t g; cudaGraphExec_t ge; + CK(cudaStreamBeginCapture(s, cudaStreamCaptureModeGlobal)); + for (int l = 0; l < L; ++l) + launch_producer_fused(rs, reinterpret_cast(db), upper, hidden, 2, dt[8 + l % 8], TBL_W, dt[2 + l], TBL_W, + dn[1 + l], K, 1, dacc, ACC_ROWS, 0, s); + CK(cudaStreamEndCapture(s, &g)); + CK(cudaGraphInstantiate(&ge, g, 0)); + for (int rep = 0; rep < 200; ++rep) { + if (rep % 5 == 0) { // register / retire / move chains: data only + const uint32_t k = uint32_t(rnd() % K); + if (k < 2) set_chain(k, int(k)); + else for (int l = 0; l < L; ++l) nodes[l].randomize_slot(k, rnd() % 2 ? HEAD_SINGLE : HEAD_FREE, 0); + for (int l = 0; l < L; ++l) up_node(1 + l, nodes[l]); + } + std::vector need(U, 0); for (int64_t r = 0; r < U; ++r) need[r] = int(rnd() % 100) < 40; + std::vector> ht(L), gt(L); + for (int l = 0; l < L; ++l) { + ht[l] = (rnd() % 10 == 0) ? std::vector(1 + 2 * TBL_N, -1) : table_with(U, need, 15); + gt[l] = random_table(U, false); + up_table(2 + l, ht[l]); up_table(8 + l, gt[l]); + } + CK(cudaGraphLaunch(ge, s)); + CK(cudaStreamSynchronize(s)); + for (int l = 0; l < L; ++l) m.fused(bbytes, hb_at(hidden), 2, U, >[l], &ht[l], nodes[l]); + n_fused += L; n_ops += L; ++replays; + } + CK(cudaGraphExecDestroy(ge)); CK(cudaGraphDestroy(g)); + } + CK(cudaDeviceSynchronize()); + + std::vector dev(CAP); + CK(cudaMemcpy(dev.data(), rs.payload_buf, CAP, cudaMemcpyDeviceToHost)); + std::vector de(TCAP); + for (uint64_t i = 0; i < TCAP; ++i) + de[i] = __atomic_load_n(&rs.publication_slots[i], __ATOMIC_ACQUIRE); + const uint64_t th = *rs.task_head, ph = *rs.payload_head, ac = *rs.actual_bytes_counter; + + int fails = 0; + uint64_t bad = 0, first = ~0ull, scores = 0, bad_scores = 0, nan_ok = 0; + double max_abs = 0, max_rel = 0; + std::vector is_score(CAP, 0); + for (auto& kv : m.score) { + for (int j = 0; j < 4; ++j) is_score[kv.first + j] = 1; + float d; std::memcpy(&d, &dev[kv.first], 4); + const double ref = kv.second.first, scale = kv.second.second; + ++scores; + if (std::isnan(ref)) { if (std::isnan(d)) ++nan_ok; else ++bad_scores; continue; } + const double err = std::fabs(double(d) - ref); + if (!(err <= 1e-5 * scale + 1e-6)) ++bad_scores; + if (err > max_abs) max_abs = err; + if (std::fabs(ref) > 1e-3 && err / std::fabs(ref) > max_rel) max_rel = err / std::fabs(ref); + } + for (uint64_t i = 0; i < CAP; ++i) + if (!is_score[i] && dev[i] != m.payload[i]) { if (!bad) first = i; ++bad; } + if (bad) { printf("FAIL payload: %lu non-score bytes differ (first at %lu)\n", (unsigned long)bad, (unsigned long)first); ++fails; } + if (bad_scores) { printf("FAIL scores: %lu of %lu outside tolerance\n", (unsigned long)bad_scores, (unsigned long)scores); ++fails; } + if (de != m.words) { printf("FAIL publication words\n"); ++fails; } + if (th != m.th || ph != m.ph || ac != m.actual) { + printf("FAIL heads: task %lu/%lu payload %lu/%lu actual %lu/%lu\n", + (unsigned long)th, (unsigned long)m.th, (unsigned long)ph, (unsigned long)m.ph, + (unsigned long)ac, (unsigned long)m.actual); ++fails; + } + if (inv_fail) { printf("FAIL batch invariance: %d configurations differ bitwise\n", inv_fail); ++fails; } + printf("ops=%d fused_launches=%d chains=%d graph_replays=%d tasks=%lu ring_wraps=%lu scores_checked=%lu (NaN %lu) " + "max_abs_err=%.3g max_rel_err=%.3g batch_invariant=%s -> %s\n", + n_ops, n_fused, chains, replays, (unsigned long)th, (unsigned long)(ph / CAP), (unsigned long)scores, + (unsigned long)nan_ok, max_abs, max_rel, inv_fail ? "no" : "yes", fails ? "FAIL" : "PASS"); + return fails ? 1 : 0; +} diff --git a/tests/native/ring/test_inline_head.cu b/tests/native/ring/test_inline_head.cu new file mode 100644 index 000000000..c4e4a752c --- /dev/null +++ b/tests/native/ring/test_inline_head.cu @@ -0,0 +1,274 @@ +// GPU equivalence test for the in-graph linear head node (producer_head / +// launch_producer_head). +// +// Script: eager launches with random hidden sizes (bf16 vector path, fp32 +// scalar path), slot counts K in [1,16], random row tables (the gather table +// contract incl. -1 = disabled, 0 = empty task, garbage, > GATHER_MAX_SEGS), +// interleaved with ordinary gather producers into one small wrapping ring; +// then a captured CUDA graph (head node + gather node) replayed 200 times +// while BOTH the row table and the weight-slot arena change between replays +// (register/retire as data, no recapture). +// +// Host model: the same packed row order as the gather contract; scores in +// double. Checks: every non-score payload byte, every task entry, task_head, +// payload_head and actual_bytes_counter EXACTLY; every score within +// 1e-5 * sum_i |w_i h_i| + 1e-6 (fp32 accumulation), NaN where the slot is +// retired (b = NaN). Reports max abs / relative error. +// +// Usage: test_inline_head (exit 0 = PASS) +// -DDMI_TEST_MUTATE_SLOT: score slot k with slot (k+1)%K's weights -> must FAIL. +#include "ring/producer.cuh" +#include "ring/publication_word.h" +#include "ring/ring_alloc.h" + +#include +#include +#include +#include +#include +#include +#include + +using namespace ring; + +#define CK(x) do { cudaError_t e_ = (x); if (e_ != cudaSuccess) { \ + fprintf(stderr, "%s:%d %s: %s\n", __FILE__, __LINE__, #x, cudaGetErrorString(e_)); \ + exit(2); } } while (0) + +static uint64_t g_rng = 0x2545F4914F6CDD1Dull; +static uint64_t rnd() { g_rng ^= g_rng << 13; g_rng ^= g_rng >> 7; g_rng ^= g_rng << 17; return g_rng; } +static int64_t rint_(int64_t lo, int64_t hi) { return lo + int64_t(rnd() % uint64_t(hi - lo + 1)); } +static float rflt() { return float(int64_t(rnd() % 2000001) - 1000000) / 1000000.f; } + +static const uint64_t CAP = 3ull * 1024 * 1024 + 4096 + 16; +static const uint64_t TCAP = 64; +static const int TBL_N = 300; +static const int64_t TBL_W = 1 + 2 * TBL_N; // table words (table_numel) +static const int MAXROWS = 700; +static const int MAXH = 4096; + +struct Model { + std::vector payload; + std::vector words; // publication word per task slot + std::map> score; // payload pos -> (ref, scale) + uint64_t th = 0, ph = 0, actual = 0; + + uint64_t pos(const TwoSpan& sp, uint64_t o) { return (o < sp.len1) ? sp.off1 + o : sp.off2 + (o - sp.len1); } + void publish(const TwoSpan&, uint64_t alloc, uint64_t total) { + words[th % TCAP] = PUBLICATION_READY | total; + th += 1; ph += alloc; actual += total; + } + // Row set and order exactly as the gather contract (serial clamp form). + static std::vector rows_of(const std::vector& t, int64_t U) { + std::vector out; + int64_t n = t[0]; if (n < 0) n = 0; if (n > int64_t(GATHER_MAX_SEGS)) n = GATHER_MAX_SEGS; + int64_t remaining = U; + for (int64_t k = 0; k < n; ++k) { + int64_t rs = t[1 + 2 * k], nr = t[2 + 2 * k]; + if (rs < 0) rs = 0; if (rs > U) rs = U; + if (nr < 0) nr = 0; if (nr > U - rs) nr = U - rs; if (nr > remaining) nr = remaining; + for (int64_t j = 0; j < nr; ++j) out.push_back(rs + j); + remaining -= nr; + } + return out; + } + void gather(const std::vector& src, uint64_t off, uint64_t upper, + const std::vector& t, uint64_t rb) { + if (t[0] == -1) return; + const int64_t U = int64_t(upper / rb); + std::vector rows = rows_of(t, U); + const uint64_t total = rows.size() * rb, a = align_up(total, PAYLOAD_ALIGN); + TwoSpan sp = payload_compute_spans(ph, CAP, a); + for (size_t r = 0; r < rows.size(); ++r) + for (uint64_t i = 0; i < rb; ++i) { + const uint64_t p = pos(sp, r * rb + i); + payload[p] = src[off + uint64_t(rows[r]) * rb + i]; + score.erase(p - (p % 4)); + } + publish(sp, a, total); + } + // h(row, i) supplied as doubles (bf16/fp32 exact values). + template + void head(H h, int64_t U, const std::vector& t, uint32_t hidden, + const std::vector& arena, uint32_t K) { + if (t[0] == -1) return; + std::vector rows = rows_of(t, U); + const uint64_t total = rows.size() * K * 4, a = align_up(total, PAYLOAD_ALIGN); + TwoSpan sp = payload_compute_spans(ph, CAP, a); + for (size_t r = 0; r < rows.size(); ++r) + for (uint32_t k = 0; k < K; ++k) { +#ifdef DMI_TEST_MUTATE_SLOT + const uint32_t kw = (k + 1) % K; +#else + const uint32_t kw = k; +#endif + double acc = 0, scale = 0; + for (uint32_t i = 0; i < hidden; ++i) { + const double v = double(arena[uint64_t(kw) * hidden + i]) * h(rows[r], i); + acc += v; scale += std::fabs(v); + } + const double b = arena[uint64_t(K) * hidden + k]; + acc += b; scale += std::isnan(b) ? 0 : std::fabs(b); + const uint64_t p = pos(sp, (r * K + k) * 4); + score[p] = {acc, scale}; + const float f = float(acc); + std::memcpy(&payload[p], &f, 4); // placeholder; compared by tolerance + } + publish(sp, a, total); + } +}; + +static std::vector random_table(int64_t U) { + std::vector t(1 + 2 * TBL_N, 0); + const int mode = int(rint_(0, 9)); + if (mode == 0) { t[0] = -1; return t; } + if (mode == 1) { t[0] = 0; return t; } + if (mode == 2) { t[0] = 1; t[1] = 0; t[2] = U; return t; } + if (mode == 3) { t[0] = 1; t[1] = rint_(0, U); t[2] = rint_(0, U); return t; } + if (mode == 4) { t[0] = TBL_N; for (int k = 0; k < TBL_N; ++k) { t[1+2*k] = rint_(0, U); t[2+2*k] = rint_(0, 2); } return t; } + if (mode == 5) { + int64_t n = rint_(2, 40); t[0] = n; + for (int k = 0; k < n; ++k) { t[1+2*k] = rint_(-5, U + 5); t[2+2*k] = rint_(-3, U); } + return t; + } + int64_t n = rint_(2, 200), r = 0, k = 0; + while (k < n && r < U) { + r += rint_(0, 3); if (r >= U) break; + int64_t len = rint_(1, 4); if (len > U - r) len = U - r; + t[1+2*k] = r; t[2+2*k] = len; r += len; ++k; + } + t[0] = k; return t; +} + +static void random_arena(std::vector& a, uint32_t K, uint32_t hidden) { + a.assign(size_t(K) * hidden + K, 0.f); + for (uint32_t k = 0; k < K; ++k) { + const int state = int(rnd() % 6); // 0 = retired slot (w = 0, b = NaN) + for (uint32_t i = 0; i < hidden; ++i) a[size_t(k) * hidden + i] = state ? rflt() / 64.f : 0.f; + a[size_t(K) * hidden + k] = state ? rflt() : NAN; + } +} + +int main() { + CK(cudaSetDevice(0)); + RingConfig cfg; cfg.task_ring_entries = TCAP; cfg.payload_ring_bytes = CAP; + AllocatedRing ar(cfg); ar.init(); CK(cudaDeviceSynchronize()); + RingState rs = ar.state(); + CK(cudaMemset(rs.payload_buf, 0xCD, CAP)); + Model m; m.payload.assign(CAP, 0xCD); m.words.assign(TCAP, 0); + + // Source rows: bf16 view and fp32 view of the same random values. + const uint64_t NSRC = uint64_t(MAXROWS) * MAXH; + std::vector<__nv_bfloat16> hb(NSRC); std::vector hf(NSRC); + for (uint64_t i = 0; i < NSRC; ++i) { hf[i] = rflt() * 4.f; hb[i] = __float2bfloat16(hf[i]); } + std::vector hbytes(NSRC * 2); + std::memcpy(hbytes.data(), hb.data(), NSRC * 2); + __nv_bfloat16* db; CK(cudaMalloc(&db, NSRC * 2)); CK(cudaMemcpy(db, hb.data(), NSRC * 2, cudaMemcpyHostToDevice)); + float* df; CK(cudaMalloc(&df, NSRC * 4)); CK(cudaMemcpy(df, hf.data(), NSRC * 4, cudaMemcpyHostToDevice)); + int64_t* dt; CK(cudaMalloc(&dt, 8 * (1 + 2 * TBL_N))); + int64_t* dt2; CK(cudaMalloc(&dt2, 8 * (1 + 2 * TBL_N))); + float* darena; CK(cudaMalloc(&darena, 4 * (16 * MAXH + 16 + 2048))); + cudaStream_t s; CK(cudaStreamCreateWithFlags(&s, cudaStreamNonBlocking)); + + int n_ops = 0, n_head = 0; + std::vector arena; + // ---- eager phase ---- + for (int it = 0; it < 400; ++it) { + const int kind = int(rnd() % 3); + if (kind == 0) { // ordinary gather producer between head nodes + const uint64_t rb = 8192; const int64_t U = rint_(1, 128); const uint64_t upper = U * rb; + std::vector t = random_table(U); + CK(cudaMemcpyAsync(dt, t.data(), 8 * t.size(), cudaMemcpyHostToDevice, s)); + launch_producer_gather(rs, reinterpret_cast(db), upper, dt, TBL_W, rb, 0, s); + m.gather(hbytes, 0, upper, t, rb); + } else { + const bool fp32 = kind == 2; + const uint32_t hidden = fp32 ? uint32_t(rint_(1, 1100)) : (rnd() % 2 ? 4096u : uint32_t(rint_(1, 512)) * 8u); + const uint32_t K = (rnd() % 3 == 0) ? 16u : uint32_t(rint_(1, 16)); + const int64_t maxU = int64_t(NSRC / hidden) < MAXROWS ? int64_t(NSRC / hidden) : MAXROWS; + const int64_t U = (rnd() % 8 == 0) ? rint_(256, maxU) : rint_(1, 48); + const uint64_t upper = uint64_t(U) * hidden * (fp32 ? 4 : 2); + std::vector t = random_table(U); + random_arena(arena, K, hidden); + CK(cudaMemcpyAsync(dt, t.data(), 8 * t.size(), cudaMemcpyHostToDevice, s)); + CK(cudaMemcpyAsync(darena, arena.data(), 4 * arena.size(), cudaMemcpyHostToDevice, s)); + launch_producer_head(rs, fp32 ? reinterpret_cast(df) : reinterpret_cast(db), + upper, dt, TBL_W, hidden, fp32 ? 4 : 2, darena, K, 24, s); + if (fp32) m.head([&](int64_t r, uint32_t i) { return double(hf[uint64_t(r) * hidden + i]); }, U, t, hidden, arena, K); + else m.head([&](int64_t r, uint32_t i) { return double(__bfloat162float(hb[uint64_t(r) * hidden + i])); }, U, t, hidden, arena, K); + ++n_head; + } + CK(cudaStreamSynchronize(s)); + ++n_ops; + } + // ---- graph phase: head node + gather node captured once; tables AND the + // weight arena are rewritten between replays (runtime register/retire). ---- + { + const uint32_t hidden = 4096, K = 8; const int64_t U = 48; + const uint64_t upper = uint64_t(U) * hidden * 2; + cudaGraph_t g; cudaGraphExec_t ge; + CK(cudaStreamBeginCapture(s, cudaStreamCaptureModeGlobal)); + launch_producer_gather(rs, reinterpret_cast(db), upper, dt2, TBL_W, hidden * 2, 0, s); + launch_producer_head(rs, reinterpret_cast(db), upper, dt, TBL_W, hidden, 2, darena, K, 24, s); + CK(cudaStreamEndCapture(s, &g)); + CK(cudaGraphInstantiate(&ge, g, 0)); + random_arena(arena, K, hidden); + for (int rep = 0; rep < 200; ++rep) { + std::vector t = random_table(U), t2 = random_table(U); + if (rep % 5 == 0) { // re-register / retire some slots: data only + std::vector fresh; random_arena(fresh, K, hidden); + const uint32_t k = uint32_t(rnd() % K); + for (uint32_t i = 0; i < hidden; ++i) arena[size_t(k) * hidden + i] = fresh[size_t(k) * hidden + i]; + arena[size_t(K) * hidden + k] = fresh[size_t(K) * hidden + k]; + } + CK(cudaMemcpyAsync(dt, t.data(), 8 * t.size(), cudaMemcpyHostToDevice, s)); + CK(cudaMemcpyAsync(dt2, t2.data(), 8 * t2.size(), cudaMemcpyHostToDevice, s)); + CK(cudaMemcpyAsync(darena, arena.data(), 4 * arena.size(), cudaMemcpyHostToDevice, s)); + CK(cudaGraphLaunch(ge, s)); + CK(cudaStreamSynchronize(s)); + m.gather(hbytes, 0, upper, t2, hidden * 2); + m.head([&](int64_t r, uint32_t i) { return double(__bfloat162float(hb[uint64_t(r) * hidden + i])); }, U, t, hidden, arena, K); + n_ops += 2; ++n_head; + } + CK(cudaGraphExecDestroy(ge)); CK(cudaGraphDestroy(g)); + } + CK(cudaDeviceSynchronize()); + + std::vector dev(CAP); + CK(cudaMemcpy(dev.data(), rs.payload_buf, CAP, cudaMemcpyDeviceToHost)); + std::vector de(TCAP); + for (uint64_t i = 0; i < TCAP; ++i) + de[i] = __atomic_load_n(&rs.publication_slots[i], __ATOMIC_ACQUIRE); + const uint64_t th = *rs.task_head, ph = *rs.payload_head, ac = *rs.actual_bytes_counter; + + int fails = 0; + uint64_t bad = 0, first = ~0ull, scores = 0, bad_scores = 0, nan_ok = 0; + double max_abs = 0, max_rel = 0; + std::vector is_score(CAP, 0); + for (auto& kv : m.score) { + for (int j = 0; j < 4; ++j) is_score[kv.first + j] = 1; + float d; std::memcpy(&d, &dev[kv.first], 4); + const double ref = kv.second.first, scale = kv.second.second; + ++scores; + if (std::isnan(ref)) { if (std::isnan(d)) ++nan_ok; else ++bad_scores; continue; } + const double err = std::fabs(double(d) - ref); + if (!(err <= 1e-5 * scale + 1e-6)) ++bad_scores; + if (err > max_abs) max_abs = err; + if (std::fabs(ref) > 1e-3 && err / std::fabs(ref) > max_rel) max_rel = err / std::fabs(ref); + } + for (uint64_t i = 0; i < CAP; ++i) + if (!is_score[i] && dev[i] != m.payload[i]) { if (!bad) first = i; ++bad; } + if (bad) { printf("FAIL payload: %lu non-score bytes differ (first at %lu)\n", (unsigned long)bad, (unsigned long)first); ++fails; } + if (bad_scores) { printf("FAIL scores: %lu of %lu outside tolerance\n", (unsigned long)bad_scores, (unsigned long)scores); ++fails; } + if (de != m.words) { printf("FAIL publication words\n"); ++fails; } + if (th != m.th || ph != m.ph || ac != m.actual) { + printf("FAIL heads: task %lu/%lu payload %lu/%lu actual %lu/%lu\n", + (unsigned long)th, (unsigned long)m.th, (unsigned long)ph, (unsigned long)m.ph, + (unsigned long)ac, (unsigned long)m.actual); ++fails; + } + printf("ops=%d head_launches=%d tasks=%lu ring_wraps=%lu live_scores_checked=%lu (retired NaN %lu) " + "max_abs_err=%.3g max_rel_err=%.3g -> %s\n", + n_ops, n_head, (unsigned long)th, (unsigned long)(ph / CAP), (unsigned long)scores, + (unsigned long)nan_ok, max_abs, max_rel, fails ? "FAIL" : "PASS"); + return fails ? 1 : 0; +} diff --git a/tests/native/ring/test_producer_equivalence.cu b/tests/native/ring/test_producer_equivalence.cu new file mode 100644 index 000000000..93e92e485 --- /dev/null +++ b/tests/native/ring/test_producer_equivalence.cu @@ -0,0 +1,353 @@ +// GPU equivalence test for the static, prefix, chunked and gather producers. +// +// Runs a fixed pseudo-random script of launches (eager, and graph replays +// whose device-side row counts and row tables change between replays) +// against a small wrapping payload ring, and checks the ENTIRE ring state -- +// every payload byte, every publication word, task_head, payload_head and +// actual_bytes_counter -- against an independent host model of the kernels' +// contract: a serial reference of the row-table clamps. With device graph +// updates it also replays planner-gated device-updatable gather nodes +// (toggle_device.h), within one graph and across two graphs. +// +// Usage: test_producer_equivalence [dump_file] +// exit 0 = all checks pass; dump_file (optional) receives the final ring +// state so two builds (e.g. before/after a kernel change) can be compared +// byte for byte with cmp. +// -DDMI_TEST_MUTATE_PLANNER gates each device node by its neighbour's row +// and must FAIL. +#include "ring/producer.cuh" +#include "ring/publication_word.h" +#include "ring/ring_alloc.h" +#include "ring/toggle_device.h" + +#include +#include +#include +#include + +using namespace ring; + +#define CK(x) do { cudaError_t e_ = (x); if (e_ != cudaSuccess) { \ + fprintf(stderr, "%s:%d %s: %s\n", __FILE__, __LINE__, #x, cudaGetErrorString(e_)); \ + exit(2); } } while (0) + +static uint64_t g_rng = 0x9E3779B97F4A7C15ull; +static uint64_t rnd() { g_rng ^= g_rng << 13; g_rng ^= g_rng >> 7; g_rng ^= g_rng << 17; return g_rng; } +static int64_t rint_(int64_t lo, int64_t hi) { return lo + int64_t(rnd() % uint64_t(hi - lo + 1)); } + +static const uint64_t SRC_BYTES = 1200ull * 8192; +static const uint64_t CAP = 24ull * 1024 * 1024 + 4096 + 16; // odd-ish, 16-aligned +static const uint64_t TCAP = 64; +static const int TBL_N = 300; // > GATHER_MAX_SEGS: exercises the clamp +static const int64_t TBL_W = 1 + 2 * TBL_N; // table words (table_numel) + +struct Model { + std::vector payload; + std::vector words; // publication word per task slot + uint64_t th = 0, ph = 0, actual = 0; + const std::vector* src; + + void put(const TwoSpan& sp, uint64_t dst_off, const uint8_t* s, uint64_t n) { + for (uint64_t i = 0; i < n; ++i) { + const uint64_t o = dst_off + i; + const uint64_t p = (o < sp.len1) ? sp.off1 + o : sp.off2 + (o - sp.len1); + payload[p] = s[i]; + } + } + void publish(const TwoSpan&, uint64_t alloc, uint64_t total) { + words[th % TCAP] = PUBLICATION_READY | total; + th += 1; ph += alloc; actual += total; + } + void stat(uint64_t off, uint64_t n) { + const uint64_t a = align_up(n, PAYLOAD_ALIGN); + TwoSpan sp = payload_compute_spans(ph, CAP, a); + put(sp, 0, src->data() + off, n); + publish(sp, a, n); + } + void prefix(uint64_t off, uint64_t upper, int64_t rc, uint64_t rb) { + if (rc < 0) rc = 0; + uint64_t n = uint64_t(rc) * rb; if (n > upper) n = upper; + const uint64_t a = align_up(n, PAYLOAD_ALIGN); + TwoSpan sp = payload_compute_spans(ph, CAP, a); + put(sp, 0, src->data() + off, n); + publish(sp, a, n); + } + void chunked(uint64_t off, uint64_t upper, const std::vector& T) { + const uint64_t K = T.size(), cin = upper / K; + uint64_t total = 0; + for (uint64_t k = 0; k < K; ++k) { + int64_t v = T[k]; if (v < 0) v = 0; if (v > int64_t(cin)) v = int64_t(cin); + total += uint64_t(v); + } + const uint64_t a = align_up(total, PAYLOAD_ALIGN); + TwoSpan sp = payload_compute_spans(ph, CAP, a); + uint64_t d = 0; + for (uint64_t k = 0; k < K; ++k) { + int64_t v = T[k]; if (v < 0) v = 0; if (v > int64_t(cin)) v = int64_t(cin); + put(sp, d, src->data() + off + k * cin, uint64_t(v)); d += uint64_t(v); + } + publish(sp, a, d); + } + // Serial reference: a single thread walking the table with a running + // `remaining` clamp (the kernel's parallel load + warp scan must match). + void gather(uint64_t off, uint64_t upper, const std::vector& t, uint64_t rb) { + if (t[0] == -1) return; + const int64_t U = rb ? int64_t(upper / rb) : 0; + int64_t n = t[0]; if (n < 0) n = 0; if (n > int64_t(GATHER_MAX_SEGS)) n = GATHER_MAX_SEGS; + int64_t remaining = U; std::vector srcr, cnt; + for (int64_t k = 0; k < n; ++k) { + int64_t rs = t[1 + 2 * k], nr = t[2 + 2 * k]; + if (rs < 0) rs = 0; if (rs > U) rs = U; + if (nr < 0) nr = 0; if (nr > U - rs) nr = U - rs; if (nr > remaining) nr = remaining; + srcr.push_back(rs); cnt.push_back(nr); remaining -= nr; + } + uint64_t rows = 0; for (int64_t c : cnt) rows += uint64_t(c); + const uint64_t total = rows * rb, a = align_up(total, PAYLOAD_ALIGN); + TwoSpan sp = payload_compute_spans(ph, CAP, a); + uint64_t d = 0; + for (size_t k = 0; k < cnt.size(); ++k) { + put(sp, d, src->data() + off + uint64_t(srcr[k]) * rb, uint64_t(cnt[k]) * rb); + d += uint64_t(cnt[k]) * rb; + } + publish(sp, a, total); + } +}; + +static std::vector random_table(int64_t U) { + std::vector t(1 + 2 * TBL_N, 0); + const int mode = int(rint_(0, 9)); + if (mode == 0) { t[0] = -1; return t; } + if (mode == 1) { t[0] = 0; return t; } + if (mode == 2) { t[0] = 1; t[1] = 0; t[2] = U; return t; } // full + if (mode == 3) { t[0] = 1; t[1] = rint_(0, U); t[2] = rint_(0, U); return t; } // one window (may clamp) + if (mode == 4) { t[0] = TBL_N; for (int k = 0; k < TBL_N; ++k) { t[1+2*k] = rint_(0, U); t[2+2*k] = rint_(0, 3); } return t; } + if (mode == 5) { // garbage: negatives, out of range, sum > U + int64_t n = rint_(2, 40); t[0] = n; + for (int k = 0; k < n; ++k) { t[1+2*k] = rint_(-5, U + 5); t[2+2*k] = rint_(-3, U); } + return t; + } + // modes 6..9: realistic fragmented selection: disjoint ascending runs + int64_t n = rint_(2, 200), r = 0, k = 0; + while (k < n && r < U) { + r += rint_(0, 3); if (r >= U) break; + int64_t len = rint_(1, 4); if (len > U - r) len = U - r; + t[1+2*k] = r; t[2+2*k] = len; r += len; ++k; + } + t[0] = k; return t; +} + +int main(int argc, char** argv) { + CK(cudaSetDevice(0)); + RingConfig cfg; cfg.task_ring_entries = TCAP; cfg.payload_ring_bytes = CAP; + AllocatedRing ar(cfg); ar.init(); CK(cudaDeviceSynchronize()); + RingState rs = ar.state(); + CK(cudaMemset(rs.payload_buf, 0xCD, CAP)); + + Model m; m.payload.assign(CAP, 0xCD); m.words.assign(TCAP, 0); + std::vector hsrc(SRC_BYTES); + for (auto& b : hsrc) b = uint8_t(rnd()); + m.src = &hsrc; + uint8_t* src; CK(cudaMalloc(&src, SRC_BYTES)); + CK(cudaMemcpy(src, hsrc.data(), SRC_BYTES, cudaMemcpyHostToDevice)); + int64_t* dvec; CK(cudaMalloc(&dvec, sizeof(int64_t) * (1 + 2 * TBL_N))); + cudaStream_t s; CK(cudaStreamCreateWithFlags(&s, cudaStreamNonBlocking)); + + const uint64_t rbs[] = {8192, 2064, 4096}; + int n_ops = 0; + // ---- eager phase ---- + for (int it = 0; it < 600; ++it) { + const uint64_t rb = rbs[rnd() % 3]; + const int64_t maxrows = int64_t(SRC_BYTES / rb) - 8; + // mostly decode-sized, sometimes prefill-sized (grid at the block cap) + const int64_t U = (rnd() % 8 == 0) ? rint_(600, maxrows < 1150 ? maxrows : 1150) : rint_(1, 300); + const uint64_t upper = uint64_t(U) * rb; + const uint64_t off = uint64_t(rint_(0, 7)) * 16; + const int kind = int(rnd() % 4); + if (off + upper > SRC_BYTES) continue; + if (kind == 0) { + uint64_t n = upper - uint64_t(rint_(0, 15)); // tail bytes + launch_producer_static(rs, src + off, n, 0, s); m.stat(off, n); + } else if (kind == 1) { + int64_t rc = rint_(-2, U + 3); + CK(cudaMemcpyAsync(dvec, &rc, 8, cudaMemcpyHostToDevice, s)); + launch_producer_prefix(rs, src + off, upper, dvec, rb, 0, s); m.prefix(off, upper, rc, rb); + } else if (kind == 2) { + const uint64_t upper_rows_bytes = upper; + const uint32_t K = uint32_t(rint_(2, 8)); + // chunk_in = upper/K must stay uint4-aligned (callers' contract) + const uint64_t upper = (upper_rows_bytes / (16ull * K)) * 16ull * K; + if (upper == 0) continue; + std::vector T(K); + for (auto& v : T) v = rint_(-16, int64_t(upper / K) + 64) & ~int64_t(15); + CK(cudaMemcpyAsync(dvec, T.data(), 8 * K, cudaMemcpyHostToDevice, s)); + launch_producer_chunked(rs, src + off, upper, dvec, K, 0, s); m.chunked(off, upper, T); + } else { + std::vector t = random_table(U); + CK(cudaMemcpyAsync(dvec, t.data(), 8 * t.size(), cudaMemcpyHostToDevice, s)); + launch_producer_gather(rs, src + off, upper, dvec, TBL_W, rb, 0, s); m.gather(off, upper, t, rb); + } + CK(cudaStreamSynchronize(s)); // dvec is rewritten next iteration + ++n_ops; + } + // ---- graph phase: grid baked by the upper bound, values re-read ---- + { + const uint64_t rb = 8192; const int64_t U = 160; const uint64_t upper = U * rb; + int64_t* dt; CK(cudaMalloc(&dt, sizeof(int64_t) * (1 + 2 * TBL_N))); + int64_t* drc; CK(cudaMalloc(&drc, sizeof(int64_t))); + cudaGraph_t g; cudaGraphExec_t ge; + CK(cudaStreamBeginCapture(s, cudaStreamCaptureModeGlobal)); + launch_producer_gather(rs, src, upper, dt, TBL_W, rb, 0, s); + launch_producer_static(rs, src + 4096, 150 * rb, 0, s); + launch_producer_prefix(rs, src + 16, upper, drc, rb, 0, s); + launch_producer_gather(rs, src + 32, upper, dt, TBL_W, rb, 0, s); + CK(cudaStreamEndCapture(s, &g)); + CK(cudaGraphInstantiate(&ge, g, 0)); + for (int rep = 0; rep < 120; ++rep) { + std::vector t = random_table(U); + int64_t rc = rint_(0, U); + CK(cudaMemcpyAsync(dt, t.data(), 8 * t.size(), cudaMemcpyHostToDevice, s)); + CK(cudaMemcpyAsync(drc, &rc, 8, cudaMemcpyHostToDevice, s)); + CK(cudaGraphLaunch(ge, s)); + CK(cudaStreamSynchronize(s)); + m.gather(0, upper, t, rb); m.stat(4096, 150 * rb); + m.prefix(16, upper, rc, rb); m.gather(32, upper, t, rb); + n_ops += 4; + } + CK(cudaGraphExecDestroy(ge)); CK(cudaGraphDestroy(g)); + } +#if DMX_HAS_DEVICE_GRAPH_UPDATE + // ---- device-toggle phase: planner-gated device-updatable gather nodes ---- + // Same contract as an early-exit launch: a node the planner disables + // (row[0] == -1) must leave the ring exactly as the -1 early exit does. + // Half of the replays force -1 rows; others use random_table (which + // itself yields -1 sometimes). 3 gather nodes + one ordinary static node. + { + const uint64_t rb = 8192; const int64_t U = 160; const uint64_t upper = U * rb; + const int NG = 3; + int64_t* dt[NG]; + for (int i = 0; i < NG; ++i) CK(cudaMalloc(&dt[i], sizeof(int64_t) * (1 + 2 * TBL_N))); + ring_py::ToggleDevPool* pool; + CK(cudaMallocManaged(&pool, sizeof(ring_py::ToggleDevPool))); + std::memset(pool, 0, sizeof(ring_py::ToggleDevPool)); + cudaGraph_t g; cudaGraphExec_t ge; + CK(cudaStreamBeginCapture(s, cudaStreamCaptureModeGlobal)); + launch_toggle_planner(pool, 0, s); + for (int i = 0; i < NG; ++i) { + CK(launch_producer_gather_device_updatable(rs, src + 48 * i, upper, dt[i], TBL_W, rb, 0, s)); + cudaStreamCaptureStatus st; const cudaGraphNode_t* deps; size_t nd; + CK(dmx_capture_info(s, &st, nullptr, nullptr, &deps, &nd)); + cudaKernelNodeAttrValue v = {}; + CK(cudaGraphKernelNodeGetAttribute(deps[nd - 1], cudaKernelNodeAttributeDeviceUpdatableKernelNode, &v)); + if (!v.deviceUpdatableKernelNode.deviceUpdatable || !v.deviceUpdatableKernelNode.devNode) { printf("FAIL devNode\n"); return 1; } +#ifdef DMI_TEST_MUTATE_PLANNER // gate each node by its NEIGHBOUR's row -> must FAIL + pool->entry[i] = ring_py::ToggleDevEntry{dt[(i + 1) % NG], v.deviceUpdatableKernelNode.devNode}; +#else + pool->entry[i] = ring_py::ToggleDevEntry{dt[i], v.deviceUpdatableKernelNode.devNode}; +#endif + pool->slot[0].n = i + 1; + if (i == 0) launch_producer_static(rs, src + 4096, 150 * rb, 0, s); + } + CK(cudaStreamEndCapture(s, &g)); + CK(cudaGraphInstantiate(&ge, g, 0)); + CK(cudaGraphUpload(ge, s)); + int disabled = 0; + for (int rep = 0; rep < 200; ++rep) { + std::vector t[NG]; + for (int i = 0; i < NG; ++i) { + t[i] = random_table(U); + if (rnd() % 2) { t[i].assign(1 + 2 * TBL_N, 0); t[i][0] = -1; } + disabled += t[i][0] == -1; + CK(cudaMemcpyAsync(dt[i], t[i].data(), 8 * t[i].size(), cudaMemcpyHostToDevice, s)); + } + CK(cudaGraphLaunch(ge, s)); + CK(cudaStreamSynchronize(s)); + m.gather(0, upper, t[0], rb); m.stat(4096, 150 * rb); + m.gather(48, upper, t[1], rb); m.gather(96, upper, t[2], rb); + n_ops += 4; + } + printf("device-toggle phase: 200 replays, %d of %d gather nodes planner-disabled\n", disabled, 200 * NG); + CK(cudaGraphExecDestroy(ge)); CK(cudaGraphDestroy(g)); + CK(cudaFree(pool)); + } + // ---- cross-graph phase: ONE planner (in graph G0) gates the producers + // of a later graph G1 too (a forward captured as consecutive segment + // graphs). ---- + { + const uint64_t rb = 4096; const int64_t U = 96; const uint64_t upper = U * rb; + int64_t* dt[3]; + for (int i = 0; i < 3; ++i) CK(cudaMalloc(&dt[i], sizeof(int64_t) * (1 + 2 * TBL_N))); + ring_py::ToggleDevPool* pool; + CK(cudaMallocManaged(&pool, sizeof(ring_py::ToggleDevPool))); + std::memset(pool, 0, sizeof(ring_py::ToggleDevPool)); + cudaGraph_t g[2]; cudaGraphExec_t ge[2]; + int e = 0; + for (int gi = 0; gi < 2; ++gi) { + CK(cudaStreamBeginCapture(s, cudaStreamCaptureModeGlobal)); + if (gi == 0) launch_toggle_planner(pool, 0, s); + for (int i = (gi == 0 ? 0 : 1); i < (gi == 0 ? 1 : 3); ++i) { + CK(launch_producer_gather_device_updatable(rs, src + 64 * i, upper, dt[i], TBL_W, rb, 0, s)); + cudaStreamCaptureStatus st; const cudaGraphNode_t* deps; size_t nd; + CK(dmx_capture_info(s, &st, nullptr, nullptr, &deps, &nd)); + cudaKernelNodeAttrValue v = {}; + CK(cudaGraphKernelNodeGetAttribute(deps[nd - 1], cudaKernelNodeAttributeDeviceUpdatableKernelNode, &v)); + pool->entry[e++] = ring_py::ToggleDevEntry{dt[i], v.deviceUpdatableKernelNode.devNode}; + pool->slot[0].n = e; + } + CK(cudaStreamEndCapture(s, &g[gi])); + CK(cudaGraphInstantiate(&ge[gi], g[gi], 0)); + CK(cudaGraphUpload(ge[gi], s)); + } + int disabled = 0; + for (int rep = 0; rep < 200; ++rep) { + std::vector t[3]; + for (int i = 0; i < 3; ++i) { + t[i] = random_table(U); + if (rnd() % 2) { t[i].assign(1 + 2 * TBL_N, 0); t[i][0] = -1; } + disabled += t[i][0] == -1; + CK(cudaMemcpyAsync(dt[i], t[i].data(), 8 * t[i].size(), cudaMemcpyHostToDevice, s)); + } + CK(cudaGraphLaunch(ge[0], s)); + CK(cudaGraphLaunch(ge[1], s)); + CK(cudaStreamSynchronize(s)); + m.gather(0, upper, t[0], rb); m.gather(64, upper, t[1], rb); m.gather(128, upper, t[2], rb); + n_ops += 3; + } + printf("cross-graph phase: 200 replays x 2 graphs, %d of 600 nodes planner-disabled, planner errors %llu\n", + disabled, pool->errors); + if (pool->errors) { printf("FAIL planner errors\n"); return 1; } + for (int gi = 0; gi < 2; ++gi) { CK(cudaGraphExecDestroy(ge[gi])); CK(cudaGraphDestroy(g[gi])); } + CK(cudaFree(pool)); + } +#endif // DMX_HAS_DEVICE_GRAPH_UPDATE + CK(cudaDeviceSynchronize()); + + std::vector dev(CAP); + CK(cudaMemcpy(dev.data(), rs.payload_buf, CAP, cudaMemcpyDeviceToHost)); + std::vector de(TCAP); + for (uint64_t i = 0; i < TCAP; ++i) + de[i] = __atomic_load_n(&rs.publication_slots[i], __ATOMIC_ACQUIRE); + const uint64_t th = *rs.task_head, ph = *rs.payload_head, ac = *rs.actual_bytes_counter; + + int fails = 0; + uint64_t bad = 0, first = ~0ull; + for (uint64_t i = 0; i < CAP; ++i) if (dev[i] != m.payload[i]) { if (!bad) first = i; ++bad; } + if (bad) { printf("FAIL payload: %lu bytes differ (first at %lu)\n", (unsigned long)bad, (unsigned long)first); ++fails; } + if (de != m.words) { printf("FAIL publication words\n"); ++fails; } + if (th != m.th || ph != m.ph || ac != m.actual) { + printf("FAIL heads: task %lu/%lu payload %lu/%lu actual %lu/%lu\n", + (unsigned long)th, (unsigned long)m.th, (unsigned long)ph, (unsigned long)m.ph, + (unsigned long)ac, (unsigned long)m.actual); ++fails; + } + printf("ops=%d tasks=%lu payload_head=%lu (ring wrapped %lu times) actual=%lu -> %s\n", + n_ops, (unsigned long)th, (unsigned long)ph, (unsigned long)(ph / CAP), (unsigned long)ac, + fails ? "FAIL" : "PASS"); + if (argc > 1) { + FILE* f = fopen(argv[1], "wb"); + if (!f) { perror("dump"); return 2; } + fwrite(dev.data(), 1, CAP, f); + fwrite(de.data(), sizeof(uint64_t), TCAP, f); + fwrite(&th, 8, 1, f); fwrite(&ph, 8, 1, f); fwrite(&ac, 8, 1, f); + fclose(f); + } + return fails ? 1 : 0; +} diff --git a/tests/native/ring/test_ring_torch_op_host.cpp b/tests/native/ring/test_ring_torch_op_host.cpp new file mode 100644 index 000000000..1e3bf1adf --- /dev/null +++ b/tests/native/ring/test_ring_torch_op_host.cpp @@ -0,0 +1,270 @@ +// Host-only harness for the ring producer ops (no CUDA toolkit or GPU). +// +// native/csrc/ring/ring_torch_op.cpp is compiled unchanged against CPU-only +// Torch and tests/native/ring/cpu_cuda_stub/ (CUDA runtime API and +// stand-ins), and linked with this file, which +// supplies a recording RingEnginePy double. Tensors are host buffers tagged +// as CUDA device 0 (from_blob with a target device), so the ops' own argument +// validation runs; nothing is ever launched. tests/test_ring_torch_op_host.py +// builds it as a shared library and drives it through the C entry points +// below, from the real dmi.node_toggle lifecycle. +#include +#include + +#include +#include +#include +#include + +#include "producer.cuh" +#include "ring_engine_py.h" +#include "ring_torch_op.h" + +// The op implementations (registered by TORCH_LIBRARY_IMPL in +// ring_torch_op.cpp; not declared in a header). +void ring_producer_impl(const at::Tensor&, const at::Tensor&, int64_t, int64_t); +void ring_producer_prefix_impl(const at::Tensor&, const at::Tensor&, + const at::Tensor&, int64_t, int64_t, int64_t); +void ring_producer_chunked_impl(const at::Tensor&, const at::Tensor&, + const at::Tensor&, int64_t, int64_t); +void ring_producer_gather_impl(const at::Tensor&, const at::Tensor&, + const at::Tensor&, int64_t, int64_t, int64_t); +void ring_producer_head_impl(const at::Tensor&, const at::Tensor&, + const at::Tensor&, const at::Tensor&, int64_t, + int64_t, int64_t); +void ring_producer_fused_impl(const at::Tensor&, const at::Tensor&, + const c10::optional&, int64_t, + const c10::optional&, + const at::Tensor&, const at::Tensor&, int64_t, + int64_t, int64_t, int64_t, int64_t); + +// --------------------------------------------------------------------------- +// RingEnginePy double: records which launcher each op reached and the capture +// nodes it registered. + +namespace { + +struct EngineRecord { + uint64_t launcher_calls = 0; + uint64_t anomalies = 0; + // graph -> (hook_type, layer_no) of each registered node + std::map>> nodes; +}; + +EngineRecord g_engine; + +} // namespace + +namespace ring_py { + +struct RingEnginePy::Impl {}; + +RingEnginePy::RingEnginePy(RingConfig, SubmitFn) : impl_(new Impl) {} +RingEnginePy::~RingEnginePy() = default; + +void RingEnginePy::hook_no_notify(uint64_t, uint64_t, uint32_t, uint64_t) { + ++g_engine.launcher_calls; +} +void RingEnginePy::hook_no_notify_prefix(uint64_t, uint64_t, uint64_t, + uint64_t, uint32_t, uint64_t) { + ++g_engine.launcher_calls; +} +void RingEnginePy::hook_no_notify_chunked(uint64_t, uint64_t, uint64_t, + uint32_t, uint32_t, uint64_t) { + ++g_engine.launcher_calls; +} +void RingEnginePy::hook_no_notify_gather(uint64_t, uint64_t, uint64_t, + int64_t, uint64_t, uint32_t, + uint64_t) { + ++g_engine.launcher_calls; +} +void RingEnginePy::hook_no_notify_head(uint64_t, uint64_t, uint64_t, int64_t, + uint32_t, uint32_t, uint64_t, uint32_t, + uint32_t, uint64_t) { + ++g_engine.launcher_calls; +} +void RingEnginePy::hook_no_notify_fused(uint64_t, uint64_t, uint32_t, + uint32_t, uint64_t, int64_t, + uint64_t, int64_t, uint64_t, + uint32_t, uint32_t, uint64_t, + uint64_t, uint32_t, uint64_t) { + ++g_engine.launcher_calls; +} +void RingEnginePy::record_no_notify(uint64_t, uint64_t, uint64_t, int32_t, + uint64_t) {} +void RingEnginePy::record_no_notify_prefix(uint64_t, uint64_t, uint64_t, + uint64_t, uint64_t, int32_t, + uint64_t) {} +void RingEnginePy::record_no_notify_chunked(uint64_t, uint64_t, uint64_t, + uint32_t, uint64_t, int32_t, + uint64_t) {} +void RingEnginePy::record_no_notify_seq_prefix_pack(uint64_t, uint64_t, + uint64_t, uint64_t, + uint32_t, uint64_t, + uint64_t, int32_t, + uint64_t) {} +void RingEnginePy::record_no_notify_segmented_pack(uint64_t, uint64_t, + uint64_t, uint64_t, + uint32_t, uint64_t, + uint64_t, int32_t, + uint64_t) {} +void RingEnginePy::register_capture_node(uint64_t graph, int hook_type, + int layer_no, uint64_t) { + g_engine.nodes[graph].emplace_back(hook_type, layer_no); +} +void RingEnginePy::note_capture_anomaly() { ++g_engine.anomalies; } + +} // namespace ring_py + +// --------------------------------------------------------------------------- +// Fixture tensors: fp32 x[ROWS, HIDDEN] and the per-op side tensors, all host +// memory tagged as CUDA device 0. + +namespace { + +constexpr int64_t ROWS = 4; +constexpr int64_t HIDDEN = 8; +constexpr int64_t ROW_BYTES = HIDDEN * 4; // fp32 rows, a 16-byte multiple +constexpr int64_t K = 2; +constexpr int64_t HOOK_TYPE = 3; +constexpr int64_t HEAD_HOOK_TYPE = 24; + +struct Fixture { + std::vector> blobs; + ring_py::RingEnginePy engine{ring_py::RingConfig{}, ring_py::SubmitFn{}}; + + at::Tensor tensor(at::IntArrayRef sizes, at::ScalarType dtype, + bool on_cuda = true) { + int64_t n = 1; + for (auto s : sizes) n *= s; + blobs.emplace_back(static_cast(n) * at::elementSize(dtype) + 16); + const at::Device device = on_cuda ? at::Device(at::kCUDA, 0) + : at::Device(at::kCPU); + return at::for_blob(blobs.back().data(), sizes) + .options(at::TensorOptions().dtype(dtype).device(device)) + .target_device(device) + .make_tensor(); + } +}; + +Fixture& fixture() { + static Fixture f; + return f; +} + +std::string g_last_error; + +// Runs one op. Returns 0 on success, 1 when the op threw, 2 for an unknown +// op name. +int run(const std::string& op) { + Fixture& f = fixture(); + at::Tensor payload = f.tensor({16}, at::kByte); + at::Tensor x = f.tensor({ROWS, HIDDEN}, at::kFloat); + at::Tensor table = f.tensor({3}, at::kLong); + const auto arena_floats = static_cast(ring::head_arena_floats( + static_cast(K), static_cast(HIDDEN))); + const auto node_words = static_cast(ring::head_node_words( + static_cast(K), static_cast(HIDDEN), false)); + at::Tensor acc = f.tensor({ROWS, static_cast(ring::HEAD_MAX_ACC)}, + at::kFloat); + const c10::optional none; + if (op == "static") { + ring_producer_impl(payload, x, HOOK_TYPE, 0); + } else if (op == "prefix") { + ring_producer_prefix_impl(payload, x, f.tensor({1}, at::kLong), + ROW_BYTES, HOOK_TYPE, 0); + } else if (op == "chunked") { + ring_producer_chunked_impl(payload, x, f.tensor({2}, at::kLong), + HOOK_TYPE, 0); + } else if (op == "gather") { + ring_producer_gather_impl(payload, x, table, ROW_BYTES, HOOK_TYPE, 0); + } else if (op == "head") { + ring_producer_head_impl(payload, x, table, + f.tensor({arena_floats}, at::kFloat), K, + HEAD_HOOK_TYPE, 0); + } else if (op == "fused") { // rows and heads in one node + ring_producer_fused_impl(payload, x, table, ROW_BYTES, table, + f.tensor({node_words}, at::kInt), acc, K, 0, + HOOK_TYPE, 0, HEAD_HOOK_TYPE); + } else if (op == "fused_rows") { + ring_producer_fused_impl(payload, x, table, ROW_BYTES, none, + f.tensor({1}, at::kInt), acc, K, 0, + HOOK_TYPE, 0, HEAD_HOOK_TYPE); + } else if (op == "fused_heads") { + ring_producer_fused_impl(payload, x, none, ROW_BYTES, table, + f.tensor({node_words}, at::kInt), acc, K, 0, + HOOK_TYPE, 0, HEAD_HOOK_TYPE); + } else if (op == "fused_empty") { // neither part: nothing launches + ring_producer_fused_impl(payload, x, none, ROW_BYTES, none, + f.tensor({1}, at::kInt), acc, K, 0, + HOOK_TYPE, 0, HEAD_HOOK_TYPE); + } else if (op == "gather_host_tensor") { // not a CUDA tensor: skipped + ring_producer_gather_impl(payload, f.tensor({ROWS, HIDDEN}, at::kFloat, + /*on_cuda=*/false), + table, ROW_BYTES, HOOK_TYPE, 0); + } else if (op == "gather_bad_table") { // refused before launching + ring_producer_gather_impl(payload, x, f.tensor({3}, at::kInt), + ROW_BYTES, HOOK_TYPE, 0); + } else { + return 2; + } + return 0; +} + +} // namespace + +extern "C" { + +// Activates the engine double (ring_set_active_engine) on first use. +int dmx_host_op(const char* op, uint64_t capturing_graph) { + Fixture& f = fixture(); + ring_set_active_engine(&f.engine); + fake_capture_graph = reinterpret_cast(capturing_graph); + int rc; + try { + rc = run(op); + g_last_error.clear(); + } catch (const std::exception& e) { + g_last_error = e.what(); + rc = 1; + } + fake_capture_graph = nullptr; + return rc; +} + +const char* dmx_host_last_error() { return g_last_error.c_str(); } + +void dmx_host_set_toggle_capture(int enabled) { + ring_set_toggle_capture(enabled != 0); +} + +uint64_t dmx_host_launch_count() { return ring_producer_launch_count(); } + +// Launcher calls the engine double received (one per launched producer). +uint64_t dmx_host_launcher_calls() { return g_engine.launcher_calls; } + +uint64_t dmx_host_anomalies() { return g_engine.anomalies; } + +// Nodes registered for `graph`; graph 0 counts every graph's. +uint64_t dmx_host_node_count(uint64_t graph) { + if (graph != 0) { + auto it = g_engine.nodes.find(graph); + return it == g_engine.nodes.end() ? 0 : it->second.size(); + } + uint64_t n = 0; + for (const auto& kv : g_engine.nodes) n += kv.second.size(); + return n; +} + +// Toggle key of the i-th node registered for `graph`: hook_type << 32 | layer. +uint64_t dmx_host_node_key(uint64_t graph, uint64_t i) { + const auto& v = g_engine.nodes.at(graph).at(i); + return (static_cast(v.first) << 32) | static_cast(v.second); +} + +void dmx_host_clear_nodes() { + g_engine.nodes.clear(); + g_engine.anomalies = 0; +} + +} // extern "C" diff --git a/tests/native/test_row_table_clamp.cpp b/tests/native/test_row_table_clamp.cpp new file mode 100644 index 000000000..2ca35a4e8 --- /dev/null +++ b/tests/native/test_row_table_clamp.cpp @@ -0,0 +1,84 @@ +// Row-table segment-count clamp (native/csrc/ring/row_table.h): the gather, +// head and fused producers read table[1 + 2k] and table[2 + 2k] for every +// k < row_table_segment_count(table[0], table_numel). Those reads must stay +// below table_numel whatever the count word holds, and the count must also +// respect GATHER_MAX_SEGS. Compiled and run on the CPU by +// tests/test_row_table_clamp.py; plain C++17, no CUDA. +#include "ring/row_table.h" + +#include +#include +#include + +namespace { + +int failures = 0; + +void expect(int64_t n_raw, int64_t numel, uint32_t want) { + const uint32_t got = ring::row_table_segment_count(n_raw, numel); + if (got != want) { + std::printf("FAIL count(n_raw=%lld, numel=%lld) = %u, want %u\n", + static_cast(n_raw), static_cast(numel), + got, want); + ++failures; + } + // Every word the loader reads is inside the table. + for (uint32_t k = 0; k < got; ++k) { + const int64_t last = 2 + 2 * static_cast(k); + if (last >= numel) { + std::printf("FAIL read table[%lld] past numel=%lld (n_raw=%lld)\n", + static_cast(last), static_cast(numel), + static_cast(n_raw)); + ++failures; + break; + } + } +} + +} // namespace + +int main() { + static_assert(ring::GATHER_MAX_SEGS == 256, "row-table capacity"); + const int64_t big = std::numeric_limits::max(); + + // One-word tables hold only the count: no segment is ever read. + expect(0, 1, 0); + expect(1, 1, 0); + expect(256, 1, 0); + expect(big, 1, 0); + // Three words hold one segment. + expect(2, 3, 1); + expect(1, 3, 1); + // Even lengths: the trailing odd word is not a whole segment. + expect(2, 2, 0); + expect(2, 4, 1); + expect(3, 6, 2); + // Odd lengths. + expect(2, 5, 2); + expect(3, 7, 3); + expect(9, 7, 3); + // Negative counts (including the -1 disable sentinel) read nothing. + expect(-1, 9, 0); + expect(-7, 9, 0); + expect(std::numeric_limits::min(), 9, 0); + // GATHER_MAX_SEGS limits: a full-width table. + expect(255, 1 + 2 * 256, 255); + expect(256, 1 + 2 * 256, 256); + expect(257, 1 + 2 * 256, 256); + expect(257, 1 + 2 * 257, 256); + expect(big, 1 + 2 * 300, 256); + // One word short of 256 segments. + expect(256, 2 * 256, 255); + // Degenerate storage (rejected by the host before launch) reads nothing. + expect(1, 0, 0); + expect(1, -3, 0); + + if (ring::row_table_capacity_segs(1) != 0 || ring::row_table_capacity_segs(3) != 1 || + ring::row_table_capacity_segs(4) != 1 || ring::row_table_capacity_segs(513) != 256 || + ring::row_table_capacity_segs(0) != 0) { + std::printf("FAIL row_table_capacity_segs\n"); + ++failures; + } + if (failures == 0) std::printf("ok\n"); + return failures == 0 ? 0 : 1; +} diff --git a/tests/test_adapter_protocol.py b/tests/test_adapter_protocol.py index 117b9c8be..4405588f7 100644 --- a/tests/test_adapter_protocol.py +++ b/tests/test_adapter_protocol.py @@ -494,3 +494,21 @@ def test_plan_step_needs_no_transport(): ctx = dataclasses.replace(_make_ctx(), q_len=2, batch=0, flattened=True) plan = BackendAdapter.plan_step(adaptor, ctx) assert (plan.total_bytes, plan.hook_count, plan.needs_eager) == (64, 1, False) + + +def test_replace_active_specs_keeps_the_transport_in_sync(): + from types import SimpleNamespace + + adaptor = _make_adaptor(0) + engine = adaptor.engine + hook_a, hook_b = object(), object() + first = [SimpleNamespace(module=hook_a, hook_type=0, layer_no=0), + SimpleNamespace(module=hook_b, hook_type=0, layer_no=1)] + adaptor.active_specs = list(first) + head = SimpleNamespace(module=hook_a, hook_type=24, layer_no=0) + adaptor.replace_active_specs([first[0], head, first[1]]) + assert adaptor.active_specs == [first[0], head, first[1]] + assert engine._ring_transport._active_specs is adaptor.active_specs + with pytest.raises(ValueError, match="installed HookPoints"): + adaptor.replace_active_specs([SimpleNamespace(module=object(), hook_type=0, layer_no=2)]) + assert engine._ring_transport._active_specs is adaptor.active_specs diff --git a/tests/test_hook_row_features.py b/tests/test_hook_row_features.py new file mode 100644 index 000000000..0d391bdc5 --- /dev/null +++ b/tests/test_hook_row_features.py @@ -0,0 +1,548 @@ +"""Gather tables and in-graph heads on a HookPoint: lifecycle and dispatch. + +Two contracts, checked with fakes only (no native extension, no CUDA): + +* Attachment lifecycle. ``configure_hook_gather`` / ``configure_hook_heads`` + arm a HookPoint for one attachment. ``uninstall_ring_hooks`` and + ``install_ring_hooks`` clear that state -- the tables, the head arena and + accumulator references, and the strip stride gather replaced -- so a + detached HookPoint releases them and a plain re-attach runs its ordinary + producer instead of a stale fused node. + +* Task agreement. For every combination of row mode (static, prefix, + chunked, gather) and heads, the host plan (``BackendAdapter.plan_step``), + the fast (capturable) path and the eager safety net agree on the tasks a + forward publishes and on their order: the row task, then the head task. + Adding heads to a HookPoint without a gather table keeps its row producer. +""" + +from __future__ import annotations + +import gc +import weakref +from types import SimpleNamespace + +import pytest +import torch + +from dmi.adapters.base import BackendAdapter +from dmi.adapters.types import StepContext +from dmi.hook_rows import ( + clear_hook_row_features, + configure_hook_gather, + configure_hook_heads, +) +from dmi.hooks.dispatch import install_ring_hooks, uninstall_ring_hooks +from dmi.hooks.point import HookPoint +from dmi.hooks.specs import ( + HOOK_TYPE_RESID_PRE, + HOOK_TYPE_RESID_PRE_HEAD, + HookSpec, + ModelShapeConfig, +) + +pytestmark = pytest.mark.cpu + +ROW, HEAD = HOOK_TYPE_RESID_PRE, HOOK_TYPE_RESID_PRE_HEAD +ROWS, HID, SLOTS = 4, 4, 2 +ROW_BYTES = HID * 4 +OPS = ("producer", "producer_prefix", "producer_chunked", "producer_gather", + "producer_head", "producer_fused") + + +class _FakeCudaTensor(torch.Tensor): + @property + def is_cuda(self) -> bool: + return True + + +def _x(): + return torch.zeros(ROWS, HID).as_subclass(_FakeCudaTensor) + + +def _table(): + return torch.zeros(1 + 2 * 4, dtype=torch.int64) + + +@pytest.fixture +def calls(monkeypatch): + """Record every producer op launch: (op, args).""" + import dmi.transport.ring as rt + + log = [] + for name in OPS: + monkeypatch.setattr( + torch.ops.ring, name, + lambda *args, _name=name: log.append((_name, args)), raising=False) + monkeypatch.setattr(rt, "_active_transport", None) + return log + + +def _published(calls): + """Tasks a list of producer launches publishes, as (hook_type, layer). + + A producer_fused launch publishes its row part when it has a gather + table and its head part when it has a head table (the kernel's g_on / + h_on predicates), rows first.""" + tasks = [] + for name, args in calls: + if name == "producer_fused": + gather, head = args[2], args[4] + hook_type, hook_id, head_type = args[9], args[10], args[11] + if gather is not None: + tasks.append((hook_type, hook_id)) + if head is not None: + tasks.append((head_type, hook_id)) + elif name in ("producer", "producer_prefix", "producer_chunked", + "producer_gather"): + tasks.append((args[-2], args[-1])) + elif name == "cpu_direct": + tasks.append(args) + return tasks + + +def _arm_heads(hook, table=None, layer=0): + arena = torch.zeros(80, dtype=torch.int32) + acc = torch.zeros(ROWS, 16) + configure_hook_heads(hook, head_type=HEAD, layer_no=layer, + table=_table() if table is None else table, + arena=arena, acc=acc, slots=SLOTS, weight_bf16=False) + return arena, acc + + +def _arm_strip(hook, mode): + if mode == "prefix": + hook._strip_tensor = torch.tensor([ROWS], dtype=torch.int64) + hook._strip_row_bytes = ROW_BYTES + elif mode == "chunked": + hook._strip_tensor = torch.tensor([ROWS * ROW_BYTES // 2] * 2, + dtype=torch.int64) + hook._strip_row_bytes = 0 + + +def _feature_fields(hook): + return (hook._gather_table, hook._head_hook_type, hook._head_hook_id, + hook._head_table, hook._head_arena, hook._head_acc, + hook._head_slots, hook._head_wbf16) + + +CLEARED = (None, None, None, None, None, None, 0, 1) + + +# --------------------------------------------------------------------------- +# Attachment lifecycle + + +def test_detach_clears_gather_and_heads_and_releases_the_tensors(calls): + hook = HookPoint() + spec = HookSpec(hook_type=ROW, layer_no=0, module=hook) + install_ring_hooks([spec], ring_payload=object()) + table = _table() + configure_hook_gather(hook, table, ROW_BYTES) + arena, acc = _arm_heads(hook) + hook.forward(_x()) + assert [c[0] for c in calls] == ["producer_fused"] + + refs = [weakref.ref(t) for t in (table, arena, acc)] + del table, arena, acc + calls.clear() # the launch log holds the args + uninstall_ring_hooks([spec]) + gc.collect() + assert _feature_fields(hook) == CLEARED + assert hook._ring_hook_type is None and hook._ring_payload is None + assert all(ref() is None for ref in refs) + + # A plain re-attach runs the ordinary producer and publishes one task. + install_ring_hooks([spec], ring_payload=object()) + calls.clear() + hook.forward(_x()) + assert [c[0] for c in calls] == ["producer"] + assert _published(calls) == [(ROW, 0)] + + +def test_disabled_hook_point_is_cleared_and_keeps_its_selection(calls): + hook = HookPoint() + hook.enabled = False + spec = HookSpec(hook_type=ROW, layer_no=0, module=hook) + install_ring_hooks([spec], ring_payload=object()) + configure_hook_gather(hook, _table(), ROW_BYTES) + _arm_heads(hook) + hook.forward(_x()) + assert calls == [] + uninstall_ring_hooks([spec]) + assert _feature_fields(hook) == CLEARED + assert hook.enabled is False # the selection is not attachment state + hook.enabled = True + install_ring_hooks([spec], ring_payload=object()) + hook.forward(_x()) + assert [c[0] for c in calls] == ["producer"] + + +def test_detach_in_finally_after_a_failing_forward_clears_the_state(calls, monkeypatch): + hook = HookPoint() + spec = HookSpec(hook_type=ROW, layer_no=0, module=hook) + install_ring_hooks([spec], ring_payload=object()) + configure_hook_gather(hook, _table(), ROW_BYTES) + _arm_heads(hook) + + def boom(*args): + raise RuntimeError("producer failed") + + monkeypatch.setattr(torch.ops.ring, "producer_fused", boom, raising=False) + with pytest.raises(RuntimeError, match="producer failed"): + try: + hook.forward(_x()) + finally: + uninstall_ring_hooks([spec]) + assert _feature_fields(hook) == CLEARED + + +def test_rejected_configuration_leaves_the_previous_state_untouched(): + hook = HookPoint() + table = _table() + configure_hook_gather(hook, table, ROW_BYTES) + with pytest.raises(ValueError): + configure_hook_gather(hook, table.to(torch.int32), ROW_BYTES) + with pytest.raises(ValueError): + configure_hook_heads(hook, head_type=HEAD, layer_no=0, table=table, + arena=torch.zeros(8, dtype=torch.int32), + acc=torch.zeros(1, 16), slots=17, weight_bf16=False) + with pytest.raises(ValueError, match="head row table"): + configure_hook_heads(hook, head_type=HEAD, layer_no=0, table=None, + arena=torch.zeros(8, dtype=torch.int32), + acc=torch.zeros(1, 16), slots=1, weight_bf16=False) + assert hook._gather_table is table and hook._strip_row_bytes == ROW_BYTES + assert hook._head_hook_type is None and hook._head_arena is None + + +def test_reinstall_without_detach_starts_a_new_attachment(calls): + hook = HookPoint() + spec = HookSpec(hook_type=ROW, layer_no=0, module=hook) + install_ring_hooks([spec], ring_payload=object()) + configure_hook_gather(hook, _table(), ROW_BYTES) + _arm_heads(hook) + install_ring_hooks([spec], ring_payload=object()) + assert _feature_fields(hook) == CLEARED + hook.forward(_x()) + assert [c[0] for c in calls] == ["producer"] + + +@pytest.mark.parametrize("mode,op", [("prefix", "producer_prefix"), + ("chunked", "producer_chunked")]) +def test_clearing_gather_restores_the_strip_mode(calls, mode, op): + hook = HookPoint() + spec = HookSpec(hook_type=ROW, layer_no=0, module=hook) + install_ring_hooks([spec], ring_payload=object()) + _arm_strip(hook, mode) + strip_rb = hook._strip_row_bytes + configure_hook_gather(hook, _table(), ROW_BYTES) + configure_hook_gather(hook, _table(), ROW_BYTES) # re-arming keeps the original + assert hook._strip_row_bytes == ROW_BYTES + configure_hook_gather(hook, None, 0) + assert hook._strip_row_bytes == strip_rb and hook._gather_table is None + hook.forward(_x()) + configure_hook_gather(hook, _table(), ROW_BYTES) + uninstall_ring_hooks([spec]) + install_ring_hooks([spec], ring_payload=object()) + hook.forward(_x()) + assert [c[0] for c in calls] == [op, op] + assert hook._strip_row_bytes == strip_rb + + +def test_clear_is_idempotent_on_plain_modules(): + module = SimpleNamespace(_gather_table=None, _strip_row_bytes=0) + clear_hook_row_features(module) + clear_hook_row_features(module) + assert module._gather_table is None and module._head_hook_type is None + assert module._strip_row_bytes == 0 + + +def test_a_head_spec_never_replaces_the_row_key(calls): + hook = HookPoint() + row = HookSpec(hook_type=ROW, layer_no=3, module=hook) + head = HookSpec(hook_type=HEAD, layer_no=3, module=hook, + dtype=torch.float32) + install_ring_hooks([row, head], ring_payload=object()) + assert (hook._ring_hook_type, hook._ring_hook_id) == (ROW, 3) + lone = HookPoint() + install_ring_hooks([HookSpec(hook_type=HEAD, layer_no=1, module=lone)], + ring_payload=object()) + assert lone._ring_hook_type is None and lone._ring_payload is not None + # Re-installing a former row HookPoint with only its head spec (no + # uninstall in between) leaves it head-only, not with the old row key. + install_ring_hooks([head], ring_payload=object()) + assert hook._ring_hook_type is None and hook._ring_payload is not None + + +# --------------------------------------------------------------------------- +# Task agreement: plan, fast path and eager safety net + + +class _Engine: + def __init__(self, enabled): + self.enabled = enabled + self.events = [] + + def is_hook_enabled(self, hook_type, layer): + return (hook_type, layer) in self.enabled + + def available_capacity(self): + return 1 << 30 + + def available_task_slots(self): + return 64 + + def flush_and_wait(self): + self.events.append("flush") + + def reserve_one(self, nbytes): + self.events.append(("reserve", nbytes)) + + +def _eager_transport(engine, calls, *, gate): + def cpu_direct(tensor, hook_type, hook_id): + calls.append(("cpu_direct", (hook_type, hook_id))) + + return SimpleNamespace( + capture_step=True, force_eager=True, _ring_engine=engine, + toggle_gate_active=gate, _gated_step=True, effective_cap=1 << 30, + submit_cpu_direct=cpu_direct) + + +def _plan(specs): + cfg = ModelShapeConfig(hidden_dim=HID, num_heads=1, num_kv_heads=1, + head_dim=HID, dtype=torch.float32, head_slots=SLOTS) + adapter = SimpleNamespace( + model_cfg=cfg, active_specs=list(specs), + transport=SimpleNamespace(_active_specs=list(specs)), + _spec_needs_eager=lambda _: False) + ctx = StepContext(model_id="m", flattened=True, req_ids=["A"], + token_ranges=[(0, ROWS)], dim0_offsets=[0], kv_offsets=[0], + q_len=ROWS, batch=0) + return BackendAdapter.plan_step(adapter, ctx) + + +def _build(mode, heads, row=True): + hook = HookPoint() + row_spec = HookSpec(hook_type=ROW, layer_no=0, module=hook, + dim0_is_actual_tokens=True) + head_spec = HookSpec(hook_type=HEAD, layer_no=0, module=hook, + dtype=torch.float32, dim0_is_actual_tokens=True) + specs = ([row_spec] if row else []) + ([head_spec] if heads else []) + install_ring_hooks(specs, ring_payload=object()) + if row: + _arm_strip(hook, mode) + if mode == "gather": + configure_hook_gather(hook, _table(), ROW_BYTES) + if heads: + _arm_heads(hook) + return hook, specs + + +FAST_OPS = { + # (row mode, heads) -> launches of the fast (capturable) path + ("static", True): ["producer", "producer_fused"], + ("prefix", True): ["producer_prefix", "producer_fused"], + ("chunked", True): ["producer_chunked", "producer_fused"], + ("gather", True): ["producer_fused"], + ("static", False): ["producer"], + ("gather", False): ["producer_gather"], +} + + +@pytest.mark.parametrize("mode,heads", sorted(FAST_OPS)) +def test_fast_path_publishes_the_planned_tasks_in_spec_order(calls, mode, heads): + hook, specs = _build(mode, heads) + hook.forward(_x()) + assert [c[0] for c in calls] == FAST_OPS[(mode, heads)] + published = _published(calls) + assert published == [(s.hook_type, s.layer_no) for s in specs] + assert _plan(specs).hook_count == len(published) + if heads and mode != "gather": + # The head-only node carries no row table: it is recorded and + # toggled under the head key alone, independent of the row node. + fused = calls[-1][1] + assert fused[2] is None and fused[4] is hook._head_table + assert (fused[9], fused[10], fused[11]) == (HEAD, 0, HEAD) + + +def test_head_only_hook_point_publishes_the_head_task(calls): + hook, specs = _build("static", True, row=False) + assert hook._ring_hook_type is None + hook.forward(_x()) + assert [c[0] for c in calls] == ["producer_fused"] + assert _published(calls) == [(HEAD, 0)] + assert _plan(specs).hook_count == 1 + + +def test_reviewed_plain_plus_heads_configuration_publishes_both_tasks(calls): + """Heads configured on an installed ordinary HookPoint, the head spec + appended with the public replace_active_specs API: the plan has the row + and the head task (64 + 16 bytes) and the forward publishes both.""" + hook = HookPoint() + spec = HookSpec(hook_type=ROW, layer_no=0, module=hook) + install_ring_hooks([spec], ring_payload=object()) + _arm_heads(hook, table=torch.tensor([1, 0, ROWS], dtype=torch.int64)) + head_spec = HookSpec(hook_type=HEAD, layer_no=0, module=hook, + dtype=torch.float32) + cfg = ModelShapeConfig(hidden_dim=HID, num_heads=1, num_kv_heads=1, + head_dim=HID, dtype=torch.float32, head_slots=1) + adapter = SimpleNamespace(model_cfg=cfg, active_specs=[spec], + transport=SimpleNamespace(_active_specs=[spec]), + _spec_needs_eager=lambda _: False) + BackendAdapter.replace_active_specs(adapter, [spec, head_spec]) + ctx = StepContext(model_id="m", flattened=True, req_ids=["A"], + token_ranges=[(0, ROWS)], dim0_offsets=[0], kv_offsets=[0], + q_len=ROWS, batch=0) + plan = BackendAdapter.plan_step(adapter, ctx) + assert plan.hook_count == 2 and plan.total_bytes == 80 + hook.forward(_x()) + assert [c[0] for c in calls] == ["producer", "producer_fused"] + assert _published(calls) == [(ROW, 0), (HEAD, 0)] + + +EAGER_ROW_OPS = {"static": "producer", "gather": "producer_gather", + "prefix": "cpu_direct", "chunked": "cpu_direct"} + + +@pytest.mark.parametrize("mode", sorted(EAGER_ROW_OPS)) +@pytest.mark.parametrize("row_on,head_on", [(True, True), (True, False), + (False, True), (False, False)]) +def test_eager_safety_net_matches_the_fast_path_and_the_toggles( + calls, monkeypatch, mode, row_on, head_on): + """With independent row/head toggle states, the eager safety net + publishes exactly the enabled subset, row first; with every key + enabled it publishes what the fast path publishes.""" + import dmi.transport.ring as rt + + hook, specs = _build(mode, True) + enabled = set() + if row_on: + enabled.add((ROW, 0)) + if head_on: + enabled.add((HEAD, 0)) + engine = _Engine(enabled) + monkeypatch.setattr(rt, "_active_transport", + _eager_transport(engine, calls, gate=True)) + hook.forward(_x()) + want = [key for key in ((ROW, 0), (HEAD, 0)) if key in enabled] + assert _published(calls) == want + names = [c[0] for c in calls] + expected = ([EAGER_ROW_OPS[mode]] if row_on else []) + ( + ["producer_fused"] if head_on else []) + assert names == expected + if head_on: + # The eager head is a head-only node (no row table). + assert calls[-1][1][2] is None + reserves = [e for e in engine.events if isinstance(e, tuple)] + assert len(reserves) == sum(1 for n in names if n != "cpu_direct") + + if row_on and head_on: + # Same tasks as the ungated fast path. + calls.clear() + monkeypatch.setattr(rt, "_active_transport", None) + hook.forward(_x()) + assert _published(calls) == want + + +# --------------------------------------------------------------------------- +# Eager head sizing: every physical source row, whatever the HookPoint rank + +SHAPES = { + "batch_seq": (2, 3, HID), # [B, S, H] + "one_batch_many_tokens": (1, 5, HID), + "one_token_per_batch": (3, 1, HID), + "packed": (6, HID), # [N, H] + "single_row": (HID,), # [H] +} + + +def _shaped_ctx(shape): + if len(shape) == 3: + b, s, _ = shape + return StepContext(model_id="m", flattened=False, + req_ids=[f"r{i}" for i in range(b)], + token_ranges=[(0, s)] * b, + dim0_offsets=list(range(b)), kv_offsets=[0] * b, + batch=b, q_len=s) + rows = 1 if len(shape) == 1 else shape[0] + return StepContext(model_id="m", flattened=True, req_ids=["A"], + token_ranges=[(0, rows)], dim0_offsets=[0], + kv_offsets=[0], q_len=rows, batch=0) + + +def _shaped_heads(shape, *, row): + """A HookPoint with full-row heads over a ``shape`` source (the table an + integration writes on an eager step), its specs and the step plan.""" + rows = 1 + for d in shape[:-1]: + rows *= d + hook = HookPoint() + row_spec = HookSpec(hook_type=ROW, layer_no=0, module=hook) + head_spec = HookSpec(hook_type=HEAD, layer_no=0, module=hook, + dtype=torch.float32) + specs = ([row_spec] if row else []) + [head_spec] + install_ring_hooks(specs, ring_payload=object()) + configure_hook_heads(hook, head_type=HEAD, layer_no=0, + table=torch.tensor([1, 0, rows], dtype=torch.int64), + arena=torch.zeros(80, dtype=torch.int32), + acc=torch.zeros(rows, 16), slots=SLOTS, + weight_bf16=False) + cfg = ModelShapeConfig(hidden_dim=HID, num_heads=1, num_kv_heads=1, + head_dim=HID, dtype=torch.float32, head_slots=SLOTS) + adapter = SimpleNamespace(model_cfg=cfg, active_specs=specs, + transport=SimpleNamespace(_active_specs=specs), + _spec_needs_eager=lambda _: False) + plan = BackendAdapter.plan_step(adapter, _shaped_ctx(shape)) + return hook, rows, plan + + +@pytest.mark.parametrize("shape", list(SHAPES.values()), ids=list(SHAPES)) +@pytest.mark.parametrize("row", [True, False], ids=["rows_and_heads", "head_only"]) +def test_eager_heads_reserve_what_the_plan_budgets_for_every_rank( + calls, monkeypatch, shape, row): + """The eager head reserves rows x K x 4 bytes with rows the product of + every dimension before hidden -- what the native launcher writes + (nbytes / (hidden x element size)) and what plan_step budgets -- not + shape[0] x K x 4.""" + import dmi.transport.ring as rt + + hook, rows, plan = _shaped_heads(shape, row=row) + engine = _Engine({(ROW, 0), (HEAD, 0)}) + monkeypatch.setattr(rt, "_active_transport", + _eager_transport(engine, calls, gate=False)) + x = torch.zeros(*shape).as_subclass(_FakeCudaTensor) + hook.forward(x) + head_bytes = rows * SLOTS * 4 + reserves = [e[1] for e in engine.events if isinstance(e, tuple)] + assert reserves == ([x.nbytes] if row else []) + [head_bytes] + # The plan budgets each task 16-byte aligned, as the ring stores it. + assert sum(-(-r // 16) * 16 for r in reserves) == plan.total_bytes + assert [c[0] for c in calls] == (["producer"] if row else []) + [ + "producer_fused"] + + +@pytest.mark.parametrize("shape", [SHAPES["batch_seq"], SHAPES["packed"]], + ids=["batch_seq", "packed"]) +def test_eager_heads_over_the_effective_capacity_are_refused( + calls, monkeypatch, shape): + """The capacity guard sees the full head output: 48 bytes for six rows + of two slots is refused under a 32-byte effective capacity before any + reservation or launch, and admitted at exactly 48.""" + import dmi.transport.ring as rt + + hook, rows, _ = _shaped_heads(shape, row=False) + head_bytes = rows * SLOTS * 4 + assert head_bytes == 48 + engine = _Engine({(HEAD, 0)}) + transport = _eager_transport(engine, calls, gate=False) + transport.effective_cap = 32 + monkeypatch.setattr(rt, "_active_transport", transport) + x = torch.zeros(*shape).as_subclass(_FakeCudaTensor) + with pytest.raises(RuntimeError, match=r"48 bytes for 6 rows x 2 slots"): + hook.forward(x) + assert engine.events == [] and calls == [] + transport.effective_cap = head_bytes + hook.forward(x) + assert engine.events == [("reserve", head_bytes)] + assert [c[0] for c in calls] == ["producer_fused"] diff --git a/tests/test_hook_rows.py b/tests/test_hook_rows.py new file mode 100644 index 000000000..30acfd1a7 --- /dev/null +++ b/tests/test_hook_rows.py @@ -0,0 +1,541 @@ +"""Request-row planning for the gather producer (dmi.hook_rows). + +The reference planners (plan_hook_rows, plan_hook_groups, HookRowPlanner) +are checked against independent enumerations of the selected token +positions; PackedDemandPlanner is checked against the reference planners; +the helpers that feed the producer (full_row_plan, encode_row_table, +segment_array, configure_hook_gather, configure_hook_heads) against their +contracts. Pure Python: no GPU or native extension. +""" +import random + +import numpy as np +import pytest + +from dmi.hook_rows import ( + EMPTY_HOOK_ROWS, + GATHER_MAX_SEGS, + HookRowPlanner, + HookRows, + PackedDemandPlanner, + canonical_token_windows, + compile_token_windows, + configure_hook_gather, + configure_hook_heads, + encode_row_table, + full_row_plan, + plan_hook_groups, + plan_hook_rows, + row_signature, + segment_array, +) + +pytestmark = pytest.mark.cpu + +H18, H20, M20 = (0, 18), (0, 20), (13, 20) + + +# --------------------------------------------------------------------------- +# Reference planners: cached membership never caches positions or rows + + +def test_cold_churn_does_not_construct_reusable_indices(monkeypatch): + p = HookRowPlanner((H18, H20), {"A": [H18], "B": [H20]}, []) + original = p._group + calls = [] + def record(ids): + calls.append(ids) + return original(ids) + monkeypatch.setattr(p, "_group", record) + for ids in (("A", "B"), ("B", "A"), ("A", "B")): + result = p.plan(ids, [(3, 4)] * 2, [0, 1], [0, 1], 2, 8) + assert result[H18].source_offsets == (ids.index("A"),) + assert not calls + result = p.plan(("A", "B"), [(4, 5)] * 2, [0, 1], [0, 1], 2, 8) + assert calls == [("A", "B")] + assert result[H20].token_ranges == ((4, 5),) + p.plan(("A", "B"), [(5, 6)] * 2, [0, 1], [0, 1], 2, 8) + assert len(calls) == 1 + + +def test_invalid_layout_does_not_publish_a_new_cache(): + p = HookRowPlanner((H18,), {"A": [H18]}, []) + p.plan(["A"], [(0, 1)], [0], [0], 1, 8) + with pytest.raises(ValueError): + p.plan(["B", "A"], [(0, 1)] * 2, [0, 0], [0, 1], 2, 8) + assert p._request_order == ("A",) and p.group_builds == 1 + result = p.plan(["A"], [(1, 2)], [0], [0], 1, 8) + assert result[H18].token_ranges == ((1, 2),) + + +def test_same_membership_advancing_positions_and_changed_chunk_size(): + p = HookRowPlanner((H18, H20), {"A": [H18], "B": [H20]}, []) + first = p.plan(["A", "B"], [(0, 63), (8, 9)], [0, 63], [0, 1], 64, 16) + second = p.plan(["A", "B"], [(63, 64), (9, 10)], [0, 1], [0, 1], 2, 16) + assert (p.group_builds, p.group_reuses) == (1, 1) + assert second[H18].token_ranges == ((63, 64),) + assert second[H20].segments == ((1, 1),) + assert second[H20].token_ranges == ((9, 10),) + assert first[H20].segments == ((63, 1),) # previous queued metadata is immutable + + +def test_expired_prefill_window_does_not_leak_into_decode_on_cache_hit(): + windows = compile_token_windows({"A": {H18: ((61, 64),)}}) + p = HookRowPlanner((H18,), {"A": [H18]}, [], windows) + first = p.plan(["A"], [(32, 63)], [0], [0], 31, 16) + tail = p.plan(["A"], [(63, 64)], [0], [0], 1, 16) + decode = p.plan(["A"], [(64, 65)], [0], [0], 1, 16) + assert first[H18].segments == ((29, 2),) + assert tail[H18].segments == ((0, 1),) + assert not decode[H18].segments and decode[H18].row_count == 0 + assert p.group_reuses == 2 + + +def test_reordering_arrival_and_snapshot_ownership(): + rows = {"A": [H18], "B": []} + p = HookRowPlanner((H18, H20), rows, [H20]) + rows["A"].append(H20) + for ids in (["A", "B"], ["B", "A"], ["new", "A", "B"], ["A"]): + n = len(ids) + plans = p.plan(ids, [(9, 10)] * n, list(range(n)), list(range(n)), n, 16) + assert plans[H18].req_ids == ("A",) + assert plans[H18].source_offsets == (ids.index("A"),) + assert plans[H20].req_ids == (("new",) if "new" in ids else ()) + assert p.group_builds == 4 + + +def test_cache_hit_still_honors_table_capacity_and_current_offsets(): + p = HookRowPlanner((H18, H20), {"A": [H18, H20], "C": [H18, H20]}, []) + args = (["A", "B", "C"], [(9, 10)] * 3, [0, 1, 2], [0, 1, 2], 3) + exact = p.plan(*args, 2) + fallback = p.plan(*args, 1) + assert exact[H18] is exact[H20] + assert not exact[H18].capacity_fallback + assert fallback[H18].capacity_fallback + assert fallback[H18].segments == ((0, 3),) + assert fallback[H18].offsets == (0, 2) + assert fallback[H18].req_ids == ("A", "C") + + +def test_generated_layouts_against_independent_position_enumeration(): + rng = random.Random(20260912) + keys = (H18, H20, M20) + ids = [f"r{i}" for i in range(9)] + for _ in range(30): + selection = {rid: frozenset(k for k in keys if rng.randrange(2)) for rid in ids[:-1]} + windows = {rid: {key: ((3, 5), (7, None)) for key in keys if rng.randrange(2)} + for rid in ids if rng.randrange(2)} + compiled = compile_token_windows(windows) + p = HookRowPlanner(keys, selection, None, compiled) + order = list(ids) + for step in range(8): + if step % 3 == 0: + rng.shuffle(order) + lengths = [rng.randrange(1, 5) for rid in order] + offsets, total = [], 0 + for n in lengths: + offsets.append(total) + total += n + ranges = [(step, step + n) for n in lengths] + pos = list(range(len(order))) + actual = p.plan(order, ranges, offsets, pos, total, 64) + assert actual == plan_hook_groups(keys, order, ranges, offsets, pos, + selection, None, total, 64, compiled) + for key in keys: + expected = [] + for rid, (lo, hi), off in zip(order, ranges, offsets): + if key not in selection.get(rid, keys): + continue + spans = windows.get(rid, {}).get(key) + expected.extend((rid, t, off + t - lo) for t in range(lo, hi) + if spans is None or any(a <= t and (b is None or t < b) + for a, b in spans)) + rows = actual[key] + got = [(rid, t, off + t - lo) for rid, (lo, hi), off + in zip(rows.req_ids, rows.token_ranges, rows.source_offsets) + for t in range(lo, hi)] + assert got == expected + + +# --------------------------------------------------------------------------- +# Token windows + + +def test_union_is_canonical_and_preserves_empty(): + assert canonical_token_windows([(8, 11), (3, 7), (6, 8), (20, None)]) == ((3, 11), (20, None)) + assert canonical_token_windows([(5, None), (7, 9), (0, 5)]) == ((0, None),) + assert canonical_token_windows([]) == () + + +def test_compiled_window_profiles_share_storage_and_own_snapshot(): + windows = {"A": {H18: ((61, 64),), H20: ((61, 64),)}} + compiled = compile_token_windows(windows) + assert compiled.by_hook[H18] is compiled.by_hook[H20] + windows["A"][H18] = () + assert compiled.by_hook[H18]["A"] == ((61, 64),) + with pytest.raises(TypeError): + compiled.by_hook[H18]["A"] = () + + +@pytest.mark.parametrize("start,length", [(0, 3), (3, 4), (509, 8), (960, 1), + (1021, 9), (2048, 64)]) +@pytest.mark.parametrize("unbounded", [False, True]) +def test_long_window_list_seeks_current_chunk(start, length, unbounded): + spans = [(i * 4, i * 4 + 1) for i in range(256)] + if unbounded: + spans.append((1100, None)) + windows = {"A": canonical_token_windows(spans)} + rows = plan_hook_rows(["A"], [(start, start + length)], [0], [3], {"A"}, + length, 128, windows) + expected = [token for token in range(start, start + length) + if any(lo <= token and (hi is None or token < hi) for lo, hi in spans)] + actual = [t for lo, hi in rows.token_ranges for t in range(lo, hi)] + assert actual == expected + assert [row for off, n in rows.segments for row in range(off, off + n)] == [t - start for t in expected] + # A later call may go backwards (e.g. a recomputed prefill chunk). A seek + # optimization must not carry a monotone cursor into the next call. + earlier = plan_hook_rows(["A"], [(0, 1)], [0], [3], {"A"}, 1, 128, windows) + assert earlier.token_ranges == ((0, 1),) + + +@pytest.mark.parametrize("selected", [(), (0,), (127,), (0, 63, 127), tuple(range(128))]) +def test_sparse_membership_keeps_packed_order_and_positions(selected): + ids = tuple(f"r{i}" for i in range(128)) + ranges = tuple((i, i + 1) for i in range(128)) + offsets = tuple(range(128)) + selection = {rid: frozenset([H18]) if i in selected else frozenset() + for i, rid in enumerate(ids)} + plans = plan_hook_groups((H18, H20), ids, ranges, offsets, offsets, + selection, frozenset(), 128, 128) + assert plans[H18].req_ids == tuple(ids[i] for i in selected) + assert plans[H18].token_ranges == tuple(ranges[i] for i in selected) + assert plans[H18].source_offsets == selected + assert plans[H20].row_count == 0 + + +def test_seeded_window_oracle(): + rng = random.Random(20260912) + for _ in range(250): + ids = [f"r{i}" for i in range(rng.randint(1, 12))] + rng.shuffle(ids) + ranges, offsets, windows, expected, total = [], [], {}, [], 0 + selected = set(rng.sample(ids, rng.randint(0, len(ids)))) + for rid in ids: + start, length = rng.randrange(12), rng.randint(1, 16) + ranges.append((start, start + length)) + offsets.append(total) + raw = [(rng.randrange(25), rng.randrange(25, 40)) for _ in range(rng.randrange(4))] + windows[rid] = canonical_token_windows(raw) + # Deliberately enumerate the original, un-reduced window union. + expected.extend((rid, token, total + token - start) + for token in range(start, start + length) + if rid in selected and any(lo <= token < hi for lo, hi in raw)) + total += length + rows = plan_hook_rows(ids, ranges, offsets, list(range(len(ids))), + selected, total, 256, windows) + source_rows = [row for off, length in rows.segments for row in range(off, off + length)] + got = [(rid, token, source_rows[off + token - start]) + for rid, (start, end), off in zip(rows.req_ids, rows.token_ranges, rows.offsets) + for token in range(start, end)] + assert got == expected + + +def test_group_planner_checks_layout_once(monkeypatch): + import dmi.hook_rows as module + checked = [] + original = module._checked_layout + def check(*args): + checked.append(1) + return original(*args) + monkeypatch.setattr(module, "_checked_layout", check) + selection = {"A": frozenset([H18]), "B": frozenset([H20])} + plans = plan_hook_groups([H18, H20], ["A", "B"], [(0, 64), (7, 8)], + [0, 64], [0, 1], selection, frozenset(), 65, 8) + assert len(checked) == 1 + assert plans[H18].row_count == 64 and plans[H20].row_count == 1 + + +@pytest.mark.parametrize("keys", [[], [H18, H20]]) +def test_group_planner_rejects_invalid_even_if_no_selected_rows(keys): + with pytest.raises(ValueError, match="overlapping"): + plan_hook_groups(keys, ["A", "B"], [(0, 3), (0, 2)], [0, 2], [0, 1], + {}, frozenset(), 5, 8) + with pytest.raises(ValueError, match="lengths"): + plan_hook_groups(keys, ["A"], [], [], [], {}, frozenset(), 0, 8) + + +def test_grouped_window_oracle_with_dynamic_replacement(): + """Random profiles, holes, chunk positions and segment fallback. + + Enumerate source rows independently instead of calling the row planner + whose assembly helper is shared with the grouped implementation. + """ + rng = random.Random(20260913) + keys = (H18, H20, (1, 18), (1, 20)) + for _ in range(300): + ids = [f"r{i}" for i in range(rng.randint(1, 12))] + rng.shuffle(ids) + ranges, offsets, total = [], [], 0 + selection, windows = {}, {} + default = rng.choice([None, frozenset(), frozenset([H18])]) + for rid in ids: + total += rng.randrange(3) # padding gaps must not be captured + length, start = rng.randint(1, 64), rng.randrange(100) + ranges.append((start, start + length)) + offsets.append(total) + total += length + if rng.randrange(3): + selection[rid] = rng.choice([None, frozenset(rng.sample(keys, rng.randrange(5)))]) + windows[rid] = {} + for key in keys: + if rng.randrange(3): + windows[rid][key] = canonical_token_windows([ + (rng.randrange(120), rng.choice([None, 180])) + for _ in range(rng.randrange(4))]) + maximum = rng.choice([1, 4, 128]) + plans = plan_hook_groups(keys, ids, ranges, offsets, list(range(len(ids))), + selection, default, total, maximum, + compile_token_windows(windows)) + for key, plan in plans.items(): + expected = [] + for pos, (rid, (start, end), off) in enumerate(zip(ids, ranges, offsets)): + selected = selection.get(rid, default) + intervals = windows[rid].get(key) + if selected is None or key in selected: + expected.extend((rid, t, off + t - start, pos) + for t in range(start, end) + if intervals is None or any(lo <= t and (hi is None or t < hi) + for lo, hi in intervals)) + copied = [i for off, count in plan.segments for i in range(off, off + count)] + actual = [(rid, t, copied[off + t - start], pos) + for rid, (start, end), off, pos in zip(plan.req_ids, plan.token_ranges, + plan.offsets, plan.orig_positions) + for t in range(start, end)] + assert actual == expected + assert plan.row_count == (total if plan.capacity_fallback else len(expected)) + + + +# --------------------------------------------------------------------------- +# PackedDemandPlanner matches the reference planner + + +def _validate(keys): + def validate(hooks): + if hooks is None: + return None + hooks = frozenset(tuple(h) for h in hooks) + unknown = hooks - set(keys) + if unknown: + raise ValueError(f"unknown HookPoint {sorted(unknown)}") + return hooks + return validate + + +def _random_step(rng, ids): + ranges, offsets, total = [], [], 0 + for _ in ids: + total += rng.randrange(3) # graph padding between requests + length, start = rng.randint(1, 40), rng.randrange(80) + ranges.append((start, start + length)) + offsets.append(total) + total += length + return ranges, offsets, total + + +def test_packed_planner_matches_reference_planner_on_random_demand(): + rng = random.Random(20261103) + keys = (H18, H20, M20, (1, 18)) + for _ in range(200): + ids = [f"r{i}" for i in range(rng.randint(1, 10))] + rows, windows = {}, {} + default = rng.choice([None, frozenset(), frozenset([H18])]) + for rid in ids: + if rng.randrange(3): + rows[rid] = rng.choice( + [None, frozenset(rng.sample(keys, rng.randrange(len(keys) + 1)))]) + if rng.randrange(2): + windows[rid] = {key: [(rng.randrange(60), rng.choice([None, 70, 120])) + for _ in range(rng.randrange(3))] + for key in keys if rng.randrange(2)} + canonical = {rid: {key: canonical_token_windows(spans) + for key, spans in per.items()} + for rid, per in windows.items()} + compiled = compile_token_windows(canonical) + planner = PackedDemandPlanner(keys, _validate(keys)) + for step in range(4): + rng.shuffle(ids) + ranges, offsets, total = _random_step(rng, ids) + positions = list(range(len(ids))) + maximum = rng.choice([1, 3, 256]) + got = planner.plan(ids, ranges, offsets, positions, total, maximum, + rows, windows, default) + want = plan_hook_groups(keys, ids, ranges, offsets, positions, + rows, default, total, maximum, compiled) + for key in keys: + assert got[key] == want[key], (step, key) + assert got[key].row_count == want[key].row_count + assert set(planner.present_keys) == { + key for key in keys if got[key] is not EMPTY_HOOK_ROWS} + + +def test_packed_planner_compiles_rules_once_per_demand_change(): + keys = (H18, H20) + planner = PackedDemandPlanner(keys, _validate(keys)) + rows = {"A": frozenset([H18]), "B": None} + windows = {"A": {H18: [(2, None)]}} + identity = object() + for step in range(5): + plans = planner.plan(["A", "B"], [(step, step + 4), (10 + step, 11 + step)], + [0, 4], [0, 1], 5, 8, rows, windows, + demand_identity=identity) + assert plans[H20].req_ids == ("B",) + assert planner.rule_compiles == 2 # one rule per request + assert (planner.group_builds, planner.group_reuses) == (1, 4) + rows = dict(rows, A=frozenset([H18, H20])) # a demand change recompiles A + planner.plan(["A", "B"], [(5, 9), (15, 16)], [0, 4], [0, 1], 5, 8, rows, + windows, demand_identity=object()) + assert planner.rule_compiles == 3 + + +def test_packed_planner_shares_equal_plans_and_exposes_native_layouts(): + keys = (H18, H20, M20) + planner = PackedDemandPlanner(keys, _validate(keys)) + plans = planner.plan(["A", "B"], [(0, 3), (7, 9)], [0, 3], [0, 1], 5, 8, + {"A": None, "B": frozenset([H18, H20])}, {}) + assert plans[H18] is plans[H20] + assert plans[M20].req_ids == ("A",) + assert plans[H18].layout() == [("A", 0, 3, 0, 0, 0), ("B", 7, 9, 3, 0, 1)] + index, lo, hi, offset, pos = plans[M20].layout_arrays() + assert index.tolist() == [0] and lo.tolist() == [0] and hi.tolist() == [3] + assert offset.tolist() == [0] and pos.tolist() == [0] + fallback = plans[H18].full_payload(5) + assert fallback.segments == ((0, 5),) and fallback.offsets == (0, 3) + + +def test_packed_planner_validates_layout_and_hooks(): + keys = (H18,) + planner = PackedDemandPlanner(keys, _validate(keys)) + with pytest.raises(ValueError, match="overlapping"): + planner.plan(["A", "B"], [(0, 3), (0, 2)], [0, 2], [0, 1], 5, 8, {}, {}) + with pytest.raises(ValueError, match="lengths"): + planner.plan(["A"], [], [], [], 0, 8, {}, {}) + with pytest.raises(ValueError, match="unknown HookPoint"): + planner.plan(["A"], [(0, 1)], [0], [0], 1, 8, {"A": [H20]}, {}) + + +# --------------------------------------------------------------------------- +# Static full-row baseline and the producer row table + + +def test_full_row_plan_selects_every_scheduled_row_with_identity_offsets(): + ids, ranges, offsets = ["A", "B"], [(5, 8), (0, 2)], [0, 4] # row 3 is padding + plan = full_row_plan(ids, ranges, offsets, [0, 1], 6) + assert plan.segments == ((0, 6),) + assert plan.row_count == 6 + assert plan.offsets == plan.source_offsets == (0, 4) + assert plan.req_ids == ("A", "B") and plan.token_ranges == ((5, 8), (0, 2)) + # Same attribution as a reference plan selecting every request. + reference = plan_hook_rows(ids, ranges, offsets, [0, 1], None, 6, 8) + assert reference.req_ids == plan.req_ids + assert reference.token_ranges == plan.token_ranges + empty = full_row_plan([], [], [], [], 0) + assert empty.segments == () and empty.row_count == 0 + + +def _full_row_slices(plan, source): + """Each request's rows as the consumer slices them from the payload the + gather producer copies for ``plan``, next to the rows the ungathered + producer's payload (the source itself) attributes to that request.""" + payload = [source[i] for start, count in plan.segments + for i in range(start, start + count)] + actual = [payload[off:off + hi - lo] + for (lo, hi), off in zip(plan.token_ranges, plan.offsets)] + expected = [source[off:off + hi - lo] + for (lo, hi), off in zip(plan.token_ranges, plan.source_offsets)] + return payload, actual, expected + + +@pytest.mark.parametrize("ids,ranges,offsets,total", [ + (["A", "B"], [(0, 2), (0, 1)], [2, 4], 8), # leading + trailing padding + (["A", "B"], [(0, 2), (3, 5)], [0, 5], 7), # interior padding + (["A", "B"], [(0, 2), (3, 5)], [0, 2], 6), # trailing padding only + (["B"], [(7, 9)], [3], 5), # filtered: first offset nonzero + (["A", "B", "C"], [(0, 1), (2, 4), (9, 10)], [1, 3, 6], 7), +]) +def test_full_row_plan_copies_the_ungathered_payload_with_physical_offsets( + ids, ranges, offsets, total): + plan = full_row_plan(ids, ranges, offsets, list(range(len(ids))), total) + assert plan.segments == ((0, total),) and plan.row_count == total + assert plan.offsets == plan.source_offsets == tuple(offsets) + assert not plan.capacity_fallback + payload, actual, expected = _full_row_slices(plan, list(range(100, 100 + total))) + assert payload == list(range(100, 100 + total)) # the whole physical tensor + assert actual == expected + + +def test_full_row_plan_without_requests_is_empty(): + for total in (0, 4): + plan = full_row_plan([], [], [], [], total) + assert plan.segments == () and plan.row_count == 0 and plan.offsets == () + assert encode_row_table(plan, 4).tolist()[0] == 0 + with pytest.raises(ValueError): + full_row_plan(["A"], [(0, 2)], [3], [0], 4) # rows past the tensor + + +def test_encode_row_table_matches_the_producer_contract(): + plan = plan_hook_rows(["A", "B"], [(0, 2), (4, 5)], [0, 3], [0, 1], None, 4, 8) + table = encode_row_table(plan, 4) + assert table.dtype == np.int64 and table.shape == (9,) + assert table.tolist() == [2, 0, 2, 3, 1, 0, 0, 0, 0] + assert encode_row_table(None, 2).tolist() == [-1, 0, 0, 0, 0] + assert encode_row_table(EMPTY_HOOK_ROWS, 1).tolist() == [0, 0, 0] + with pytest.raises(ValueError, match="more segments"): + encode_row_table(plan, 1) + + +def test_row_signature_and_segment_array_agree_across_planners(): + ids, ranges, offsets, pos = ["A", "B", "C"], [(0, 2), (4, 6), (9, 10)], [0, 2, 4], [0, 1, 2] + reference = plan_hook_rows(ids, ranges, offsets, pos, {"A", "C"}, 5, 8) + planner = PackedDemandPlanner((H18,), _validate((H18,))) + packed = planner.plan(ids, ranges, offsets, pos, 5, 8, + {"A": frozenset([H18]), "C": frozenset([H18])}, {})[H18] + assert segment_array(reference).tolist() == [[0, 2], [4, 1]] + assert segment_array(packed).tolist() == segment_array(reference).tolist() + assert row_signature(packed) == row_signature(packed) + other = plan_hook_rows(ids, ranges, offsets, pos, {"B"}, 5, 8) + assert row_signature(other) != row_signature(reference) + assert segment_array(EMPTY_HOOK_ROWS).shape == (0, 2) + assert GATHER_MAX_SEGS == 256 + + +def test_configure_hook_gather_and_heads_validate_their_inputs(): + import torch + from types import SimpleNamespace + + hook = SimpleNamespace(_gather_table=None, _strip_row_bytes=0) + table = torch.zeros(1 + 2 * 4, dtype=torch.int64) + configure_hook_gather(hook, table, 64) + assert hook._gather_table is table and hook._strip_row_bytes == 64 + with pytest.raises(ValueError, match="int64"): + configure_hook_gather(hook, table.to(torch.int32), 64) + with pytest.raises(ValueError, match="int64"): + configure_hook_gather(hook, table.reshape(3, 3), 64) + with pytest.raises(ValueError, match="positive row stride"): + configure_hook_gather(hook, table, 0) + configure_hook_gather(hook, None, 0) + assert hook._gather_table is None + + arena = torch.zeros(80, dtype=torch.int32) + acc = torch.zeros(4, 16) + configure_hook_heads(hook, head_type=24, layer_no=3, table=table, arena=arena, + acc=acc, slots=2, weight_bf16=True) + assert (hook._head_hook_type, hook._head_hook_id, hook._head_slots, + hook._head_wbf16) == (24, 3, 2, 1) + assert hook._head_arena is arena and hook._head_acc is acc + with pytest.raises(ValueError, match="1..16"): + configure_hook_heads(hook, head_type=24, layer_no=3, table=table, + arena=arena, acc=acc, slots=17, weight_bf16=False) + with pytest.raises(ValueError, match="int32"): + configure_hook_heads(hook, head_type=24, layer_no=3, table=table, + arena=arena.float(), acc=acc, slots=2, weight_bf16=False) diff --git a/tests/test_inline_head_eager.py b/tests/test_inline_head_eager.py new file mode 100644 index 000000000..e5b84e1e7 --- /dev/null +++ b/tests/test_inline_head_eager.py @@ -0,0 +1,57 @@ +"""The eager safety net of in-graph heads respects the ring's effective +capacity (min of payload ring and pinned staging): an entry the drain can +never move is refused before anything is reserved. Fakes only.""" +from types import SimpleNamespace + +import pytest +import torch + +pytestmark = pytest.mark.cpu + + +class Engine: + def __init__(self, available): + self.available = available + self.calls = [] + + def available_capacity(self): + return self.available + + def available_task_slots(self): + return 8 + + def flush_and_wait(self): + self.calls.append("flush") + + def reserve_one(self, nbytes): + self.calls.append(("reserve", nbytes)) + + +def _hook(slots): + from dmi.hooks.point import HookPoint + + hook = HookPoint() + hook._head_hook_type, hook._head_hook_id, hook._head_slots = 24, 0, slots + return hook + + +def test_head_output_above_the_effective_capacity_is_refused(monkeypatch): + launches = [] + monkeypatch.setattr(torch.ops.ring, "producer_fused", + lambda *args: launches.append(args), raising=False) + engine = Engine(available=1 << 20) + transport = SimpleNamespace(effective_cap=256) + x = torch.zeros(8, 4) # 8 rows x 16 slots x 4 B = 512 B + with pytest.raises(RuntimeError, match="effective capacity"): + _hook(16)._eager_head(transport, engine, None, x) + assert engine.calls == [] and launches == [] + + +def test_head_output_within_capacity_flushes_when_the_ring_is_full(monkeypatch): + launches = [] + monkeypatch.setattr(torch.ops.ring, "producer_fused", + lambda *args: launches.append(args), raising=False) + engine = Engine(available=64) + transport = SimpleNamespace(effective_cap=1 << 20) + _hook(4)._eager_head(transport, engine, None, torch.zeros(8, 4)) + assert engine.calls == ["flush", ("reserve", 128)] and len(launches) == 1 diff --git a/tests/test_node_toggle.py b/tests/test_node_toggle.py index a067ed0ae..3bb959f00 100644 --- a/tests/test_node_toggle.py +++ b/tests/test_node_toggle.py @@ -267,6 +267,20 @@ def test_ungated_steps_use_the_full_active_set(): assert transport.gated_step_requires_registry() is True +def test_table_gate_makes_every_step_follow_the_toggle(): + specs = _specs() + transport = _transport(FakeToggleEngine(_keys(specs)), specs) + transport.set_active_hooks_lazy([(HOOK_TYPE_RESID_PRE, 0)]) + transport.set_gated_step(False) + transport.set_prefill_table_gate(True) + assert _keys(transport.specs_for_step()) == [(HOOK_TYPE_RESID_PRE, 0)] + assert transport.gated_step_requires_registry() is False + transport.set_gated_step(True) + assert transport.gated_step_requires_registry() is False + transport.set_prefill_table_gate(False) + assert transport.gated_step_requires_registry() is True + + def test_clear_resets_registry_and_gate_together(): specs = _specs() engine = FakeToggleEngine(_keys(specs)) diff --git a/tests/test_per_hook_gather_gpu.py b/tests/test_per_hook_gather_gpu.py new file mode 100644 index 000000000..9de245a71 --- /dev/null +++ b/tests/test_per_hook_gather_gpu.py @@ -0,0 +1,604 @@ +"""Per-HookPoint request-row gather and in-graph heads on a real legacy ring. + +HookPoints are armed in gather mode (a per-HookPoint row table) and, for the +heads, with a fused-node arena; plans from dmi.hook_rows feed both the +producer tables and the per-hook request layouts in the metadata. An +in-process consumer checks that every delivered slice is attributed to the +right request and token range and carries the right rows (or scores). + +Needs CUDA and the full native backend. Memory use is a few MiB. +""" +from __future__ import annotations + +import time + +import numpy as np +import pytest +import torch + +from tests._requirements import require_cuda, require_native_backend + +pytestmark = [ + pytest.mark.gpu, + pytest.mark.native_backend, + require_cuda(), + require_native_backend(), +] + +HID = 16 +MAX_SEGS = 8 + + +class Ring: + def __init__(self, payload=8 << 20): + from dmi.hooks.specs import ModelShapeConfig + from dmi.transport import native as ne + from dmi.transport import ring as rt + + self.slices = [] + cfg = ne.RingConfig() + cfg.payload_ring_bytes = payload + cfg.pinned_staging_bytes = payload + cfg.drain_poll_timeout_us = 200 + + def consume(model_id, shard_rank, req_id, act_name, layer_no, + start, end, tensor): + self.slices.append((act_name, int(layer_no), req_id, int(start), + int(end), tensor.clone())) + + self.engine = ne.RingEngine(cfg, consume) + self.engine.init(0) + self.engine.start() + self.transport = rt.RingTransport(self.engine) + rt.activate(self.transport) + self.transport.set_model_cfg(ModelShapeConfig( + hidden_dim=HID, num_heads=2, num_kv_heads=2, head_dim=HID // 2, + dtype=torch.float32, head_slots=0)) + self.payload = self.engine.payload_tensor() + + def drain(self, expected): + torch.cuda.synchronize() + self.engine.flush_and_wait() + deadline = time.monotonic() + 5.0 + while len(self.slices) < expected and time.monotonic() < deadline: + time.sleep(0.005) + time.sleep(0.02) + got = list(self.slices) + self.slices.clear() + return got + + def close(self): + from dmi.transport import ring as rt + torch.cuda.synchronize() + self.engine.flush_and_wait() + rt.deactivate() + self.engine.stop() + + +@pytest.fixture +def ring(): + r = Ring() + try: + yield r + finally: + r.close() + + +def _hook(ring, hook_type, layer, table, row_bytes=HID * 4): + from dmi.hook_rows import configure_hook_gather + from dmi.hooks.point import HookPoint + + hook = HookPoint() + hook._ring_hook_type = hook_type + hook._ring_hook_id = layer + hook._ring_payload = ring.payload + configure_hook_gather(hook, table, row_bytes) + return hook + + +def _spec(hook_type, layer): + from dmi.hooks.specs import HookSpec + return HookSpec(hook_type=hook_type, module=None, layer_no=layer, + dim0_is_actual_tokens=True) + + +def _write_table(table, plan): + from dmi.hook_rows import encode_row_table + table.copy_(torch.from_numpy(encode_row_table(plan, MAX_SEGS))) + + +def _expected(x, plan, act, layer): + """Consumer slices a correct gather must produce for `plan`.""" + out = [] + for rid, (lo, hi), source in zip(plan.req_ids, plan.token_ranges, + plan.source_offsets): + out.append((act, layer, rid, lo, hi, x[source:source + (hi - lo)].cpu())) + return out + + +def _compare(got, want): + key = lambda s: (s[0], s[1], s[2], s[3]) + got, want = sorted(got, key=key), sorted(want, key=key) + assert [key(s) for s in got] == [key(s) for s in want] + for g, w in zip(got, want): + assert torch.equal(g[5], w[5]), key(g) + + +def _step_layout(): + # Three requests packed into 9 rows with a padding row at 4. + req_ids = ["a", "b", "c"] + ranges = [(10, 14), (0, 3), (7, 8)] + offsets = [0, 5, 8] + return req_ids, ranges, offsets, [0, 1, 2], 9 + + +def test_two_hooks_with_independent_selections_eager_and_graph(ring): + from dmi.hook_rows import plan_hook_rows + from dmi.hooks.specs import HOOK_TYPE_RESID_PRE + + ht = HOOK_TYPE_RESID_PRE + req_ids, ranges, offsets, pos, total = _step_layout() + x = torch.arange(total * HID, dtype=torch.float32, device="cuda").reshape(total, HID) + tables = [torch.zeros(1 + 2 * MAX_SEGS, dtype=torch.int64, device="cuda") + for _ in range(2)] + hooks = [_hook(ring, ht, layer, tables[layer]) for layer in range(2)] + ring.transport._active_specs = [_spec(ht, 0), _spec(ht, 1)] + + selections = [({"a", "c"}, {"b"}), ({"b"}, set()), (None, {"a", "b", "c"})] + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for hook in hooks: + hook(x) + ring.drain(0) + for mode in ("eager", "graph"): + for sel0, sel1 in selections: + plans = [plan_hook_rows(req_ids, ranges, offsets, pos, sel, total, MAX_SEGS) + for sel in (sel0, sel1)] + for table, plan in zip(tables, plans): + _write_table(table, plan) + ring.transport.set_step_context( + model_id="m", req_ids=req_ids, token_ranges=ranges, + dim0_offsets=offsets, flattened=True, + hook_rows={(ht, 0): plans[0], (ht, 1): plans[1]}) + ring.transport.pre_push_all_metas(batch=0, q_len=total, kv_dim=0) + if mode == "eager": + for hook in hooks: + hook(x) + else: + graph.replay() + want = (_expected(x, plans[0], "blocks.hook_resid_pre", 0) + + _expected(x, plans[1], "blocks.hook_resid_pre", 1)) + _compare(ring.drain(len(want)), want) + + +def test_disabled_rows_publish_nothing_and_skip_eager_launches(ring): + from dmi.hook_rows import plan_hook_rows + from dmi.hooks.specs import HOOK_TYPE_RESID_PRE + + ht = HOOK_TYPE_RESID_PRE + req_ids, ranges, offsets, pos, total = _step_layout() + x = torch.randn(total, HID, device="cuda") + host = torch.zeros(2, 1 + 2 * MAX_SEGS, dtype=torch.int64).pin_memory() + dev = torch.zeros_like(host, device="cuda") + hooks = [_hook(ring, ht, layer, dev[layer]) for layer in range(2)] + ring.engine.set_gather_host_mirror(dev.data_ptr(), host.data_ptr(), + host.stride(0) * 8, 2) + plan = plan_hook_rows(req_ids, ranges, offsets, pos, {"b"}, total, MAX_SEGS) + from dmi.hook_rows import encode_row_table + host[0] = torch.from_numpy(encode_row_table(plan, MAX_SEGS)) + host[1] = torch.from_numpy(encode_row_table(None, MAX_SEGS)) # disabled + dev.copy_(host, non_blocking=True) + # Only the enabled hook has metadata: a disabled row has none. + ring.transport._active_specs = [_spec(ht, 0)] + ring.transport.set_step_context(model_id="m", req_ids=req_ids, + token_ranges=ranges, dim0_offsets=offsets, + flattened=True, hook_rows={(ht, 0): plan}) + ring.transport.pre_push_all_metas(batch=0, q_len=total, kv_dim=0) + skips = ring.engine.gather_host_skips() + for hook in hooks: + hook(x) + _compare(ring.drain(1), _expected(x, plan, "blocks.hook_resid_pre", 0)) + assert ring.engine.gather_host_skips() == skips + 1 + assert ring.engine.available_task_slots() == ring.engine.task_cap() + ring.engine.set_gather_host_mirror(0, 0, 0, 0) + + +def test_gather_nodes_are_toggle_managed(ring): + """A gather producer captured inside the toggle window records its node, + so disabling its key disables it (otherwise it would fire while the + metadata gate excluded it).""" + from dmi.hook_rows import plan_hook_rows + from dmi.hooks.specs import HOOK_TYPE_RESID_PRE + + ht = HOOK_TYPE_RESID_PRE + req_ids, ranges, offsets, pos, total = _step_layout() + x = torch.randn(total, HID, device="cuda") + tables = [torch.zeros(1 + 2 * MAX_SEGS, dtype=torch.int64, device="cuda") + for _ in range(3)] + hooks = [_hook(ring, ht, layer, tables[layer]) for layer in range(3)] + ring.transport._active_specs = [_spec(ht, layer) for layer in range(3)] + ring.engine.enable_toggle_capture(True) + ring.engine.set_null_mode(True) + graph = torch.cuda.CUDAGraph(keep_graph=True) + with torch.cuda.graph(graph): + for hook in hooks: + hook(x) + graph.instantiate() + ring.engine.set_null_mode(False) + ring.engine.enable_toggle_capture(False) + ring.engine.bind_graph_exec(graph.raw_cuda_graph(), graph.raw_cuda_graph_exec()) + assert ring.engine.toggle_node_count() == 3 + assert ring.engine.capture_anomaly_count() == 0 + + plan = plan_hook_rows(req_ids, ranges, offsets, pos, {"a", "b"}, total, MAX_SEGS) + for table in tables: + _write_table(table, plan) + ring.transport.set_active_hooks_lazy([(ht, 2)]) + ring.transport.set_step_context( + model_id="m", req_ids=req_ids, token_ranges=ranges, dim0_offsets=offsets, + flattened=True, hook_rows={(ht, layer): plan for layer in range(3)}) + ring.transport.pre_push_all_metas(batch=0, q_len=total, kv_dim=0) + assert ring.transport.ensure_graph_current(graph.raw_cuda_graph()) == 0 + graph.replay() + _compare(ring.drain(2), _expected(x, plan, "blocks.hook_resid_pre", 2)) + assert ring.engine.available_task_slots() == ring.engine.task_cap() + + +def test_table_gate_follows_the_toggle_on_ungated_steps(ring): + """With the table gate armed, an ungated (eager) step still publishes + only the enabled hooks: disabled hooks' tables hold -1 and the eager + safety net skips them.""" + from dmi.hook_rows import encode_row_table, full_row_plan + from dmi.hooks.specs import HOOK_TYPE_RESID_PRE + + ht = HOOK_TYPE_RESID_PRE + req_ids, ranges, offsets, pos, total = _step_layout() + x = torch.randn(total, HID, device="cuda") + tables = [torch.zeros(1 + 2 * MAX_SEGS, dtype=torch.int64, device="cuda") + for _ in range(2)] + hooks = [_hook(ring, ht, layer, tables[layer]) for layer in range(2)] + ring.transport._active_specs = [_spec(ht, 0), _spec(ht, 1)] + ring.engine.enable_toggle_capture(True) + ring.engine.set_null_mode(True) + graph = torch.cuda.CUDAGraph(keep_graph=True) + with torch.cuda.graph(graph): + for hook in hooks: + hook(x) + graph.instantiate() + ring.engine.set_null_mode(False) + ring.engine.enable_toggle_capture(False) + ring.engine.bind_graph_exec(graph.raw_cuda_graph(), graph.raw_cuda_graph_exec()) + ring.transport.set_active_hooks_lazy([(ht, 1)]) + ring.transport.set_prefill_table_gate(True) + ring.transport.set_gated_step(False) + try: + plan = full_row_plan(req_ids, ranges, offsets, pos, total) + tables[0].copy_(torch.from_numpy(encode_row_table(None, MAX_SEGS))) + _write_table(tables[1], plan) + assert [(s.hook_type, s.layer_no) for s in ring.transport.specs_for_step()] == [(ht, 1)] + ring.transport.set_step_context( + model_id="m", req_ids=req_ids, token_ranges=ranges, + dim0_offsets=offsets, flattened=True, hook_rows={(ht, 1): plan}) + ring.transport.pre_push_all_metas(batch=0, q_len=total, kv_dim=0) + ring.transport.force_eager = True + for hook in hooks: + hook(x) + ring.transport.force_eager = False + _compare(ring.drain(3), _expected(x, plan, "blocks.hook_resid_pre", 1)) + assert ring.transport.gated_step_requires_registry() is False + finally: + ring.transport.force_eager = False + ring.transport.set_prefill_table_gate(False) + ring.transport.set_gated_step(True) + + +def _fused_arena(weights, bias, modes, acc_cols): + """int32 node arena: modes, accumulator columns, fp32 biases, fp32 weights.""" + K, hidden = weights.shape + words = (64 + K * hidden + 3) & ~3 + arena = np.zeros(words, dtype=np.int32) + arena[0:K] = modes + arena[16:16 + K] = acc_cols + arena[32:32 + K] = np.asarray(bias, dtype=np.float32).view(np.int32) + arena[64:64 + K * hidden] = weights.astype(np.float32).reshape(-1).view(np.int32) + return torch.from_numpy(arena).cuda() + + +def test_fused_heads_score_selected_rows_and_chain_across_layers(ring): + from dmi.hook_rows import configure_hook_heads, plan_hook_rows + from dmi.hooks.specs import (HOOK_TYPE_RESID_PRE, HOOK_TYPE_RESID_PRE_HEAD, + ModelShapeConfig) + + ht, hh = HOOK_TYPE_RESID_PRE, HOOK_TYPE_RESID_PRE_HEAD + K = 3 + ring.transport.set_model_cfg(ModelShapeConfig( + hidden_dim=HID, num_heads=2, num_kv_heads=2, head_dim=HID // 2, + dtype=torch.float32, head_slots=K)) + req_ids, ranges, offsets, pos, total = _step_layout() + rng = np.random.default_rng(7) + xs = [torch.from_numpy(rng.standard_normal((total, HID)).astype(np.float32)).cuda() + for _ in range(2)] + w = [rng.standard_normal((K, HID)).astype(np.float32) for _ in range(2)] + bias = np.array([0.5, -1.0, 2.0], dtype=np.float32) + # Slot 0: single-layer head at layer 0. Slot 1: a linear head over + # layers 0 and 1 (FIRST at 0, LAST at 1, accumulator column 3). + # Slot 2: free. + arenas = [_fused_arena(w[0], bias, [1, 2, 0], [0, 3, 0]), + _fused_arena(w[1], bias, [0, 4, 0], [0, 3, 0])] + acc = torch.zeros(64, 16, dtype=torch.float32, device="cuda") + head_tables = [torch.zeros(1 + 2 * MAX_SEGS, dtype=torch.int64, device="cuda") + for _ in range(2)] + row_tables = [torch.zeros(1 + 2 * MAX_SEGS, dtype=torch.int64, device="cuda") + for _ in range(2)] + hooks = [] + for layer in range(2): + hook = _hook(ring, ht, layer, row_tables[layer]) + configure_hook_heads(hook, head_type=hh, layer_no=layer, + table=head_tables[layer], arena=arenas[layer], + acc=acc, slots=K, weight_bf16=False) + hooks.append(hook) + rows_plan = plan_hook_rows(req_ids, ranges, offsets, pos, {"b"}, total, MAX_SEGS) + head_plan = plan_hook_rows(req_ids, ranges, offsets, pos, {"a", "c"}, total, MAX_SEGS) + for layer in range(2): + _write_table(row_tables[layer], rows_plan) + _write_table(head_tables[layer], head_plan) + ring.transport._active_specs = [_spec(ht, 0), _spec(hh, 0), _spec(ht, 1), _spec(hh, 1)] + hook_rows = {} + for layer in range(2): + hook_rows[(ht, layer)] = rows_plan + hook_rows[(hh, layer)] = head_plan + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for hook, x in zip(hooks, xs): + hook(x) + ring.drain(0) + for mode in ("eager", "graph"): + ring.transport.set_step_context( + model_id="m", req_ids=req_ids, token_ranges=ranges, + dim0_offsets=offsets, flattened=True, hook_rows=hook_rows) + ring.transport.pre_push_all_metas(batch=0, q_len=total, kv_dim=0) + if mode == "eager": + for hook, x in zip(hooks, xs): + hook(x) + else: + graph.replay() + got = ring.drain(2 + 2 * len(head_plan.req_ids)) + rows = [s for s in got if s[0] == "blocks.hook_resid_pre"] + _compare(rows, _expected(xs[0], rows_plan, "blocks.hook_resid_pre", 0) + + _expected(xs[1], rows_plan, "blocks.hook_resid_pre", 1)) + heads = {(s[1], s[2]): s[5] for s in got if s[0] == "blocks.hook_resid_pre_head"} + assert sorted(heads) == [(0, "a"), (0, "c"), (1, "a"), (1, "c")] + for rid, (lo, hi), source in zip(head_plan.req_ids, head_plan.token_ranges, + head_plan.source_offsets): + h0 = xs[0][source:source + hi - lo].double().cpu() + h1 = xs[1][source:source + hi - lo].double().cpu() + w0 = torch.from_numpy(w[0]).double() + w1 = torch.from_numpy(w[1]).double() + layer0, layer1 = heads[(0, rid)], heads[(1, rid)] + assert layer0.shape == (hi - lo, K) + torch.testing.assert_close(layer0[:, 0].double(), h0 @ w0[0] + 0.5, + rtol=1e-5, atol=1e-5) + assert torch.isnan(layer0[:, 1]).all() and torch.isnan(layer0[:, 2]).all() + chained = h0 @ w0[1] + h1 @ w1[1] - 1.0 + torch.testing.assert_close(layer1[:, 1].double(), chained, + rtol=1e-5, atol=1e-5) + assert torch.isnan(layer1[:, 0]).all() and torch.isnan(layer1[:, 2]).all() + assert ring.engine.head_launches() >= 2 + + +def test_original_positions_select_final_logits_rows(ring): + """A filtered request list keeps each request's position in the + unfiltered step, so final_logits rows stay attributed correctly.""" + from dmi.hooks.point import HookPoint + from dmi.hooks.specs import HOOK_TYPE_FINAL_LOGITS, HookSpec, ModelShapeConfig + + vocab = 8 + ring.transport.set_model_cfg(ModelShapeConfig( + hidden_dim=HID, num_heads=2, num_kv_heads=2, head_dim=HID // 2, + dtype=torch.float32, vocab_size=vocab)) + logits = torch.arange(4 * vocab, dtype=torch.float32, device="cuda").reshape(4, vocab) + hook = HookPoint() + hook._ring_hook_type = HOOK_TYPE_FINAL_LOGITS + hook._ring_hook_id = -1 + hook._ring_payload = ring.payload + ring.transport._active_specs = [HookSpec(hook_type=HOOK_TYPE_FINAL_LOGITS, + module=None, layer_no=-1)] + # Requests 1 and 3 of four survive the filter. + ring.transport.set_step_context(model_id="m", req_ids=["q1", "q3"], + token_ranges=[(0, 5), (0, 2)], + dim0_offsets=[3, 9], flattened=True, + orig_positions=[1, 3]) + ring.transport.pre_push_all_metas(batch=0, q_len=11, kv_dim=0, logits_to_keep=4) + hook(logits) + got = {s[2]: s[5] for s in ring.drain(2)} + assert torch.equal(got["q1"], logits[1:2].cpu()) + assert torch.equal(got["q3"], logits[3:4].cpu()) + + +def test_native_layout_parser_accepts_rows_and_arrays(): + from dmi.transport import native as ne + + rows = [("b", 0, 3, 0, 0, 1), ("a", 10, 14, 3, 0, 0)] + assert ne.parse_hook_requests(["a", "b"], rows) == rows + arrays = (np.array([1, 0]), np.array([0, 10]), np.array([3, 14]), + np.array([0, 3]), np.array([1, 0])) + assert ne.parse_hook_requests(["a", "b"], arrays) == rows + with pytest.raises(ValueError, match="outside the step"): + ne.parse_hook_requests(["a"], tuple(np.array([v]) for v in (2, 0, 1, 0, 0))) + with pytest.raises(ValueError, match="five fields"): + ne.parse_hook_requests(["a"], (np.array([0]),)) + with pytest.raises(ValueError, match="range or offset"): + ne.parse_hook_requests(["a"], [("a", 4, 2, 0, 0, 0)]) + + +def test_device_toggle_gates_gather_producers_of_excluded_captures( + ring, monkeypatch): + """In device-toggle mode a capture excluded from the host registry + (e.g. a partial-model segment graph) captures its gather producers as + device-updatable nodes behind one planner node: each node follows its + own table row every replay, with no host SetEnabled call.""" + from dmi import node_toggle + from dmi.hook_rows import encode_row_table, plan_hook_rows + from dmi.hooks.specs import HOOK_TYPE_RESID_PRE + + monkeypatch.setattr(torch.cuda, "CUDAGraph", torch.cuda.CUDAGraph) + monkeypatch.setattr(node_toggle, "_PATCHED_CLASS", None) + node_toggle.begin_capture_window(ring.engine, exclude_capture=lambda: True, + device_toggle=True) + if not ring.engine.toggle_device_mode(): + node_toggle.end_capture_window(ring.engine) + pytest.skip("device graph updates unavailable") + ht = HOOK_TYPE_RESID_PRE + req_ids, ranges, offsets, pos, total = _step_layout() + x = torch.randn(total, HID, device="cuda") + tables = [torch.zeros(1 + 2 * MAX_SEGS, dtype=torch.int64, device="cuda") + for _ in range(3)] + hooks = [_hook(ring, ht, layer, tables[layer]) for layer in range(3)] + ring.engine.set_null_mode(True) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for hook in hooks: + hook(x) + graph.instantiate() + ring.engine.set_null_mode(False) + node_toggle.end_capture_window(ring.engine) + assert graph._dmi_toggle_excluded + assert ring.engine.toggle_device_node_count() == 3 + assert ring.engine.toggle_node_count() == 0 # nothing host-toggled + assert ring.engine.capture_anomaly_count() == 0 + + plan = plan_hook_rows(req_ids, ranges, offsets, pos, {"a", "c"}, total, MAX_SEGS) + for enabled in ([0, 2], [1], [], [0, 1, 2]): + for layer, table in enumerate(tables): + table.copy_(torch.from_numpy(encode_row_table( + plan if layer in enabled else None, MAX_SEGS))) + ring.transport._active_specs = [_spec(ht, layer) for layer in enabled] + ring.transport.set_step_context( + model_id="m", req_ids=req_ids, token_ranges=ranges, + dim0_offsets=offsets, flattened=True, + hook_rows={(ht, layer): plan for layer in enabled}) + ring.transport.pre_push_all_metas(batch=0, q_len=total, kv_dim=0) + graph.replay() + want = [] + for layer in enabled: + want += _expected(x, plan, "blocks.hook_resid_pre", layer) + _compare(ring.drain(len(want)), want) + assert ring.engine.toggle_device_errors() == 0 + + +def _misaligned(rows, hidden, seed): + """A contiguous [rows, hidden] fp32 view at storage offset 1 (4-byte, not + 16-byte, aligned base).""" + gen = torch.Generator().manual_seed(seed) + buf = torch.randn(1 + rows * hidden, generator=gen).cuda() + view = buf[1:1 + rows * hidden].view(rows, hidden) + assert view.is_contiguous() and view.data_ptr() % 16 != 0 + return view + + +def test_heads_score_tensors_with_misaligned_bases(ring): + """hidden % 8 == 0 alone must not select 16-byte vector loads: a view at + an odd storage offset takes the scalar path, with the same scores.""" + from dmi.hook_rows import full_row_plan + from dmi.hooks.specs import (HOOK_TYPE_RESID_PRE, HOOK_TYPE_RESID_PRE_HEAD, + ModelShapeConfig) + + hid, K, rows = 64, 2, 5 + ht, hh = HOOK_TYPE_RESID_PRE, HOOK_TYPE_RESID_PRE_HEAD + ring.transport.set_model_cfg(ModelShapeConfig( + hidden_dim=hid, num_heads=2, num_kv_heads=2, head_dim=hid // 2, + dtype=torch.float32, head_slots=K)) + req_ids, ranges, offsets, pos = ["a"], [(0, rows)], [0], [0] + plan = full_row_plan(req_ids, ranges, offsets, pos, rows) + table = torch.zeros(1 + 2 * MAX_SEGS, dtype=torch.int64, device="cuda") + _write_table(table, plan) + x = _misaligned(rows, hid, 1) + rng = np.random.default_rng(3) + w = rng.standard_normal((K, hid)).astype(np.float32) + bias = np.array([0.25, -0.5], dtype=np.float32) + want = x.double().cpu() @ torch.from_numpy(w).double().T + torch.from_numpy(bias).double() + + # Standalone head node: fp32 arena (weights, biases, split scratch), also + # at a misaligned base. + from dmi.transport import native as ne + floats = K * hid + K + 16 * 8 * K + arena_buf = torch.zeros(1 + floats, dtype=torch.float32, device="cuda") + arena = arena_buf[1:] + arena[:K * hid] = torch.from_numpy(w.reshape(-1)).cuda() + arena[K * hid:K * hid + K] = torch.from_numpy(bias).cuda() + ring.transport._active_specs = [_spec(hh, 0)] + ring.transport.set_step_context(model_id="m", req_ids=req_ids, token_ranges=ranges, + dim0_offsets=offsets, flattened=True, + hook_rows={(hh, 0): plan}) + ring.transport.pre_push_all_metas(batch=0, q_len=rows, kv_dim=0) + ring.engine.reserve_one(rows * K * 4) + torch.ops.ring.producer_head(ring.payload, x, table, arena, K, hh, 0) + got = ring.drain(1) + assert len(got) == 1 + torch.testing.assert_close(got[0][5].double(), want, rtol=1e-5, atol=1e-5) + + # Fused node: rows and heads from the misaligned view. + node = _fused_arena(w, bias, [1, 1], [0, 0]) + acc = torch.zeros(16, 16, dtype=torch.float32, device="cuda") + head_table = table.clone() + hook = _hook(ring, ht, 0, table, row_bytes=hid * 4) + from dmi.hook_rows import configure_hook_heads + configure_hook_heads(hook, head_type=hh, layer_no=0, table=head_table, + arena=node, acc=acc, slots=K, weight_bf16=False) + ring.transport._active_specs = [_spec(ht, 0), _spec(hh, 0)] + ring.transport.set_step_context(model_id="m", req_ids=req_ids, token_ranges=ranges, + dim0_offsets=offsets, flattened=True, + hook_rows={(ht, 0): plan, (hh, 0): plan}) + ring.transport.pre_push_all_metas(batch=0, q_len=rows, kv_dim=0) + align = lambda n: (n + 15) & ~15 + ring.engine.prepare_step(align(rows * hid * 4) + align(rows * K * 4), 2) + hook(x) + got = {s[0]: s[5] for s in ring.drain(2)} + torch.testing.assert_close(got["blocks.hook_resid_pre"], x.cpu()) + torch.testing.assert_close(got["blocks.hook_resid_pre_head"].double(), want, + rtol=1e-5, atol=1e-5) + torch.cuda.synchronize() # no misaligned-address fault poisoned the context + + +@pytest.mark.parametrize("op", ["gather", "head", "fused"]) +@pytest.mark.parametrize("bad", ["int32", "strided", "host", "empty"]) +def test_producers_refuse_malformed_row_tables(ring, op, bad): + from dmi.hooks.specs import HOOK_TYPE_RESID_PRE, HOOK_TYPE_RESID_PRE_HEAD + + x = torch.zeros(4, HID, device="cuda") + good = torch.zeros(1 + 2 * MAX_SEGS, dtype=torch.int64, device="cuda") + table = {"int32": good.to(torch.int32), "strided": torch.zeros( + 2 * (1 + 2 * MAX_SEGS), dtype=torch.int64, device="cuda")[::2], + "host": good.cpu(), "empty": good[:0]}[bad] + K = 2 + with pytest.raises(RuntimeError, match="row table"): + if op == "gather": + torch.ops.ring.producer_gather(ring.payload, x, table, HID * 4, + HOOK_TYPE_RESID_PRE, 0) + elif op == "head": + arena = torch.zeros(K * HID + K + 16 * 8 * K, device="cuda") + torch.ops.ring.producer_head(ring.payload, x, table, arena, K, + HOOK_TYPE_RESID_PRE_HEAD, 0) + else: + node = torch.zeros((64 + K * HID + 3) & ~3, dtype=torch.int32, device="cuda") + acc = torch.zeros(4, 16, device="cuda") + torch.ops.ring.producer_fused(ring.payload, x, None, 0, table, node, acc, + K, 0, HOOK_TYPE_RESID_PRE, 0, + HOOK_TYPE_RESID_PRE_HEAD) + + +def test_head_arena_and_accumulator_must_share_the_tensor_device(ring): + from dmi.hooks.specs import HOOK_TYPE_RESID_PRE, HOOK_TYPE_RESID_PRE_HEAD + + K = 2 + x = torch.zeros(4, HID, device="cuda") + table = torch.zeros(1 + 2 * MAX_SEGS, dtype=torch.int64, device="cuda") + with pytest.raises(RuntimeError, match="arena"): + torch.ops.ring.producer_head(ring.payload, x, table, + torch.zeros(K * HID + K + 16 * 8 * K), K, + HOOK_TYPE_RESID_PRE_HEAD, 0) + node = torch.zeros((64 + K * HID + 3) & ~3, dtype=torch.int32, device="cuda") + with pytest.raises(RuntimeError, match="acc"): + torch.ops.ring.producer_fused(ring.payload, x, None, 0, table, node, + torch.zeros(4, 16), K, 0, + HOOK_TYPE_RESID_PRE, 0, HOOK_TYPE_RESID_PRE_HEAD) diff --git a/tests/test_ring_torch_op_host.py b/tests/test_ring_torch_op_host.py new file mode 100644 index 000000000..7f4a75eb7 --- /dev/null +++ b/tests/test_ring_torch_op_host.py @@ -0,0 +1,306 @@ +"""Host-only tests of the ring producer ops' launch counter (no CUDA toolkit +or GPU). + +native/csrc/ring/ring_torch_op.cpp is compiled unchanged against CPU-only +Torch and tests/native/ring/cpu_cuda_stub/, with the RingEnginePy double in +tests/native/ring/test_ring_torch_op_host.cpp, into a shared library. A fresh +interpreter loads it (the ops register the ``ring`` Torch library, which must +not collide with a native extension loaded by another test) and: + +* calls every producer op and reads the real ``ring_producer_launch_count``; +* drives the real ``dmi.node_toggle`` capture lifecycle with an engine whose + producer *is* the compiled op, so whether a capture counts as carrying + producers comes from the op's own counter, not from a test double. + +Device-side behavior (the uploaded exec, node state on the GPU) is covered by +the GPU tests. +""" + +from __future__ import annotations + +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path + +import pytest + +ROOT = Path(__file__).resolve().parents[1] +RING_DIR = ROOT / "native" / "csrc" / "ring" +TESTS_DIR = ROOT / "tests" / "native" / "ring" + +NEW_OPS = ("gather", "head", "fused", "fused_rows", "fused_heads") +ALL_OPS = ("static", "prefix", "chunked") + NEW_OPS + +_DRIVER = r""" +import ctypes, json, runpy, sys +import pytest +import torch + +lib = ctypes.CDLL(sys.argv[1]) +for name in ("dmx_host_launch_count", "dmx_host_launcher_calls", + "dmx_host_anomalies", "dmx_host_node_count", "dmx_host_node_key"): + getattr(lib, name).restype = ctypes.c_uint64 +lib.dmx_host_node_count.argtypes = [ctypes.c_uint64] +lib.dmx_host_node_key.argtypes = [ctypes.c_uint64, ctypes.c_uint64] +lib.dmx_host_op.argtypes = [ctypes.c_char_p, ctypes.c_uint64] +lib.dmx_host_last_error.restype = ctypes.c_char_p +ops = json.loads(sys.argv[3]) +out = {"ops": {}, "captured": {}, "excluded": {}, "late": {}} + + +def run(op, graph=0): + c0, l0 = lib.dmx_host_launch_count(), lib.dmx_host_launcher_calls() + rc = lib.dmx_host_op(op.encode(), graph) + return {"rc": rc, "counter_delta": lib.dmx_host_launch_count() - c0, + "launcher_calls": lib.dmx_host_launcher_calls() - l0, + "error": lib.dmx_host_last_error().decode().splitlines()[:1]} + + +for op in ops + ["fused_empty", "gather_host_tensor", "gather_bad_table"]: + out["ops"][op] = run(op) + +# Under stream capture with node recording on. +lib.dmx_host_set_toggle_capture(1) +for i, op in enumerate(ops): + lib.dmx_host_clear_nodes() + graph = 0x7000 + i + res = run(op, graph) + res["keys"] = [divmod(lib.dmx_host_node_key(graph, j), 1 << 32) + for j in range(lib.dmx_host_node_count(graph))] + res["anomalies"] = lib.dmx_host_anomalies() + out["captured"][op] = res +lib.dmx_host_set_toggle_capture(0) +lib.dmx_host_clear_nodes() + +# The real node_toggle lifecycle, the compiled op as the producer. The ring +# ops are this library's; a locally built native extension must not register +# them a second time. +import dmi.transport.native as native + + +def _no_extension(): + raise ImportError("the host harness provides the ring ops") + + +native._load_extension = _no_extension +fx = runpy.run_path(sys.argv[2]) +from dmi import node_toggle as nt +from dmi.transport import ring + + +class NativeOpEngine(fx["LifecycleEngine"]): + def __init__(self, op): + super().__init__() + self.op = op + + def produce(self, layer=0, *, counted=True): + lib.dmx_host_set_toggle_capture(int(bool(self.capture))) + raw = self.capturing.raw if self.capturing is not None else 0 + try: + res = run(self.op, raw) + finally: + lib.dmx_host_set_toggle_capture(0) + assert res["rc"] == 0, res + + def producer_launch_count(self): + return lib.dmx_host_launch_count() + + def toggle_node_count(self): + return lib.dmx_host_node_count(0) + + def toggle_node_count_for_graph(self, raw): + return lib.dmx_host_node_count(raw) + + def clear_toggle_registry(self): + super().clear_toggle_registry() + lib.dmx_host_clear_nodes() + + def is_graph_ready(self, raw): + return lib.dmx_host_node_count(raw) > 0 and raw in self.bound + + +def capture(engine): + graph = torch.cuda.CUDAGraph() + graph.capture_begin() + engine.capturing = graph + engine.produce() + engine.capturing = None + graph.capture_end() + return graph + + +for op in ops: + # Excluded capture in device-toggle mode: no node is recorded, so only + # the launch counter can tell that it holds a producer that needs the + # upload before its first replay. + with pytest.MonkeyPatch.context() as mp: + mp.setattr(torch.cuda, "CUDAGraph", fx["LifecycleGraph"]) + fx["_reset_patch_state"](mp, nt) + mp.setattr(ring, "_active_transport", None) + engine = NativeOpEngine(op) + nt.begin_capture_window(engine, exclude_capture=lambda: True, + device_toggle=True) + graph = capture(engine) + trace_at_capture_end = list(graph.trace) + uploads_before_replay = list(engine.uploads) + nt.end_capture_window(engine) + nt.set_replay_guard(True) + mp.setattr(ring, "_active_transport", + fx["GatedTransport"](engine, requires_registry=False)) + graph.replay() + out["excluded"][op] = { + "producers": graph._dmi_producers, + "trace_at_capture_end": trace_at_capture_end, + "uploaded_before_replay": uploads_before_replay == [ + (graph.raw, graph.exec_ptr)], + "replays": graph.replays, + } + nt._WINDOW.engine = None + + # A capture after the window closed: no node, no bind. The guard must + # refuse its replay. + with pytest.MonkeyPatch.context() as mp: + mp.setattr(torch.cuda, "CUDAGraph", fx["LifecycleGraph"]) + fx["_reset_patch_state"](mp, nt) + mp.setattr(ring, "_active_transport", None) + engine = NativeOpEngine("static") + nt.begin_capture_window(engine) + monitored = capture(engine) + nt.bind_graph(engine, monitored) + nt.end_capture_window(engine) + nt.set_replay_guard(True) + engine.op = op + late = capture(engine) + mp.setattr(ring, "_active_transport", fx["GatedTransport"](engine)) + error = None + try: + late.replay() + except RuntimeError as exc: + error = str(exc) + monitored.replay() + out["late"][op] = {"producers": late._dmi_producers, + "replays": late.replays, "error": error, + "monitored_replays": monitored.replays} + nt._WINDOW.engine = None + +out["cuda_initialized"] = torch.cuda.is_initialized() +print("RESULT " + json.dumps(out)) +""" + + +def _torch_root() -> Path: + import torch + + return Path(torch.__file__).resolve().parent + + +@pytest.fixture(scope="module") +def op_results(tmp_path_factory): + compiler = shutil.which("g++") or shutil.which("c++") + if compiler is None: + pytest.skip("a C++20 compiler is required") + import torch + + torch_root = _torch_root() + build_dir = tmp_path_factory.mktemp("ring_torch_op_host") + library = build_dir / "libring_torch_op_host.so" + build = subprocess.run( + [compiler, "-std=c++20", "-O0", "-fPIC", "-shared", + "-Wl,--no-undefined", + f"-D_GLIBCXX_USE_CXX11_ABI={int(torch.compiled_with_cxx11_abi())}", + # The stub directory comes first: it shadows the CUDA runtime and + # , so no CUDA toolkit header is used. + f"-I{TESTS_DIR / 'cpu_cuda_stub'}", f"-I{RING_DIR}", + f"-I{RING_DIR.parent}", f"-I{torch_root / 'include'}", + f"-I{torch_root / 'include' / 'torch' / 'csrc' / 'api' / 'include'}", + str(RING_DIR / "ring_torch_op.cpp"), + str(TESTS_DIR / "test_ring_torch_op_host.cpp"), + f"-L{torch_root / 'lib'}", f"-Wl,-rpath,{torch_root / 'lib'}", + "-ltorch_cpu", "-lc10", "-o", str(library)], + capture_output=True, text=True, check=False, + ) + assert build.returncode == 0, build.stdout + build.stderr + env = dict(os.environ) + env["CUDA_VISIBLE_DEVICES"] = "" + env["PYTHONPATH"] = os.pathsep.join( + [str(ROOT / "src")] + ([env["PYTHONPATH"]] if env.get("PYTHONPATH") else [])) + run = subprocess.run( + [sys.executable, "-c", _DRIVER, str(library), + str(ROOT / "tests" / "test_node_toggle.py"), json.dumps(list(ALL_OPS))], + capture_output=True, text=True, check=False, timeout=300, env=env, + cwd=build_dir, + ) + assert run.returncode == 0, run.stdout + run.stderr + line = [ln for ln in run.stdout.splitlines() if ln.startswith("RESULT ")] + assert line, run.stdout + run.stderr + results = json.loads(line[-1][len("RESULT "):]) + assert results["cuda_initialized"] is False + return results + + +@pytest.mark.cpu +@pytest.mark.parametrize("op", ALL_OPS) +def test_every_producer_op_counts_its_launch_once(op_results, op): + res = op_results["ops"][op] + assert res["rc"] == 0, res + assert res["launcher_calls"] == 1 + assert res["counter_delta"] == 1 + + +@pytest.mark.cpu +@pytest.mark.parametrize("op", ["fused_empty", "gather_host_tensor"]) +def test_an_op_that_launches_nothing_counts_nothing(op_results, op): + res = op_results["ops"][op] + assert res["rc"] == 0 + assert (res["launcher_calls"], res["counter_delta"]) == (0, 0) + + +@pytest.mark.cpu +def test_a_refused_op_counts_nothing(op_results): + res = op_results["ops"]["gather_bad_table"] + assert res["rc"] == 1 and "row table" in res["error"][0] + assert (res["launcher_calls"], res["counter_delta"]) == (0, 0) + + +@pytest.mark.cpu +@pytest.mark.parametrize("op,keys,anomalies", [ + ("static", [[3, 0]], 0), + ("prefix", [[3, 0]], 0), + ("chunked", [], 1), # not toggle-managed: fails closed + ("gather", [[3, 0]], 0), + ("head", [[24, 0]], 0), + ("fused", [[3, 0], [24, 0]], 0), # one node, both keys + ("fused_rows", [[3, 0]], 0), + ("fused_heads", [[24, 0]], 0), +]) +def test_a_captured_op_counts_one_launch_and_records_its_keys( + op_results, op, keys, anomalies): + res = op_results["captured"][op] + assert res["rc"] == 0 and res["counter_delta"] == 1 + assert res["keys"] == keys + assert res["anomalies"] == anomalies + + +@pytest.mark.cpu +@pytest.mark.parametrize("op", ALL_OPS) +def test_an_excluded_device_toggle_capture_is_uploaded_before_it_replays( + op_results, op): + res = op_results["excluded"][op] + assert res["producers"] is True + assert res["trace_at_capture_end"] == [ + "capture_begin", "capture_end", "instantiate"] + assert res["uploaded_before_replay"] is True + assert res["replays"] == 1 + + +@pytest.mark.cpu +@pytest.mark.parametrize("op", ALL_OPS) +def test_a_late_unbound_capture_fails_closed(op_results, op): + res = op_results["late"][op] + assert res["producers"] is True + assert res["replays"] == 0 + assert "not registered and bound" in res["error"] + assert res["monitored_replays"] == 1 diff --git a/tests/test_row_table_clamp.py b/tests/test_row_table_clamp.py new file mode 100644 index 000000000..78f50980c --- /dev/null +++ b/tests/test_row_table_clamp.py @@ -0,0 +1,83 @@ +"""Producer row tables are read only inside their physical storage. + +A row table is ``[n, src0, len0, ...]`` int64. The kernels re-read ``n`` +every replay, so the host never reads it (no per-launch sync, nothing baked +into a graph); it passes the table's element count instead, and the device +loader clamps the segment count by both ``GATHER_MAX_SEGS`` and that +capacity. A one-word table whose count word is 1 used to read ``table[1]`` +and ``table[2]`` past its allocation. + +The clamp itself (``native/csrc/ring/row_table.h``) is plain C++ shared by +host and device code, so it is compiled and run here on the CPU; the static +checks pin the wiring from the torch ops to the shared loader. +""" + +from __future__ import annotations + +import re +import shutil +import subprocess +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.cpu + +ROOT = Path(__file__).resolve().parents[1] +RING = ROOT / "native" / "csrc" / "ring" + + +def test_segment_count_is_bounded_by_table_storage(tmp_path): + compiler = shutil.which("g++") or shutil.which("c++") + if compiler is None: + pytest.skip("a C++17 compiler is required") + source = ROOT / "tests" / "native" / "test_row_table_clamp.cpp" + executable = tmp_path / "test_row_table_clamp" + built = subprocess.run( + [compiler, "-std=c++17", "-O0", "-Wall", "-Wextra", "-Werror", + f"-I{ROOT / 'native' / 'csrc'}", str(source), "-o", str(executable)], + capture_output=True, text=True, check=False) + assert built.returncode == 0, built.stdout + built.stderr + ran = subprocess.run([str(executable)], capture_output=True, text=True, + check=False, timeout=60) + assert ran.returncode == 0, ran.stdout + ran.stderr + assert ran.stdout.strip() == "ok" + + +def test_python_and_native_segment_limits_agree(): + from dmi.hook_rows import GATHER_MAX_SEGS + + header = (RING / "row_table.h").read_text() + match = re.search(r"GATHER_MAX_SEGS\s*=\s*(\d+)", header) + assert match and int(match.group(1)) == GATHER_MAX_SEGS + + +def test_every_table_reader_receives_the_table_capacity(): + producer = (RING / "producer.cu").read_text() + loader = producer[producer.index("load_row_table(const int64_t* table"):] + loader = loader[:loader.index("__syncthreads")] + # The count comes from the shared clamp, bounded by table_numel. + assert "row_table_segment_count(table[0], table_numel)" in loader + assert "GATHER_MAX_SEGS) n_raw" not in loader + # Each kernel passes its own table's capacity to the loader. + assert producer.count("load_row_table(seg_table_dev_ptr, table_numel,") == 2 + assert "load_row_table(gtab, gtab_numel," in producer + assert "load_row_table(htab, htab_numel," in producer + + op = (RING / "ring_torch_op.cpp").read_text() + for impl, needles in ( + ("ring_producer_gather_impl", ["seg_table.numel()"]), + ("ring_producer_head_impl", ["seg_table.numel()"]), + ("ring_producer_fused_impl", ["gather_table->numel()", + "head_table->numel()"]), + ): + body = op[op.index(f"void {impl}("):] + body = body[:body.index("\n}\n")] + for needle in needles: + assert needle in body, (impl, needle) + # The count word is never read on the host (no sync, nothing baked + # into a captured graph). + assert ".item" not in body and "[0]" not in body, impl + assert "check_row_table(" in body, impl + # A table must at least hold its count word. + assert "table.numel() >= 1" in op