[MISC] Fix Tensor Parallelism for Quantized Mamba Models with n_groups=1

[vadiklyutiy]

Summary

Enable tensor parallelism (TP > 1) for quantized hybrid Mamba models (e.g., Falcon-H1R-7B with FP8) when n_groups=1.

Root Cause

Custom weight loaders for group replication were only implemented for non-quantized layers. This PR extends support to quantized layers by leveraging the weight_loader property on ModelWeightParameter (extends BasevLLMParameter).

Test

vllm serve tiiuae/Falcon-H1R-7B -tp 2 --quantization fp8

Previously failed, now works

Validation (lm_eval)

lm_eval --model local-chat-completions \
  --model_args model=tiiuae/Falcon-H1R-7B,base_url=http://localhost:8000/v1/chat/completions,num_concurrent=250 \
  --tasks gsm8k --apply_chat_template --num_fewshot 5

Results with TP=2 (this PR):

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.3283	±	0.0129
		strict-match	5	exact_match	↑	0.0963	±	0.0081

Results with TP=1 (baseline):

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.3154	±	0.0128
		strict-match	5	exact_match	↑	0.1031	±	0.0084

Results are consistent within statistical error margins, confirming correctness.

Related Issues

• Fixes: [Bug]: Quantization In MambaMixer2 Not Supported when Tensor Parallel is enabled #14618
• Related PRs: Enforce that TP > 1 is not supported for Mamba2 if Quantization is Enabled. #14617, [Mamba] Support TP>1 with quantization for mamba2 mixer in case n_groups % tp_size == 0 #24593

[gemini-code-assist]

Code Review

This pull request enables tensor parallelism for quantized Mamba models with n_groups=1 by extending custom weight loader support to quantized layers. The changes are well-structured and address the issue effectively. The removal of the restrictive assertion and the new logic for applying the mamba_v2_sharded_weight_loader to both quantized and non-quantized weights are correct. The implementation correctly distinguishes between BasevLLMParameter subclasses and standard torch.nn.Parameter to set the weight loader. The changes look good and improve support for quantized Mamba models.

[vadiklyutiy]

cc @tomeras91

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{Bugbot Autofix is OFF. To automatically fix reported issues with Cloud Agents, enable Autofix in the Cursor dashboard.}

Comment @cursor review or bugbot run to trigger another review on this PR

[tomeras91]

Nice!

[vadiklyutiy]

want to cover with wider testing. Add ready for that

[vadiklyutiy]

@tomeras91 @tdoublep
CI passed, review comment fixed, failed Falcon-H1R-7B is fixed
Could take a look?

[tomeras91]

LGTM
As a final validation, can you double check that NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 still works and produces similar GSM8K results on main and on this PR?

[vadiklyutiy]

LGTM As a final validation, can you double check that NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 still works and produces similar GSM8K results on main and on this PR?

For NVIDIA-Nemotron-3-Nano-30B-A3B-FP8

Baseline

| Tasks | Version | Filter           | n-shot | Metric      |   | Value  |   | Stderr |
|-------|--------:|------------------|-------:|-------------|---|-------:|---|-------:|
| gsm8k |       3 | flexible-extract |      5 | exact_match | ↑ | 0.3169 | ± | 0.0128 |
|       |         | strict-match     |      5 | exact_match | ↑ | 0.4792 | ± | 0.0138 |

With PR

| Tasks | Version | Filter           | n-shot | Metric      |   | Value  |   | Stderr |
|-------|--------:|------------------|-------:|-------------|---|-------:|---|-------:|
| gsm8k |       3 | flexible-extract |      5 | exact_match | ↑ | 0.3017 | ± | 0.0126 |
|       |         | strict-match     |      5 | exact_match | ↑ | 0.4693 | ± | 0.0137 |