Repository: microsoft/onnxruntime
File: onnxruntime/core/mlas/lib/qlutgemm.cpp (C++)

Code before:
    // TODO(vraspar): other way to do it
    auto lut_buffer = std::make_unique<int8_t[]>(lut_buffer_size);
    memset(lut_buffer.get(), 0, lut_buffer_size);

    int8_t* qlut = reinterpret_cast<int8_t*>(lut_buffer.get());
    float* lut_scales = reinterpret_cast<float*>(qlut + lut_size_bytes);                  // after lut
    float* lut_biases = reinterpret_cast<float*>(lut_scales + lut_scales_size * M);       // after scales

    const auto* a_float = reinterpret_cast<const float*>(A);  // Activation data

    // const int num_groups = static_cast<int>(K / BlkLen);

    // Iterate over M (batch dimension)
    // Each iteration processes one row of the activation matrix.
    // NOTE: This loop is intentionally serialized. Previous attempts to parallelize
    // using MlasTrySimpleParallel caused flaky test failures (race conditions)
    // when M > 1 (e.g., Batch32 case). Since GenerateLUT is lightweight,
    // serial execution ensures correctness with negligible performance impact.
    // TODO(vraspar): Ideally we have to do block parallelism here

    for (size_t ine11 = 0; ine11 < static_cast<size_t>(M); ine11++) {
        const size_t row_offset = ine11 * K;
        // Use 16-byte stride to match the buffer allocation and kernel expectation.
        const size_t lut_offset = ine11 * K * 16;
        const size_t scale_bias_offset = ine11 * lut_scales_size;

        // Call the dispatch function for this row
        // ggml_tmac_mul_mat_task_init
        Dispatch->GenerateLUT(
            const_cast<float*>(a_float + row_offset),  // Input activation for this row
            qlut + lut_offset,                         // Output LUT for this row
            lut_scales + scale_bias_offset,            // Scales for this row
            lut_biases + scale_bias_offset,            // Biases for this row
            M,
            K,
            N,
            tmac_params.act_group_size
        );
    }

    // all relevant LUT's have been generated
    // equivalent of lut_mul_mat's ggml_backend_tmac_mul_mat function ggml_barrier line

    const size_t n_tiles_num = tmac_params.n_tiles_num;
    assert(N % n_tiles_num == 0);

    const size_t bits = tmac_params.bits;

    // Pre-calculate sizes for offset calculations
    const size_t w_size = N * K * bits / 8;
    const size_t w_chunk_size = w_size / n_tiles_num;

Review comment (performance):
sorry, If I am misunderstanding, but this change will pad lut buffer with extra zeros. It seems inefficient to do that. 

If I understand correctly, GenerateLUT_avx2 writes at stride 128 bytes, but the tbl_g4_int8_float_update_impl reads at stride 256 bytes.

This causes the tbl_g4_int8_float_update_impl to read correct data at offset 0 but  read data intended for a different K iteration at offset 256

I wonder if modifying `GenerateLUT_avx2` to write at stride `kfactor * (1 << g) = 256` bytes, matching what `tbl_g4_int8_float_update_impl` expects would be more better solution
```
// Instead of:
lut_ctor_g4_int8_impl(..., &qlut[k_outer_1 * act_group_size * 4], ...);

// can try:
lut_ctor_g4_int8_impl(..., &qlut[k_outer_1 * kfactor * 16], ...); 

```

I ran this test 100s of time, but I can't seem to reproduce this locally on my machine

Code after:
     *
     * Currently, we only support FP32, add FP16 support later which requires conversion buffer
     *
     * LUT Buffer for FP32 : K * M * 4 * sizeof(uint8_t) bytes + lut_scale_size * m * 2 * sizeof(float) bytes  + allignment
     *
     */

    // n_tiles_num = m * bits / bm;

    // TODO(vraspar): support other bitwidths
    // For T-MAC, kernel properties (bm, n_tiles_num) are primarily driven by the number of output features (N).
    // Initialization during packing (LutGemmPackQuantBDataSize) uses N as the major dimension,
    // so we must match that here to ensure consistent weight tiling.
    MlasInitLutGemmKernelConfig(N, K, 2, BlkLen, HasZeroPoint);
    const MlasTMACKernelParams& tmac_params = MlasGetLutGemmKernelParams(N, K, 2, BlkLen, HasZeroPoint);
    const size_t lut_scales_size = K / tmac_params.act_group_size;
    const size_t lut_size_bytes = static_cast<size_t>(M) * static_cast<size_t>(K) * 4;
    size_t lut_buffer_size = CalculateLutBufferSize(N, K, M, tmac_params);

    // make buffer of lut_buffer_size bytes
    // TODO(vraspar): other way to do it
    auto lut_buffer = std::make_unique<int8_t[]>(lut_buffer_size);
    memset(lut_buffer.get(), 0, lut_buffer_size);

    int8_t* qlut = reinterpret_cast<int8_t*>(lut_buffer.get());
    float* lut_scales = reinterpret_cast<float*>(qlut + lut_size_bytes);                  // after lut
    float* lut_biases = reinterpret_cast<float*>(lut_scales + lut_scales_size * M);       // after scales

    const auto* a_float = reinterpret_cast<const float*>(A);  // Activation data

    // const int num_groups = static_cast<int>(K / BlkLen);

    // Iterate over M (batch dimension)
    // Each iteration processes one row of the activation matrix.
    // NOTE: This loop is intentionally serialized. Previous attempts to parallelize
    // using MlasTrySimpleParallel caused flaky test failures (race conditions)
    // when M > 1 (e.g., Batch32 case). Since GenerateLUT is lightweight,
    // serial execution ensures correctness with negligible performance impact.
    // TODO(vraspar): Ideally we have to do block parallelism here

    for (size_t ine11 = 0; ine11 < static_cast<size_t>(M); ine11++) {
        const size_t row_offset = ine11 * K;
        // Call the LUT generation kernel for this activation row.
        // We use a 4-byte stride (per activation) for the LUT entries to satisfy
        // the memory layout requirements of the computation kernel.
        const size_t lut_offset = ine11 * K * 4;
        const size_t scale_bias_offset = ine11 * lut_scales_size;

        Dispatch->GenerateLUT(
            const_cast<float*>(a_float + row_offset),  // Input activation for this row
            qlut + lut_offset,                         // Output LUT for this row