[MXFP4] Support for linear layers + compressed-tensors integration

[dsikka]

Purpose

• Update scheme name from a16 to a4 to be agnostic for activation quantization
• Extend flashinfer_mm_fp4 and flashinfer_scaled_fp4_mm to take in a configurable block_size and boolean flag use_nvfp4 to enable the mxfp4 linear forward pass based on https://github.com/flashinfer-ai/flashinfer/blob/393e83ea8497ff9fb9ad61e170b89797a6b682a3/flashinfer/gemm/gemm_base.py#L5511
• Follow the pattern of linear kernels to allow selection between marlin and flashinfer cutlass based on env overrides and platform

Test Plan

• Added basic smoke tests
• LM Eval validation

LM-Eval

lm_eval \
  --model vllm \
  --model_args pretrained="nm-testing/Meta-Llama-3-8B-Instruct-MXFP4-GPTQ",dtype=auto,add_bos_token=True,max_model_len=4096,max_gen_toks=1024,tensor_parallel_size=1 \
  --tasks gsm8k_cot_llama \
  --fewshot_as_multiturn \
  --apply_chat_template \
  --num_fewshot 8 \
  --batch_size auto

flashinfer

|     Tasks     |Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|---------------|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k_cot_llama|      4|flexible-extract|     8|exact_match|↑  |0.6892|±  |0.0127|
|               |       |strict-match    |     8|exact_match|↑  |0.6846|±  |0.0128|

marlin (VLLM_MXFP4_USE_MARLIN=1)

|     Tasks     |Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|---------------|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k_cot_llama|      4|flexible-extract|     8|exact_match|↑  |0.7604|±  |0.0118|
|               |       |strict-match    |     8|exact_match|↑  |0.7551|±  |0.0118|

dense (meta-llama/Meta-Llama-3-8B-Instruct)

|     Tasks     |Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|---------------|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k_cot_llama|      4|flexible-extract|     8|exact_match|↑  |0.7998|±  | 0.011|
|               |       |strict-match    |     8|exact_match|↑  |0.7991|±  | 0.011|

[gemini-code-assist]

Code Review

This pull request renames the MXFP4 quantization scheme to CompressedTensorsW4A4Mxfp4 and introduces support for true W4A4 quantization using FlashInfer on SM100+ devices. The changes include new utility functions for FlashInfer FP4 operations and logic to handle activation quantization. A critical issue was identified in the weight scale processing where swizzle_mxfp4_scales might cause a RuntimeError during reshaping if the output feature size is not a multiple of 128 due to internal padding.

[gemini-code-assist]

The swizzle_mxfp4_scales function pads the N dimension to the nearest multiple of 128. If N (the output feature size) is not a multiple of 128, the total number of elements in the swizzled tensor will be padded_N * padded_scale_cols, which is not necessarily divisible by N. This will cause a RuntimeError during the .reshape(N, -1) call. Even if it were divisible, the resulting 2D tensor would have misaligned scale data because of the padding introduced during swizzling. You should ensure that the scale tensor's shape is compatible with what the FlashInfer kernel expects, which likely involves keeping the padded dimensions or ensuring the kernel handles the original N correctly with the swizzled layout.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[dsikka]

@yewentao256 @mgoin

[yewentao256]

Thanks for the work!

[kylesayrs]

Support looks good, was able to verify locally

[mgoin]

LGTM, thanks!

[mergify]

Hi @dsikka, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[yewentao256]

Sorry to block for a while, please take a look at my previous comment

[dsikka]

Sorry to block for a while, please take a look at my previous comment

Sorry to block for a while, please take a look at my previous comment

@yewentao256 Please take a look at the latest commits. This has been addressed to use the padded_N for the reshape

Dismiss request change as already solved

[yewentao256]

Thanks for the work! A small update