[ROCm] Apply FP8 weights padding to values not divisible by 512 bytes on ROCm #13231
Conversation
Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
… strides Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add 🚀 |
| and (weight.stride(-2) * weight.element_size()) % 512 == 0): | ||
| num_pad = 256 // weight.element_size() | ||
| weight = F.pad(weight, (0, num_pad), "constant", 0)[..., :-num_pad] | ||
| torch.cuda.empty_cache() |
There was a problem hiding this comment.
is empty_cache really necessary here?
There was a problem hiding this comment.
Without it there is a possibility of having double the memory allocated, depending on the allocator behavior
|
Thanks for contributing! 🙏🏻 |
Co-authored-by: Michael Goin <mgoin64@gmail.com> Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
6106325 to
f3da192
Compare
| @@ -477,7 +477,7 @@ def w8a8_block_fp8_matmul( | |||
| assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1] | |||
| M = A.numel() // A.shape[-1] | |||
|
|
|||
| assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2 | |||
| assert B.ndim == 2 and Bs.ndim == 2 | |||
There was a problem hiding this comment.
Are we sure this is okay?
There was a problem hiding this comment.
The kernel works just fine with a padded non-contiguous tensor. And in any scenario other than with padding it should be contiguous already, so no existing workflow is supposed to break.
There was a problem hiding this comment.
One other option is just to call weight.contiguous() after we pad it in process_weights_after_loading?
There was a problem hiding this comment.
This would remove the padding, reverting the F.pad action
There was a problem hiding this comment.
sorry, that was a dumb comment by me
There was a problem hiding this comment.
@gshtras I agree contiguous here was overly strict. But should we still check that the stride is 1 for the last dimension? B.stride(-1) == 1?
|
Nice work! |
Performance improvement for ROCm working around the hardware limitation.
One requirement for this is for w8a8_block_fp8_matmul to support the non-contiguous weights, which it seems to already do, so the leftover assertion is obsolete.
While maintaining the same correctness, this shows the following latency improvement on ROCm:
amd/Llama-3.1-8B-Instruct-FP8-KV bs=64 in=512 out=512 tp=1:
5.95s -> 5.7s (4%)
amd/Llama-3.1-70B-Instruct-FP8-KV bs=64 4in=512 out=512 tp=1:
25.6s -> 24.3s (5%)
deepseek-ai/DeepSeek-R1 bs=64 in=256 out=256 tp=8:
26.1s -> 24.9 (5%)