Fail to compile `T.copy` for fp8 on sm89

[bingps]

When running python examples/deepseek_v32/fp8_lighting_indexer.py on sm89, it reports error like

/tmp/tmpts7207id.cu(44): error: identifier "make_fp8_e4_16_t" is undefined
        condval = make_fp8_e4_16_t(fp8_e4_t(0x0p+0f ), fp8_e4_t(0x0p+0f ), fp8_e4_t(0x0p+0f ), fp8_e4_t(0x0p+0f ), fp8_e4_t(0x0p+0f ), fp8_e4_t(0x0p+0f ), fp8_e4_t(0x0p+0f ), fp8_e4_t(0x0p+0f ), fp8_e4_t(0x0p+0f ), fp8_e4_t(0x0p+0f ), fp8_e4_t(0x0p+0f ), fp8_e4_t(0x0p+0f ), fp8_e4_t(0x0p+0f ), fp8_e4_t(0x0p+0f ), fp8_e4_t(0x0p+0f ), fp8_e4_t(0x0p+0f ));

this is because the T.copy(IndexQ[seq_len_i * heads, 0], index_q_shared) will be lowed to CUDA like

  #pragma unroll
  for (int i = 0; i < 2; ++i) {
    fp8_e4_16_t condval;
    if ((((((((int)blockIdx.x) * 128) + (i * 64)) + (((int)threadIdx.x) >> 2)) < (seq_len * 32)) && ((((((int)blockIdx.x) * 128) + (i * 64)) + (((int)threadIdx.x) >> 2)) < (seq_len * 32)))) {
      condval = *(fp8_e4_16_t*)(IndexQ + (((((int64_t)((int)blockIdx.x)) * (int64_t)8192) + (((int64_t)i) * (int64_t)4096)) + (((int64_t)((int)threadIdx.x)) * (int64_t)16)));
    } else {
      condval = make_fp8_e4_16_t(fp8_e4_t(0x0p+0f/*0.000000e+00*/), fp8_e4_t(0x0p+0f/*0.000000e+00*/), fp8_e4_t(0x0p+0f/*0.000000e+00*/), fp8_e4_t(0x0p+0f/*0.000000e+00*/), fp8_e4_t(0x0p+0f/*0.000000e+00*/), fp8_e4_t(0x0p+0f/*0.000000e+00*/), fp8_e4_t(0x0p+0f/*0.000000e+00*/), fp8_e4_t(0x0p+0f/*0.000000e+00*/), fp8_e4_t(0x0p+0f/*0.000000e+00*/), fp8_e4_t(0x0p+0f/*0.000000e+00*/), fp8_e4_t(0x0p+0f/*0.000000e+00*/), fp8_e4_t(0x0p+0f/*0.000000e+00*/), fp8_e4_t(0x0p+0f/*0.000000e+00*/), fp8_e4_t(0x0p+0f/*0.000000e+00*/), fp8_e4_t(0x0p+0f/*0.000000e+00*/), fp8_e4_t(0x0p+0f/*0.000000e+00*/));
    }
    *(fp8_e4_16_t*)(((fp8_e4_t*)buf_dyn_shmem) + (((((i * 4096) + ((((int)threadIdx.x) >> 2) * 64)) + (((((((int)threadIdx.x) & 31) >> 4) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 32)) + (((((((int)threadIdx.x) & 15) >> 3) + (((int)threadIdx.x) & 1)) & 1) * 16)) + 49152)) = condval;
  }

and make_fp8_e4_16_t is not defined.

I made a temp fix in bingps@2115e1f by adding definition at src/tl_templates/cuda/cuda_fp8.h, and it passes the tests of fp8_lighting_indexer.py.

Are such manual fixes required? Or there are some config that auto-generates these make-functions disabled on sm89 🤔

[bingps]

I tried to isolate this error in a unit test like

import tilelang
from tilelang import language as T
import torch

tilelang.disable_cache()

@tilelang.jit()
def copy_shared(seq_len, heads, dim):
    shape = [seq_len * heads, dim]
    dtype = "float8_e4m3"
    threads = 512

    block = 128 // heads

    @T.prim_func
    def main(A: T.Tensor(shape, dtype)):
        with T.Kernel(T.ceildiv(seq_len, block), threads=threads) as bx:
            A_shared = T.alloc_shared([dim], dtype)
            seq_len_i = bx * block
            T.copy(A[seq_len_i * heads, 0], A_shared)
    
    return main

def test_sm89():
    copy_shared_kernel = copy_shared(4096, 32, 64)
    A = torch.randn((4096 * 32, 64), dtype=torch.bfloat16, device='cuda').to(torch.float8_e4m3fn)
    copy_shared_kernel(A)

    # _kernel = tilelang.compile(copy_shared_kernel, execution_backend="cython")
    # print(_kernel.get_kernel_source())

but it can be successfully compiled and run. Its CUDA is like

extern "C" __global__ void __launch_bounds__(512, 1) main_kernel(fp8_e4_t* __restrict__ A) {
  extern __shared__ __align__(1024) fp8_e4_t A_shared[];
  if (((int)threadIdx.x) < 64) {
    #pragma unroll
    for (int i = 0; i < 1; ++i) {
      A_shared[((int)threadIdx.x)] = A[((((int)blockIdx.x) * 8192) + ((int)threadIdx.x))];
    }
  }
}

which does not involve the vector fp8 type.

[LeiWang1999]

likely can be resolved via pr #967 but have to fix the ci.