vortex_torch.abs.tensor

Pure-metadata tensor type used throughout the vortex graph + compiler.

A vTensor does not carry any real storage. It is a small descriptor that records exactly what the rest of the system needs to reason about a tensor at graph-construction and codegen time:

  • shape — tuple of ints, the real logical shape.

  • padded_shape — tuple of ints, each entry rounded up to the next

    power of two (only shape[1] / shape[2] may differ from shape; shape[0] is the leading axis and is not subject to Triton block-shape pow2 constraints).

  • dtype — torch.dtype

  • device — torch.device / str / None

  • _format — FORMAT (BATCHED / RAGGED / PAGED)

  • tensor_id — int, the graph-level identity used by the compiler

Why padded_shape: Triton’s block-shape constexprs (tl.zeros, tl.arange, tl.reshape, tl.make_block_ptr.block_shape, etc.) must each be a power of two. Real models (e.g. Qwen3-14B with num_attention_heads // num_key_value_heads == 5) carry tensors whose shape[1] is not a power of two. The compiler emits tile sizes from padded_shape while keeping memory-addressing math (strides, per-row offsets, divisors in Mean) anchored to shape; codegens add load/store masks that select the real-shape lanes whenever padded_shape != shape. When shape is already pow2, padded_shape == shape and no masking is emitted, so power-of-two models pay no overhead.

It also exposes dim() for parity with torch.Tensor so existing profile-time validation code (assert x.dim() == 3) keeps working.

There is intentionally no torch op support, no __torch_function__ override, no parent torch.Tensor class. vTensor is just metadata. Real tensors used by the runtime/execute path stay as plain torch.Tensor instances.

Functions

as_vtensor([x, _format, tensor_id, shape, ...])

Build a vTensor.

Classes

FORMAT(*values)

Tensor storage/layout format.

vTensor([shape, dtype, device, _format, ...])

Pure-metadata virtual tensor.
class FORMAT(*values)[source]

Bases: Enum

Tensor storage/layout format.

BATCHED

Standard dense batched tensors (e.g., [B, N, D]).

RAGGED

Ragged tensors with variable-length sequences or elements per batch.

PAGED

Paged tensors used for large or streaming data split into pages/chunks.

BATCHED = 0
RAGGED = 1
PAGED = 2
class vTensor(shape=(), dtype=torch.bfloat16, device=None, _format=FORMAT.BATCHED, tensor_id=-1, padded_shape=None)[source]

Bases: object

Pure-metadata virtual tensor.

Carries the descriptor fields used by the graph builder, the compiler, and the codegen layer. It does not own any GPU / CPU memory and intentionally cannot participate in torch ops — code that wants to compute on real data should hold a torch.Tensor separately and use the vTensor only for graph bookkeeping.

Parameters:

  • shape (tuple)

  • dtype (torch.dtype)

  • device (torch.device | str | None)

  • _format (FORMAT)

  • tensor_id (int)

  • padded_shape (tuple)

shape: tuple
padded_shape: tuple
dtype: torch.dtype
device: torch.device | str | None
tensor_id: int
dim()[source]

Number of dimensions; mirrors torch.Tensor.dim().

Return type:

int

property ndim: int
numel()[source]
Return type:

int

size(dim=None)[source]

Mirror of torch.Tensor.size().

Parameters:

dim (int | None)

needs_padding()[source]

True iff padded_shape != shape — i.e. at least one inner dim is not already a power of two and codegen must emit load/store masks.

Return type:

bool

as_vtensor(x=None, _format=FORMAT.BATCHED, tensor_id=-1, *, shape=None, dtype=None, device=None)[source]

Build a vTensor.

Three calling styles, all returning a fresh vTensor (or, for an existing vTensor, the same object re-tagged):

  1. Re-tag an existing vTensor — as_vtensor(vt, fmt, tid) overwrites vt._format and vt.tensor_id in place and returns vt. Useful when the caller wants to add an existing tensor descriptor to the graph under a fresh id. padded_shape is preserved.

  2. Extract metadata from a torch.Tensor — as_vtensor(real, fmt, tid) reads shape, dtype, device from real and returns a brand-new vTensor. padded_shape is derived from shape. The original tensor is not retained — vTensor is pure metadata.

  3. Direct construction by kwargs — as_vtensor(_format=fmt, tensor_id=tid, shape=..., dtype=..., device=...). Use this when no real torch tensor is available (the common case once the compile path is fully virtualized).

Parameters:

  • x (Any)

  • _format (FORMAT)

  • tensor_id (int)

  • shape (Sequence[int] | None)

  • dtype (torch.dtype | None)

  • device (torch.device | str | None)

Return type:

vTensor