[0cc4m]

Overview

Handle src1 precision in the Vulkan backend MUL_MAT and MUL_MAT_ID implementation, by scaling the values into the fp16 range, similar to how it is handled by Metal (#29029).

I also added the GGML helper functions and switch cases discussed in #30105 (comment).

Supercedes #30105 and #30544

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES

[jeffbolznv]

I understand that ggml-metal is already doing this scaling+f16 hack, but it doesn't seem like a legitimate implementation of PREC_F32. There are obvious cases where it breaks down, surely somebody will eventually find a real-world case where it doesn't work. IMO ggml should either require a legitimate fp32-based implementation of PREC_F32 and accept that it will be slow, or disallow F32 as a src prec and move this amax+scale trick into the graph. CC @ggerganov for his thoughts.

[0cc4m]

I had the same concern and discussed it with Georgi before, my understanding was that this is a performance-compromise to avoid fp16 overflow, which was the actual problem that this src prec flag was meant to solve. And the overflow is avoided by this path. But I agree, it's not "true" FP32 precision.

[ggerganov]

Technically, I agree that the amax is a bit of a hack, so thanks for pushing on this - let's clean it up.

For the Metal backend I have a few refactors in the pipe to enable proper BF16 support for src1 in GGML_OP_MUL_MAT and GGML_OP_MUL_MAT_ID. After these land, we will replace the amax trick with these.

I think what we have to do in general is:

• Change the graph to use ggml_prec_set_src(experts, GGML_PREC_BF16, 1);. This should allow the backend to use BF16 which is also correct in the cases of Mistral4 and Mellum.
• When the backend sees GGML_PREC_BF16 for src1 there are 2 options:
  • Use a src1 BF16 capable kernel (if available)
  • Use a src1 F32 kernel

The amax trick could be a stopgap until the above can be satisfied and if it make sense in terms of complexity/maintenance. But we should likely mark all such tricks with TODOs so that we know we have to make them correct eventually.

Does this sound OK?

[0cc4m]

I'm not sure about the hardware support in Metal devices, but while bfloat16 is better and I certainly want to support that, for Vulkan we have a lot of devices that have good performant fp16 support, but not bfloat16, and falling back to fp32 will be very slow in comparison. Possibly some of this can be cushioned with some kind of bf16 emulation just for shmem storage + fp32 math shader, but I think this will be painful for older devices.

Maybe that's fine as long as it just affects a limited amount of huge models, like Mistral 4.

[Green-Sky]

It also effects smaller models like mellum2.1 . It also effects cuda. #30704

[ggerganov]

I'm not sure about the hardware support in Metal devices

Pre-M3 devices do not have hardware for BF16 - it is emulated and hence a bit slower than F16.

[0cc4m]

I'm not sure about the hardware support in Metal devices

Pre-M3 devices do not have hardware for BF16 - it is emulated and hence a bit slower than F16.

I'm most worried about RDNA3, RADV does not expose bfloat16 support without an env var (because of numerical inconsistencies in the hardware), so falling back to fp32 means disabling tensor cores and all scalar fp16 dual issue instructions, so it's gonna be a lot slower. For older devices it would just be about losing the dual issue and possibly having to store twice as much data in shared memory.

[ggerganov]

Yeah, I'm not 100% sure what is the right call here. The fundamental issue is that using F16 is by definition not enough for LLMs because all models are trained in BF16. So there is never a guarantee that F16 will not break for some specific activation values. We've only observed Mistral4 and Mellum, but in theory any model can exhibit such issues. So from correctness perspective, all compute has to be done in BF16. But on the other hand the performance loss on a lot of consumer hardware is significant - we can't ignore that too.

The approach that we can take is to relax the API requirements a bit. I guess we can declare in the ggml_prec API that when a hardware BF16 support is not available, the implementation is allowed to use best-effort approximation and this way we permit the F16 + amax trick for GGML_PREC_BF16 when we can't use BF16? The benefit of this would be that at least the ggml API is correct in terms of what it computes, though it becomes more vague. Likely better than being technically incorrect.

[0cc4m]

bfloat16 is also problematic for precision, since it has fewer bits for the mantissa, so I'm not even sure which path is closer to actual fp32.

[jeffbolznv]

The graph builder could check for bf16 support in the backend and choose between using PREC_BF16 vs amax+PREC_F16.

[IMbackK]

The graph builder could check for bf16 support in the backend and choose between using PREC_BF16 vs amax+PREC_F16.

Support is not enough, and graph builder would need to make decisions about performance too (non trivial eg. mi100 has bf16 matrix instructions but they are half rate compared to fp16, amd hw in general dosent have packed valu bf16 instructions for its fp16 equivalents etc)
Using bf16 just because there is support hits slow paths, always using amax+f16 precludes using bf16 at all if its found necessary in some case.

[ggerganov]

@IMbackK I think Jeff's point is that the decision which path to take (fast and not-exact or slow and exact) should be lifted up to the user code. This way the ggml supports both options and it's up to the user application to figure out which option to use when. How exactly (via CLI args, some extra logic, etc.) not yet defined. Btw, we have the llama_prec_policy recently introduced which is rather rudimentary atm, but can probably grow to handle all this.

[jeffbolznv]

I think Jeff's point is that the decision which path to take (fast and not-exact or slow and exact) should be lifted up to the user code.

Right, and the user code can make the choice on a per-model basis, based on correctness and performance testing.

bf16 is also half speed on Geforce (well, bf16 is only supported with fp32 accumulators, and fp32 accumulators are half speed, but presumably fp32 accumulators are required for these models so maybe no real slowdown?). Also I don't think token gen will really be penalized by using bf16. codex tells me the bf16 work is maybe 20% of the math in these models. So I don't think half speed for that fraction of the workload is catastrophic. Maybe for GPUs that need to fall back from tensor core to non-tensor core (rdna3) the situation would be worse.

[ggerganov]

We'll lift the amax trick to the graph for the specified models as a stopgap. After that we'll work to add auto-detection of which backends have support for the necessary src1 precision (i.e. GGML_PREC_BF16) and avoid applying the extra scales to those layers.

[0cc4m]

How do you intend to handle bf16-capable hardware that would prefer fp16 for performance reasons?

[ggerganov]

How do you intend to handle bf16-capable hardware that would prefer fp16 for performance reasons?

Not sure yet - this will be part of the second step. For now I'd like to close the gap (e.g. #30733) and in the meantime we can figure out what to do for fast/slow BF16 hardware.

[ggerganov]

It seems doing the amax trick at the graph level is not so simple - injecting an argmax op before the down projection is not compatible with the tenor parallelism in the Meta backend because we need the full rows to compute the max. And with tensor parallelism enabled, the rows are split across the devices. So not sure if this is going to work.

[ggerganov]

The tensor parallel issue aside, the prompt processing performance drops significantly due to the ggml_argmax on the activations. With Metal at -ub 2048 it's ~2x slower.

@jeffbolznv Let me know if you have any suggestions - here is the WIP (981f0a6)

[jeffbolznv]

I tried it on Vulkan and the argmax was very slow, like many times slower. Maybe top_k would be faster? For vulkan at least we have a fast path for k=1. I had codex put this together and perf was much better, but I haven't validated that it's doing the right thing.

not compatible with the tenor parallelism in the Meta backend because we need the full rows to compute the max

Is it really splitting the inner dimension of the src matrices? I'd have expected it to split dim1 of the result.

[jeffbolznv]

When I dumped the perf logger stats for top_k I see this:

TOP_K K=1 (14680064,1,1,1): 297 x 681.71 us = 202468 us

Maybe the tensor is incorrectly being treated as flattened and that's why it's so slow? Or it could just mean codex gave me bad code to run.

[0cc4m]

Okay, I reduced it to only the Vulkan precision handling without scaling.

[jeffbolznv]

My comments are just nitpicks, fine to merge if you want.