What content is inside workspace_buffer and how should I choose a proper buffer size?

[Tomorrowdawn]

As document described, each handler has a initializing parameter 'workspace_buffer'. However, the search result gives me zero information about this buffer. I would like to ask what's the usage of this buffer and how should I allocate a proper size in advance?

[yzh119]

Sorry for the late reply.

These workspace are prepared for two kind of buffers:

1. integer buffers (metadata): after we tile the attention computation on query length dimension or k/v length dimension, we need some metadata for indexing the page-table after partition.
2. float buffers (intermediate attention output), we use the split-k trick to increase the parallelism for small problem size, in this case, we need some buffers on global memory to store the intermediate attention results (partial attention output and logsumexp factor, see https://docs.flashinfer.ai/tutorials/recursive_attention.html) corresponding to each kv-chunk, and we need another round of reduction to get the complete attention result.