submission 688762

#!POPCORN leaderboard amd-mxfp4-mm
#!POPCORN gpu MI355X
"""
Crossover of prog_0013 + prog_0015 + prog_0017: best config per shape.

Shape 1 (M=4, K=512, N=2880): BSM=8, BSN=32 from prog_0017 (8.32µs).
  90 blocks vs 45 with "any". num_warps=2 for tiny tile.

Shape 2 (M=16, K=7168, N=2112): KSPLIT=14 from prog_0013 (12.9µs).
  KSPLIT=16 (prog_0015) was worse at 14.9µs. Keep stages=2, wpe=4.

Shape 3 (M=32, K=512, N=4096): BSM=16 from prog_0015 (7.18µs).
  256 blocks. waves_per_eu=3 (between 2 and 4).

Shape 4 (M=32, K=512, N=2880): BSM=16 from prog_0015 (7.44µs).
  180 blocks. waves_per_eu=3.

Shape 5 (M=64, K=2048, N=7168): KSPLIT=2 from prog_0013 (13.5µs).
  KSPLIT=4 (prog_0015) was worse at 17.0µs. BSM=16, BSN=256, wpe=4.

Shape 6: CK ASM (unchanged, ~13µs).
"""
import torch, os, json, re, importlib, inspect, math
from task import input_t, output_t
from torch.utils.cpp_extension import load_inline

os.environ['PYTORCH_ROCM_ARCH'] = 'gfx950'

from aiter.ops.triton.gemm.basic.gemm_a16wfp4 import gemm_a16wfp4_preshuffle
from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH
from aiter.ops.triton.utils._triton import arch_info
import triton

# EVOLVE-BLOCK-START
_cfgs={"N=2880-K=512":{"M_LEQ_4":{"BLOCK_SIZE_M":8,"BLOCK_SIZE_N":32,"BLOCK_SIZE_K":512,"GROUP_SIZE_M":1,"NUM_KSPLIT":1,"num_warps":2,"num_stages":1,"waves_per_eu":2,"matrix_instr_nonkdim":16,"cache_modifier":".cg"},"M_LEQ_32":{"BLOCK_SIZE_M":16,"BLOCK_SIZE_N":32,"BLOCK_SIZE_K":512,"GROUP_SIZE_M":1,"NUM_KSPLIT":1,"num_warps":4,"num_stages":1,"waves_per_eu":3,"matrix_instr_nonkdim":16,"cache_modifier":".cg"},"any":{"BLOCK_SIZE_M":32,"BLOCK_SIZE_N":64,"BLOCK_SIZE_K":512,"GROUP_SIZE_M":1,"NUM_KSPLIT":1,"num_warps":8,"num_stages":1,"waves_per_eu":2,"matrix_instr_nonkdim":16,"cache_modifier":None}},"N=4096-K=512":{"M_LEQ_32":{"BLOCK_SIZE_M":16,"BLOCK_SIZE_N":32,"BLOCK_SIZE_K":512,"GROUP_SIZE_M":1,"NUM_KSPLIT":1,"num_warps":4,"num_stages":1,"waves_per_eu":3,"matrix_instr_nonkdim":16,"cache_modifier":".cg"},"any":{"BLOCK_SIZE_M":32,"BLOCK_SIZE_N":64,"BLOCK_SIZE_K":512,"GROUP_SIZE_M":1,"NUM_KSPLIT":1,"num_warps":8,"num_stages":1,"waves_per_eu":2,"matrix_instr_nonkdim":16,"cache_modifier":None}},"N=2112-K=7168":{"M_LEQ_16":{"BLOCK_SIZE_M":16,"BLOCK_SIZE_N":128,"BLOCK_SIZE_K":256,"GROUP_SIZE_M":1,"NUM_KSPLIT":14,"num_warps":4,"num_stages":2,"waves_per_eu":4,"matrix_instr_nonkdim":16,"cache_modifier":".cg"},"any":{"BLOCK_SIZE_M":32,"BLOCK_SIZE_N":128,"BLOCK_SIZE_K":256,"GROUP_SIZE_M":4,"NUM_KSPLIT":1,"num_warps":8,"num_stages":2,"waves_per_eu":4,"matrix_instr_nonkdim":16,"cache_modifier":None}},"N=7168-K=2048":{"M_LEQ_64":{"BLOCK_SIZE_M":16,"BLOCK_SIZE_N":256,"BLOCK_SIZE_K":256,"GROUP_SIZE_M":1,"NUM_KSPLIT":2,"num_warps":8,"num_stages":2,"waves_per_eu":4,"matrix_instr_nonkdim":16,"cache_modifier":".cg"},"any":{"BLOCK_SIZE_M":32,"BLOCK_SIZE_N":256,"BLOCK_SIZE_K":256,"GROUP_SIZE_M":4,"NUM_KSPLIT":1,"num_warps":8,"num_stages":2,"waves_per_eu":4,"matrix_instr_nonkdim":16,"cache_modifier":None}},"N=3072-K=1536":{"M_LEQ_64":{"BLOCK_SIZE_M":16,"BLOCK_SIZE_N":128,"BLOCK_SIZE_K":256,"GROUP_SIZE_M":1,"NUM_KSPLIT":3,"num_warps":4,"num_stages":1,"waves_per_eu":2,"matrix_instr_nonkdim":16,"cache_modifier":".cg"},"M_LEQ_256":{"BLOCK_SIZE_M":32,"BLOCK_SIZE_N":256,"BLOCK_SIZE_K":256,"GROUP_SIZE_M":4,"NUM_KSPLIT":2,"num_warps":8,"num_stages":2,"waves_per_eu":4,"matrix_instr_nonkdim":32,"cache_modifier":".cg"},"any":{"BLOCK_SIZE_M":32,"BLOCK_SIZE_N":128,"BLOCK_SIZE_K":256,"GROUP_SIZE_M":4,"NUM_KSPLIT":1,"num_warps":8,"num_stages":2,"waves_per_eu":4,"matrix_instr_nonkdim":16,"cache_modifier":None}}}
try:_dev=arch_info.get_arch()
except:_dev="gfx950"
_cd=f"{AITER_TRITON_CONFIGS_PATH}/gemm";os.makedirs(_cd,exist_ok=True)
for _sk,_cfg in _cfgs.items():
    with open(f"{_cd}/{_dev}-GEMM-A16WFP4_PRESHUFFLED-{_sk}.json","w") as f:json.dump(_cfg,f)

