For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo struct

BF16TokenFormat

struct BF16TokenFormat[output_layout: TensorLayout, //, _hid_dim: Int, _top_k: Int, _alignment: Int = Int(0)]

Token format that transmits the full hidden state in BFloat16.

Packs hid_dim BF16 elements directly into the wire buffer without quantization. Used when bandwidth headroom is sufficient or when the model does not apply EP quantization.

Parameters

• ​output_layout (TensorLayout): Layout of the output TileTensor that receives decoded tokens.
• ​_hid_dim (Int): Hidden dimension (number of BF16 elements per token).
• ​_top_k (Int): Number of experts each token is routed to.
• ​_alignment (Int): Override for the byte alignment of the wire buffer; 0 selects get_device_alignment().

Fields

• ​output_tokens (BF16TokenFormat[_hid_dim, _top_k, _alignment].TensorType):

Implemented traits

AnyType, Copyable, Deinitable, DevicePassable, ImplicitlyCopyable, Movable, RegisterPassable, TokenFormat, TrivialRegisterPassable

comptime members

alignment

comptime alignment = _alignment if _alignment.__bool__() else get_device_alignment()

device_type

comptime device_type = BF16TokenFormat[_hid_dim, _top_k, _alignment]

dispatch_smem_size

comptime dispatch_smem_size = 0

dispatch_wait_tile_shape

comptime dispatch_wait_tile_shape = Tuple(Int(128), Int(1))

hid_dim

comptime hid_dim = _hid_dim

TensorType

comptime TensorType = TileTensor[DType.bfloat16, output_layout, MutUntrackedOrigin]

top_k

comptime top_k = _top_k

Methods

__init__

def __init__(output_tokens: TileTensor[DType.bfloat16, output_layout, address_space=output_tokens.address_space, linear_idx_type=output_tokens.linear_idx_type]) -> Self

get_type_name

static def get_type_name() -> String

Returns:

String

token_size

static def token_size() -> Int

Returns:

Int

copy_token_to_send_buf

static def copy_token_to_send_buf[src_type: DType, block_size: Int, buf_addr_space: AddressSpace = AddressSpace.GENERIC](buf_p: Pointer[UInt8, address_space=buf_addr_space], src_p: Pointer[Scalar[src_type], address_space=src_p.address_space], input_scale: Float32)

copy_msg_to_output_tensor

def copy_msg_to_output_tensor[buf_addr_space: AddressSpace = AddressSpace.GENERIC](self, buf_p: Pointer[UInt8, address_space=buf_addr_space], token_index: Int, expert_slot: Int = Int(0), expert_start: Int = Int(0))