Fix BMM style MoE export in fp8_pc_pt recipe

[Edwardf0t1]

What does this PR do?

Type of change: Bug fix

Overview: The Llama-4-Scout-17B-16E-Instruct model uses Llama4TextExperts, which stores expert weights in a BMM (batch matrix multiply) layout: (num_experts, input_dim, output_dim). This is different from standard MoE models. The FP8_PC_PT (FP8 per-channel per-token) quantization code didn't handle this layout properly, causing shape mismatches.

Usage

python3 hf_ptq.py --pyt_ckpt_path /home/scratch.omniml_data_2/models/Llama-4-Scout-17B-16E-Instruct --qformat fp8_pc_pt --export_path /home/scratch.omniml_data_2/zhiyuc/checkpoints/llama4-scout-fp8_pc_pt --trust_remote_code

Testing

Before your PR is "Ready for review"

• Make sure you read and follow Contributor guidelines and your commits are signed.
• Is this change backward compatible?: Yes
• Did you write any new necessary tests?: No
• Did you add or update any necessary documentation?: No
• Did you update Changelog?: No

Additional Information

[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.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 74.45%. Comparing base (a703e22) to head (5078573).
⚠️ Report is 5 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #515   +/-   ##
=======================================
  Coverage   74.45%   74.45%           
=======================================
  Files         182      182           
  Lines       18250    18250           
=======================================
  Hits        13588    13588           
  Misses       4662     4662           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[cjluo-nv]

Could you share what models have you tested?

[Edwardf0t1]

Could you share what models have you tested?

Llama4 scout is tested

[meenchen]

LGTM. Will defer to @cjluo-nv for approval.

[cjluo-nv]

can we just do assert here instead for simplification? Same for line 794