NVIDIA / NVIDIA/cutlass

[BUG] e4m3, int8, bf16 pytorch emitter not working

Open
#1,663 4 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

? - Needs Triage bug inactive-30d inactive-90d
Dominant language
C++
Stars
10.5k
Forks
2.1k
Avg merge
3d 11h
Merged PRs (30d)
7

Description

I am attempting to emit pytorch code but unfortunately it does not work for fp8, bf16, and int8. I have tried to patch the converter type dict https://github.com/OrenLeung/cutlass/commit/6d619c964eb8b9c150a5f97891849d33f6ee8b64

This patch fixed the initial set of issue but unfortunately I run into a deeper issue when the emitter tries to compile the kernel.

New Error Issue
c /workspace/cutlass/examples/python/gemm_mod_fp8_kernel.cu -o gemm_mod_fp8_kernel.cuda.o 
/workspace/cutlass/examples/python/gemm_mod_fp8_kernel.cu(100): error: a value of type "cutlass::float_e4m3_t *" cannot be used to initialize an entity of type "const float *"
          D,
          ^

Reprod

import torch
import cutlass

# FP8 CUTLASS GEMM plan
plan_fp8 = cutlass.op.Gemm(
    element=torch.float8_e4m3fn,
    element_accumulator=torch.float32,
    element_D=torch.float32,
    layout_A=cutlass.LayoutType.RowMajor,
    layout_B=cutlass.LayoutType.ColumnMajor,
    layout_C=cutlass.LayoutType.ColumnMajor)
op_fp8 = plan_fp8.construct()

# Generate the PyTorch module for the FP8 GEMM operation
mod_fp8 = cutlass.emit.pytorch(op_fp8, name='gemm_mod_fp8', cc=plan_fp8.cc, jit=True)

# BF16 CUTLASS GEMM plan
plan_bf16 = cutlass.op.Gemm(
    element=torch.bfloat16,
    element_accumulator=torch.float32,
    element_D=torch.float32,
    layout_A=cutlass.LayoutType.RowMajor,
    layout_B=cutlass.LayoutType.ColumnMajor,
    layout_C=cutlass.LayoutType.ColumnMajor)
op_bf16 = plan_bf16.construct()

# Generate the PyTorch module for the BF16 GEMM operation
mod_bf16 = cutlass.emit.pytorch(op_bf16, name='gemm_mod_bf16', cc=plan_bf16.cc, jit=True)

plan = cutlass.op.Gemm(
    element=cutlass.DataType.s8,
    element_accumulator=cutlass.DataType.s32,
    element_D=cutlass.DataType.s32,
    layout=cutlass.LayoutType.RowMajor)
op = plan.construct()

# Generate the PyTorch module for the GEMM operation
mod = cutlass.emit.pytorch(op, name='gemm_mod', cc=plan.cc, jit=True)

plan = cutlass.op.Gemm(
    element=cutlass.DataType.s8,
    element_accumulator=cutlass.DataType.s32,
    element_D=cutlass.DataType.s32,
    layout=cutlass.LayoutType.RowMajor)
op = plan.construct()

# Generate the PyTorch module for the GEMM operation
mod = cutlass.emit.pytorch(op, name='gemm_mod', cc=plan.cc, jit=True)

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the provided cutlass.emit.pytorch reproducer and inspect the generated examples/python/gemm_mod_fp8_kernel.cu, especially the failing call near line 100. Re-run the FP8, BF16, and int8 cases after investigating the emitter's generated types; done means each case emits and compiles without the reported pointer-type error.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.