[CPU][Brgemm] add support for int8 brgemm

[Valentine233]

For INT8 SDPA kernel usage, we add support for INT8 Brgemm.

cc @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10 @voznesenskym @penguinwu @EikanWang @Guobing-Chen @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @aakhundov

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/143384

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Failures because of deprecated version of actions/download-artifact: v3

✅ You can merge normally! (1 Unrelated Failure)

As of commit ea8db8b with merge base 9631d1a ():

UNSTABLE - The following job failed but was likely due to flakiness present on trunk and has been marked as unstable:

• pull / linux-jammy-py3-clang12-executorch / test (executorch, 1, 1, lf.linux.2xlarge) (gh) (#144480)
  backends/xnnpack/test/ops/test_conv1d.py::TestConv1d::test_qs8_conv1d_batchnorm_seq

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[jgong5]

Can you elaborate on the usage of batch_size? How is it to be used by SDPA?

[Valentine233]

Can you elaborate on the usage of batch_size? How is it to be used by SDPA?

The batch_size here refers to the batch parameter in BRGemm (Batch-Reduced Gemm). For the second Gemm in INT8 SDPA, i.e. the product of attention [q_len, kv_len] and value [kv_len, head_size], we do blocking on kv_len to split it into several kv_blocks. So, we can take advantage of parameter batch_size to merge several Gemms of [q_len, kv_block_size] and [kv_block_size, head_size] into one BRGemm with batch_size=kv_block_num, with the reduced dimension of K.

[jgong5]

Can you elaborate on the usage of batch_size? How is it to be used by SDPA?

The batch_size here refers to the batch parameter in BRGemm (Batch-Reduced Gemm). For the second Gemm in INT8 SDPA, i.e. the product of attention [q_len, kv_len] and value [kv_len, head_size], we do blocking on kv_len to split it into several kv_blocks. So, we can take advantage of parameter batch_size to merge several Gemms of [q_len, kv_block_size] and [kv_block_size, head_size] into one BRGemm with batch_size=kv_block_num, with the reduced dimension of K.

In that case, why not making the kv_block_size larger and invoke a single small gemm instead of making a brgemm call?

[Valentine233]

Can you elaborate on the usage of batch_size? How is it to be used by SDPA?

The batch_size here refers to the batch parameter in BRGemm (Batch-Reduced Gemm). For the second Gemm in INT8 SDPA, i.e. the product of attention [q_len, kv_len] and value [kv_len, head_size], we do blocking on kv_len to split it into several kv_blocks. So, we can take advantage of parameter batch_size to merge several Gemms of [q_len, kv_block_size] and [kv_block_size, head_size] into one BRGemm with batch_size=kv_block_num, with the reduced dimension of K.

In that case, why not making the kv_block_size larger and invoke a single small gemm instead of making a brgemm call?

To efficiently do the softmax, which is before the gemm of A and V, we create the buffer shape for A as [kv_block_num, q_block_size, kv_block_size] instead of [q_block_size, kv_block_num, kv_block_size]. With this kind of shape, we can do softmax-related calculations on a small contiguous block [q_block_size, kv_block_size], which shows good perf. As the kv_block_num is the outermost dimension, we need to make a brgemm call with batch size greater than 1.

[jgong5]

To efficiently do the softmax, which is before the gemm of A and V, we create the buffer shape for A as [kv_block_num, q_block_size, kv_block_size] instead of [q_block_size, kv_block_num, kv_block_size]. With this kind of shape, we can do softmax-related calculations on a small contiguous block [q_block_size, kv_block_size], which shows good perf. As the kv_block_num is the outermost dimension, we need to make a brgemm call with batch size greater than 1.

Can you elaborate the necessity of [kv_block_num, q_block_size, kv_block_size] instead of [q_block_size, kv_block_num, kv_block_size] to make softmax computation more efficient? Why can't we make the output layout of softmax [q_block_size, kv_len] instead? The reason why I am asking these questions is that I'd like to understand the necessity of supporting "batch-reduced" gemm here. My assumption was that the "batch-reduced" semantics is only necessary for computations like convolutions where the data window to be reduced is non-contiguous by nature.

[Valentine233]

To efficiently do the softmax, which is before the gemm of A and V, we create the buffer shape for A as [kv_block_num, q_block_size, kv_block_size] instead of [q_block_size, kv_block_num, kv_block_size]. With this kind of shape, we can do softmax-related calculations on a small contiguous block [q_block_size, kv_block_size], which shows good perf. As the kv_block_num is the outermost dimension, we need to make a brgemm call with batch size greater than 1.

Can you elaborate the necessity of [kv_block_num, q_block_size, kv_block_size] instead of [q_block_size, kv_block_num, kv_block_size] to make softmax computation more efficient? Why can't we make the output layout of softmax [q_block_size, kv_len] instead? The reason why I am asking these questions is that I'd like to understand the necessity of supporting "batch-reduced" gemm here. My assumption was that the "batch-reduced" semantics is only necessary for computations like convolutions where the data window to be reduced is non-contiguous by nature.

Thanks, Jiong. We assume kv_len is pretty large here, e.g. 9216 from SD. First of all, we need to do kv blocking, because operating on the whole large [q_block_size, kv_len] is memory inefficient. Then we have two options for the layout of attention. If we make attention be [q_block_size, kv_block_num, kv_block_size], for each [q_block_size, kv_block_size] the stride is kv_len. If we use [kv_block_num, q_block_size, kv_block_size] instead, the stride is kv_block_size. The latter one is more efficient for the calculations of softmax, because each block of [q_block_size, kv_block_size] is contiguous, and we don't need to load much unused buffer in memory. This is also proved by the experiment.

[jgong5]

Thanks, Jiong. We assume kv_len is pretty large here, e.g. 9216 from SD. First of all, we need to do kv blocking, because operating on the whole large [q_block_size, kv_len] is memory inefficient. Then we have two options for the layout of attention. If we make attention be [q_block_size, kv_block_num, kv_block_size], for each [q_block_size, kv_block_size] the stride is kv_len. If we use [kv_block_num, q_block_size, kv_block_size] instead, the stride is kv_block_size. The latter one is more efficient for the calculations of softmax, because each block of [q_block_size, kv_block_size] is contiguous, and we don't need to load much unused buffer in memory. This is also proved by the experiment.

Are you referring to the input to the softmax or the output of it, or both?

[Valentine233]

Thanks, Jiong. We assume kv_len is pretty large here, e.g. 9216 from SD. First of all, we need to do kv blocking, because operating on the whole large [q_block_size, kv_len] is memory inefficient. Then we have two options for the layout of attention. If we make attention be [q_block_size, kv_block_num, kv_block_size], for each [q_block_size, kv_block_size] the stride is kv_len. If we use [kv_block_num, q_block_size, kv_block_size] instead, the stride is kv_block_size. The latter one is more efficient for the calculations of softmax, because each block of [q_block_size, kv_block_size] is contiguous, and we don't need to load much unused buffer in memory. This is also proved by the experiment.

Are you referring to the input to the softmax or the output of it, or both?

Both the input and output of softmax.

[jgong5]

Thanks, Jiong. We assume kv_len is pretty large here, e.g. 9216 from SD. First of all, we need to do kv blocking, because operating on the whole large [q_block_size, kv_len] is memory inefficient. Then we have two options for the layout of attention. If we make attention be [q_block_size, kv_block_num, kv_block_size], for each [q_block_size, kv_block_size] the stride is kv_len. If we use [kv_block_num, q_block_size, kv_block_size] instead, the stride is kv_block_size. The latter one is more efficient for the calculations of softmax, because each block of [q_block_size, kv_block_size] is contiguous, and we don't need to load much unused buffer in memory. This is also proved by the experiment.

Are you referring to the input to the softmax or the output of it, or both?

Both the input and output of softmax.

Can we make the layout of input and output different? The input can be blocked and the output can be in 2D and then we don't need batch-reduce semantics.

[Valentine233]

Thanks, Jiong. We assume kv_len is pretty large here, e.g. 9216 from SD. First of all, we need to do kv blocking, because operating on the whole large [q_block_size, kv_len] is memory inefficient. Then we have two options for the layout of attention. If we make attention be [q_block_size, kv_block_num, kv_block_size], for each [q_block_size, kv_block_size] the stride is kv_len. If we use [kv_block_num, q_block_size, kv_block_size] instead, the stride is kv_block_size. The latter one is more efficient for the calculations of softmax, because each block of [q_block_size, kv_block_size] is contiguous, and we don't need to load much unused buffer in memory. This is also proved by the experiment.

Are you referring to the input to the softmax or the output of it, or both?

Both the input and output of softmax.

Can we make the layout of input and output different? The input can be blocked and the output can be in 2D and then we don't need batch-reduce semantics.

Thanks, I will have a try to make the output layout different.

[Valentine233]

As the layout change may impact the kernel perf, I would continue this PR after the perf is confirmed.

[Valentine233]

As the layout change may impact the kernel perf, I would continue this PR after the perf is confirmed.

@jgong5 Hi, I have confirmed the perf, and remove the support of batch_size in brgemm. Please help review again, thanks!

[Valentine233]

@pytorchbot merge

[pytorchmergebot]

Merge failed

Reason: Approvers from one of the following sets are needed:

• superuser (pytorch/metamates)
• Core Reviewers (mruberry, lezcano, Skylion007, ngimel, peterbell10, ...)
• Core Maintainers (soumith, gchanan, ezyang, dzhulgakov, malfet, ...)

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

[Valentine233]

@peterbell10 @ezyang Could you help take a look? Thanks!

[ezyang]

okey dokey but you sure you don't want tests?

[ezyang]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here