[BUGFIX] Fix accuracy regression for NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4 with TP>1

[vadiklyutiy]

Purpose

Fix accuracy regression for NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4 with TP>1.

Alternative to #34151.

What happened

PR #33257 (a372f3f40) was a squash-merge of two logical changes:

1. Adding TP support for quantized Mamba models with n_groups=1 (e.g., Falcon-H1R-7B with FP8) — correct and needed.
2. Unifying both code paths (n_groups % tp_size == 0 and n_groups == 1) to always use ColumnParallelLinear with a custom mamba_v2_sharded_weight_loader — incorrect for quantized models.

The second change (d5d6d0b88) replaced MergedColumnParallelLinear with ColumnParallelLinear + custom weight loader for ALL cases, but only overrode the weight_loader on in_proj.weight. For quantized models (NVFP4, FP8), there are also scale parameters (weight_scale, input_scale, weight_scale_2) that still use ColumnParallelLinear's default contiguous sharding.

The mamba in_proj has a composite weight layout [gate, intermediate, B_groups, C_groups, dt_heads]. The custom mamba_loader shards each component separately, but the scale parameters do simple contiguous sharding, causing weight-scale misalignment — the dequantization scale at row N corresponds to a different model component than the weight at row N.

MergedColumnParallelLinear handles this correctly because it knows the per-component output sizes and shards ALL parameters (weights and scales) accordingly.

Fix

Revert d5d6d0b88 ("Unify MambaMixer2 TP sharding to use custom weight loader"), the last commit from #33257. This restores MergedColumnParallelLinear for the n_groups % tp_size == 0 case while preserving the n_groups == 1 quantized TP support.

Test Results

Model: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4, TP=8, gsm8k 5-shot (limit=500), 8×B200

After fix (3 runs):

Run	flexible-extract	strict-match	Status
1	0.430	0.760	GOOD
2	0.416	0.746	GOOD
3	0.424	0.730	GOOD

3/3 runs correct.

[gemini-code-assist]

Code Review

This pull request addresses an accuracy regression for quantized Mamba models with tensor parallelism greater than one. The root cause was a previous change that incorrectly unified weight loading logic, leading to a misalignment between weights and their quantization scales. This fix correctly reverts the problematic part of that change by restoring the use of MergedColumnParallelLinear when the number of groups is divisible by the tensor-parallel size. This is the right approach, as MergedColumnParallelLinear correctly handles the sharding of all parameters, including quantization scales, resolving the issue. The special handling for n_groups=1, which was already correct, is preserved. The change is well-justified, clearly explained, and effectively fixes the bug.