[Bug Fix] Illegal memory access, FP8 Llama 3.1 405b

[LucasWilkinson]

Co-authored-by: Tyler Michael Smith tyler@neuralmagic.com

Should resolve: #6767, #6689 (comment) etc.

Currently the Sm90ColOrScalarBroadcast could read off the end of the scales buffer when the batch size is not a multiple of gemm tile size. Added a predicate tensor and used copy_if to prevent this.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which consists a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of default ones by unblocking the steps in your fast-check build on Buildkite UI.

Once the PR is approved and ready to go, please make sure to run full CI as it is required to merge (or just use auto-merge).

To run full CI, you can do one of these:

• Comment /ready on the PR
• Add ready label to the PR
• Enable auto-merge.

🚀

[LucasWilkinson]

Whoever merges this can you please add Tyler Michael Smith (tlrmchlsmth) tyler@neuralmagic.com as a co-author? Thanks!

[LucasWilkinson]

/ready

[comaniac]

Thanks for the fix! Does this problem reproducible on your end now? If so can we add a unit test? It not then I'm good with CI green.

[mgoin]

Amazing effort!!

[LucasWilkinson]

@comaniac

Does this problem reproducible on your end now?

Unfortunately only when we run Llama 405b server on 8xH100 and hit it with a bunch of requests (granted fairly consistently), but not something I think we can get up and running on the CI (although I'm not that spun up on it).

[comaniac]

LGTM!