Skip to main content

EmbedOverlay

Struct EmbedOverlay 

Source
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

Source

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.

Source

pub fn ctx(&self) -> &Arc<CudaContext>

The context rows live in.

Source

pub fn resident_in(&self, e: &Engine) -> bool

true iff e can dereference rows — i.e. e runs in the same CUDA context.

Source

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.

Source

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.

Source

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).

Source

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.

Auto Trait Implementations§

Blanket Implementations§

Source§

impl<T> Any for T
where T: 'static + ?Sized,

Source§

fn type_id(&self) -> TypeId

Gets the TypeId of self. Read more
Source§

impl<T> Borrow<T> for T
where T: ?Sized,

Source§

fn borrow(&self) -> &T

Immutably borrows from an owned value. Read more
Source§

impl<T> BorrowMut<T> for T
where T: ?Sized,

Source§

fn borrow_mut(&mut self) -> &mut T

Mutably borrows from an owned value. Read more
Source§

impl<T> From<T> for T

Source§

fn from(t: T) -> T

Returns the argument unchanged.

Source§

impl<T, U> Into<U> for T
where U: From<T>,

Source§

fn into(self) -> U

Calls U::from(self).

That is, this conversion is whatever the implementation of From<T> for U chooses to do.

Source§

impl<T> Same for T

Source§

type Output = T

Should always be Self
Source§

impl<T, U> TryFrom<U> for T
where U: Into<T>,

Source§

type Error = !

The type returned in the event of a conversion error.
Source§

fn try_from(value: U) -> Result<T, !>

Performs the conversion.
Source§

impl<T, U> TryInto<U> for T
where U: TryFrom<T>,

Source§

type Error = <U as TryFrom<T>>::Error

The type returned in the event of a conversion error.
Source§

fn try_into(self) -> Result<U, <U as TryFrom<T>>::Error>

Performs the conversion.