[Bug]: Cutlass NVFP4 MoE produces the wrong results

[nzmora-nvidia]

System Info

Blackwell

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

pytest tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_trtllm_moe.py::test_trtllm_fused_moe_nvfp4[8-otype1-2-2-2688-1856-1]

Expected behavior

No errors with reference

actual behavior

There are errors compared to the reference. It can be seen that the error has a structure (about every ~32 elements).

diff.abs()>=2e-1=tensor([[
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, True, True, True, True, True, True, True,
True, False, True, True, True, True, True, False, True, True,
True, True, False, True, True, True, False, True, False, True,
True, False, True, True, False, False, False, False, True, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, True, True, True, True,
True, True, True, False, True, True, True, True, False, True,
True, True, True, True, True, False, True, True, False, True,
True, False, True, True, True, True, True, True, False, False,
False, False, False, False, False, False, True, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
True, True, False, False, True, True, True, True, True, False,
True, True, True, True, False, False, False, True, True, True,
True, False, True, True, True, True, False, True, False, True,
True, True, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, True, False, True, True, False, True,
True, False, True, False, True, False, True, True, False, True,
True, True, True, False, True, False, False, True, False, False,
True, True, True, False, True, False, False, False, True, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, True,
False, False, True, True, False, False, False, False, True, True,
False, True, True, False, False, False, False, True, True, True,
False, True, False, True, True, True, False, False, True, True,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, True, False, True, True, True, True, True, True,
True, True, True, True, True, True, True, True, False, True,
True, True, True, True, False, False, True, True, True, True,
False, True, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, True, False, False, False,
False, False, False, False, False, False, True, False, True, True,
True, True, True, True, True, True, True, True, True, True,
True, True, True, True, False, True, False, True, True, True,
True, True, False, True, False, False, True, True, False, False,
False, False, False, False, False, False, False, False, False, True,
False, False, False, False, True, False, False, False, False, False,
False, False, True, False, False, False, False, False, False, False,
True, False, False, True, True, True, True, True, True, True,
False, True, False, True, True, True, True, False, True, True,
True, True, False, True, True, True, True, False, True, True,
True, True, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, True, False, False, False, True, True,
True, True, True, False, True, True, True, True, False, True,
True, False, True, True, True, False, True, False, True, True,
True, True, True, False, True, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, True, False, True, False, True, True, True, True, True,
False, True, False, True, True, False, False, True, True, False,
True, True, False, False, True, True, True, True, True, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, True, True, True, False, True, True, True, False,
True, True, True, False, True, False, True, True, True, True,
False, True, False, False, True, True, False, True, True, True,
True, True, True, False, False, True, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, True, False, True, False,
True, True, False, True, True, False, True, True, False, True,
False, False, True, True, False, False, True, False, True, True,
True, True, False, True, True, True, True, True, False, False,
False, False, False, False, False, False, False, False, False, True,
False, True, False, False, True, False, True, False, False, False,
False, False, False, False, False, False, False, False, False, False,
True, True, True, False, True, True, True, True, False, False,
True, False, False, True, False, False, False, True, True, True,
True, False, True, True, True, True, True, False, True, True,
False, True, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, True, False, True, True, True, False,
True, True, True, True, True, True, False, True, True, True,
True, True, True, True, True, False, True, True, False, True,
True, True, False, True, True, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, True, True,
True, True, False, True, True, False, True, False, False, True,
False, True, True, False, True, False, False, True, True, True,
True, True, False, True, False, True, False, True, True, True,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, True, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, True, True, False, True, True, False, False, True,
False, True, True, True, True, False, False, True, False, True,
False, True, False, True, True, True, False, False, True, True,
True, True, True, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, True,
False, False, False, False, False, True, False, True, True, False,
True, True, True, True, True, False, False, False, True, True,
True, True, True, False, False, False, True, True, False, True,
False, True, True, False, False, False, True, True, False, False,
False, False, False, False, False, False, True, False, False, False,
True, False, False, False, False, False, False, False, False, True,
False, False, False, False, False, False, False, False, False, False,
True, False, True, False, True, False, False, False, True, True,
True, True, True, True, True, True, False, True, False, False,
True, True, True, True, True, True, True, False, True, True,
True, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
True, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, True, True, True, True,
True, False, False, True, True, False, False, True, True, True,
True, False, True, False, True, True, True, False, False, True,
True, True, True, False, False, False, False, False, False, False,
False, False, False, False, False, False, True, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, True, True,
True, True, True, False, True, False, True, True, True, True,
True, False, True, True, True, True, True, True, False, True,
True, False, True, True, False, True, True, True, True, False,
False, False, False, False, True, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, True, True, True, True, True, True, False,
True, True, True, True, True, True, True, True, False, True,
False, True, True, True, True, False, True, True, False, True,
True, True, True, False, False, False, False, False, True, False,
False, False, True, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, True, False, True, True, True, True,
True, True, True, False, True, False, True, True, True, True,
True, True, True, True, True, False, True, False, False, False,
True, False, True, True, True, True, False, False, False, True,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
True, False, True, False, True, True, True, False, True, False,
True, False, True, True, True, True, True, True, True, False,
False, True, False, True, True, True, True, True, False, True,
True, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, True, False, True, True, True, True,
True, True, False, True, False, True, True, True, False, True,
False, True, False, True, False, True, True, True, True, False,
True, True, False, False, True, True, False, False, False, False,
False, True, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, True,
True, True, True, True, True, True, False, True, True, True,
False, False, True, False, False, False, True, True, True, True,
True, True, True, True, False, True, True, True, False, True,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, True, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, True, False, True, True, False, False, False, False,
True, True, True, False, True, True, True, True, False, True,
True, False, True, True, False, False, False, True, True, True,
True, True, True, False, False, False, False, False, False, False,
False, False, True, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, True, False, False, True, False, True, False,
True, False, True, True, False, True, False, True, True, True,
False, True, True, True, True, True, True, False, True, True,
True, False, True, False, False, True, False, True, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
True, False, True, False, True, True, False, False, True, True,
False, True, True, True, True, True, True, True, True, True,
True, True, False, True, False, True, True, False, True, False,
True, True, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, True, False, False,
False, False, False, False, False, True, True, True, False, False,
True, False, True, True, True, False, True, True, True, False,
True, True, True, True, True, True, True, True, False, True,
True, False, True, True, False, True, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, True,
False, False, False, False, False, False, False, False, True, True,
True, True, False, True, False, True, True, False, False, False,
True, True, True, False, False, True, False, False, True, False,
True, False, True, True, True, True, True, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, True, False, False, True, True, True, True, False,
True, False, True, True, False, False, True, True, True, False,
True, True, True, True, False, True, False, True, True, False,
True, True, True, True, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, True, False, True, False,
True, True, False, True, True, True, True, False, True, True,
True, True, True, True, False, True, True, True, False, True,
True, True, False, True, False, True, True, False, False, False,
False, False, False, False, False, False, False, False, False, False,
True, False, False, False, True, True, False, False, True, False,
False, False, False, False, False, False, False, False, False, False,
False, True, False, True, True, True, True, True, False, True,
True, True, True, True, True, False, True, True, True, True,
False, True, False, True, True, True, True, True, True, True,
True, True, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, True,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, True, True, True, True, True, True,
False, True, True, True, True, True, True, False, True, True,
True, True, True, False, False, True, True, True, False, True,
True, False, True, True, True, False, False, False, False, False,
False, False, False, False, True, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, True,
True, True, True, False, True, True, True, False, True, False,
True, False, True, False, True, True, False, True, True, True,
True, True, False, True, True, True, True, True, True, True,
False, False, False, False, False, False, False, False, False, True,
False, True, False, False, False, False, False, False, False, False,
False, False, False, False, False, True, False, False, False, False,
False, False, False, False, False, False, True, False, True, True,
True, False, True, True, True, True, True, False, False, True,
True, True, True, True, False, True, True, True, True, False,
False, True, False, True, False, False, False, True, False, False,
True, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, True, False, True, True,
False, False, True, True, True, True, True, True, True, True,
True, False, True, True, True, True, False, True, True, False,
False, True, True, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
True, True, True, False, True, False, True, True, True, False,
True, False, True, True, False, True, True, True, False, True,
False, True, False, False, True, True, True, True, True, True,
True, True, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, True, False, False, True, True, False, True, True, True,
True, True, False, True, True, True, True, False, True, True,
True, False, True, True, True, True, True, True, True, False,
True, True, False, True, True, True, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, True, False,
True, True, True, True, True, True, True, True, True, True,
True, True, False, True, False, True, True, True, True, False,
True, True, False, True, True, True, True, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False, False, False,
False, False, True, False, False, False, False, False, False, False,
False, False, True, True, False, False, True, False, True, True,
True, False, False, True, True, False, True, False, True, True,
True, True, False, False, False, True, False, True, True, True,
True, True, False, True, False, False, False, False, True, False,
False, False, False, False, True, False, True, False, False, False,
False, False, False, False, False, False, False, False, True, False,
False, False, False, False, False, False, True, False, True, True,
True, True, False, False, True, True, True, True, False, False,
False, True, False, True, True, False, True, True, False, True,
True, True, True, True, True, True, True, False]],
device='cuda:0')
TEST RAISED ERROR: Tensor-likes are not close!

Mismatched elements: 486 / 2688 (18.1%)

additional notes

Currently generates the wrong results when the number of rows in MoE FC1 weights is not divisible by 128

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.

[coderabbitai]

📝 CodeRabbit Plan Mode

Generate an implementation plan and prompts that you can use with your favorite coding agent.

• Create Plan

Examples

• Example 1
• Example 2

---

🧪 Issue enrichment is currently in open beta.

You can configure auto-planning by selecting labels in the issue_enrichment configuration.

To disable automatic issue enrichment, add the following to your .coderabbit.yaml:

issue_enrichment:
  auto_enrich:
    enabled: false

💬 Have feedback or questions? Drop into our discord!

[tcherckez-nvidia]

Got a fix for the test and implementation, but I want to see full model run before moving forward, so blocked on #10358