# This script provides a template for using load_inline to run a HIP kernel for
import os 
os.environ['PYTORCH_ROCM_ARCH'] = 'gfx942'
import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t
import triton
from copy import deepcopy


CPP_WRAPPER = """
void fp8_mm(torch::Tensor a, torch::Tensor b, torch::Tensor a_scale, torch::Tensor b_scale, torch::Tensor c);
"""

CUDA_SRC = R""" 
// File: fp8_mm.cpp

#include <torch/extension.h>
#include <hip/hip_runtime.h>
#include <hip/amd_detail/amd_hip_fp8.h>    // __hip_fp8_e4m3_fnuz, operator float()
#include <hip/amd_detail/amd_hip_bf16.h>   // __hip_bfloat16, __float2bfloat16, __bfloat162float
#include <hip/hip_fp16.h>

// ----------------------------------------
// 1. 辅助宏与模板
// ----------------------------------------
template <int N>
__device__ __host__ inline int cdiv(int x) {
    return (x + N - 1) / N;
}
template <int M, int N>
constexpr __device__ __host__ inline int cdiv() {
    return (M + N - 1) / N;
}

template <int BM, int BN, int mode = 0>
__device__ __host__ inline int2 get_mn_begin(int block_idx, int m, int n) {
    int num_m_block = (m + BM - 1) / BM;
    int num_n_block = (n + BN - 1) / BN;

    if constexpr (mode == 0) {
        int m_idx = block_idx % num_m_block;
        int n_idx = block_idx / num_m_block;
        return make_int2(m_idx * BM, n_idx * BN);
    } else {
        int m_idx = block_idx / num_n_block;
        int n_idx = block_idx % num_n_block;
        return make_int2(m_idx * BM, n_idx * BN);
    }
}

// MFMA wrapper
struct MFMA_FP8_16x16x32 {
    static constexpr int M = 16;
    static constexpr int N = 16;
    static constexpr int K = 32;

    using Areg = __hip_fp8_e4m3_fnuz[8];
    using Breg = __hip_fp8_e4m3_fnuz[8];
    using Creg = float[4];

    __device__ static void load_a(const __hip_fp8_e4m3_fnuz* __restrict__ a_mem,
                                  int stride_m, int stride_k,
                                  Areg& a_reg, int lane_id) {
        constexpr int vec_size = 8;
        int k_i = (lane_id / 16) * vec_size;
        int m_i = (lane_id % 16);
        for (int i = 0; i < vec_size; i++) {
            a_reg[i] = a_mem[m_i * stride_m + (k_i + i) * stride_k];
        }
    }

    __device__ static void load_b(const __hip_fp8_e4m3_fnuz* __restrict__ b_mem,
                                  int stride_n, int stride_k,
                                  Breg& b_reg, int lane_id) {
        constexpr int vec_size = 8;
        int k_i = (lane_id / 16) * vec_size;
        int n_i = (lane_id % 16);
        for (int i = 0; i < vec_size; i++) {
            b_reg[i] = b_mem[n_i * stride_n + (k_i + i) * stride_k];
        }
    }

    __device__ static void mma(const Areg& a_reg,
                               const Breg& b_reg,
                               Creg& c_reg) {
        using floatx4 = __attribute__((__vector_size__(4 * sizeof(float)))) float;
        *(reinterpret_cast<floatx4*>(c_reg)) =
            __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8(
                *(reinterpret_cast<const long*>(&a_reg)),
                *(reinterpret_cast<const long*>(&b_reg)),
                *(reinterpret_cast<floatx4*>(c_reg)),
                0, 0, 0
            );
    }

    __device__ static void store_c(__hip_bfloat16* __restrict__ c_mem,
                                   int stride_m, int stride_n,
                                   const Creg& c_reg, int lane_id) {
        int n_i = lane_id % 16;
        int m_i = (lane_id / 16) * 4;
        constexpr int vec_size = 4;
        for (int i = 0; i < vec_size; ++i) {
            c_mem[(m_i + i) * stride_m + n_i * stride_n] = __float2bfloat16(c_reg[i]);
        }
    }
};

template <int WM_, int WN_, int WK_, typename INST>
struct CTA {
    static constexpr int WM = WM_;
    static constexpr int WN = WN_;
    static constexpr int WK = WK_;
    static constexpr int INST_M = WM_ * INST::M;
    static constexpr int INST_N = WN_ * INST::N;
    static constexpr int INST_K = WK_ * INST::K;

    __device__ __host__ inline int3 get_offset(int warp_idx,
                                               int inst_m, int inst_n, int inst_k) const {
        int warp_mn    = warp_idx % (WM * WN);
        int warp_m_idx = warp_mn / WN;
        int warp_n_idx = warp_mn % WN;
        int warp_k_idx = warp_idx / (WM * WN);

        int off_m = inst_m + warp_m_idx * INST::M;
        int off_n = inst_n + warp_n_idx * INST::N;
        int off_k = inst_k + warp_k_idx * INST::K;
        return make_int3(off_m, off_n, off_k);
    }
};

constexpr int SCALE_BLOCK_SIZE = 128;

template <int BM, int BK>
__device__ inline void load_a_scale_to_smem(
    const float* __restrict__ a_scale,
    float* __restrict__ smem_a_scale,
    int m_begin,
    int m,
    int base_tile_k
) {
    constexpr int local_num_tile_k = (BK + SCALE_BLOCK_SIZE - 1) / SCALE_BLOCK_SIZE;
    int idx = threadIdx.x;
    int total = BM * local_num_tile_k;
    if (idx < total) {
        int row = idx % BM;
        int p   = idx / BM;
        int global_row = m_begin + row;
        int tile_k = base_tile_k + p;
        smem_a_scale[p * BM + row] =
            a_scale[tile_k * m + global_row];
            //a_scale[ global_row * global_num_tile_k + tile_k ];
    }
}

template <int BN, int BK>
__device__ inline void load_b_scale_to_smem(
    const float* __restrict__ b_scale,
    float* __restrict__ smem_b_scale,
    int tile_n_begin,
    int n_tiles,
    int global_num_tile_n,
    int base_tile_k
) {
    constexpr int local_num_tile_k = (BK + SCALE_BLOCK_SIZE - 1) / SCALE_BLOCK_SIZE;
    int idx = threadIdx.x;
    int total = n_tiles * local_num_tile_k;
    if (idx < total) {
        int ntile = idx / local_num_tile_k;
        int p     = idx % local_num_tile_k;
        int global_ntile = tile_n_begin + ntile;
        int tile_k = base_tile_k + p;
        smem_b_scale[ ntile * local_num_tile_k + p ] =
            b_scale[tile_k * global_num_tile_n + global_ntile];
            //b_scale[ global_ntile * global_num_tile_k + tile_k ];
    }
}

template <int BM, int BN, int BK, int WM, int WN, int WK>
__global__ void custom_kernel(
    const __hip_fp8_e4m3_fnuz* __restrict__ a,
    const __hip_fp8_e4m3_fnuz* __restrict__ b,
    const float* __restrict__ a_scale,
    const float* __restrict__ b_scale,
    __hip_bfloat16* __restrict__ c,
    int m, int n, int k
) {
    using Inst = MFMA_FP8_16x16x32;
    using Cta  = CTA<WM, WN, WK, Inst>;

    constexpr int INST_M = WM * Inst::M;
    constexpr int INST_N = WN * Inst::N;
    constexpr int INST_K = WK * Inst::K;
    constexpr int M_CNT = BM / INST_M;
    constexpr int N_CNT = BN / INST_N;
    constexpr int K_CNT = BK / INST_K;

    int global_num_tile_k = (k + SCALE_BLOCK_SIZE - 1) / SCALE_BLOCK_SIZE;
    int local_num_tile_k  = (BK + SCALE_BLOCK_SIZE - 1) / SCALE_BLOCK_SIZE;
    int n_tiles_bound     = (BN + SCALE_BLOCK_SIZE - 1) / SCALE_BLOCK_SIZE;

    int2 mn_begin = get_mn_begin<BM, BN>(blockIdx.x, m, n);
    int m_begin = mn_begin.x;
    int n_begin = mn_begin.y;

    int lane_id = threadIdx.x % warpSize;
    int warp_id = threadIdx.x / warpSize;

    typename Inst::Areg a_reg[M_CNT][K_CNT];
    typename Inst::Breg b_reg[N_CNT][K_CNT];
    typename Inst::Creg c_reg[M_CNT][N_CNT];
    #pragma unroll
    for (int im = 0; im < M_CNT; ++im) {
        #pragma unroll
        for (int in = 0; in < N_CNT; ++in) {
            #pragma unroll
            for (int t = 0; t < 4; ++t) {
                c_reg[im][in][t] = 0.0f;
            }
        }
    }

    extern __shared__ float shared_mem[];
    float* smem_a_scale = shared_mem;
    float* smem_b_scale = shared_mem + BM * local_num_tile_k;

    for (int k_begin = 0; k_begin < k; k_begin += BK) {
        int base_tile_k = k_begin / SCALE_BLOCK_SIZE;

        int tile_n_begin = n_begin / SCALE_BLOCK_SIZE;
        int tile_n_end = (n_begin + BN - 1) / SCALE_BLOCK_SIZE;
        int n_tiles = tile_n_end - tile_n_begin + 1;

        load_a_scale_to_smem< BM, BK >(
            a_scale,
            smem_a_scale,
            m_begin,
            m,
            base_tile_k
        );
        
        load_b_scale_to_smem< BN, BK >(
            b_scale,
            smem_b_scale,
            tile_n_begin,
            n_tiles,
            (n + SCALE_BLOCK_SIZE - 1) / SCALE_BLOCK_SIZE,
            base_tile_k
        );
        __syncthreads();
        #pragma unroll
        for (int ik = 0; ik < K_CNT; ++ik) {
            int sub_k_begin = k_begin + ik * INST_K;
            #pragma unroll
            for (int im = 0; im < M_CNT; ++im) {
                #pragma unroll
                for (int in = 0; in < N_CNT; ++in) {
                    int3 local_off = Cta().get_offset(
                        warp_id,
                        im * INST_M,
                        in * INST_N,
                        sub_k_begin
                    );
                    int off_m = m_begin + local_off.x;
                    int off_n = n_begin + local_off.y;
                    int off_k = local_off.z;
                    if (off_m >= m || off_n >= n || off_k >= k) {
                        continue;
                    }

                    Inst::load_a(
                        a + off_m + off_k * m,
                        1, m,
                        a_reg[im][ik],
                        lane_id
                    );
                    Inst::load_b(
                        b + off_n + off_k * n,
                        1, n,
                        b_reg[in][ik],
                        lane_id
                    );

                    float c_tmp[4] = {0.0f, 0.0f, 0.0f, 0.0f};
                    using floatx4 = __attribute__((__vector_size__(4 * sizeof(float)))) float;
                    *(reinterpret_cast<floatx4*>(c_tmp)) =
                        __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8(
                            *(reinterpret_cast<const long*>(&a_reg[im][ik])),
                            *(reinterpret_cast<const long*>(&b_reg[in][ik])),
                            *(reinterpret_cast<floatx4*>(c_tmp)),
                            0, 0, 0
                        );

                    int p_local = (off_k - k_begin) / SCALE_BLOCK_SIZE;
                    int row_in_cta = local_off.x + (lane_id/16) * 4;
                    float a_s[4];
                    a_s[0] = smem_a_scale[ p_local * BM + row_in_cta ];
                    a_s[1] = smem_a_scale[ p_local * BM + row_in_cta + 1];
                    a_s[2] = smem_a_scale[ p_local * BM + row_in_cta + 2];
                    a_s[3] = smem_a_scale[ p_local * BM + row_in_cta + 3];
                    
                    int this_tile_n = off_n / SCALE_BLOCK_SIZE;
                    int ntile_idx   = this_tile_n - tile_n_begin;
                    float b_s = smem_b_scale[ ntile_idx * local_num_tile_k + p_local ];
                    #pragma unroll
                    for (int t = 0; t < 4; ++t) {
                        c_reg[im][in][t] += c_tmp[t] * a_s[t] * b_s;
                    }
                }
            }
        }
        __syncthreads();
    }
    #pragma unroll
    for (int im = 0; im < M_CNT; ++im) {
        #pragma unroll
        for (int in = 0; in < N_CNT; ++in) {
            int3 local_off = Cta().get_offset(
                warp_id,
                im * INST_M,
                in * INST_N,
                0
            );
            int off_m = m_begin + local_off.x;
            int off_n = n_begin + local_off.y;
            if (off_m < m && off_n < n) {
                Inst::store_c(
                    c + off_m * n + off_n,
                    n, 1,
                    c_reg[im][in],
                    lane_id
                );
            }
        }
    }
}

// Wrapper function
void fp8_mm(torch::Tensor a,
            torch::Tensor b,
            torch::Tensor a_scale,
            torch::Tensor b_scale,
            torch::Tensor c) {
    
    int64_t m = a.size(0);
    int64_t k = a.size(1);
    int64_t n = b.size(0);

    int64_t num_tile_k = (k + SCALE_BLOCK_SIZE - 1) / SCALE_BLOCK_SIZE;
    int64_t num_tile_n = (n + SCALE_BLOCK_SIZE - 1) / SCALE_BLOCK_SIZE;

    constexpr int BM = 128;
    constexpr int BN = 64;
    constexpr int BK = 64;
    constexpr int WM = 8;
    constexpr int WN = 1;
    constexpr int WK = 1;

    dim3 blockDim(warpSize * WM * WN * WK, 1, 1);
    int grid_dim_m = (m + BM - 1) / BM;
    int grid_dim_n = (n + BN - 1) / BN;
    dim3 gridDim(grid_dim_m * grid_dim_n, 1, 1);

    constexpr int local_num_tile_k = (BK + SCALE_BLOCK_SIZE - 1) / SCALE_BLOCK_SIZE;
    constexpr int n_tiles_bound    = (BN + SCALE_BLOCK_SIZE - 1) / SCALE_BLOCK_SIZE;
    size_t sharedBytes = (BM * local_num_tile_k + n_tiles_bound * local_num_tile_k) * sizeof(float);

    hipLaunchKernelGGL(
        HIP_KERNEL_NAME(custom_kernel<BM, BN, BK, WM, WN, WK>),
        gridDim, blockDim, sharedBytes, 0,
        (__hip_fp8_e4m3_fnuz*)a.data_ptr(), (__hip_fp8_e4m3_fnuz*)b.data_ptr(), 
        a_scale.data_ptr<float>(), b_scale.data_ptr<float>(), (__hip_bfloat16*)c.data_ptr(),
        (int)m, (int)n, (int)k
    );
    hipDeviceSynchronize();
}
"""

