DecodeBenchmark.Config(
        model=Transformer.Config(
        │   channels_in=512,
        │   channels_out=256,
        │   num_layers=8,
        │   proj_in=Embedding.Config(
        │       channels_in=256,
        │       channels_out=512,
        │       padding_idx=None,
        │       device=None,
        │       dtype=None,
        │       lookup_dtype=None,
        │       shard=None,
        │       depth_index=(),
        │       init_weight=functools.partial(priml.model.init.truncated_normal, std=0.02)
        │   ),
        │   block=[
        │   │   TransformerBlock.Config(
        │   │       channels_in=512,
        │   │       channels_out=512,
        │   │       attn=Attention.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           num_heads=8,
        │   │           channels_head=64,
        │   │           num_heads_kv=2,
        │   │           bias=False,
        │   │           dropout=0.0,
        │   │           causal=True,
        │   │           rope=RoPE.Config(
        │   │               channels_head=64,
        │   │               reduction_mode='cat',
        │   │               frequencies=HuggingFaceFrequencies.Config(base=10000.0),
        │   │               dtype=torch.float32
        │   │           ),
        │   │           rotate=None,
        │   │           norm_qk=None,
        │   │           share_qk_norm=True,
        │   │           norm_out=None,
        │   │           split_qkv_projection=False,
        │   │           proj_qkv=None,
        │   │           proj_out=None,
        │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │           depth_index=((0, 8),),
        │   │           attn_kernel=SdpaFused.Config()
        │   │       ),
        │   │       ffn=SwiGLU.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           channels_hidden=1_536,
        │   │           expansion=2.6666666666666665,
        │   │           round_to=256,
        │   │           gate=True,
        │   │           bias=False,
        │   │           split_gate_projection=False,
        │   │           norm=None,
        │   │           act=priml.model.swiglu.silu,
        │   │           depth_index=((0, 8),),
        │   │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │           init_weight_out=torch.nn.init.zeros_,
        │   │           up_proj=None,
        │   │           middle=None,
        │   │           down_proj=None,
        │   │           shard='colwise'
        │   │       ),
        │   │       norm1=RMSNorm.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           eps=1e-06,
        │   │           elementwise_affine=False,
        │   │           device=None,
        │   │           dtype=None
        │   │       ),
        │   │       norm2=RMSNorm.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           eps=1e-06,
        │   │           elementwise_affine=False,
        │   │           device=None,
        │   │           dtype=None
        │   │       ),
        │   │       prenorm=True,
        │   │       checkpoint=False,
        │   │       depth_index=((0, 8),)
        │   │   ),
        │   │   TransformerBlock.Config(
        │   │       channels_in=512,
        │   │       channels_out=512,
        │   │       attn=Attention.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           num_heads=8,
        │   │           channels_head=64,
        │   │           num_heads_kv=2,
        │   │           bias=False,
        │   │           dropout=0.0,
        │   │           causal=True,
        │   │           rope=RoPE.Config(
        │   │               channels_head=64,
        │   │               reduction_mode='cat',
        │   │               frequencies=HuggingFaceFrequencies.Config(base=10000.0),
        │   │               dtype=torch.float32
        │   │           ),
        │   │           rotate=None,
        │   │           norm_qk=None,
        │   │           share_qk_norm=True,
        │   │           norm_out=None,
        │   │           split_qkv_projection=False,
        │   │           proj_qkv=None,
        │   │           proj_out=None,
        │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │           depth_index=((1, 8),),
        │   │           attn_kernel=SdpaFused.Config()
        │   │       ),
        │   │       ffn=SwiGLU.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           channels_hidden=1_536,
        │   │           expansion=2.6666666666666665,
        │   │           round_to=256,
        │   │           gate=True,
        │   │           bias=False,
        │   │           split_gate_projection=False,
        │   │           norm=None,
        │   │           act=priml.model.swiglu.silu,
        │   │           depth_index=((1, 8),),
        │   │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │           init_weight_out=torch.nn.init.zeros_,
        │   │           up_proj=None,
        │   │           middle=None,
        │   │           down_proj=None,
        │   │           shard='colwise'
        │   │       ),
        │   │       norm1=RMSNorm.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           eps=1e-06,
        │   │           elementwise_affine=False,
        │   │           device=None,
        │   │           dtype=None
        │   │       ),
        │   │       norm2=RMSNorm.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           eps=1e-06,
        │   │           elementwise_affine=False,
        │   │           device=None,
        │   │           dtype=None
        │   │       ),
        │   │       prenorm=True,
        │   │       checkpoint=False,
        │   │       depth_index=((1, 8),)
        │   │   ),
        │   │   TransformerBlock.Config(
        │   │       channels_in=512,
        │   │       channels_out=512,
        │   │       attn=Attention.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           num_heads=8,
        │   │           channels_head=64,
        │   │           num_heads_kv=2,
        │   │           bias=False,
        │   │           dropout=0.0,
        │   │           causal=True,
        │   │           rope=RoPE.Config(
        │   │               channels_head=64,
        │   │               reduction_mode='cat',
        │   │               frequencies=HuggingFaceFrequencies.Config(base=10000.0),
        │   │               dtype=torch.float32
        │   │           ),
        │   │           rotate=None,
        │   │           norm_qk=None,
        │   │           share_qk_norm=True,
        │   │           norm_out=None,
        │   │           split_qkv_projection=False,
        │   │           proj_qkv=None,
        │   │           proj_out=None,
        │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │           depth_index=((2, 8),),
        │   │           attn_kernel=SdpaFused.Config()
        │   │       ),
        │   │       ffn=SwiGLU.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           channels_hidden=1_536,
        │   │           expansion=2.6666666666666665,
        │   │           round_to=256,
        │   │           gate=True,
        │   │           bias=False,
        │   │           split_gate_projection=False,
        │   │           norm=None,
        │   │           act=priml.model.swiglu.silu,
        │   │           depth_index=((2, 8),),
        │   │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │           init_weight_out=torch.nn.init.zeros_,
        │   │           up_proj=None,
        │   │           middle=None,
        │   │           down_proj=None,
        │   │           shard='colwise'
        │   │       ),
        │   │       norm1=RMSNorm.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           eps=1e-06,
        │   │           elementwise_affine=False,
        │   │           device=None,
        │   │           dtype=None
        │   │       ),
        │   │       norm2=RMSNorm.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           eps=1e-06,
        │   │           elementwise_affine=False,
        │   │           device=None,
        │   │           dtype=None
        │   │       ),
        │   │       prenorm=True,
        │   │       checkpoint=False,
        │   │       depth_index=((2, 8),)
        │   │   ),
        │   │   TransformerBlock.Config(
        │   │       channels_in=512,
        │   │       channels_out=512,
        │   │       attn=Attention.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           num_heads=8,
        │   │           channels_head=64,
        │   │           num_heads_kv=2,
        │   │           bias=False,
        │   │           dropout=0.0,
        │   │           causal=True,
        │   │           rope=RoPE.Config(
        │   │               channels_head=64,
        │   │               reduction_mode='cat',
        │   │               frequencies=HuggingFaceFrequencies.Config(base=10000.0),
        │   │               dtype=torch.float32
        │   │           ),
        │   │           rotate=None,
        │   │           norm_qk=None,
        │   │           share_qk_norm=True,
        │   │           norm_out=None,
        │   │           split_qkv_projection=False,
        │   │           proj_qkv=None,
        │   │           proj_out=None,
        │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │           depth_index=((3, 8),),
        │   │           attn_kernel=SdpaFused.Config()
        │   │       ),
        │   │       ffn=SwiGLU.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           channels_hidden=1_536,
        │   │           expansion=2.6666666666666665,
        │   │           round_to=256,
        │   │           gate=True,
        │   │           bias=False,
        │   │           split_gate_projection=False,
        │   │           norm=None,
        │   │           act=priml.model.swiglu.silu,
        │   │           depth_index=((3, 8),),
        │   │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │           init_weight_out=torch.nn.init.zeros_,
        │   │           up_proj=None,
        │   │           middle=None,
        │   │           down_proj=None,
        │   │           shard='colwise'
        │   │       ),
        │   │       norm1=RMSNorm.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           eps=1e-06,
        │   │           elementwise_affine=False,
        │   │           device=None,
        │   │           dtype=None
        │   │       ),
        │   │       norm2=RMSNorm.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           eps=1e-06,
        │   │           elementwise_affine=False,
        │   │           device=None,
        │   │           dtype=None
        │   │       ),
        │   │       prenorm=True,
        │   │       checkpoint=False,
        │   │       depth_index=((3, 8),)
        │   │   ),
        │   │   TransformerBlock.Config(
        │   │       channels_in=512,
        │   │       channels_out=512,
        │   │       attn=Attention.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           num_heads=8,
        │   │           channels_head=64,
        │   │           num_heads_kv=2,
        │   │           bias=False,
        │   │           dropout=0.0,
        │   │           causal=True,
        │   │           rope=RoPE.Config(
        │   │               channels_head=64,
        │   │               reduction_mode='cat',
        │   │               frequencies=HuggingFaceFrequencies.Config(base=10000.0),
        │   │               dtype=torch.float32
        │   │           ),
        │   │           rotate=None,
        │   │           norm_qk=None,
        │   │           share_qk_norm=True,
        │   │           norm_out=None,
        │   │           split_qkv_projection=False,
        │   │           proj_qkv=None,
        │   │           proj_out=None,
        │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │           depth_index=((4, 8),),
        │   │           attn_kernel=SdpaFused.Config()
        │   │       ),
        │   │       ffn=SwiGLU.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           channels_hidden=1_536,
        │   │           expansion=2.6666666666666665,
        │   │           round_to=256,
        │   │           gate=True,
        │   │           bias=False,
        │   │           split_gate_projection=False,
        │   │           norm=None,
        │   │           act=priml.model.swiglu.silu,
        │   │           depth_index=((4, 8),),
        │   │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │           init_weight_out=torch.nn.init.zeros_,
        │   │           up_proj=None,
        │   │           middle=None,
        │   │           down_proj=None,
        │   │           shard='colwise'
        │   │       ),
        │   │       norm1=RMSNorm.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           eps=1e-06,
        │   │           elementwise_affine=False,
        │   │           device=None,
        │   │           dtype=None
        │   │       ),
        │   │       norm2=RMSNorm.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           eps=1e-06,
        │   │           elementwise_affine=False,
        │   │           device=None,
        │   │           dtype=None
        │   │       ),
        │   │       prenorm=True,
        │   │       checkpoint=False,
        │   │       depth_index=((4, 8),)
        │   │   ),
        │   │   TransformerBlock.Config(
        │   │       channels_in=512,
        │   │       channels_out=512,
        │   │       attn=Attention.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           num_heads=8,
        │   │           channels_head=64,
        │   │           num_heads_kv=2,
        │   │           bias=False,
        │   │           dropout=0.0,
        │   │           causal=True,
        │   │           rope=RoPE.Config(
        │   │               channels_head=64,
        │   │               reduction_mode='cat',
        │   │               frequencies=HuggingFaceFrequencies.Config(base=10000.0),
        │   │               dtype=torch.float32
        │   │           ),
        │   │           rotate=None,
        │   │           norm_qk=None,
        │   │           share_qk_norm=True,
        │   │           norm_out=None,
        │   │           split_qkv_projection=False,
        │   │           proj_qkv=None,
        │   │           proj_out=None,
        │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │           depth_index=((5, 8),),
        │   │           attn_kernel=SdpaFused.Config()
        │   │       ),
        │   │       ffn=SwiGLU.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           channels_hidden=1_536,
        │   │           expansion=2.6666666666666665,
        │   │           round_to=256,
        │   │           gate=True,
        │   │           bias=False,
        │   │           split_gate_projection=False,
        │   │           norm=None,
        │   │           act=priml.model.swiglu.silu,
        │   │           depth_index=((5, 8),),
        │   │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │           init_weight_out=torch.nn.init.zeros_,
        │   │           up_proj=None,
        │   │           middle=None,
        │   │           down_proj=None,
        │   │           shard='colwise'
        │   │       ),
        │   │       norm1=RMSNorm.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           eps=1e-06,
        │   │           elementwise_affine=False,
        │   │           device=None,
        │   │           dtype=None
        │   │       ),
        │   │       norm2=RMSNorm.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           eps=1e-06,
        │   │           elementwise_affine=False,
        │   │           device=None,
        │   │           dtype=None
        │   │       ),
        │   │       prenorm=True,
        │   │       checkpoint=False,
        │   │       depth_index=((5, 8),)
        │   │   ),
        │   │   TransformerBlock.Config(
        │   │       channels_in=512,
        │   │       channels_out=512,
        │   │       attn=Attention.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           num_heads=8,
        │   │           channels_head=64,
        │   │           num_heads_kv=2,
        │   │           bias=False,
        │   │           dropout=0.0,
        │   │           causal=True,
        │   │           rope=RoPE.Config(
        │   │               channels_head=64,
        │   │               reduction_mode='cat',
        │   │               frequencies=HuggingFaceFrequencies.Config(base=10000.0),
        │   │               dtype=torch.float32
        │   │           ),
        │   │           rotate=None,
        │   │           norm_qk=None,
        │   │           share_qk_norm=True,
        │   │           norm_out=None,
        │   │           split_qkv_projection=False,
        │   │           proj_qkv=None,
        │   │           proj_out=None,
        │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │           depth_index=((6, 8),),
        │   │           attn_kernel=SdpaFused.Config()
        │   │       ),
        │   │       ffn=SwiGLU.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           channels_hidden=1_536,
        │   │           expansion=2.6666666666666665,
        │   │           round_to=256,
        │   │           gate=True,
        │   │           bias=False,
        │   │           split_gate_projection=False,
        │   │           norm=None,
        │   │           act=priml.model.swiglu.silu,
        │   │           depth_index=((6, 8),),
        │   │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │           init_weight_out=torch.nn.init.zeros_,
        │   │           up_proj=None,
        │   │           middle=None,
        │   │           down_proj=None,
        │   │           shard='colwise'
        │   │       ),
        │   │       norm1=RMSNorm.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           eps=1e-06,
        │   │           elementwise_affine=False,
        │   │           device=None,
        │   │           dtype=None
        │   │       ),
        │   │       norm2=RMSNorm.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           eps=1e-06,
        │   │           elementwise_affine=False,
        │   │           device=None,
        │   │           dtype=None
        │   │       ),
        │   │       prenorm=True,
        │   │       checkpoint=False,
        │   │       depth_index=((6, 8),)
        │   │   ),
        │   │   TransformerBlock.Config(
        │   │       channels_in=512,
        │   │       channels_out=512,
        │   │       attn=Attention.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           num_heads=8,
        │   │           channels_head=64,
        │   │           num_heads_kv=2,
        │   │           bias=False,
        │   │           dropout=0.0,
        │   │           causal=True,
        │   │           rope=RoPE.Config(
        │   │               channels_head=64,
        │   │               reduction_mode='cat',
        │   │               frequencies=HuggingFaceFrequencies.Config(base=10000.0),
        │   │               dtype=torch.float32
        │   │           ),
        │   │           rotate=None,
        │   │           norm_qk=None,
        │   │           share_qk_norm=True,
        │   │           norm_out=None,
        │   │           split_qkv_projection=False,
        │   │           proj_qkv=None,
        │   │           proj_out=None,
        │   │           init_weight=priml.model.init.kaiming_uniform,
        │   │           depth_index=((7, 8),),
        │   │           attn_kernel=SdpaFused.Config()
        │   │       ),
        │   │       ffn=SwiGLU.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           channels_hidden=1_536,
        │   │           expansion=2.6666666666666665,
        │   │           round_to=256,
        │   │           gate=True,
        │   │           bias=False,
        │   │           split_gate_projection=False,
        │   │           norm=None,
        │   │           act=priml.model.swiglu.silu,
        │   │           depth_index=((7, 8),),
        │   │           init_weight=priml.model.init.unit_fan_in_uniform,
        │   │           init_weight_out=torch.nn.init.zeros_,
        │   │           up_proj=None,
        │   │           middle=None,
        │   │           down_proj=None,
        │   │           shard='colwise'
        │   │       ),
        │   │       norm1=RMSNorm.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           eps=1e-06,
        │   │           elementwise_affine=False,
        │   │           device=None,
        │   │           dtype=None
        │   │       ),
        │   │       norm2=RMSNorm.Config(
        │   │           channels_in=512,
        │   │           channels_out=512,
        │   │           eps=1e-06,
        │   │           elementwise_affine=False,
        │   │           device=None,
        │   │           dtype=None
        │   │       ),
        │   │       prenorm=True,
        │   │       checkpoint=False,
        │   │       depth_index=((7, 8),)
        │   │   )
        │   ],
        │   proj_out=Sequential.Config(
        │       channels_in=512,
        │       channels_out=256,
        │       elements=[
        │           RMSNorm.Config(
        │               channels_in=512,
        │               channels_out=512,
        │               eps=1e-06,
        │               elementwise_affine=False,
        │               device=None,
        │               dtype=None
        │           ),
        │           Linear.Config(
        │               channels_in=512,
        │               channels_out=256,
        │               bias=False,
        │               depth_index=(),
        │               device=None,
        │               dtype=None,
        │               init_weight=priml.model.init.kaiming_uniform,
        │               init_bias=torch.nn.init.zeros_,
        │               linear_fn=<built-in function linear>,
        │               shard=None
        │           )
        │       ],
        │       repeat=1,
        │       depth_index=()
        │   )
        ),
        allocate_cache=priml.model.attention.kvcache.alloc_layer_cache,
        prompts=[
            'The quick brown fox jumps over the lazy dog.',
            'Machine learning is a subset of artificial intelligence that',
            'The mitochondria is the powerhouse of the cell',
            'The speed of light is approximately 299792458 meters per second',
            'Photosynthesis is the process by which plants convert',
            'Quantum entanglement is a phenomenon where two particles',
            'To be or not to be, that is the question',
            'Water boils at 100 degrees Celsius at sea level'
        ],
        prompt_len=24,
        num_decode_steps=4_072,
        num_warmup=1,
        num_repeats=3,
        dtype=torch.bfloat16,
        device='cuda',
        seed=None,
        roofline_device=''
    )
