Skip to content

[Core] Add per-hook request-row gather, PackedDemandPlanner, and fused inline projection heads #172

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zhaoyboo merged 1 commit into from
Nov 5, 2026
Merged

[Core] Add per-hook request-row gather, PackedDemandPlanner, and fused inline projection heads #172

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
179 changes: 179 additions & 0 deletions docs/request-row-gather.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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.
146 changes: 142 additions & 4 deletions native/csrc/bindings.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,12 @@
#include <torch/extension.h>
#endif
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <pybind11/stl.h>
#include <chrono>
#include <climits>
#include <memory>
#include <unordered_map>
#include <cmath>
#include <limits>
#include <type_traits>
Expand Down Expand Up @@ -259,6 +263,83 @@ std::vector<ring::RecordDescriptor> CopyRecordDescriptors(
}
#endif

#ifndef DMI_HOST_ONLY
using HookRowsPtr = std::shared_ptr<const std::vector<ring_py::RequestMeta>>;

// 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<ring_py::RequestMeta>& requests) {
auto parsed = std::make_shared<std::vector<ring_py::RequestMeta>>();
if (py::isinstance<py::tuple>(rows_obj)) {
auto fields = py::reinterpret_borrow<py::tuple>(rows_obj);
if (py::len(fields) != 5)
throw py::value_error("hook request arrays need five fields");
using I64 = py::array_t<int64_t, py::array::c_style | py::array::forcecast>;
I64 idx = fields[0].cast<I64>(), lo = fields[1].cast<I64>();
I64 hi = fields[2].cast<I64>(), off = fields[3].cast<I64>();
I64 pos = fields[4].cast<I64>();
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<int64_t>(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<int32_t>(v);
};
parsed->reserve(static_cast<size_t>(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<size_t>(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<py::list>();
parsed->reserve(py::len(rows));
for (auto row : rows) {
auto fields = row.cast<py::tuple>();
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<std::string>();
rm.start_token = fields[1].cast<int32_t>();
rm.end_token = fields[2].cast<int32_t>();
rm.dim0_offset = fields[3].cast<int64_t>();
rm.kv_offset = fields[4].cast<int32_t>();
rm.orig_position = fields[5].cast<int32_t>();
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) {
Expand Down Expand Up @@ -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<size_t>(py::len(hook_types_py));
if (py::len(hook_requests_py) &&
static_cast<size_t>(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<ring_py::StepContext>();
auto* ctx = owned_ctx.get();
ctx->model_id = model_id;
ctx->tp_rank = tp_rank;
ctx->dp_rank = dp_rank;
Expand All @@ -914,12 +1002,20 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
rm.dim0_offset = py::cast<int64_t>(dim0_offsets_py[i]);
if (i < static_cast<size_t>(py::len(kv_offsets_py)))
rm.kv_offset = py::cast<int32_t>(kv_offsets_py[i]);
// Only when the integration filtered the request list
// (empty list = unfiltered; -1 = use the list index).
if (i < static_cast<size_t>(py::len(orig_positions_py)))
rm.orig_position = py::cast<int32_t>(orig_positions_py[i]);
ctx->requests.push_back(std::move(rm));
}
// Build per-hook metas
size_t n = static_cast<size_t>(py::len(hook_types_py));
std::vector<ring_py::TensorMeta> 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<PyObject*, HookRowsPtr> layouts;
for (size_t i = 0; i < n; ++i) {
ring_py::TensorMeta meta;
meta.hook_type = py::cast<int>(hook_types_py[i]);
Expand All @@ -930,11 +1026,22 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
py::list shape = shapes_py[i].cast<py::list>();
for (auto d : shape)
meta.shape.push_back(py::cast<int64_t>(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"),
Expand All @@ -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"),
Expand Down Expand Up @@ -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<ring_py::RequestMeta> requests;
requests.reserve(py::len(req_ids_py));
for (auto rid : req_ids_py) {
ring_py::RequestMeta rm;
rm.req_id = rid.cast<std::string>();
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
}
Loading
Loading