Flash Attention Shared Memory constraint

[AntonOresten]

Hi! Thanks for making this.
I understand this package is still experimental, but wanted to document this constraint, and also ask if anyone knows whether e.g. head dim of 128 would be difficult to implement?

julia> x = CUDA.rand(64, 4096, 4, 4); Jjama3.NNop.flash_attention(x, x, x; causal=true);

julia> x = CUDA.rand(128, 4096, 4, 4); NNop.flash_attention(x, x, x; causal=true);
ERROR: Failed to find groupsize for Flash Attention that satisfies Shared Memory constraint.
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] #flash_attention_groupsize#8
   @ ~/.julia/packages/NNop/cRsoL/src/attention.jl:186 [inlined]
 [3] flash_attention_groupsize
   @ ~/.julia/packages/NNop/cRsoL/src/attention.jl:175 [inlined]
 [4] _flash_attention(q::CuArray{Float32, 4, CUDA.DeviceMemory}, k::CuArray{Float32, 4, CUDA.DeviceMemory}, v::CuArray{Float32, 4, CUDA.DeviceMemory}; causal::Bool)
   @ NNop ~/.julia/packages/NNop/cRsoL/src/attention.jl:127
 [5] _flash_attention
   @ ~/.julia/packages/NNop/cRsoL/src/attention.jl:113 [inlined]
 [6] #flash_attention#18
   @ ~/.julia/packages/NNop/cRsoL/src/attention_crc.jl:5 [inlined]
 [7] top-level scope
   @ REPL[121]:1

[pxl-th]

Hi. There's no real obstacle for 128 head dim support, only how much shmem is available.
For groupsize estimation we use backward kernel for both forward & backward kernels.

How much shmem do you have?

julia> NNop.flash_attention_shmem_bwd(Float32; emb_dim=128, groupsize=64, qk_fp16=false) |> Base.format_bytes
"144.000 KiB"
julia> NNop.flash_attention_shmem_bwd(Float32; emb_dim=128, groupsize=64, qk_fp16=true) |> Base.format_bytes
"112.000 KiB"
julia> NNop.flash_attention_shmem_bwd(Float32; emb_dim=128, groupsize=32, qk_fp16=true) |> Base.format_bytes
"52.000 KiB"
julia> NNop.flash_attention_shmem_bwd(Float32; emb_dim=128, groupsize=32, qk_fp16=false) |> Base.format_bytes
"68.000 KiB"

[AntonOresten]

Ah, that's probably it. I'm using an RTX 6000 Ada (48 GB), but I'm not sure about shared memory.
Does NNop.flash_attention_shmem_fwd and NNop.flash_attention_shmem_bwd tell me how much would be needed? How is groupsize decided?

Edit:

julia> Int64(CUDA.attribute(first(CUDA.devices()), CUDA.CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK))
49152

julia> Int64(CUDA.attribute(first(CUDA.devices()), CUDA.CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN))
101376

[AntonOresten]

Okay.. I can do 128 if I use fp16, which is reasonable. But fp32 may still work with smaller groupsize?

julia> x = CUDA.rand(Float16, 128, 1, 1, 1);

julia> NNop.flash_attention(x, x, x, causal=true);

[AntonOresten]

Wanted to see if it was as easy as switching to CUDA.CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN in NNopCUDAExt and setting qk_fp16s = (false, true) in flash_attention_groupsize:

fp32 flash attention with head dim 128

julia> x = CUDA.rand(Float32, 128, 10, 4, 4);

julia> NNop.flash_attention(x, x, x; causal=true)
ERROR: Failed to compile PTX code (ptxas exited with code 255)
Invocation arguments: --generate-line-info --verbose --gpu-name sm_89 --output-file /tmp/jl_lO3UjaHMX6.cubin /tmp/jl_RgDk9Whgdw.ptx
ptxas error   : Entry function '_Z25gpu__flash_attention_fwd_16CompilerMetadataI11DynamicSize12DynamicCheckv16CartesianIndicesILi3E5TupleI5OneToI5Int64ES6_S6_EE7NDRangeILi3ES0_10StaticSizeI10_32__1__1_ES8_vEE4TypeI12FATileConfigILi32ELi128ELi32ELi8ELi4ELb0ELb0ELb0EEESE_ISF_ILi32ELi32ELi128ELi16ELi8ELb0ELb1ELb1EEE13CuDeviceArrayI7Float32Li4ELi1EESK_ISL_Li3ELi1EESN_SM_SM_SM_SL_3ValILi128EESO_ILi1EESO_ILb0EESO_ILb1EE' uses too much shared data (0xd000 bytes, 0xc000 max)
ptxas info    : 217 bytes gmem, 48 bytes cmem[4]
ptxas info    : Compiling entry function '_Z25gpu__flash_attention_fwd_16CompilerMetadataI11DynamicSize12DynamicCheckv16CartesianIndicesILi3E5TupleI5OneToI5Int64ES6_S6_EE7NDRangeILi3ES0_10StaticSizeI10_32__1__1_ES8_vEE4TypeI12FATileConfigILi32ELi128ELi32ELi8ELi4ELb0ELb0ELb0EEESE_ISF_ILi32ELi32ELi128ELi16ELi8ELb0ELb1ELb1EEE13CuDeviceArrayI7Float32Li4ELi1EESK_ISL_Li3ELi1EESN_SM_SM_SM_SL_3ValILi128EESO_ILi1EESO_ILb0EESO_ILb1EE' for 'sm_89'
ptxas info    : Function properties for _Z25gpu__flash_attention_fwd_16CompilerMetadataI11DynamicSize12DynamicCheckv16CartesianIndicesILi3E5TupleI5OneToI5Int64ES6_S6_EE7NDRangeILi3ES0_10StaticSizeI10_32__1__1_ES8_vEE4TypeI12FATileConfigILi32ELi128ELi32ELi8ELi4ELb0ELb0ELb0EEESE_ISF_ILi32ELi32ELi128ELi16ELi8ELb0ELb1ELb1EEE13CuDeviceArrayI7Float32Li4ELi1EESK_ISL_Li3ELi1EESN_SM_SM_SM_SL_3ValILi128EESO_ILi1EESO_ILb0EESO_ILb1EE
    32 bytes stack frame, 0 bytes spill stores, 0 bytes spill loads
