[FEATURE] How to generate new kernel config for matrix (3072, 3072) ?

[llcnt]

Hey!
I am trying to use flute on a model with matrices of shape (3072, 3072).
Some already existing templates (the one that validates the function is_template_supported here if I understood well) can work for such matrices. However the compute time is not very good (ie inference takes long). I guess it is because the templates pre-loaded in TEMPLATE_CONFIGS do not suit well for shape (3072, 3072).
Am I missing something? Should I re-run the codegen_utils.py file with different hyper-parameters?

Thanks for your help :)

[HanGuo97]

Thank you for trying it out! I have a couple of quick clarification questions to help me understand the situation:

1. Which GPU are you using?
2. Do you have an idea of how “not very good” it is?
3. Are you using the latest FLUTE PyPI wheel? (Earlier versions have hard-coded dispatch, while more recent versions have an autotune.)

[llcnt]

Thanks for your quick reply!

1. I have tested flute on 2 GPUS (A10G and L40S). I know they have not the same number of SMs than the GPUs you have reported results on, but it works with a smaller model (it works well and I obtain a x2 speedup on both GPU when I use a llama-1B (that contains matrices of shapes (2048, 2048) and (2048, 512)) ). To be precise: on my GPUs inference with matrices of shapes (2048, 2048) takes as long as original model after quantization, but if I compile the forward pass (using torch.compile) I obtain some speedup at inference. While quantizing a matrix of shape (3072, 3072), the inference is slower (about x2). I manage to recover the same speed as the non-quantized version when I use again torch.compile, but there are no additional speedups. I am not sure, but I guess it is maybe because the quantized matrix without compilation is slower than the default matrix, no ?
2. The quality (ie the output) is not bad, but the inference takes as long as the default (ie 16 bits) model;
3. I have built FLUTE from source, and yes I am using the maybe_tune_and_repack function to enable autotune. But still, if I understood correctly, the kernels templates tested in autotune will be limited by the one pre-defined in the codegen_utils.py file, right?

Thx :)

[HanGuo97]

Thank you for the thorough information; it’s truly helpful!

I have limited knowledge about the GPUs you mentioned, but I’ve heard good things regarding their performance in inference tasks.

From my understanding, torch.compile doesn’t go that much inside custom kernels. Therefore, your observation that it could make a notable difference suggests that non-kernel factors might be the cause. I recommend starting there before experimenting with new kernel templates.

Could you elaborate on the framework you use to run the model? For the HF model, we observed in a prior experiment that several quantization kernels (including FLUTE) did not lead to a improvement, which puzzled us, although we couldn’t pinpoint the reason. In contrast, results seem more consistent with GPT-Fast and vLLM. My (perhaps uninformed) assumption is that frameworks not optimized for performance may introduce additional overheads.

[llcnt]

Thank you for the thorough information; it’s truly helpful!

I have limited knowledge about the GPUs you mentioned, but I’ve heard good things regarding their performance in inference tasks.

From my understanding, torch.compile doesn’t go that much inside custom kernels. Therefore, your observation that it could make a notable difference suggests that non-kernel factors might be the cause. I recommend starting there before experimenting with new kernel templates.

Could you elaborate on the framework you use to run the model? For the HF model, we observed in a prior experiment that several quantization kernels (including FLUTE) did not lead to a improvement, which puzzled us, although we couldn’t pinpoint the reason. In contrast, results seem more consistent with GPT-Fast and vLLM. My (perhaps uninformed) assumption is that frameworks not optimized for performance may introduce additional overheads.

I am using diffusers lib from HF: to be precise I am replacing all nn.linear layers in a Flux model with Flute quantized one.
Here are the results/differences I found wrt LLMs:

• Compared to LLMs, the bottleneck is not only due to computing qkv. Computing the scale dot product attention is also very consuming. For small stablediffusion models it is manageable, but for larger models (like Flux) the matrices involved are very large (eg [1, 24, 4608, 128] for a batch of 1) and computations involve GEMM kernels (the problem is here similar to the one you would expect while computing sdpa on LLM with a very large sequence length)
• Compared to LLMs, computing qkv matrices also involves more computation (in general). The hidden states are typically very large (eg [1, 4096, 3072] for a batch of size 1, while q matrix is of shape [3072, 3072]), so the computation also involved large GEMM kernels
• Currently diffusers do not benefit from larger batch sizes(cf this issue). I guess this can be explained by the above 2 points. What is still unclear to me is that FLUTE should also provide speedups for GEMM, right?
• Maybe the above first two points could explain the differences you obtained with GPT-Fast and vLLM ? At least could be interesting to have a look at how the qkv and the sdpa are coded there ;)

I am planning to run experiments on just running one forward with:

• a matrix quantized with FLUTE of shape [3072, 3072] and a vector [1, 3072] to test GEMV
• a matrix quantized with FLUTE of shape [3072, 3072] and a matrix [4608, 3072] to test GEMM

I will keep you updated :)

