Fix MoE expert decompression for non-32-divisible bit widths

[KKothuri]

What does this PR do?

DecompressExperts.convert (in integrations/compressed_tensors.py) rebuilds each MoE expert's unpacked input dimension as packed_cols * (32 // num_bits). compressed-tensors uses dense bit-packing, so this reconstruction is only exact when num_bits divides 32 (i.e. 4-bit and 8-bit). For 3/5/6/7-bit weights it under-counts.

Concretely, a 3-bit down_proj stored as weight_packed of shape (2048, 48) reconstructs an input dim of 48 * (32 // 3) = 48 * 10 = 480, whereas the true dim is 512. The experts then decompress to the wrong shape and loading fails with:

RuntimeError: contraction dimension of mat_a and mat_b must match

from torch._grouped_mm. 4-bit and 8-bit checkpoints are unaffected because 8 and 4 divide 32 evenly.

Fix

The checkpoint already stores an exact per-expert weight_shape ((out-dim, in-dim)), and it is already routed into convert via the .weight_shape$ source pattern in quantizer_compressed_tensors.py. This PR uses that stored shape, and falls back to the packed-tensor reconstruction only when it is missing/empty (e.g. left empty on most ranks under TP/EP sharding — the case the previous code was written for).

Testing

Verified on a 3-bit (W3A16) AWQ-quantized MoE checkpoint (GlmMoeDsaForCausalLM):

• Before: loading crashes with the _grouped_mm contraction error above.
• After: loads successfully; perplexity 1.0002 vs. 1.0000 for the bf16 base on the same text.

4-bit and 8-bit checkpoints load identically before and after (the stored shape matches the old reconstruction for those bit widths).

AI assistance was used for this change; every line was reviewed and the tests above were run.

Before submitting

• Did you read the contributor guideline?
• Did you make sure to update the documentation with your changes? (n/a — bug fix, no API change)
• Did you write any new necessary tests? (happy to add a regression test if maintainers prefer)

Who can review?

Quantization / compressed-tensors: @SunMarc @MekkCyber

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 29320067190:2
Result: success | Jobs: 15 | Tests: 169,249 | Failures: 0 | Duration: 13h 32m

[Rocketknight1]

Not sure how significant it is but could be real, cc @SunMarc for quants maybe?

[ErenAta16]

The fallback ordering is right: prefer the checkpoint's stored weight_shape when present, since it's exact for any bit width, and only reconstruct from pack_factor when it's missing or empty, which is the TP/EP-sharded case where weight_shape isn't populated on most ranks. That's a real gap in the previous unconditional-reconstruction path: pack_factor = 32 // num_bits only recovers the true in-dim exactly when num_bits divides 32, so 3/5/6/7-bit packing was silently under-counting the unpacked dimension. The fix keeps the original reconstruction as a fallback rather than replacing it outright, so the sharded-checkpoint case that logic was originally written for still works.