[ Kernel ] AWQ Fused MoE

[robertgshaw2-redhat]

SUMMARY:

• Picks up No executable after building vllm from source with CPU support #6259 2761 to support AWQ MoE models via a fused kernel, after the refactor I did in [ Misc ] Refactor MoE to isolate Fp8 From Mixtral #5970. The kernels for this PR were developed by @chu-tianxiang
• Adds AWQMoEmethod, supporting loading AutoAWQ models
• Refactors FusedMoE.weight_loader, to enable loading AWQ models, which have transposed weights (input_dim, output_dim) on disk. Fp16 and Fp8 models have share (input_dim, output_dim). This required more complex logic for handling indexing in the TP case and MergedColumn case
• Refactors expert_params_mapping, which was overfit to fp16 and fp8 checkpoints. This required renaming the scale parameters in fp8 which to better match the state dicts that we create in autofp8, limiting the amount of remapping we need to do in the model files
• Updates layers to use fused_topk/grouped_topk and fused_experts, rather than calling fused_moe directly, such that we can reuse the logic across fp16, fp8, and awq

Latency Benchmarking with 2xA100 80GB:

Mixtral Fused MoE with AWQ:
> python benchmarks/benchmark_latency.py --model casperhansen/mixtral-instruct-awq --tensor-parallel-size 2 --input-len 512 --output-len 128 --max-model-len 1024 --quantization awq --batch-size 1
Avg latency: 1.3650233147976298 seconds
10% percentile latency: 1.3638953405432404 seconds
25% percentile latency: 1.3643284866120666 seconds
50% percentile latency: 1.3648834400810301 seconds
75% percentile latency: 1.3656865148805082 seconds
90% percentile latency: 1.3661401799879969 seconds
99% percentile latency: 1.366987234679982 seconds

Mixtral Fp16:
> python benchmarks/benchmark_latency.py --model mistralai/Mixtral-8x7B-Instruct-v0.1 --tensor-parallel-size 2 --input-len 512 --output-len 128 --max-model-len 1024 --batch-size 1
Avg latency: 1.470517885716011 seconds
10% percentile latency: 1.4687837794423104 seconds
25% percentile latency: 1.4693051476497203 seconds
50% percentile latency: 1.4700921354815364 seconds
75% percentile latency: 1.4719352358952165 seconds
90% percentile latency: 1.4728145461529494 seconds
99% percentile latency: 1.4740593627933414 seconds

Followup work:

• Deepseek is not supported yet, illegal memory access with TechxGenus/DeepSeek-V2-Lite-Chat-AWQ

[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 trigger fastcheck CI to run, which consists only a small and essential subset of tests to quickly catch errors with the flexibility to run extra individual tests on top (you can do this by unblocking test steps in the Buildkite run).

Full CI run is still required to merge this PR so once the PR is ready to go, please make sure to run it. If you need all test signals in between PR commits, you can trigger full CI as well.

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

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

🚀

[AlpinDale]

Thanks for doing this.

[dsikka]

/ready

[mgoin]

This looks to be shaping up well! I appreciate the refactoring of fused_moe. It would be more clear if it was separated out first, but I think we can land it within this PR if needed.

[mgoin]

This seems a bit high and it is worth commenting how it was calibrated (what model, benchmark, GPU used)

[dsikka]

@robertgshaw2-neuralmagic do we know why this is 1024 specifically?

[dsikka]

@mgoin can't resolve but addressed all but one comment

[dsikka]

Latency Benchmarking with Two 82 GB A100s:

Mixtral Fused MoE with AWQ:
Avg latency: 1.3650233147976298 seconds
10% percentile latency: 1.3638953405432404 seconds
25% percentile latency: 1.3643284866120666 seconds
50% percentile latency: 1.3648834400810301 seconds
75% percentile latency: 1.3656865148805082 seconds
90% percentile latency: 1.3661401799879969 seconds
99% percentile latency: 1.366987234679982 seconds

Mixtral Fp16:
Avg latency: 1.470517885716011 seconds
10% percentile latency: 1.4687837794423104 seconds
25% percentile latency: 1.4693051476497203 seconds
50% percentile latency: 1.4700921354815364 seconds
75% percentile latency: 1.4719352358952165 seconds
90% percentile latency: 1.4728145461529494 seconds
99% percentile latency: 1.4740593627933414 seconds

[mgoin]

LGTM! We should add a model test, but considering there isn't a small mixtral awq to test, maybe just including a lm-eval large test would be sufficient.

[dsikka]

Note: Splitting this PR into two separate PRs.
PR 1/2: #7334

[github-actions]

This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this pull request should remain open. Thank you!

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @robertgshaw2-neuralmagic.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork