e2e demonstration for bigger models

[jiwonsong-dev]

Hi, thank you for great work and efforts.

Current kernels seem to support only dimensions of 7B models with hidden dimension 4096.
How can I extend it for larger models like Llama-30B or 65B?
It returns an error when I just add template instances for larger dimension.

Thank you.

[happierpig]

Hi @jiwonsong-dev,

Thanks for your interest in this project!

For supporting different input shapes, GEMM kernel can be used without any performance tuning by changing

Atom/kernels/src/GEMM/bench_dense_layer_gemm_i4_o16.cu

Line 69 in 7e3618b

.add_int64_axis("hidden_dim", {4096});

. Attention kernels can be used as different models are naturally supported by FlashInfer. However, the REORDER and RMS_NORM kernels' design are coupled with the shape. The blockDim are hard coded (

Atom/kernels/include/Reorder/Reorder.cuh

Line 217 in 7e3618b

dim3 blocks(4, 32);

) and the last block are forced to quantize the INT8 outliers (

Atom/kernels/include/Reorder/Reorder.cuh

Line 171 in 7e3618b

if(ty == bdy - 1){

). Some kernel efforts are needed for support larger shapes.