vortex_torch.cache.fill

Classes

Fill([alpha])

In-place page-wise fill dispatcher.
class vortex_torch.cache.fill.Fill(alpha=0.0)[source]

Bases: vOp

In-place page-wise fill dispatcher.

This operator performs an in-place, format-aware fill over a batched tensor. The input is treated as a rank-3 tensor

\[X \in \mathbb{R}^{B \times D_0 \times D_1},\]

where:

  • \(B\) is a batch-like axis (for example, batch * heads),

  • \(D_0, D_1\) encode per-page/per-token features.

Using per-position metadata stored in loc together with layout information in Context, the implementation identifies positions where a page has reached its logical end (e.g. an end-of-page or end-of-sequence condition) and fills the corresponding tiles with a scalar value alpha:

\[X[b, d_0, d_1] \leftarrow \alpha \quad \text{for all positions } (b, d_0, d_1) \text{ that lie past the end-of-page mark.}\]

All modifications are performed in-place; the tensor format (PAGED, RAGGED, etc.) determines how page boundaries are interpreted and which elements are affected.

Key properties

  • Dispatch is keyed only by the input tensor format x._format.

  • The operator is purely in-place: no output buffer is allocated.

  • The same vTensor is returned from profile().

_impl_map

Dispatch table keyed by x_format. Each entry maps to a callable implementation that performs the in-place fill.

Type:

Dict[FORMAT, Callable]

alpha

Scalar fill value written into selected positions.

Type:

float

impl

The resolved implementation selected during profile().

Type:

Optional[Callable]

profile(x, loc, ctx)[source]

Validate inputs and select the in-place implementation.

Since the operation is in-place, no output buffer is allocated and this method simply returns the input vTensor unmodified.

Parameters:

  • x (vTensor) – Input tensor to be modified in-place. Expected logical shape [B, D_0, D_1].

  • loc (torch.Tensor) – Auxiliary tensor carrying page/position metadata used to detect end-of-page locations.

  • ctx (Context) – Execution context providing layout information and any additional metadata needed by the implementation.

Returns:

The same object as x, returned as the output view.

Return type:

vTensor

Raises:

AssertionError – If x is not a vTensor, if its rank is not 3, if loc is not a torch.Tensor, or if no implementation is registered for x._format.

execute(x, loc, ctx)[source]

Execute the in-place page-wise fill and return the input tensor.

The selected implementation examines page/position metadata in loc and context, and overwrites elements in x with the scalar value alpha once an end-of-page condition is detected.

Parameters:

  • x (torch.Tensor) – Input tensor to be modified in-place. Must be compatible with the vTensor validated in profile().

  • loc (torch.Tensor) – Auxiliary tensor carrying page/position metadata.

  • ctx (Context) – Execution context forwarded to the implementation.

Returns:

The same tensor instance x, after in-place filling.

Return type:

torch.Tensor

Raises:

AssertionError – If profile() has not been called and no implementation is available.

Parameters:

alpha (float)