Add support to modelopt quantization of Mixtral model

[yueshen2016]

Add vllm support to modelopt fp8 quantization of Mixtral8x7b model

[github-actions]

👋 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 fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

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 ready label to the PR or enable auto-merge.

🚀

[mgoin]

This seems to add support for kv cache quantization to mixtral. However there is no MoE layer registered for modelopt to use, so what happens when the MoE layers are quantized?

vllm/vllm/model_executor/layers/quantization/modelopt.py

Lines 222 to 231 in 463bbb1

	def get_quant_method(self, layer: torch.nn.Module,
	prefix: str) -> Optional["QuantizeMethodBase"]:
	from vllm.attention.layer import Attention # Avoid circular import
	if isinstance(layer, LinearBase):
	if is_layer_skipped(prefix, self.exclude_modules):
	return UnquantizedLinearMethod()
	return ModelOptNvFp4LinearMethod(self)
	elif isinstance(layer, Attention):
	return ModelOptFp8KVCacheMethod(self)
	return None

[yueshen2016]

This seems to add support for kv cache quantization to mixtral. However there is no MoE layer registered for modelopt to use, so what happens when the MoE layers are quantized?

vllm/vllm/model_executor/layers/quantization/modelopt.py

Lines 222 to 231 in 463bbb1

	def get_quant_method(self, layer: torch.nn.Module,
	prefix: str) -> Optional["QuantizeMethodBase"]:
	from vllm.attention.layer import Attention # Avoid circular import
	if isinstance(layer, LinearBase):
	if is_layer_skipped(prefix, self.exclude_modules):
	return UnquantizedLinearMethod()
	return ModelOptNvFp4LinearMethod(self)
	elif isinstance(layer, Attention):
	return ModelOptFp8KVCacheMethod(self)
	return None

Hi @mgoin , the purpose of this PR is to do key remapping:
model.layers.X.self_attn.k_proj.k_scale -> model.layers.X.self_attn.attn.k_scale
model.layers.X.self_attn.v_proj.v_scale -> model.layers.X.self_attn.attn.v_scale

As for the moe layer, for each expert, a input_scale and weight_scale is added to it.

[pavanimajety]

@mgoin This is mixtral_quant.py which doesn't use FusedMoE layer. It is simply an MLP layer: here

And the switch between architectures happens here

edit: just fixed the links to vllm-project

[pavanimajety]

LGTM.
@mgoin This is a python change, don't see why the docker is failing. Can we restart the CI?

[pavanimajety]

@simon-mo / @mgoin Can you PTAL?

[mgoin]

I think this PR was made before we made changes to the docker image. Can you please merge with main?