[jiqing-feng]

What

Fixes a regression from #49548 (Quantization cleanpup). That PR replaced GptOssMLP with Mxfp4GptOssMLP. The new class dropped @use_kernel_forward_from_hub("MegaBlocksMoeMLP"), so use_kernels=True stopped swapping in the megablocks MoE kernel for MXFP4 GPT-OSS:

• XPU silently falls back to the slower Triton path.
• CPU fails on the first forward: Triton has no CPU driver.

This PR adds the decorator back to Mxfp4GptOssMLP. It also adds a slow test that loads openai/gpt-oss-20b with use_kernels=True, checks that every MLP runs the megablocks kernel, and checks the generated text. The test is skipped on devices other than CPU and XPU.

use_kernels=False is unchanged, since kernelize only runs when use_kernels=True. CUDA is also unchanged: use_kernels=True dequantizes there, so Mxfp4GptOssMLP is never created.

Perf

openai/gpt-oss-20b on one Intel Arc Pro B60 (BMG), bf16, greedy decoding:

	Speedup vs main
1st token	4.27x
2nd+ token	1.90x

import time

import torch

from transformers import AutoModelForCausalLM, AutoTokenizer


model_id = "openai/gpt-oss-20b"
device = "xpu"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map=device, dtype=torch.bfloat16, use_kernels=True)
print("MLP forward:", model.model.layers[0].mlp.forward.__qualname__)

inputs = tokenizer(["The capital of France is"], return_tensors="pt").to(device)


def timed_generate(new_tokens):
    torch.xpu.synchronize()
    start = time.perf_counter()
    model.generate(**inputs, max_new_tokens=new_tokens, min_new_tokens=new_tokens, do_sample=False)
    torch.xpu.synchronize()
    return time.perf_counter() - start


timed_generate(8)  # warmup
first = min(timed_generate(1) for _ in range(3))
total = min(timed_generate(64) for _ in range(3))
print(f"1st token: {first * 1000:.2f} ms, 2nd+ token: {(total - first) / 63 * 1000:.2f} ms")

Tests

RUN_SLOW=1 pytest tests/quantization/mxfp4/test_mxfp4.py::Mxfp4KernelsModelTest

The new test fails on main and passes with this PR, on both XPU and CPU.

@ArthurZucker could you please review? Thanks!

[jiqing-feng]

Also bumped the outdated flash-attn2 kernel version in benchmark_v2 (2d28f92).

is_fa2_or_kernel_available() still pinned kernels-community/flash-attn2 to version=1, while the model now loads v4 through FLASH_ATTN_KERNEL_VERSIONS. v1 has no XPU build, so flash_attention_2 benchmarks on XPU silently fell back to SDPA. The check now uses the same version as the model.

[SunMarc]

Thanks a lot !

[github-actions]

Thank you for your contribution 🤗!

CI Security Gate — automatic approval blocked

This PR was not automatically approved for CI because the security gate failed.

Possible reasons:

• The PR touches 50 or more files — only PRs with fewer than 50 changed files are automatically approved
• A changed file is outside the allowed directories (src/, tests/, docs/, utils/), has a disallowed extension (only .py, .txt, .md permitted outside tests/ and docs/), or is not .md/.yml inside docs/ — this covers files the PR deletes or renames, not only the ones it edits
• A new security issue was detected in the changed Python files (Bandit check) — a high-severity one, shell execution, exec/eval, a URL fetch, or disabled TLS verification
• The PR touches a path this repository protects from untrusted PRs, such as the file that decides who reviews it — a maintainer must make that change in a separate PR

See the workflow run for the exact violations.

A maintainer can review and manually approve CI if a finding is a false positive.

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: mxfp4

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.