Repository navigation
[Core] Add per-hook request-row gather, PackedDemandPlanner, and fused inline projection heads - #172
Conversation
|
Independent audit of #172 at Evidence below separates executed CPU reproductions from static CUDA-path analysis. No GPU tests, CUDA initialization, inference, profiler, or performance benchmarks were run. All GPU verification remains deferred until the owner's explicit instruction. The scripts use the real checked-out Python implementation with producer/graph boundary doubles as identified below; they do not establish an observed GPU failure. G172-1 — [P1] Clear gather/head attachment state on detach and ordinary reattachCode: fused dispatch, new retained head configuration, and existing uninstall path. Trigger: Install a row HookPoint, configure gather plus heads, detach it through Actual, reproduced on CPU: The reattached HookPoint still dispatches Expected: A plain new attachment uses the ordinary producer and releases the prior attachment's feature resources. A retained enabled head table makes the native fused publication logic include an extra task without corresponding new metadata/reservation; retained row selection can also affect the row payload. That downstream consequence follows from the code and still needs GPU delivery validation. Fix: Give gather/head configuration an explicit attachment lifetime and clear it in detach/reinstall, including tensor references and restoration of the previous strip configuration. Cover detach/reattach after enabled, disabled and exceptional configurations. See the shared executable reproduction below ( G172-2 — [P1] Adding heads to a plain HookPoint removes its original row publicationCode: accepted public head configuration, unconditional fused dispatch, native enable predicates, and task publication count. Trigger: Call Actual, reproduced on CPU: For four fp32 rows with hidden width 4 and one head slot, the real host planner returns 2 tasks / 80 bytes (64 row bytes + 16 head bytes). Before configuring heads, forward calls Expected: Both the original row task and the head task, in matching metadata order. Alternatively, reject an unsupported configuration before capture/reservation. The current mismatch breaks reservation/metadata/publication agreement; actual GPU misattribution or drain behavior remains untested. The eager safety branch dispatches the row separately and thus disagrees with the fused fast path. Fix: Preserve static/prefix row dispatch when no gather table is installed, or reject that configuration and provide an explicit full-row initialization path. Cover plain+heads, prefix+heads, gather+heads and independent row/head on/off states. See the shared reproduction ( G172-3 — [P2]
|
4362f6f to
dc06342
Compare
2323c99 to
4e234b4
Compare
|
@Samfisheryu thanks for the audit. All five findings are addressed. G172-1..4 are in dc06342, rebased onto the revised #171. T172-4 is fixed in #171 (4e234b4), because the upload helper lives in G172-1 — gather/head state belongs to one attachment
G172-2 — heads never replace the row publicationDispatch now follows the spec order (row task, then head task):
The eager safety net makes the same launches in the same order, under the same per-key gates. For every row mode × heads combination, G172-3 —
|
|
Independent follow-up audit at No GPU tests, CUDA initialization, inference, GPU queries, sanitizers, or performance benchmarks were run. Device correctness and performance verification remain deferred until the owner's explicit instruction. Producer publication sizes below are derived from the native code; they are not observed GPU output. Repaired findings positively verifiedThe targeted
T172-4 — [P1, still open] New producer ops do not update the counter used to decide upload and replay protectionThe revised helper does instantiate/upload when it recognizes a producer. However, the actual new ops do not increment The inherited helper returns the counter comparison directly, so an available but unchanged counter suppresses the node-count fallback. Consequently:
These are two consequences of one missing-counter wiring defect, not two separate findings. The mandatory device-node upload consequence specifically applies to the excluded device-updatable gather path. Executed evidence: An independent CPU harness extracted and compiled the unchanged six C++ op bodies with ATen/engine API doubles, retaining their validation, engine, recording and counter calls. Results were identical at #172 and #174: Feeding those measured counter deltas into the real The CPU graph model demonstrates API order and classification, not CUDA device execution. The original execution records are retained with the audit; the copyable wiring check below exposes the missing connection directly. Why the added test passes: Small reproducible source-wiring check (CPU only; complements the executed C++/lifecycle evidence above)Run from this revision's checkout: python - <<'PY'
from pathlib import Path
s = Path('native/csrc/ring/ring_torch_op.cpp').read_text()
for label, fn in [('static', 'ring_producer_impl'),
('prefix', 'ring_producer_prefix_impl'),
('chunked', 'ring_producer_chunked_impl'),
('gather', 'ring_producer_gather_impl'),
('head', 'ring_producer_head_impl'),
('fused', 'ring_producer_fused_impl')]:
body = s[s.index('void ' + fn + '('):]
body = body[:body.index('\n}\n')]
print(label, 'counter_increment=' + str('g_producer_launches.fetch_add' in body))
PYOutput: Fix direction: Count all actual producer launches at a common boundary, or wire every new op and test that connection against the real wrapper. A partial counter must not override evidence that a graph contains registered producer nodes. Re-run excluded gather upload and late/unbound guard cases using the actual op counter behavior, then perform the deferred device tests. G172-R2-1 — [P1, newly found] Batched eager heads reserve fewer bytes than the producer writesCode: Trigger: A batched fp32 HookPoint Executed CPU result: The real planner budgets row 96 + head 48 = 144 bytes. Real metadata shapes are The original #172 Copyable minimal CPU reproduction using the real public configuration, planner and HookPoint forwardRun from the reviewed checkout with Python, PyTorch and the project Python dependencies. No native build is needed. The fake payload exposes only dispatch attributes; all allocated tensors are CPU tensors. Producer calls and engine reservations are recorded instead of executing CUDA. env CUDA_VISIBLE_DEVICES='' OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 PYTHONNOUSERSITE=1 nice -n 15 python - <<'PY'
import json, sys, types
from pathlib import Path
sys.path.insert(0, str(Path.cwd() / 'src'))
import torch
def forbid_cuda(*a, **kw):
raise AssertionError('CUDA initialization forbidden')
torch.cuda._lazy_init = forbid_cuda
from dmi.hook_rows import configure_hook_heads
from dmi.hooks.point import HookPoint
from dmi.hooks.dispatch import install_ring_hooks
from dmi.hooks.specs import HookSpec, ModelShapeConfig
from dmi.adapters.base import BackendAdapter
from dmi.adapters.types import StepContext
import dmi.transport
rt = types.ModuleType('dmi.transport.ring')
sys.modules[rt.__name__] = rt
dmi.transport.ring = rt
calls, reserved = [], []
for name in ('producer', 'producer_fused'):
setattr(torch.ops.ring, name,
lambda *a, _name=name: calls.append(_name))
engine = types.SimpleNamespace(
available_capacity=lambda: 128, available_task_slots=lambda: 16,
flush_and_wait=lambda: None, reserve_one=reserved.append)
transport = types.SimpleNamespace(
capture_step=True, force_eager=True, _ring_engine=engine,
toggle_gate_active=False, effective_cap=128)
rt._active_transport = transport
class Payload:
is_cuda = True
shape = (2, 3, 4)
nbytes = 96
def contiguous(self): return self
x = Payload()
hook = HookPoint()
row = HookSpec(hook_type=0, layer_no=0, module=hook)
head = HookSpec(hook_type=24, layer_no=0, module=hook, dtype=torch.float32)
def arm():
# Full physical source: six rows, two fp32 head slots; all actual tensors CPU.
configure_hook_heads(hook, head_type=24, layer_no=0,
table=torch.tensor([1, 0, 6], dtype=torch.int64),
arena=torch.zeros(72, dtype=torch.int32),
acc=torch.zeros(6, 16), slots=2, weight_bf16=False)
install_ring_hooks([row], ring_payload=object())
arm()
cfg = ModelShapeConfig(hidden_dim=4, num_heads=1, num_kv_heads=1,
head_dim=4, dtype=torch.float32, head_slots=2)
adapter = types.SimpleNamespace(model_cfg=cfg, active_specs=[row],
transport=transport, _spec_needs_eager=lambda _: False)
BackendAdapter.replace_active_specs(adapter, [row, head])
ctx = StepContext(model_id='m', flattened=False, req_ids=['A', 'B'],
token_ranges=[(0, 3), (0, 3)], dim0_offsets=[0, 1], kv_offsets=[0, 0],
batch=2, q_len=3)
plan = BackendAdapter.plan_step(adapter, ctx)
hook.forward(x)
print(json.dumps({'plan_tasks': plan.hook_count, 'plan_bytes': plan.total_bytes,
'reserve_one_calls': reserved, 'producer_calls': calls,
'native_head_bytes_from_size_contract': (x.nbytes // (x.shape[-1]*4))*2*4}))
assert plan.total_bytes == 144 and reserved == [96, 16]
# A 48-byte head also bypasses the explicit 32-byte effective-capacity guard.
install_ring_hooks([head], ring_payload=object())
arm()
transport.effective_cap = 32
reserved.clear(); calls.clear()
hook.forward(x)
assert reserved == [16] and calls == ['producer_fused']
print(json.dumps({'head_only_effective_cap': 32, 'reserve_one_calls': reserved,
'native_head_bytes_from_size_contract': 48, 'raised': False}))
assert not torch.cuda.is_initialized()
assert 'dmi.transport.native' not in sys.modules
PYObserved: The 128-byte capacity in the first case is smaller than the 144-byte step plan, so this is precisely a situation where the adapter can select the per-hook eager safety path. The script directly selects that path to isolate the faulty calculation. The short script above was independently executed on this revision. The larger cross-revision matrix and original-version control are retained with the audit. Fix direction: Derive the eager head's row count from all dimensions preceding hidden, consistently with the native launcher, before both capacity checking and reservation. Alternatively, if heads intentionally support only packed 2-D inputs, reject batched configuration before publishing metadata or reserving, and align the shape/documentation contract. Add |
dc06342 to
22a6482
Compare
4e234b4 to
541b10d
Compare
|
@Samfisheryu thanks for the re-audit. T172-4 and G172-R2-1 are fixed in 22a6482. It is rebased on the revised #171, which carries the helper half of T172-4. T172-4: the launch counter now covers every producer op
G172-R2-1: batched eager heads
The targeted files pass on this branch: node toggle, toggle registry, hook row features, hook rows, row table clamp and the new op test, 178 tests. All ring test binaries link. CPU only: no GPU tests, CUDA initialization or benchmarks were run. Device and sanitizer validation stays deferred until the owner gives the go-ahead. |
|
Independent third review of R2-172-counter / T172-4: I verified the real C++ op bodies, the common counter helper and the real Python lifecycle together. All six producer families now increment the counter once, including gather/head/fused. A fused producer still records both keys while counting one launch. The independent new/previous-version controls show: Recorded node evidence also remains positive when the counter has not changed. These observations establish corrected host call ordering and classification; device execution remains pending. R2-172-batched / G172-R2-1: the original batched case now agrees across the actual planner, metadata and eager forward: The independent matrix passed 10 mode/path combinations, 16 row/head gate combinations, 14 shape combinations and 16 capacity combinations at both this head and the full stack. It includes packed, batched, head-only, one-dimensional and empty-row inputs. The same runner retains the previous revision's contrary results. The targeted gather/features/table-capacity files passed 93 CPU tests. The combined toggle/registry/real-op host files passed 85 CPU tests. The earlier four gather findings remain resolved. The latest workflow is successful. No additional confirmed issue was found in this revision. GPU task delivery, numerical checks, actual device upload/replay and performance are still pending the owner's explicit instruction; no GPU work was performed. |
Samfisheryu
left a comment
Samfisheryu
left a comment
There was a problem hiding this comment.
Approved at 22a6482275d38c7bc0cb2de49eb236c74e65c13f following the completed independent code and CPU/host review. The issues raised for this PR have been addressed and independently rechecked. The reviewed commit's CI is successful.
Review record: #172 (comment)
GPU execution and end-to-end performance benchmarks remain a separate follow-up, pending the owner's explicit instruction as agreed. This approval records the completed code/CPU review and does not claim those measurements.
541b10d to
ffbedf5
Compare
…d 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.
22a6482 to
6d7e741
Compare
Depends on #171.
Summary
Each HookPoint can copy only the rows it needs -- some requests, or absolute
token windows of them -- chosen per step through a device row table that
captured graphs re-read on every replay. Optional fused projection heads score
selected rows inside the graph so only K fp32 scores leave the GPU.
Implementation details
table,
-1disables the producer, host-mirror skip of eager launches,device-updatable gather nodes gated by a planner); 8-wide head loads only
from 16-byte aligned bases (views at odd storage offsets take the scalar
path); per-hook request layouts and
orig_positionin metadata;parse_hook_requests.ring::producer_gather/head/fusedvalidate every row table (device,contiguity, int64, a segment count) and the head arena and accumulator
devices, raising instead of skipping a reserved producer.
dmi.hook_rows(plan_hook_rows,HookRowPlanner,PackedDemandPlanner,full_row_plan,encode_row_table,row_signature,segment_array,configure_hook_gather,configure_hook_heads);StepContext.hook_rows/orig_positions;RingTransport.set_prefill_table_gate;BackendAdapter.replace_active_specs;head hook types
resid_pre_head,final_ln_head. The eager head pathrefuses an output larger than the ring's effective capacity before
reserving.
docs/request-row-gather.md.Compatibility and scope
Opt-in per HookPoint.
Validation
test_producer_equivalence,test_inline_head,test_fused_head(each
DMI_TEST_MUTATE_*mutant fails as required).test_hook_rows.py(planners against independent enumerations),test_per_hook_gather_gpu.py(22 passed, including misaligned-baseheads, which fault with a misaligned address without the alignment check,
and malformed-table and device-mismatch rejection),
test_inline_head_eager.py(2 passed).