import os
os.environ["CXX"] = "clang++"

module = load_inline(
    name='fp8_mm',
    cpp_sources=[CPP_WRAPPER],
    cuda_sources=[CUDA_SRC],
    functions=['fp8_mm'],
    verbose=True,
    extra_cuda_cflags=["--offload-arch=gfx942", "-std=c++20"],
)


def custom_kernel(data: input_t) -> output_t:
    a, b, a_scale, b_scale, c = data
    module.fp8_mm(a, b, a_scale, b_scale, c)
    return c


def ref_kernel(data: input_t) -> output_t:
    """
    Highly inefficient torch reference implementation of FP8 GEMM.
    You can use this as a reference / starting template for your implementation.
    """
    # c: [m, n] is pre-allocated memory to help remove allocation overhead.
    a, b, a_scale, b_scale, c = data

    # a is M x K in column-major order, we convert here for simplicity.
    a = a.contiguous()
    a_scale = a_scale.contiguous()
    b_scale = b_scale.contiguous()

    # constants
    m = a.shape[0]
    n = b.shape[0]
    k = a.shape[1]
    block_shape_n = 128
    block_shape_k = 128
    scale_n = b_scale.shape[0]
    scale_k = b_scale.shape[1]

    # Apply blockwise scaling to input 'a'
    a_scale = a_scale.unsqueeze(-1).repeat(1, 1, block_shape_k)  # Shape: [m, scale_k, block_shape_k]
    a_scale = a_scale.reshape(m, scale_k * block_shape_k) 
    a_scale = a_scale[:, :k]

    # Dequantize 'a', in your implementation you should do this at the end.
    a = a.to(a_scale.dtype) * a_scale 

    # Apply blockwise scaling to input 'b'
    b_scale = (
        b_scale.view(-1, 1)
        .repeat(1, block_shape_n * block_shape_k)
        .view(scale_n, scale_k, block_shape_n, block_shape_k)
        .permute(0, 2, 1, 3)  # Reorder dimensions: [scale_n, blk_n, scale_k, blk_k]
        .reshape(scale_n * block_shape_n, scale_k * block_shape_k)
    )
    b_scale = b_scale[:n, :k]

    # Dequantize 'b', in your implementation you should do this at the end.
    b = b.to(b_scale.dtype) * b_scale 

    # Compute FP8 GEMM and write to 'c'. 
    c[...] = (a @ b.T).to(torch.bfloat16)
    return c

