[BUG] add_stub" not implemented for 'Float8_e4m3fn, Float4_e2m1fn_x2

Which component has the problem?

CuTe DSL

Bug Report

Describe the bug
Simple low precision arithmetic not working.

torch

import torch
def add_precision_sweep ():

    precisions = [torch.float16, torch.bfloat16, torch.float8_e4m3fn, torch.float4_e2m1fn_x2 ]
    a_float = torch.randn(4, 4, dtype=torch.float32, device="cuda")
    b_float = torch.randn(4, 4, dtype=torch.float32, device="cuda")

    for p in precisions:
        a_p = a_float.view(p)
        b_p = b_float.view(p)
        try:
            c_p = a_p + b_p
        except Exception as ex:
            print (ex)

add_precision_sweep()

torch+dlpack+cutedsl:

import torch
import cutlass
import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack


def naive_elementwise_add_kernel(
    gA: cute.Tensor,  # Input tensor A
    gB: cute.Tensor,  # Input tensor B
    gC: cute.Tensor,  # Output tensor C = A + B
):
    # Step 1: Get thread indices
    # ------------------------
    # CUDA threads are organized in a 3D grid of thread blocks
    # Here we only use the x-dimension for simplicity
    tidx, _, _ = cute.arch.thread_idx()  # Thread index within block (0 to bdim-1)
    bidx, _, _ = cute.arch.block_idx()  # Block index in grid (0 to grid_dim-1)
    bdim, _, _ = cute.arch.block_dim()  # Number of threads per block

    # Calculate global thread index
    # This gives each thread a unique ID across all blocks
    thread_idx = bidx * bdim + tidx  # Global thread ID

    # Step 2: Map thread index to tensor coordinates
    # -------------------------------------------
    # Each thread will process one element of the input tensors
    m, n = gA.shape  # Get tensor dimensions (M rows × N columns)

    # Convert linear thread index to 2D coordinates:
    # - ni: column index (0 to n-1)
    # - mi: row index (0 to m-1)
    ni = thread_idx % n  # Column index (faster varying dimension)
    mi = thread_idx // n  # Row index (slower varying dimension)

    # Step 3: Load and process data
    # ---------------------------
    # Load values from input tensors
    # The tensor layout automatically handles the conversion from
    # logical indices (mi, ni) to physical memory addresses
    a_val = gA[mi, ni]  # Load element from tensor A
    b_val = gB[mi, ni]  # Load element from tensor B

    # Step 4: Store result
    # ------------------
    # Write the sum back to the output tensor
    gC[mi, ni] = a_val + b_val

@cute.jit  # Just-in-time compilation decorator
def naive_elementwise_add(
    mA: cute.Tensor,  # Input tensor A
    mB: cute.Tensor,  # Input tensor B
    mC: cute.Tensor,  # Output tensor C
):
    # Configure kernel launch parameters
    # --------------------------------
    # Choose number of threads per block
    # 256 is a common choice as it:
    # - Allows good occupancy on most GPUs
    # - Is a multiple of 32 (warp size)
    # - Provides enough threads for latency hiding
    num_threads_per_block = 256

    # Get input dimensions
    m, n = mA.shape  # Matrix dimensions (M rows × N columns)

    # Create kernel instance
    kernel = naive_elementwise_add_kernel(mA, mB, mC)

    # Launch kernel with calculated grid dimensions
    # -------------------------------------------
    # Grid size calculation:
    # - Total elements: m * n
    # - Blocks needed: ceil(total_elements / threads_per_block)
    # - Using integer division here assumes m * n is multiple of threads_per_block
    kernel.launch(
        grid=((m * n) // num_threads_per_block, 1, 1),  # Number of blocks in x,y,z
        block=(num_threads_per_block, 1, 1),  # Threads per block in x,y,z
    )

def run_naive_elementwise_add():
    # Test Setup
    # ----------
    # Define test dimensions
    M, N = 16384, 8192  # Using large matrices to measure performance

    for input_dtype in [torch.float16, torch.float8_e4m3fn, torch.float8_e4m3fnuz, torch.float4_e2m1fn_x2]:

        # Create test data on GPU
        # ----------------------
        # Using float16 (half precision) for:
        # - Reduced memory bandwidth requirements
        # - Better performance on modern GPUs
        a = torch.randn( (M, N), device="cuda", dtype=torch.float16).view(input_dtype)  # Random input A
        b = torch.randn( (M, N), device="cuda", dtype=torch.float16).view(input_dtype)  # Random input B
        c = torch.zeros( (M, N), device="cuda", dtype=torch.float16).view(input_dtype)  # Random input C
        #c = torch.zeros(M, N, device="cuda", dtype=torch.float16)  # Output buffer

        # Calculate total elements for bandwidth calculations
        num_elements = sum([a.numel(), b.numel(), c.numel()])

        # Convert PyTorch tensors to CuTe tensors
        # -------------------------------------
        # from_dlpack creates CuTe tensor views of PyTorch tensors
        # assumed_align=16 ensures proper memory alignment for vectorized access
        a_ = from_dlpack(a, assumed_align=16)  # CuTe tensor A
        b_ = from_dlpack(b, assumed_align=16)  # CuTe tensor B
        c_ = from_dlpack(c, assumed_align=16)  # CuTe tensor C

        try:

            # Compile the kernel for the specific input types
            naive_elementwise_add_ = cute.compile(naive_elementwise_add, a_, b_, c_)

            # Run the kernel
            naive_elementwise_add_(a_, b_, c_)

            # Verify Results
            # -------------
            # Compare our kernel output with PyTorch's native implementation
            torch.testing.assert_close(c, a + b)  # Raises error if results don't match
        except Exception as ex:
            print (ex)
    print (f" completed run_naive_elementwise_add")

run_naive_elementwise_add()

output:

lib/python3.10/site-packages/nvidia_cutlass_dsl/python_packages/cutlass/base_dsl/typing.py":846:0): see current operation: %42 = "arith.addi"(%39, %41) <{overflowFlags = #arith.overflow<none>}> : (f8E4M3FN, f8E4M3FN) -> f8E4M3FN
module 'cutlass.base_dsl.typing' has no attribute 'Float8E4M3FNUZ'
Sub-byte scalar dereference not supported for type Float4E2M1FN
 completed run_naive_elementwise_add

Expected behavior
Work

Environment details (please complete the following information):
nvidia-cutlass 4.2.0.0
nvidia-cutlass-dsl 4.3.2
torch 2.10.0.dev20251031+cu130

Additional context
DGX Spark

Start by running the provided CuTe DSL reproducer and inspect lib/python3.10/site-packages/nvidia_cutlass_dsl/python_packages/cutlass/base_dsl/typing.py around the reported operation at line 846. Trace how Float8_e4m3fn, Float8_e4m3fnuz, and Float4_e2m1fn_x2 are represented during cute.compile and scalar access. Done means the shown arithmetic and elementwise kernel compile and run for the listed types without the reported errors.