For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/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, DevicePassable, ImplicitlyCopyable, ImplicitlyDeletable, 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, _safe=False], src_p: Pointer[Scalar[src_type], address_space=src_p.address_space, _safe=False], 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, _safe=False], token_index: Int, expert_slot: Int = Int(0), expert_start: Int = Int(0))

Parameters​

Fields​

Implemented traits​

comptime members​

alignment​

device_type​

dispatch_smem_size​

dispatch_wait_tile_shape​

hid_dim​

TensorType​

top_k​

Methods​

__init__​

get_type_name​

token_size​

copy_token_to_send_buf​

copy_msg_to_output_tensor​