# Warmup all shapes via normal path
for _m,_n,_k in [(4,2880,512),(16,2112,7168),(32,4096,512),(32,2880,512),(64,7168,2048),(256,3072,1536)]:
    _A=torch.randn((_m,_k),dtype=torch.bfloat16,device="cuda")
    _Bw=torch.zeros((_n//16,(_k//2)*16),dtype=torch.uint8,device="cuda")
    _Bws=torch.zeros((_n//32,_k),dtype=torch.uint8,device="cuda")
    try: gemm_a16wfp4_preshuffle(_A,_Bw,_Bws,prequant=True,dtype=torch.bfloat16)
    except: pass
torch.cuda.synchronize()

# Get the Heuristics kernel object for direct calling
print("=== Setting up fast wrapper ===")
mod = importlib.import_module('aiter.ops.triton._triton_kernels.gemm.basic.gemm_a16wfp4')
_triton_kernel = getattr(mod, '_gemm_a16wfp4_preshuffle_kernel', None)
if _triton_kernel:
    print(f"  Got kernel: type={type(_triton_kernel).__name__}")
else:
    print("  WARNING: kernel not found")

# Pre-compute configs for each shape, including KSPLIT>1 shapes.
_fast_configs = {}

# Helper to compute splitk params matching aiter's get_splitk logic
def _get_splitk(K_packed, BSK, NS):
    SB = triton.cdiv((2 * triton.cdiv(K_packed, NS)), BSK) * BSK
    while NS > 1 and BSK > 16:
        if (K_packed % (SB // 2) == 0 and SB % BSK == 0 and K_packed % (BSK // 2) == 0):
            break
        elif K_packed % (SB / 2) != 0 and NS > 1:
            NS = NS // 2
        elif SB % BSK != 0:
            if NS > 1: NS = NS // 2
            elif BSK > 16: BSK = BSK // 2
        SB = triton.cdiv((2 * triton.cdiv(K_packed, NS)), BSK) * BSK
    return SB, BSK, NS

_shape_cfgs = [
    # Shape 1: BSM=8, BSN=32 (from prog_0017, 8.32µs) — 90 blocks
    (4,  2880, 512,  "N=2880-K=512", "M_LEQ_4"),
    # Shape 2: KSPLIT=14 (from prog_0013, 12.9µs) — 238 blocks
    (16, 2112, 7168, "N=2112-K=7168", "M_LEQ_16"),
    # Shape 3: BSM=16 (from prog_0015, 7.18µs) — 256 blocks, wpe=3
    (32, 4096, 512,  "N=4096-K=512", "M_LEQ_32"),
    # Shape 4: BSM=16 (from prog_0015, 7.44µs) — 180 blocks, wpe=3
    (32, 2880, 512,  "N=2880-K=512", "M_LEQ_32"),
    # Shape 5: KSPLIT=2 (from prog_0013, 13.5µs) — 224 blocks, wpe=4
    (64, 7168, 2048, "N=7168-K=2048", "M_LEQ_64"),
    # Shape 6 handled by CK ASM.
]

for (m, n, k, cfg_key, m_key) in _shape_cfgs:
    cfg = _cfgs[cfg_key][m_key]
    BSM = cfg['BLOCK_SIZE_M']
    BSN = cfg['BLOCK_SIZE_N']
    BSK = cfg['BLOCK_SIZE_K']
    KSPLIT = cfg['NUM_KSPLIT']
    N_k = n
    K_k = k // 2
    grid_mn = math.ceil(N_k / BSN) * math.ceil(m / BSM)

    if KSPLIT > 1:
        # Compute splitk params
        sb, bk, ns = _get_splitk(K_k, BSK, KSPLIT)
        # Intermediate float32 buffer: [KSPLIT, m, n]
        y_split = torch.empty(ns, m, n, dtype=torch.float32, device="cuda")
        y_final = torch.empty(m, n, dtype=torch.bfloat16, device="cuda")
        _fast_configs[(m, n, k)] = {
            'cfg': cfg, 'BSM': BSM, 'BSN': BSN, 'BSK': bk, 'N_k': N_k, 'K_k': K_k,
            'grid_x': ns * grid_mn, 'grid_mn': grid_mn,
            'y': y_split, 'y_final': y_final,
            'ksplit': ns, 'splitk_block_size': sb,
            'is_splitk': True,
        }
        print(f"  M={m} N={n} K={k}: BSM={BSM} BSN={BSN} KSPLIT={ns} SB={sb} grid={ns*grid_mn}")
    else:
        y = torch.empty(m, n, dtype=torch.bfloat16, device="cuda")
        _fast_configs[(m, n, k)] = {
            'cfg': cfg, 'BSM': BSM, 'BSN': BSN, 'N_k': N_k, 'K_k': K_k,
            'grid_x': grid_mn, 'grid_mn': grid_mn, 'y': y,
            'is_splitk': False,
        }
        print(f"  M={m} N={n} K={k}: BSM={BSM} BSN={BSN} grid={grid_mn}")

# Build C++ module for CK ASM + fast reduce
_ck_hip = r"""
#include <hip/hip_runtime.h>
#include <torch/extension.h>

// General reduce: sum ks float partitions -> bf16.
// src layout: [ks, M, N] (or effectively ks chunks of stride_k floats)
extern "C" __global__ void fast_reduce_ks(
    const float* __restrict__ src, unsigned short* __restrict__ dst,
    int MN, int stride_k, int ks) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= MN) return;
    float sum = src[idx];
    for (int i = 1; i < ks; i++) sum += src[idx + i*stride_k];
    unsigned int u = __float_as_uint(sum);
    u = u + (((u >> 16) & 1) + 0x7FFF);
    dst[idx] = (unsigned short)(u >> 16);
}

static torch::Tensor g_ro; static int g_rm=0,g_rn=0;
torch::Tensor fused_reduce(torch::Tensor src, int M, int N, int KS) {
    if(M!=g_rm||N!=g_rn){g_ro=torch::empty({M,N},torch::dtype(torch::kBFloat16).device(src.device()));g_rm=M;g_rn=N;}
    int MN=M*N;
    // Use general ks kernel — handles any KSPLIT value correctly
    hipLaunchKernelGGL(fast_reduce_ks,dim3((MN+255)/256),dim3(256),0,0,
        src.data_ptr<float>(),(unsigned short*)g_ro.data_ptr(),MN,(int)src.stride(0),KS);
    return g_ro;
}

extern "C" __global__ void fused_quant_shuffle(
    const unsigned short* __restrict__ A, unsigned char* __restrict__ A_q,
    unsigned char* __restrict__ A_scale_sh, int M, int K, int K_half, int K_groups, int M_pad) {
    int gid = blockIdx.x * blockDim.x + threadIdx.x;
    if (gid >= M_pad * K_groups) return;
    int m = gid / K_groups; int g = gid % K_groups;
    float vals[32]; float amax = 0.0f;
    if (m < M) {
        const uint4* row4 = (const uint4*)(A + m * K + g * 32);
        #pragma unroll
        for (int v=0;v<4;v++){uint4 c=row4[v];unsigned int w[4]={c.x,c.y,c.z,c.w};
            #pragma unroll
            for(int j=0;j<4;j++){float lo=__uint_as_float((w[j]&0xFFFFu)<<16);float hi=__uint_as_float(w[j]&0xFFFF0000u);
                vals[v*8+j*2]=lo;vals[v*8+j*2+1]=hi;amax=fmaxf(amax,fmaxf(fabsf(lo),fabsf(hi)));}}
    } else { for(int i=0;i<32;i++) vals[i]=0.0f; }
    unsigned int ai=__float_as_uint(amax);unsigned int ar=(ai+0x200000u)&0xFF800000u;
    int eb=(ar>>23)&0xFF;int su=eb-129;int sb=su+127;if(sb<0)sb=0;
    int m0=m/32;int m1=(m&31)/16;int m2=m&15;int g0=g/8;int g1=(g&7)/4;int g2=g&3;
    int kg8=K_groups/8;
    A_scale_sh[m0*(kg8*256)+g0*256+g2*64+m2*4+g1*2+m1]=(unsigned char)sb;
    int qe=127-su;if(qe<1)qe=0;if(qe>254)qe=254;
    float qs=__uint_as_float((unsigned int)qe<<23);if(amax==0.0f)qs=0.0f;
    const unsigned int dmi=149u<<23;float dmf=__uint_as_float(dmi);
    const int vta=((int)(1-127)<<23)+(1<<21)-1;
    auto fp4=[&](float v)->unsigned char{float qf=v*qs;unsigned int qx=__float_as_uint(qf);
        unsigned int s=qx&0x80000000u;qx^=s;float qp=__uint_as_float(qx);unsigned char r;
        if(qp>=6.0f)r=0x7;else if(qp<1.0f)r=(unsigned char)((__float_as_uint(qp+dmf)-dmi)&0xFF);
        else{unsigned int mo=(qx>>22)&1;r=(unsigned char)((((unsigned int)((int)qx+vta)+mo)>>22)&0xFF);}
        return(r&0x7)|((unsigned char)(s>>28)&0x8);};
    unsigned char packed[16];
    #pragma unroll
    for(int i=0;i<16;i++) packed[i]=fp4(vals[2*i])|(fp4(vals[2*i+1])<<4);
    uint4* out4=(uint4*)(A_q+m*K_half+g*16);*out4=*((uint4*)packed);
}
struct __attribute__((packed)) CKArgs {
    void* ptr_D;char _p0[8];void* ptr_C;char _p1[8];void* ptr_A;char _p2[8];void* ptr_B;char _p3[8];
    float alpha;char _p4[12];float beta;char _p5[12];
    unsigned int stride_D0;char _p6[12];unsigned int stride_D1;char _p7[12];
    unsigned int stride_C0;char _p8[12];unsigned int stride_C1;char _p9[12];
    unsigned int stride_A0;char _p10[12];unsigned int stride_A1;char _p11[12];
    unsigned int stride_B0;char _p12[12];unsigned int stride_B1;char _p13[12];
    unsigned int M;char _p14[12];unsigned int N;char _p15[12];unsigned int K;char _p16[12];
    void* ptr_ScaleA;char _p17[8];void* ptr_ScaleB;char _p18[8];
    unsigned int stride_ScaleA0;char _p19[12];unsigned int stride_ScaleA1;char _p20[12];
    unsigned int stride_ScaleB0;char _p21[12];unsigned int stride_ScaleB1;char _p22[12];
    int log2_k_split;
};
static hipModule_t g_ck_mod=nullptr;static hipFunction_t g_ck_fn=nullptr;
static torch::Tensor g_aq,g_ash,g_ck_out;static int g_ck_Mp=0,g_ck_K=0,g_ck_N=0;
bool init_ck(const std::string& co,const std::string& fn){
    if(g_ck_fn)return true;
    if(hipModuleLoad(&g_ck_mod,co.c_str())!=hipSuccess)return false;
    if(hipModuleGetFunction(&g_ck_fn,g_ck_mod,fn.c_str())!=hipSuccess)return false;
    return true;}
torch::Tensor ck_dispatch(torch::Tensor A,torch::Tensor Bs,torch::Tensor Bss,int Nv){
    if(!g_ck_fn)return torch::Tensor();
    int M=A.size(0),Ke=A.size(1),Mp=((M+31)/32)*32,Kh=Ke/2,Kg=Ke/32;
    if(Mp!=g_ck_Mp||Ke!=g_ck_K){g_aq=torch::empty({Mp,Kh},torch::dtype(torch::kUInt8).device(A.device()));
        g_ash=torch::empty({Mp,Kg},torch::dtype(torch::kUInt8).device(A.device()));g_ck_Mp=Mp;g_ck_K=Ke;}
    if(Nv!=g_ck_N){g_ck_out=torch::empty({Mp,Nv},torch::dtype(torch::kBFloat16).device(A.device()));g_ck_N=Nv;}
    int tg=Mp*Kg;
    hipLaunchKernelGGL(fused_quant_shuffle,dim3((tg+127)/128),dim3(128),0,0,
        (const unsigned short*)A.data_ptr(),g_aq.data_ptr<unsigned char>(),
        g_ash.data_ptr<unsigned char>(),M,Ke,Kh,Kg,Mp);
    CKArgs ka;memset(&ka,0,sizeof(ka));
    ka.ptr_D=g_ck_out.data_ptr();ka.ptr_C=g_ck_out.data_ptr();
    ka.ptr_A=g_aq.data_ptr<unsigned char>();ka.ptr_B=Bs.data_ptr();
    ka.alpha=1.0f;ka.beta=0.0f;ka.stride_D0=Nv;ka.stride_D1=1;ka.stride_C0=Nv;ka.stride_C1=1;
    ka.stride_A0=Ke;ka.stride_A1=1;ka.stride_B0=Bs.stride(0)*2;ka.stride_B1=1;
    ka.M=Mp;ka.N=Nv;ka.K=Ke;ka.ptr_ScaleA=g_ash.data_ptr<unsigned char>();ka.ptr_ScaleB=Bss.data_ptr();
    ka.stride_ScaleA0=Kg;ka.stride_ScaleA1=1;ka.stride_ScaleB0=Bss.stride(0);ka.stride_ScaleB1=1;
    ka.log2_k_split=0;size_t asz=sizeof(ka);void* cfg[]={(void*)0x01,&ka,(void*)0x02,&asz,(void*)0x03};
    hipModuleLaunchKernel(g_ck_fn,(Nv+127)/128,(Mp+31)/32,1,256,1,1,0,0,nullptr,cfg);
    return g_ck_out.slice(0,0,M);}
"""
_ck_cpp = r"""
#include <torch/extension.h>
bool init_ck(const std::string&, const std::string&);
torch::Tensor ck_dispatch(torch::Tensor,torch::Tensor,torch::Tensor,int);
torch::Tensor fused_reduce(torch::Tensor,int,int,int);
"""
_mod = None
try:
    _mod = load_inline(name='ck591h',cpp_sources=_ck_cpp,cuda_sources=_ck_hip,
        functions=['init_ck','ck_dispatch','fused_reduce'],verbose=False,
        extra_cuda_cflags=['-O3','-w','-mcumode','--offload-arch=gfx950'])
    _ck_ok = _mod.init_ck("/home/runner/aiter/hsa/gfx950/f4gemm/f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128.co",
        "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E")
    if _ck_ok:
        _w=torch.randn(256,1536,dtype=torch.bfloat16,device="cuda")
        _mod.ck_dispatch(_w,torch.zeros(3072,768,dtype=torch.uint8,device="cuda"),
            torch.zeros(3072,48,dtype=torch.uint8,device="cuda"),3072)
        torch.cuda.synchronize();del _w
        print("CK ASM OK")
except Exception as e:
    _ck_ok = False; print(f"CK: {e}")

_orig_ps = gemm_a16wfp4_preshuffle
_use_fr = 'skip_reduce' in inspect.signature(_orig_ps).parameters
print(f"skip_reduce: {_use_fr}")

torch.cuda.empty_cache()

_ps_ck=None;_ps_cw=None;_ps_cs=None

@torch.inference_mode()
def custom_kernel(data: input_t) -> output_t:
    global _ps_ck,_ps_cw,_ps_cs
    A=data[0];B_shuffle=data[3];B_scale_sh=data[4]
    m,k=A.shape;n=data[1].shape[0]

    # Shape 6: CK ASM
    if _ck_ok and m==256 and k==1536:
        return _mod.ck_dispatch(A, B_shuffle, B_scale_sh, n)

    dp=B_shuffle.data_ptr()
    if dp!=_ps_ck:
        _ps_ck=dp
        _ps_cw=B_shuffle.view(torch.uint8).reshape(n//16,(k//2)*16)
        _ps_cs=B_scale_sh.view(torch.uint8)[:n,:].contiguous().reshape(n//32,k)

    # Direct Triton dispatch for all pre-configured shapes
    fc = _fast_configs.get((m,n,k))
    if fc and _triton_kernel:
        N_k = fc['N_k']
        K_k = fc['K_k']
        grid_x = fc['grid_x']
        cfg = fc['cfg']

        if fc['is_splitk']:
            # KSPLIT>1: dispatch to float32 intermediate, then reduce
            y_split = fc['y']  # pre-allocated [KSPLIT, m, n] float32
            ksplit = fc['ksplit']
            sb = fc['splitk_block_size']
            full_cfg = dict(cfg)
            full_cfg['SPLITK_BLOCK_SIZE'] = sb
            full_cfg['EVEN_K'] = (K_k % (fc['BSK'] // 2)) == 0
            full_cfg['GRID_MN'] = fc['grid_mn']
            _triton_kernel[(grid_x,)](
                A, _ps_cw, y_split, _ps_cs,
                m, N_k, K_k,
                A.stride(0), A.stride(1),
                _ps_cw.stride(0), _ps_cw.stride(1),
                y_split.stride(0),  # stride_ck = M*N (stride between k-split planes)
                y_split.stride(1), y_split.stride(2),
                _ps_cs.stride(0), _ps_cs.stride(1),
                PREQUANT=True,
                **full_cfg,
            )
            # Reduce float32 partials to bf16
            return _mod.fused_reduce(y_split, m, n, ksplit)
        else:
            # KSPLIT=1: direct to bf16 output
            y = fc['y']
            full_cfg = dict(cfg)
            full_cfg['SPLITK_BLOCK_SIZE'] = k  # full K, KSPLIT=1
            full_cfg['EVEN_K'] = (K_k % (cfg['BLOCK_SIZE_K'] // 2)) == 0
            full_cfg['GRID_MN'] = fc['grid_mn']
            _triton_kernel[(grid_x,)](
                A, _ps_cw, y, _ps_cs,
                m, N_k, K_k,
                A.stride(0), A.stride(1),
                _ps_cw.stride(0), _ps_cw.stride(1),
                0,  # stride_ck unused for KSPLIT=1
                y.stride(0), y.stride(1),
                _ps_cs.stride(0), _ps_cs.stride(1),
                PREQUANT=True,
                **full_cfg,
            )
            return y

    # Fallback: use aiter wrapper (should not be reached for known shapes)
    if _use_fr and _mod:
        result = _orig_ps(A,_ps_cw,_ps_cs,prequant=True,dtype=torch.bfloat16,skip_reduce=True)
        if result.dim() == 3:
            ks, rm, rn = result.shape
            return _mod.fused_reduce(result, rm, rn, ks)
        return result

    return _orig_ps(A,_ps_cw,_ps_cs,prequant=True,dtype=torch.bfloat16)
# EVOLVE-BLOCK-END