# m,n,k = 256, 256, 128
m,n,k = 1024, 1024, 1024
scale_n = triton.cdiv(n, 128)
scale_k = triton.cdiv(k, 128)
a = torch.randn((k, m), dtype=torch.bfloat16, device="cuda").to(torch.float8_e4m3fnuz)
b = torch.randn((k, n), dtype=torch.bfloat16, device="cuda").to(torch.float8_e4m3fnuz)
a_scale = torch.randn([scale_k, m], dtype=torch.float32, device="cuda")
b_scale = torch.randn([scale_k, scale_n], dtype=torch.float32, device="cuda")
a_scale = torch.ones([scale_k, m], dtype=torch.float32, device="cuda")
# for i in range(m):
#     a_scale[:,i] = torch.ones((scale_k), ) * (i/m)
# for i in range(scale_k):
#     a_scale[:,i] = torch.ones((scale_k), ) * (i/m)
c = torch.zeros((m, n), dtype=torch.bfloat16, device="cuda")
data = (a.T, b.T, a_scale.T, b_scale.T, c)
copy_data = deepcopy(data)
out = custom_kernel(data)
ref_out = ref_kernel(copy_data)
diff = torch.abs(out - ref_out)
# print(f"diff shape: {diff.shape}")
for i in range(m):
    print(f"i={i}: {diff[i]}")
