[model, quant] fix: Preserve quantized expert scale dimensions

[yaoyu-33]

Problem

AutoBridge.export_hf_weights_quant() can silently drop a valid singleton block-grid dimension from MoE expert scale tensors. The trigger is a selected expert weight whose block quantizer returns a scale shaped [block_rows, block_cols, 1] with either block-grid axis equal to one, which is valid when an expert shard dimension equals the caller-selected block size.

For grouped EP=2 export, two local scales shaped [1, 2, 1] are staged as [2, 1, 2, 1], but the current helper emits [2, 2, 1]. Grouped accumulation then preserves that malformed geometry in the Hugging Face *_scale_inv tensor. Per-expert and EP=1 paths can lose the same quantizer-owned axis.

This affects the quantize-before-gather workflow introduced by #2737. The adjacent fix in #5731 established that ordinary EP gathering must remove only its own staging axis, but the quantization-specific scale helper retained the unqualified squeeze.

Root cause and fix

gather_from_ep_ranks_scale() stages tensors with a leading EP-owned dimension, then calls unqualified squeeze() and re-adds only a trailing dimension. That removes every singleton axis, including valid block-grid axes owned by the quantizer.

Use squeeze(0) so only a size-one staging axis is removed. For grouped EP greater than one, the leading EP axis remains available to _accumulate_grouped_export() as designed.

Regression evidence

The focused test executes the real helper with a deterministic two-rank gather and a [1, 2, 1] local scale.

Fail before:

uv run --active --no-sync --no-project python -m pytest \
  tests/unit_tests/models/test_param_mapping.py::test_ep_scale_gather_preserves_singleton_block_grid_dimensions -q
1 failed — expected (2, 1, 2, 1), got (2, 2, 1)

Pass after, unchanged test:

uv run --active --no-sync --no-project python -m pytest \
  tests/unit_tests/models/test_param_mapping.py::test_ep_scale_gather_preserves_singleton_block_grid_dimensions -q
1 passed

Adjacent validation:

uv run --active --no-sync --no-project python -m pytest \
  tests/unit_tests/models/test_param_mapping.py::test_ep_gather_preserves_expert_scale_singleton_dimensions \
  tests/unit_tests/models/test_param_mapping.py::TestFusedExpertMapping \
  tests/unit_tests/models/test_param_mapping.py::TestFusedGatedExpertMapping -q
7 passed

uv run --active --no-sync --no-project python -m pytest \
  tests/unit_tests/quantization/test_quant_bridge.py::test_quantized_stream_accumulates_grouped_experts -q
1 passed

git diff --check
uv run --no-sync pre-commit run --all-files

Both repository checks passed. The focused CPU environment used the repository's pinned MCore revision and disabled only unrelated optional CUDA/Triton imports.

Scope

This changes one internal tensor-axis operation and adds one focused CPU regression test. It does not change conversion APIs, quantization formats, model registrations, dependencies, workflows, the lockfile, or Megatron-Core.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yaoyu-33]

/ok to test a0886f0

[claude]

Light review — LGTM with one doc nit

The fix is correct. .squeeze() collapsed every singleton dimension of the FP8 scale tensor, destroying legitimate block-grid dimensions; .squeeze(0) removes only the staging dimension added by unsqueeze(0), matching the sibling gather_from_ep_ranks. The added unit test correctly locks in the (2, 1, 2, 1) shape and value ordering.

• Nit: the gather_from_ep_ranks_scale docstring (line 827) is now stale — it still claims the method adds an extra unsqueeze before returning, which this PR removed. Worth updating so the two now-identical staging paths are described accurately.

Suggested test cases

• test_ep_scale_gather_preserves_singleton_block_grid_dimensions (added in this PR)
• test_ep_gather_preserves_expert_scale_singleton_dimensions (existing sibling, confirm still passes)

No perf tests impacted.

[yaoyu-33]

/ok to test 5374f9f