[HanGuo97]

Understood; knowing it’s a diffusion model is very helpful. Here are a couple of thoughts:

FLUTE is designed to work with GEMM from the very beginning. In the GEMM literature, the terms M, N, K are commonly referenced (as in [M, K] x [K, N] -> [M, N]). When dealing with matrices larger than 2D, it’s often necessary to reshape them to determine the M, N, K values.

FLUTE was designed with M in the range of [1, ..., 16] in mind, with M = 1 corresponding to the GEMV case. For very large values of M, there are fundamental challenges in achieving speedup with weight-only quantization, and FLUTE has not been optimized for such cases, so performance may decrease.

The experiments you suggested are sensible, but it may also be worthwhile to conduct an experiment using a different quantization kernel- perhaps bitsandbytes for M = 1 or marlin and machete for M > 1. Do any of these options accelerate your workload? This could help us determine if we should focus on the kernel or other parts of the model.

[llcnt]

I ran several experiments based on your recommendations! Interestingly I obtained similar differences in my results compared to what you had with vllm vs gptfast:

• in the diffusion model case I found that there are noticeable differences in between the code from huggingface (in the diffusers lib), and the initial code from BlackforestLab if we consider the Flux model (the code from BFL is much faster!)

I ran 2 basic experiments of matrix multiplication (average over 100 runs with 5 warmup runs before, on a L40S GPU) with M=1 and M=4000 (and several values for K=N), for benchmarking FLUTE vs torch (the default f16) vs torch-fp8 (with torch._scaled_mm):

• in the GEMV case, surprisingly, the original torch implementation seems to be the fastest;
• in the GEMM case the torch-fp8 is the fastest. This is something expected, and already put in practice eg here
  
• I dont understand why in all cases FLUTE does not perform well, even in the GEMV case :( Is it because of my GPU?

[llcnt]

Edit: I have enlarged the matrix and now FLUTE is way faster than the default torch implementation :)

However torch-fp8 is still better than FLUTE :(

• Could you explain why FLUTE is competitive only for very large matrices ? Do you think it is related to the GPU I tested it on ? (I was wondering if it is related to the size of the L2/cached memory available on L40S GPUs)
• Do you think it is doable to extend FLUTE to large GEMM cases ?

[HanGuo97]

Thanks for the experiments!

• I’m not surprised that FLUTE performed poorly with M = 4000. It’s primarily designed to work with M < 16, although it should manage with M < 64. There are fundamental reasons why achieving speedups with large M, N, and K is challenging. I can provide more details if you’re interested, but it relates to compute-boundness versus memory-boundness.
• Does it make sense for FP8 to be slower than the original PyTorch MM in GEMV?

Could you clarify the data type of C = mm(A, B)? Are A, B, C all fp16?

[HanGuo97]

Thanks again for sharing the experimental results—I’ve spent a bit more time thinking through them, and I had a few questions and thoughts:

• In the GEMV cases, PyTorch’s fp8 matmul appears to be more than 2× slower than fp16. If you're using fp16 x fp16 -> fp16 for the PyTorch matmul and fp8 x fp8 -> fp8 for the FP8 case, the latter should save both on compute and memory bandwidth. So I’m not quite sure why it ends up being significantly slower.

• Also in the GEMV cases, when you scale up both K and N by 4×, the FP8 matmul runtime actually decreases. Since both memory traffic and FLOPs increase, that’s a bit puzzling.

On a related note—regarding speeding up large GEMMs—it might be helpful to clarify the two types of quantization. Suppose we’re computing C = mm(A, W):

• Weight-only quantization quantizes just the weights (W), while the activations (A) remain in their original precision. This means you'll need to do type conversion at runtime, and the compute still happens in the activation's precision.
• Weight-activation quantization quantizes both A and W, enabling the matmul to potentially run entirely in the lower precision format.

For small GEMMs, the runtime is generally memory-bound, so both approaches can be fast. But for large GEMMs, where compute dominates, only weight-activation quantization is likely to yield significant speedups.

Hope that helps clarify where I’m coming from!

[llcnt]

Maybe another process was running at the same time on my GPU, I have just re-launched the experiment, and obtained :

Still, for flute, the GEMVcase 2048 seems to be slower than 3072 (0.028 vs 0.026) ?!
Here are the precise numbers ([GEMV, GEMM] for torch, torch-Fp8, FLUTE):

[HanGuo97]

Thanks! The M=1 with a relatively large N, K now seems slightly more reasonable to me. That said, I’m still a bit surprised by the factor of 2x slower, especially for FP8, with small K and N.

Could you clarify the data type of C = matmul(A, B)? Are A, B, and C all FP16 for PyTorch and FLUTE, but FP8 for FP8 PyTorch?

One small request: I’d love to see if you could include the NF4 kernel from bitsandbytes. It should be fast for M=1. Since it’s a more mature and stable library, and it is spiritually close to FLUTE, its profiling numbers would provide me with a lot of information.