python eval.py leaderboard /tmp/tmpbz92rhto
[92mSuccessfully preprocessed all matching files.[0m

/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip [skipped, already hipified]
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/main.cpp -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/main.cpp [skipped, no changes]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 0
[1/1] clang++ hip.cuda.o main.o -shared -L/usr/local/lib/python3.10/dist-packages/torch/lib -lc10 -lc10_hip -ltorch_cpu -ltorch_hip -ltorch -ltorch_python -L/opt/rocm-6.3.1/lib -lamdhip64 -o fp8_mm.so
i=0: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=1: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=2: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=3: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=4: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=5: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=6: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=7: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=8: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=9: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=10: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=11: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=12: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=13: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=14: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=15: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=16: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=17: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=18: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=19: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=20: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=21: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=22: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=23: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=24: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=25: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=26: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=27: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=28: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=29: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=30: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=31: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=32: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=33: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=34: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=35: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=36: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=37: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=38: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=39: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=40: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=41: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=42: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=43: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=44: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=45: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=46: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=47: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=48: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=49: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=50: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=51: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=52: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=53: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=54: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=55: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=56: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=57: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=58: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=59: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=60: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=61: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=62: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=63: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=64: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=65: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=66: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=67: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=68: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=69: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=70: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=71: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=72: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=73: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=74: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=75: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=76: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=77: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=78: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=79: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=80: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=81: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=82: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=83: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=84: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=85: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=86: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=87: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=88: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=89: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=90: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=91: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=92: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=93: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=94: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=95: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=96: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=97: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=98: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=99: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=100: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=101: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=102: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=103: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=104: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=105: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=106: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=107: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=108: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=109: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=110: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=111: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=112: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=113: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=114: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=115: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=116: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=117: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=118: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=119: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=120: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=121: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=122: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=123: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=124: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=125: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=126: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=127: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=128: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=129: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=130: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=131: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=132: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=133: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=134: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=135: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=136: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=137: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=138: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=139: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=140: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=141: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=142: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=143: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=144: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=145: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=146: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=147: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=148: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=149: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=150: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=151: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=152: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=153: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=154: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=155: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=156: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=157: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=158: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=159: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=160: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=161: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=162: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=163: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=164: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=165: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=166: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=167: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=168: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=169: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=170: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=171: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=172: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=173: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=174: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=175: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=176: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=177: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=178: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=179: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=180: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=181: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=182: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=183: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
i=184: tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
[...] 840 lines omitted