[BUG] rs_op_selector does not allow to use mixed FP8 types for A and B

[ex-rzr]

Describe the bug
rs_op_selector does not allow to use mixed FP8 types for A and B.
There is a static assert that prevents this (it is placed correctly in ss_op_selector, each branch has its own check if needed).
See https://github.com/NVIDIA/cutlass/blob/main/include/cute/arch/mma_sm90.hpp#L857

Steps/Code to reproduce bug

cute::make_tiled_mma(cute::GMMA::rs_op_selector<cutlass::float_e4m3_t, cutlass::float_e5m2_t, float, ...

cutlass/include/cute/arch/mma_sm90.hpp(856): error: static assertion failed with "ElementA and ElementB must be the same type for this config."
      static_assert(is_same_v<ElementA, ElementB>, "ElementA and ElementB must be the same type for this config.");
      ^
          detected during:
            instantiation of "auto cute::GMMA::rs_op_selector<ElementA,ElementB,ElementC,TileShape_MNK,MajorA,MajorB,Args...>() [with ElementA=cutlass::float_e4m3_t, ElementB=cutlass::float_e5m2_t, ElementC=float, ..."

Expected behavior
The same behavior as in ss_op_selector.

[hwu36]

@IonThruster

[jayhshah]

I also noticed that the option of FP8 operand and FP16 accumulator is excluded in ss_op_selector and rs_op_selector, even though this is accommodated by the PTX ISA and otherwise has MMA_Traits structs defined for it (F16E4M3E4M3 etc.). This is done by the following code:

// FP16 accumulator
  if constexpr (is_same_v<ElementC, half_t>) {
    static_assert(is_same_v<ElementA, half_t>, "Element types for AB must be half if ElementC is half.");
    static_assert(is_same_v<ElementB, half_t>, "Element types for AB must be half if ElementC is half.");

Was this intentional or is this an oversight?

[thakkarV]

hi @ex-rzr , we support mixed fp8 types in the RS op selector now:

cutlass/include/cute/arch/mma_sm90.hpp

Line 996 in acba5be

else if constexpr (is_same_v<ElementA, float_e4m3_t> && is_same_v<ElementB, float_e5m2_t>) {

Is this what you wanted? Can we close the issue?

[ex-rzr]

@thakkarV
Nothing has changed since I opened this issue.
There is a static assert before the line you pointed at:

cutlass/include/cute/arch/mma_sm90.hpp

Line 857 in acba5be

static_assert(is_same_v<ElementA, ElementB>, "ElementA and ElementB must be the same type for this config.");

[thakkarV]

sorry about that. I noticed the assert after replying earlier to the thread. This will get fixed with 3.4 tagging in the next few days.

[thakkarV]

Fixed witb 3.4. Closing.