[Feature][Quantization]Support NZ for A5 MX quantization

[lijiahang226]

What this PR does / why we need it?

This PR adds FRACTAL_NZ support for A5 MX quantization paths (W8A8 MXFP8 and W4A4 MXFP4).

• Extends maybe_trans_nz to forward customize_dtype and input_dtype to npu_format_cast, so FP8/FP4 MX weights can be converted to FRACTAL_NZ with the correct dtype hints.
• Updates W8A8 MXFP8 linear/MoE weight processing:
  • When NZ is enabled, transposed weights are made contiguous and converted to FRACTAL_NZ; weight scales are kept contiguous.
  • When NZ is disabled, the previous non-contiguous transpose layout is preserved.
• Updates W4A4 MXFP4 linear weight processing with the same NZ-enabled/disabled behavior.
• Avoids double NZ conversion for layers that will be processed by fused SFA preprocess (MLAPO/PROLOG_V3): such layers are marked at init time and quant methods skip their own NZ conversion.
• Refactors SFA preprocess type resolution so the fused preprocess decision can be queried before submodule weight processing.

Why needed: A5 requires FRACTAL_NZ for these MX quantized weights. The previous code always kept the non-contiguous transpose layout, which is not compatible with NZ-enabled execution on A5. This change also keeps the old behavior when NZ is disabled, avoiding regressions.

Does this PR introduce any user-facing change?

No new user-facing API. Behavior is affected only when weight_nz_mode enables NZ on supported hardware; with NZ disabled, the existing MX layout behavior is preserved.

How was this patch tested?

• Added/updated unit tests:

  • tests/ut/test_utils.py
  • tests/ut/quantization/methods/test_w4a4_mxfp4.py
  • tests/ut/quantization/methods/test_w8a8_mxfp8.py
  • tests/ut/ops/test_mla.py
  • tests/ut/attention/test_sfa_v1.py

• vLLM main: vllm-project/vllm@ba07e4a

[github-actions]

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

• A PR should do only one thing, smaller PRs enable faster reviews.
• Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
• Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces FRACTAL_NZ support for A5 MX quantization paths, specifically for W8A8 MXFP8 and W4A4 MXFP4 models. The changes ensure that weights are correctly converted to the FRACTAL_NZ format required by A5 hardware, while maintaining backward compatibility by preserving the previous layout when NZ is disabled. Additionally, the PR refactors SFA preprocess type resolution and implements mechanisms to avoid redundant conversions for layers already handled by fused SFA processes.

Highlights

• FRACTAL_NZ Support: Added FRACTAL_NZ support for A5 MX quantization (W8A8 MXFP8 and W4A4 MXFP4) to ensure hardware compatibility.
• Weight Processing Updates: Updated linear and MoE weight processing to handle FRACTAL_NZ conversion, while preserving existing layout behavior when NZ is disabled.
• Fused SFA Optimization: Refactored SFA preprocess type resolution and added logic to prevent redundant NZ conversions for layers managed by fused SFA (MLAPO/PROLOG_V3).
• Utility Enhancement: Extended the maybe_trans_nz utility to support customize_dtype and input_dtype for correct conversion hints.

New Features

🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

Suggested PR Title:

[Attention][Ops][Feature] Refactor fused preprocess resolution and support NZ-disabled layouts in quantization

Suggested PR Summary:

### What this PR does / why we need it?

This PR refactors the fused preprocess type resolution and improves weight transformation handling under different quantization schemes (W4A4 MXFP4 and W8A8 MXFP8) when FRACTAL_NZ format conversion is disabled or managed externally.

Key changes include:
- **Fused Preprocess Refactoring**: Extracted `_fused_preprocess_type()` in `sfa_v1.py` to determine the active preprocess type (PROLOG_V3 or MLAPO) and avoid redundant checks.
- **Managed Preprocess Marking**: In `mla.py`, layers managed by fused preprocess are marked with `_fused_preprocess_managed = True` to prevent quantization methods from performing duplicate NZ conversions.
- **NZ-Disabled Layout Support**: Updated W4A4 MXFP4 and W8A8 MXFP8 quantization methods to check `_should_trans_nz`. If NZ conversion is disabled, they preserve the pre-NZ non-contiguous transpose layout.
- **Utility Enhancements**: Updated `maybe_trans_nz` to support passing `customize_dtype` and `input_dtype` to `npu_format_cast`.
- **Robustness**: Replaced `.view` with `.reshape` in MXFP8 MoE weight restoration to avoid runtime errors on non-contiguous tensors.
- **Testing**: Added extensive unit tests covering the new NZ-disabled layouts, utility parameters, and preprocess marking.

### Does this PR introduce _any_ user-facing change?

No.

### How was this patch tested?

Tested via newly added and updated unit tests in:
- `tests/ut/attention/test_sfa_v1.py`
- `tests/ut/ops/test_mla.py`
- `tests/ut/quantization/methods/test_w4a4_mxfp4.py`
- `tests/ut/quantization/methods/test_w8a8_mxfp8.py`
- `tests/ut/test_utils.py`

I have no further feedback to provide as the changes are well-implemented and thoroughly tested.