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
Pack
struct Pack[MaskType: MHAMask, SchedulerType: MHATileScheduler, ValidLengthType: OptionalPointer, SinkType: OptionalPointer, KVRowOffsetsType: OptionalPointer, MaxSeqLenType: OptionallyStaticInt, PartitionType: MHAPartitionScheme]
Bundles MHA kernel parameters into a single device-passable struct.
Parameters
- MaskType (
MHAMask): Mask type applied to the attention score tiles. - SchedulerType (
MHATileScheduler): Tile scheduler that assigns work to CTAs. - ValidLengthType (
OptionalPointer): Optional pointer type for the per-batch valid-length tensor. - SinkType (
OptionalPointer): Optional pointer type for the attention-sink weights. - KVRowOffsetsType (
OptionalPointer): Optional pointer type for the KV row-offsets tensor. - MaxSeqLenType (
OptionallyStaticInt): Type of the maximum sequence length, which may be static or dynamic. - PartitionType (
MHAPartitionScheme): KV-cache partitioning scheme.
Fields
- mask (
MaskType): - scheduler (
SchedulerType): - valid_length (
ValidLengthType): - sink_weights (
SinkType): - kv_input_row_offsets (
KVRowOffsetsType): - max_seq_len (
MaxSeqLenType): - partition (
PartitionType):
Implemented traits
AnyType,
Copyable,
Deinitable,
DevicePassable,
ImplicitlyCopyable,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
device_type
comptime device_type = Pack[MaskType, SchedulerType, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType]
Methods
__init__
def __init__(mask: MaskType, scheduler: SchedulerType, valid_length: ValidLengthType, sink_weights: SinkType, kv_input_row_offsets: KVRowOffsetsType, max_seq_len: MaxSeqLenType, partition: PartitionType) -> Self