From 6d7e7418a37c6517e04fc9b8ddd1e4584294a3ea Mon Sep 17 00:00:00 2001 From: zhaoyboo Date: Tue, 3 Nov 2026 16:05:09 -0500 Subject: [PATCH] [Core] Add per-hook request-row gather, PackedDemandPlanner, and fused inline projection heads Let every HookPoint copy only the rows it needs, chosen per HookPoint and per step through a per-HookPoint row table that captured CUDA graphs re-read at every replay, and let a HookPoint compute linear projection scores inside the graph so only the scores leave the GPU. Native: - producer_gather: packed copy of a row table's (source, rows) runs; -1 disables the producer for the step (no reservation, task or payload). The table is loaded by all threads and the destination prefix is one warp's __shfl_up_sync scan; a single run takes a flat whole-grid copy. - producer_head / producer_fused: K fp32 affine scores per selected row from a stable weight arena. The fused node runs the row producer and the heads in one kernel with separate tables and toggle keys, supports linear heads over several layers (FIRST/MID/LAST partial sums) and uses a fixed reduction order, bit-identical across batch compositions. - The new kernels publish through the existing publication-word protocol (block release fences, last-block acquire, system-scope release store). - Gather, head and fused producers record their kernel nodes in the toggle capture window; a fused node carries both keys. A gather producer that recorded no node could not be disabled with its key. - Device-side toggling: in device-toggle mode, gather producers of captures outside the toggle window are captured as device-updatable nodes behind one planner node per captured forward. - A pinned host mirror of the row tables lets uncaptured launches of disabled rows be skipped. - Metadata carries optional per-hook request layouts (tuples or int64 arrays, parsed once per shared layout) and each request's original position; the consumer slices final_logits by it. - HOOK_TYPE_RESID_PRE_HEAD / HOOK_TYPE_FINAL_LN_HEAD score rows. Python: - dmi.hook_rows: reference planners (plan_hook_rows, plan_hook_groups, HookRowPlanner) with absolute token windows and table-capacity fallback; PackedDemandPlanner for declared per-request demand (rules compiled once per demand change, NumPy per step, plans equal to the reference planner's); full_row_plan (static full-row baseline) and encode_row_table. - StepContext/RingTransport carry hook_rows and orig_positions; reservation and metadata use each planned hook's row count; the per-HookPoint table gate lets the per-step set follow the toggle on every step. - HookPoint gather mode and fused heads, with their eager safety-net paths. Attachment lifecycle, dispatch and table bounds: - Gather tables and in-graph heads belong to one attachment. install_ring_hooks and uninstall_ring_hooks clear them (dmi.hook_rows.clear_hook_row_features): the row table, head table, arena and accumulator references are released and the strip stride gather replaced is restored, so a detached HookPoint drops those tensors and a plain re-attach runs its ordinary producer. Installing also drops the previous row key; a head spec binds the payload but never sets the row key. - Heads on a HookPoint without a gather table keep its row publication: with a row spec, a gather table and heads one fused node publishes both; with a static, prefix or chunked row producer and heads the row producer launches first and a head-only fused node follows; heads alone launch the head-only node. The eager safety net makes the same launches in the same order, so the host plan, the fast path and the eager path agree on the tasks (row, then head) and their toggle keys. - full_row_plan copies the whole physical tensor (0, total_rows) with the step's physical offsets, the ungathered producer's payload including leading, interior and trailing padding. - The gather, head and fused ops pass table.numel() to the kernels (no host read of the count word), and the shared loader bounds the segment count by GATHER_MAX_SEGS and by the table's storage (row_table.h), so a count larger than the table never reads past it. --- docs/request-row-gather.md | 179 ++++ native/csrc/bindings.cpp | 146 ++- native/csrc/ring/p2p_thread.cpp | 17 +- native/csrc/ring/producer.cu | 803 +++++++++++++++++ native/csrc/ring/producer.cuh | 143 +++ native/csrc/ring/ring_engine_py.cu | 201 +++++ native/csrc/ring/ring_engine_py.h | 55 ++ native/csrc/ring/ring_torch_op.cpp | 213 ++++- native/csrc/ring/row_table.h | 47 + native/csrc/ring/tensor_meta.h | 22 +- native/csrc/ring/toggle_registry.cu | 16 +- native/csrc/ring/toggle_registry.h | 10 +- src/dmi/adapters/base.py | 21 + src/dmi/adapters/types.py | 15 +- src/dmi/configuration/architecture.py | 4 +- src/dmi/hook_rows.py | 849 ++++++++++++++++++ src/dmi/hooks/catalog.py | 4 + src/dmi/hooks/dispatch.py | 51 +- src/dmi/hooks/point.py | 265 ++++-- src/dmi/hooks/specs.py | 19 + src/dmi/transport/native.py | 1 + src/dmi/transport/ring.py | 109 ++- tests/native/ring/Makefile | 29 +- .../cpu_cuda_stub/ATen/cuda/CUDAContext.h | 24 + .../native/ring/cpu_cuda_stub/cuda_runtime.h | 40 +- tests/native/ring/test_fused_head.cu | 459 ++++++++++ tests/native/ring/test_inline_head.cu | 274 ++++++ .../native/ring/test_producer_equivalence.cu | 353 ++++++++ tests/native/ring/test_ring_torch_op_host.cpp | 270 ++++++ tests/native/test_row_table_clamp.cpp | 84 ++ tests/test_adapter_protocol.py | 18 + tests/test_hook_row_features.py | 548 +++++++++++ tests/test_hook_rows.py | 541 +++++++++++ tests/test_inline_head_eager.py | 57 ++ tests/test_node_toggle.py | 14 + tests/test_per_hook_gather_gpu.py | 604 +++++++++++++ tests/test_ring_torch_op_host.py | 306 +++++++ tests/test_row_table_clamp.py | 83 ++ 38 files changed, 6790 insertions(+), 104 deletions(-) create mode 100644 docs/request-row-gather.md create mode 100644 native/csrc/ring/row_table.h create mode 100644 src/dmi/hook_rows.py create mode 100644 tests/native/ring/cpu_cuda_stub/ATen/cuda/CUDAContext.h create mode 100644 tests/native/ring/test_fused_head.cu create mode 100644 tests/native/ring/test_inline_head.cu create mode 100644 tests/native/ring/test_producer_equivalence.cu create mode 100644 tests/native/ring/test_ring_torch_op_host.cpp create mode 100644 tests/native/test_row_table_clamp.cpp create mode 100644 tests/test_hook_row_features.py create mode 100644 tests/test_hook_rows.py create mode 100644 tests/test_inline_head_eager.py create mode 100644 tests/test_per_hook_gather_gpu.py create mode 100644 tests/test_ring_torch_op_host.py create mode 100644 tests/test_row_table_clamp.py 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