| #include <torch/all.h> |
| #include <torch/library.h> |
|
|
| #include <ATen/cuda/CUDAContext.h> |
| #include <c10/cuda/CUDAGuard.h> |
| #include <cuda_bf16.h> |
|
|
| #include "quantize_fp4_sfa_bf16.cuh" |
|
|
| void ada_layer_norm_bf16( |
| const __nv_bfloat16* x, const __nv_bfloat16* scale, |
| const __nv_bfloat16* shift, __nv_bfloat16* out, |
| int seq_len, int dim, float eps, cudaStream_t stream); |
| void layer_norm_no_affine_bf16( |
| const __nv_bfloat16* x, __nv_bfloat16* out, int seq_len, int dim, float eps, |
| cudaStream_t stream); |
|
|
| namespace { |
|
|
| void reference_adaln_fp4( |
| const torch::Tensor& x, const torch::Tensor& scale, |
| const torch::Tensor& shift, double eps, torch::Tensor& packed, |
| torch::Tensor& sf) { |
| c10::cuda::CUDAGuard guard(x.device()); |
| auto stream = at::cuda::getCurrentCUDAStream(x.get_device()).stream(); |
| auto norm = torch::empty_like(x); |
| ada_layer_norm_bf16( |
| static_cast<const __nv_bfloat16*>(x.data_ptr()), |
| static_cast<const __nv_bfloat16*>(scale.data_ptr()), |
| static_cast<const __nv_bfloat16*>(shift.data_ptr()), |
| static_cast<__nv_bfloat16*>(norm.data_ptr()), |
| static_cast<int>(x.size(0)), static_cast<int>(x.size(1)), |
| static_cast<float>(eps), stream); |
| const int rc = flash_rt::fp4::quantize_fp4_dynamic_sfa_bf16_vec( |
| norm.data_ptr(), packed.data_ptr(), sf.data_ptr(), |
| static_cast<int>(x.size(0)), static_cast<int>(x.size(1)), false, stream); |
| TORCH_CHECK(rc == 0, "SM110 staged AdaLN-FP4 reference failed: rc=", rc); |
| } |
|
|
| void reference_ln_fp4( |
| const torch::Tensor& x, double eps, torch::Tensor& packed, |
| torch::Tensor& sf) { |
| c10::cuda::CUDAGuard guard(x.device()); |
| auto stream = at::cuda::getCurrentCUDAStream(x.get_device()).stream(); |
| auto norm = torch::empty_like(x); |
| layer_norm_no_affine_bf16( |
| static_cast<const __nv_bfloat16*>(x.data_ptr()), |
| static_cast<__nv_bfloat16*>(norm.data_ptr()), |
| static_cast<int>(x.size(0)), |
| static_cast<int>(x.size(1)), static_cast<float>(eps), stream); |
| const int rc = flash_rt::fp4::quantize_fp4_dynamic_sfa_bf16_vec( |
| norm.data_ptr(), packed.data_ptr(), sf.data_ptr(), |
| static_cast<int>(x.size(0)), static_cast<int>(x.size(1)), false, stream); |
| TORCH_CHECK(rc == 0, "SM110 staged LayerNorm-FP4 reference failed: rc=", rc); |
| } |
|
|
| } |
|
|
| TORCH_LIBRARY_FRAGMENT(adaptive_layernorm_producers_test, ops) { |
| ops.def("_reference_adaln_fp4(Tensor x, Tensor scale, Tensor shift, float eps, Tensor! packed, Tensor! sf) -> ()"); |
| ops.def("_reference_ln_fp4(Tensor x, float eps, Tensor! packed, Tensor! sf) -> ()"); |
| ops.impl("_reference_adaln_fp4", torch::kCUDA, &reference_adaln_fp4); |
| ops.impl("_reference_ln_fp4", torch::kCUDA, &reference_ln_fp4); |
| } |
|
|
