add custom op wrapping mxfp8 dim1 cast cuda kernel

[danielvegamyhre]

Stacked PRs:

• integrate mxfp8 dim1 cast kernel choice enum into MXLinear #2554
• modify triton_scale_swizzle kernel to accept column major inputs #2553
• add wrapper with dtensor handling for mxfp8 dim1 cast kernel #2552
• register custom sharding for mxfp8 dim1 cast cuda kernel #2551
• ->add custom op wrapping mxfp8 dim1 cast cuda kernel #2550

---

add custom op wrapping mxfp8 dim1 cast cuda kernel

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2550

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 4b0038f with merge base c011bad ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[drisspg]

Nit I hate this pattern it would be cool to create a decorator that
@requires_triton

or something that will raise this error if not avail

[danielvegamyhre]

agree this is not ideal. i like the decorator idea, will try it

[danielvegamyhre]

Hmm it seems like with a decorator approach there are problems with torch custom op internals trying to infer the schema.

../.conda/envs/ao/lib/python3.13/site-packages/torch/_library/infer_schema.py:69: in error_fn
    raise ValueError(f"infer_schema(func): {what} Got func with signature {sig})")
E   ValueError: infer_schema(func): Parameter func must have a type annotation. Got func with signature (func))

I solved this by reversing the order of the custom op decorator and new decorator so the signature is preserved:

@requires_sm100
@torch.library.custom_op("torchao::mxfp8_quantize_cuda", mutates_args=())
def mxfp8_quantize_cuda(
...

However, I then get an error when trying to register a fake for the custom op:

torchao/prototype/mx_formats/kernels.py:1770: in <module>
    @mxfp8_quantize_cuda.register_fake
     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E   AttributeError: 'function' object has no attribute 'register_fake'

I inspected the object returned by the decorator and managed to figure out how to get it to work, but it is ugly and think it will probably have side effects.... :

# ... in the decorator ... 
   return decorator.__closure__[0].cell_contents

Here's the full working decorator: https://www.internalfb.com/phabricator/paste/view/P1870964511

IMO I am not confident enough in this implementation to include it atm (I think it probably won't work correctly on non-sm100) and prefer the predictability of the statement (for now).

[drisspg]

cc @zou3519 Do you think that multiple decorators should be supported / supportable?

[vkuzo]

would be good to keep the code inside this file consistent, I'd vote for a separate PR fixing all the kernels with a better fallback syntax

[zou3519]

What's the goal? To error out when the operator gets executed and we are < sm100?

What's wrong with:

@torch.library.custom_op("torchao::mxfp8_quantize_cuda", mutates_args=())
def mxfp8_quantize_cuda(...):
    if not_is_sm100(...):
        raise RuntimeError(...)

The multiple decorators can work though. My rec is:

@torch.library.custom_op("torchao::mxfp8_quantize_cuda", mutates_args=())
@requires_sm100
def mxfp8_quantize_cuda(

but because requires_sm100 wraps a function, it should also propagate the type annotations to the wrapped function. That'll let torch.library.custom_op read from the type annotations.

[vkuzo]

would be good to use the scaling mode enum in as many places as possible, and convert to string right before the call into the c++ wrapper

[danielvegamyhre]

Makes sense, done

[danielvegamyhre]

Nevermind, I remembered why I had to use a string in the first place, it seems custom ops don't place nicely with enums:

torchao/prototype/mx_formats/kernels.py:1746: in <module>
    @torch.library.custom_op("torchao::mxfp8_quantize_cuda", mutates_args=())
     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
../.conda/envs/ao/lib/python3.13/site-packages/torch/_library/custom_ops.py:148: in inner
    schema_str = torch.library.infer_schema(fn, mutates_args=mutates_args)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
../.conda/envs/ao/lib/python3.13/site-packages/torch/_library/infer_schema.py:127: in infer_schema
    if annotation_type.__origin__ is tuple:
       ^^^^^^^^^^^^^^^^^^^^^^^^^^
E   AttributeError: type object 'ScaleCalculationMode' has no attribute '__origin__'

[zou3519]

yeah, you can't pass an enum to a custom op.