ptxas info    : Used 254 registers, used 1 barriers, 32 bytes cumulative stack size, 53248 bytes smem, 740 bytes cmem[0], 8 bytes cmem[2]
ptxas info    : Compile time = 292.440 ms
ptxas info    : Function properties for gpu_report_exception
    0 bytes stack frame, 0 bytes spill stores, 0 bytes spill loads
ptxas info    : Function properties for gpu_signal_exception
    0 bytes stack frame, 0 bytes spill stores, 0 bytes spill loads
If you think this is a bug, please file an issue and attach /tmp/jl_RgDk9Whgdw.ptx

The reason I don't really want to switch to Float16 is because even though at least the forward pass seems to work, the rest of my code seems to break somewhere, and generations are garbage.

One solution is to add groupsize 16 to flash_attention_groupsize, albeit with a performance penalty.

[AntonOresten]

Does this problem possibly go deeper, to KernelAbstractions.jl's @localmem being subject to a static 48 KiB memory limit or something? CUDA.jl at least has CuDynamicSharedArray.

[pxl-th]

Does this problem possibly go deeper, to KernelAbstractions.jl's @LocalMem being subject to a static 48 KiB memory limit or something?

No, I have 64 KiB and can use all of it.
What does NNop._shared_memory(CUDABackend(), 1) shows for you?

[pxl-th]

The reason I don't really want to switch to Float16 is because even though at least the forward pass seems to work, the rest of my code seems to break somewhere, and generations are garbage.

We can also have FP32 inputs, but store data in FP16 shmem. But that may lead to precision loss.
I can work on some config options for that.

[AntonOresten]

No, I have 64 KiB and can use all of it. What does NNop._shared_memory(CUDABackend(), 1) shows for you?

The code that I ran above came from the CUDA extension already, but here it is:

julia> NNop._shared_memory(CUDABackend(), 1)
0x000000000000c000

julia> NNop._shared_memory(CUDABackend(), 1) |> Int
49152

And as we saw, calling with CUDA.CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN showed:

julia> Int64(CUDA.attribute(first(CUDA.devices()), CUDA.CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN))
101376

but I don't know what this means or if it can even be used.

It the card should have at least 128 KiB of L1 cache, but it seems it can't all be used.

[AntonOresten]

We can also have FP32 inputs, but store data in FP16 shmem. But that may lead to precision loss.

Should such conversion be up to the user beforehand, or is doing that as part of the kernel much more efficient?

[pxl-th]

The code that I ran above came from the CUDA extension already

Ah... my bad, didn't notice.

To use more than 48 KiB of shmem on CUDA devices we have to switch to dynamic shmem (more info):

Compute capability 7.x devices allow a single thread block to address the full capacity of shared memory: 96 KB on Volta, 64 KB on Turing. Kernels relying on shared memory allocations over 48 KB per block are architecture-specific, as such they must use dynamic shared memory (rather than statically sized arrays) and require an explicit opt-in using cudaFuncSetAttribute()

But it is not yet supported in KernelAbstractions (there's a task for it though JuliaGPU/KernelAbstractions.jl#533).
So might need some further work...

[AntonOresten]

Ah, bummer. How difficult do you reckon it might be to implement dynamic shared memory in KA? I'm not seeing an open issue. I guess support for that could happen before 1.0.

I might try to port to a CUDA.jl-based kernel, although I guess there are already CUDA reference implementations.🤣

Does AMD simply support higher static shmem? Is 48 KB just a limitation of CUDA?

[pxl-th]

How difficult do you reckon it might be to implement dynamic shared memory in KA?

Should be fairly straightforward, just forwarding calls to respective backends. (CC @vchuravy)

Does AMD simply support higher static shmem? Is 48 KB just a limitation of CUDA?

I think all recent AMD GPUs simply have 64 KiB shmem size.

[AntonOresten]

Seems groupsize of 16 is allowed since #11, which is better than using naive because of the memory benefits. Dynamic shared memory might still be nice though. I see there's an old PR for it that was closed: JuliaGPU/KernelAbstractions.jl#11

Either that, or better support forBFloat16, cause maybe Float32 with 128 headdim is a bit insane.

I'd be happy to contribute to each of these in any way I can.

[pxl-th]

Yes, this is a temporary solution :)
I think BFloat16 should just work, no? Need to test though