pub struct EmbedOverlay {
pub rows: Arc<CudaSlice<f32>>,
pub spans: Vec<(usize, usize, usize)>,
/* private fields */
}Expand description
Mixed-embedding prime overlay: image embeddings that replace <|image_pad|> token
embeddings at prompt-relative positions during prime_cache_overlaid. rows holds all
images’ merger outputs concatenated ([total_rows, n_embd]); each span is
(prompt_pos, row_off, n_rows) — rows [row_off, row_off+n_rows) land at prompt
positions [prompt_pos, prompt_pos+n_rows). Spans must not overlap.
ROW RESIDENCY IS CARRIED, NOT ASSUMED (lane/glm53-vision-ppn, 2026-09-01). rows is a
device pointer, and a device pointer is only meaningful inside ONE CUDA context. The
overlay therefore records the context its rows live in, and every consumer checks it. The
pre-lane code proxied that invariant with “the overlay was built on the primary engine AND
stage 0 IS the primary engine” (std::ptr::eq on &Engine), which is both too weak (two
Engines can share one context — CudaContext::new retains the device’s PRIMARY context,
so per-stage Engines on one device are one address space) and too strong (it refused the
real 3-card serving shape, where the worker’s primary engine follows the LAST pp stage and
stage 0’s intake engine is a different device). Context identity is the exact condition.
Fields§
§rows: Arc<CudaSlice<f32>>Arc SO A WINDOW REALLY ALIASES (lane/glm53-vision-ppn follow-up, memra-next#23).
This field was a bare CudaSlice<f32> and window() cloned it, on the belief — stated
in two comments here and one on Engine::clone_dtod — that CudaSlice::clone() bumps a
refcount. It does not: in the LOCKED cudarc 0.19.8 (Cargo.lock), impl Clone for CudaSlice is
try_clone().unwrap() -> stream.clone_dtod(self), i.e. a full device allocation plus
D2D copy of every row, and an unwrap that PANICS in the GPU worker thread. Stated
precisely, because the first version of this comment overstated it: the panic is CAUGHT
(worker.rs wraps run in catch_unwind, marks the worker dead and respawns, reaching
exit(EXIT_WORKER_UNRECOVERABLE) only once respawns are exhausted), so the cost is every
IN-FLIGHT SESSION on the box plus a respawn, not necessarily immediate process death —
still an unacceptable outcome for an allocation failure that the caller is already shaped
to propagate. window() runs once per prefill tick AND once per chunk of a chunked ppN
prime, so the old code also paid several whole-buffer copies per multi-chunk image prompt
while the docs claimed zero.
With an Arc, a window is a refcount bump for real, and the buffer is freed once when
the last window drops.
spans: Vec<(usize, usize, usize)>Implementations§
Source§impl EmbedOverlay
impl EmbedOverlay
Sourcepub fn new(
e: &Engine,
rows: CudaSlice<f32>,
spans: Vec<(usize, usize, usize)>,
) -> Result<Self, Box<dyn Error>>
pub fn new( e: &Engine, rows: CudaSlice<f32>, spans: Vec<(usize, usize, usize)>, ) -> Result<Self, Box<dyn Error>>
Wrap rows built on e, taking residency FROM THE SLICE and refusing if that disagrees
with e.
WHY NOT JUST RECORD e.ctx() (memra-next#24). Allocation goes through
Engine::stream() -> Gpu::stream(), which returns the THREAD-LOCAL stream override
when one is pushed — and that override is not keyed to an engine. So a caller that
allocates inside a PpNRt::enter(s) scope gets a buffer in STAGE s’s context while
believing it used e; recording e.ctx() would then hand require_resident a label
that vouches for a foreign pointer, which is this lane’s own hazard with the label
inverted. CudaSlice::context() is the ground truth, so it is what gets recorded, and a
disagreement with e is a REFUSAL rather than a silently relabelled buffer.
Sourcepub fn ctx(&self) -> &Arc<CudaContext> ⓘ
pub fn ctx(&self) -> &Arc<CudaContext> ⓘ
The context rows live in.
Sourcepub fn resident_in(&self, e: &Engine) -> bool
pub fn resident_in(&self, e: &Engine) -> bool
true iff e can dereference rows — i.e. e runs in the same CUDA context.
Sourcepub fn require_resident(&self, e: &Engine) -> Result<(), Box<dyn Error>>
pub fn require_resident(&self, e: &Engine) -> Result<(), Box<dyn Error>>
The residency law as a REFUSAL, shared by every consumer of rows (the common splice,
and gemma4’s masked-prefill arm, whose splice arithmetic differs deliberately). One
message, one law: a site that reads rows calls this first.
Sourcepub fn new_published(
tower: &Engine,
intake: &Engine,
mode: OverlayPublish,
rows: CudaSlice<f32>,
spans: Vec<(usize, usize, usize)>,
) -> Result<Self, Box<dyn Error>>
pub fn new_published( tower: &Engine, intake: &Engine, mode: OverlayPublish, rows: CudaSlice<f32>, spans: Vec<(usize, usize, usize)>, ) -> Result<Self, Box<dyn Error>>
Build an overlay whose rows are resident in the engine that will CONSUME them.
tower is the engine the vision tower ran on (where rows currently live); intake
is the engine that owns embedding intake for the serving placement
(HybridModel::vision_intake_engine — the primary engine on a single-device or
streams-off shape, pp stage 0’s engine under a per-stage-stream ppN split).
PUBLICATION IS A HOST BOUNCE, deliberately. tower.dtoh drains the tower’s stream at
a HOST boundary and intake.htod + one synchronize place the bytes on the intake
stream, so the ordering argument needs no event plumbing and no P2P capability: it
holds on every placement, including MEMRA_PP_HOST_BOUNCE=1 boxes with no peer
access. The cost is ONE round trip of [total_rows, n_embd] f32 per session (a few
MiB; ~5 MiB for a 256-row 5120-wide image) against a tower that already host-bounces
q/k/v and the merger input in EVERY block — a peer D2D twin is a named follow-up, not
a correctness question. Nothing here is per token: decode never touches an overlay.
THAT COST SENTENCE IS ONLY TRUE SINCE memra-next#23. When this was first written,
window() deep-copied the whole rows buffer on every prefill tick and every prime
chunk (see the rows field), so a multi-chunk image prompt paid several D2D copies that
the sentence did not mention. rows is now an Arc and a window is a refcount bump, so
the publication really is the only copy on this path.
The bytes are moved, never transformed: f32 through the host is bit-exact, which is
what makes MEMRA_VISION_OVERLAY_PUBLISH=force a byte-identity gate arm.
CALL OUTSIDE ANY pp STAGE SCOPE. The ambient stream override is THREAD-LOCAL and applies
to every engine on the thread, so inside PpNRt::enter(s) this upload would bind to
stage s’s stream — which belongs to another context whenever the stages differ. The
serving caller (build_vision_overlay, at the first prefill tick) and the gate arms both
run outside stage scopes.
THE FREE IS ORDERED TOO — but not by what an earlier version of this comment claimed,
and this is the sentence a future lane will lean on, so it is worth stating exactly.
Published rows are freed on THEIR ALLOCATION STREAM (the intake engine’s ambient stream)
— free_async on cudarc’s async-alloc branch, and synchronize + free_sync on the
other, which likewise touches only the allocation stream — while the splice reads them on
pp stage 0’s STAGE stream:
two streams, one context. What orders the reads before the free is (a) the pipeline
chain itself — stage 0 feeds stage 1 feeds … feeds the last stage — and (b) the
host-synchronizing dtoh every prime call ends with in hyper_prime_tail, which
retires all of it before any later host code, including the drop, runs.
It is NOT publish_all_to: that orders the CALLER behind stage compute (the other
direction), and it is MEMRA_PP_EXIT_PUBLISH-gated, so leaning on it would make this
argument evaporate at =0. And there is no implicit safety net underneath: cudarc event
tracking is DISABLED in Engine::new unless MEMRA_EVT=1, so a drop carries no read
guard — the manual argument above is all there is. (Corrected by the memra-next peer
review; the accrace lane, MEMRA_PP_EXIT_PUBLISH, is what happens when an ordering
claim in a comment is merely assumed.)
mode is PASSED IN, not read from the environment here (memra-next#25). The door is a
FAMILY-AGNOSTIC correctness input — every vision family’s overlay comes through this
function — so an unrecognized value has to refuse at BOOT, once, for all of them. When
this read the env itself, the only boot-time validation was glm5-scoped, so a typo’d
door on a gemma/qwen/step37 deployment booted clean and then 500’d mid-prefill on the
first image request: exactly the failure this lane removed for glm5, still live for the
others. Resolve with overlay_publish_mode() at startup and thread the value.
Sourcepub fn window(&self, off: usize, len: usize) -> Option<EmbedOverlay>
pub fn window(&self, off: usize, len: usize) -> Option<EmbedOverlay>
Sub-window for a prime call covering prompt-relative [off, off+len): spans clipped
and rebased so the callee sees call-relative positions (the serve prefill tick primes
a prompt across multiple prime_cache_overlaid calls). rows is an Arc clone — a
refcount bump, which since memra-next#23 is TRUE rather than merely documented. None =
no image rows in this window (caller may prime plain).
Sourcepub fn splice_into(
&self,
e: &Engine,
embedded: &mut CudaSlice<f32>,
chunk_off: usize,
t: usize,
n_embd: usize,
) -> Result<(), Box<dyn Error>>
pub fn splice_into( &self, e: &Engine, embedded: &mut CudaSlice<f32>, chunk_off: usize, t: usize, n_embd: usize, ) -> Result<(), Box<dyn Error>>
The mixed-embedding splice itself: image rows overwrite placeholder-token embeddings
inside a prime call’s prompt-relative window [chunk_off, chunk_off + t), BEFORE any
downstream transform (stream expansion, trunk walk). ONE implementation shared by the
single-engine hyper walk and the ppN stage-0 intake (lane/glm5-vision-default-on) so
the splice point cannot drift between arms. embedded is the [t, n_embd] token
embedding buffer of this call.
RESIDENCY LAW, checked HERE — at the copy, for every arm (serial chunk walk, hyper
walk, ppN stage-0 intake), rather than at one call site’s placement assumption: rows
must live in e’s CUDA context. A mismatch is a REFUSAL. Publish the overlay into the
consuming engine with EmbedOverlay::new_published instead of relaxing this.