From 4b06e5976937d28d8fe8c8c7876c3b7c2d0fec2e Mon Sep 17 00:00:00 2001 From: STwangyingrui Date: Fri, 4 Sep 2026 10:07:50 +0000 Subject: [PATCH 1/6] feat(kernel): add SM120 FP8 GEMM with FP16 accumulation --- lightx2v_kernel/CMakeLists.txt | 1 + lightx2v_kernel/csrc/common_extension.cc | 44 + .../fp8_f16_accum_scaled_mm_kernels_sm120.cu | 801 ++++++++++++++++++ lightx2v_kernel/include/lightx2v_kernel_ops.h | 33 + .../python/lightx2v_kernel/gemm.py | 18 + 5 files changed, 897 insertions(+) create mode 100644 lightx2v_kernel/csrc/gemm/fp8_f16_accum_scaled_mm_kernels_sm120.cu diff --git a/lightx2v_kernel/CMakeLists.txt b/lightx2v_kernel/CMakeLists.txt index 369f4da9c..990a1b53a 100644 --- a/lightx2v_kernel/CMakeLists.txt +++ b/lightx2v_kernel/CMakeLists.txt @@ -82,6 +82,7 @@ list(APPEND LIGHTX2V_KERNEL_CUDA_FLAGS set(SOURCES + "csrc/gemm/fp8_f16_accum_scaled_mm_kernels_sm120.cu" "csrc/gemm/nvfp4_scaled_mm_kernels_sm120.cu" "csrc/gemm/nvfp4_quant_kernels_sm120.cu" "csrc/gemm/mxfp4_quant_kernels_sm120.cu" diff --git a/lightx2v_kernel/csrc/common_extension.cc b/lightx2v_kernel/csrc/common_extension.cc index 3bfdc746f..2061d95d2 100644 --- a/lightx2v_kernel/csrc/common_extension.cc +++ b/lightx2v_kernel/csrc/common_extension.cc @@ -6,6 +6,50 @@ TORCH_LIBRARY_FRAGMENT(lightx2v_kernel, m) { + m.def( + "cutlass_scaled_fp8_mm_f16_accum_sm120(Tensor activation, Tensor weight, Tensor activation_scale, " + "Tensor weight_scale, ScalarType out_dtype, Tensor? bias=None) -> Tensor"); + m.impl( + "cutlass_scaled_fp8_mm_f16_accum_sm120", + torch::kCUDA, + &cutlass_scaled_fp8_mm_f16_accum_sm120); + + m.def( + "cutlass_scaled_fp8_mm_f16_accum_with_config_sm120(Tensor activation, Tensor weight, " + "Tensor activation_scale, Tensor weight_scale, ScalarType out_dtype, Tensor? bias, int config_id) -> Tensor"); + m.impl( + "cutlass_scaled_fp8_mm_f16_accum_with_config_sm120", + torch::kCUDA, + &cutlass_scaled_fp8_mm_f16_accum_with_config_sm120); + + m.def("fp8_f16_accum_autotune_cache_abi_sm120() -> int"); + m.impl( + "fp8_f16_accum_autotune_cache_abi_sm120", + &fp8_f16_accum_autotune_cache_abi_sm120); + m.def("fp8_f16_accum_autotune_configs_sm120() -> str[]"); + m.impl( + "fp8_f16_accum_autotune_configs_sm120", + &fp8_f16_accum_autotune_configs_sm120); + m.def( + "set_fp8_f16_accum_autotune_config_sm120(int device_index, int m, int n, int k, " + "ScalarType out_dtype, bool has_bias, int config_id) -> ()"); + m.impl( + "set_fp8_f16_accum_autotune_config_sm120", + &set_fp8_f16_accum_autotune_config_sm120); + m.def("set_fp8_f16_accum_autotune_enabled_sm120(bool enabled) -> ()"); + m.impl( + "set_fp8_f16_accum_autotune_enabled_sm120", + &set_fp8_f16_accum_autotune_enabled_sm120); + m.def( + "get_fp8_f16_accum_autotune_cache_sm120(int device_index) -> Tensor"); + m.impl( + "get_fp8_f16_accum_autotune_cache_sm120", + &get_fp8_f16_accum_autotune_cache_sm120); + m.def("clear_fp8_f16_accum_autotune_cache_sm120(int device_index=-1) -> ()"); + m.impl( + "clear_fp8_f16_accum_autotune_cache_sm120", + &clear_fp8_f16_accum_autotune_cache_sm120); + m.def( "cutlass_scaled_nvfp4_mm_sm120(Tensor! out, Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, Tensor " "alpha, Tensor? bias) -> ()"); diff --git a/lightx2v_kernel/csrc/gemm/fp8_f16_accum_scaled_mm_kernels_sm120.cu b/lightx2v_kernel/csrc/gemm/fp8_f16_accum_scaled_mm_kernels_sm120.cu new file mode 100644 index 000000000..115c00814 --- /dev/null +++ b/lightx2v_kernel/csrc/gemm/fp8_f16_accum_scaled_mm_kernels_sm120.cu @@ -0,0 +1,801 @@ +// SM120 FP8 GEMM with FP16 accumulation. + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cutlass/epilogue/fusion/operations.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/util/packed_stride.hpp" + +using namespace cute; + +namespace { + +template < + typename TileShape_, + typename ElementD_ = cutlass::bfloat16_t, + bool FuseBias_ = false> +struct GemmDefinition { + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementD = ElementD_; + using ElementC = void; + using ElementAccumulator = cutlass::half_t; + using TileShape = TileShape_; + static constexpr bool FuseBias = FuseBias_; + using ClusterShape = Shape<_1, _1, _1>; + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::RowMajor; + using LayoutD = cutlass::layout::RowMajor; + + static constexpr int AlignmentAB = 16; + static constexpr int AlignmentD = 16 / sizeof(ElementD); + + using Accum = cutlass::epilogue::fusion::Sm90AccFetch; + using ScaleA = cutlass::epilogue::fusion::Sm90ColBroadcast< + 0, + TileShape, + float, + float, + Stride, Int<0>, Int<0>>>; + using ScaleB = cutlass::epilogue::fusion::Sm90RowBroadcast< + 0, + TileShape, + float, + float, + Stride, Int<1>, Int<0>>>; + using Multiply = cutlass::epilogue::fusion::Sm90Compute< + cutlass::multiplies, + float, + float, + cutlass::FloatRoundStyle::round_to_nearest>; + using MultiplyOutput = cutlass::epilogue::fusion::Sm90Compute< + cutlass::multiplies, + ElementD, + float, + cutlass::FloatRoundStyle::round_to_nearest>; + using AddBias = cutlass::epilogue::fusion::Sm90Compute< + cutlass::plus, + ElementD, + float, + cutlass::FloatRoundStyle::round_to_nearest>; + using Bias = cutlass::epilogue::fusion::Sm90RowBroadcast< + 0, + TileShape, + ElementD, + float, + Stride, Int<1>, Int<0>>, + AlignmentD>; + using ScaleBAccum = + cutlass::epilogue::fusion::Sm90EVT; + using ScaledEVT = cutlass::epilogue::fusion::Sm90EVT< + MultiplyOutput, + ScaleA, + ScaleBAccum>; + using OutputEVT = cutlass::epilogue::fusion::Sm90EVT< + AddBias, + ScaledEVT, + Bias>; + using EpilogueEVT = std::conditional_t; + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm120, + cutlass::arch::OpClassTensorOp, + TileShape, + ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, + float, + ElementC, + LayoutC, + AlignmentD, + ElementD, + LayoutD, + AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto, + EpilogueEVT>::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm120, + cutlass::arch::OpClassTensorOp, + ElementA, + LayoutA, + AlignmentAB, + ElementB, + LayoutB, + AlignmentAB, + ElementAccumulator, + TileShape, + ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast( + sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + static typename EpilogueEVT::Arguments prepare_epilogue( + float* activation_scale, + float* weight_scale, + ElementD const* bias) { + typename ScaleA::Arguments activation_arguments{activation_scale}; + typename ScaleB::Arguments weight_arguments{weight_scale}; + typename ScaleBAccum::Arguments scaled_accumulator{ + weight_arguments, + {}, + {}, + }; + typename ScaledEVT::Arguments scaled_output{ + activation_arguments, + scaled_accumulator, + {}, + }; + if constexpr (FuseBias) { + typename Bias::Arguments bias_arguments{bias}; + return typename OutputEVT::Arguments{ + scaled_output, + bias_arguments, + {}, + }; + } else { + return scaled_output; + } + } +}; +using NarrowGemm = GemmDefinition>; +using WideGemm = GemmDefinition>; +using NarrowGemmFp16 = GemmDefinition, cutlass::half_t>; +using WideGemmFp16 = GemmDefinition, cutlass::half_t>; + +using NarrowGemmWithBias = + GemmDefinition, cutlass::bfloat16_t, true>; +using WideGemmWithBias = + GemmDefinition, cutlass::bfloat16_t, true>; +using NarrowGemmFp16WithBias = + GemmDefinition, cutlass::half_t, true>; +using WideGemmFp16WithBias = + GemmDefinition, cutlass::half_t, true>; + +// Bump the ABI whenever a config mapping or candidate implementation changes. +constexpr int64_t kAutotuneCacheAbi = 1; +constexpr int64_t kFallbackConfigId = 0; + +struct KernelConfig { + bool wide_tile; + int swizzle; + char const* name; +}; + +constexpr std::array kKernelConfigs = {{ + {false, 1, "tile_128x128x64_swizzle_1"}, + {false, 2, "tile_128x128x64_swizzle_2"}, + {false, 4, "tile_128x128x64_swizzle_4"}, + {false, 8, "tile_128x128x64_swizzle_8"}, + {true, 1, "tile_128x256x64_swizzle_1"}, + {true, 2, "tile_128x256x64_swizzle_2"}, + {true, 4, "tile_128x256x64_swizzle_4"}, + {true, 8, "tile_128x256x64_swizzle_8"}, +}}; + +KernelConfig const& kernel_config(int64_t config_id) { + TORCH_CHECK( + config_id >= 0 && + config_id < static_cast(kKernelConfigs.size()), + "FP8-F16 GEMM config_id must be in [0, ", + kKernelConfigs.size(), + "), got ", + config_id); + return kKernelConfigs[config_id]; +} + +struct AutotuneKey { + int device_index; + int32_t m; + int32_t n; + int32_t k; + torch::ScalarType output_dtype; + bool has_bias; + + bool operator==(AutotuneKey const& other) const { + return device_index == other.device_index && m == other.m && + n == other.n && k == other.k && + output_dtype == other.output_dtype && has_bias == other.has_bias; + } +}; + +struct AutotuneKeyHash { + size_t operator()(AutotuneKey const& key) const { + size_t value = std::hash{}(key.device_index); + value = value * 31 + std::hash{}(key.m); + value = value * 31 + std::hash{}(key.n); + value = value * 31 + std::hash{}(key.k); + value = value * 31 + std::hash{}(static_cast(key.output_dtype)); + return value * 31 + std::hash{}(key.has_bias); + } +}; + +using AutotuneCache = + std::unordered_map; + +AutotuneCache& autotune_cache() { + static AutotuneCache cache; + return cache; +} + +std::shared_mutex& autotune_cache_mutex() { + static std::shared_mutex mutex; + return mutex; +} + +std::mutex& autotune_measurement_mutex() { + static std::mutex mutex; + return mutex; +} + +std::atomic& autotune_enabled() { + static std::atomic enabled{false}; + return enabled; +} + +std::optional cached_config_id(AutotuneKey const& key) { + std::shared_lock lock(autotune_cache_mutex()); + auto entry = autotune_cache().find(key); + if (entry == autotune_cache().end()) { + return std::nullopt; + } + return entry->second; +} + +void cache_config(AutotuneKey const& key, int64_t config_id) { + std::unique_lock lock(autotune_cache_mutex()); + autotune_cache()[key] = config_id; +} + +template +void launch( + torch::Tensor output, + torch::Tensor activation, + torch::Tensor weight, + torch::Tensor activation_scale, + torch::Tensor weight_scale, + c10::optional const& bias, + int swizzle) { + using Gemm = typename Definition::Gemm; + using GemmKernel = typename Definition::GemmKernel; + using ElementA = typename Definition::ElementA; + using ElementB = typename Definition::ElementB; + using ElementD = typename Definition::ElementD; + using ElementC = typename Definition::ElementC; + using StrideA = typename GemmKernel::StrideA; + using StrideB = typename GemmKernel::StrideB; + using StrideC = typename GemmKernel::StrideC; + using StrideD = typename GemmKernel::StrideD; + + int32_t m = activation.size(0); + int32_t k = activation.size(1); + int32_t n = weight.size(1); + StrideA stride_a = make_stride( + int64_t(activation.stride(0)), Int<1>{}, int64_t(0)); + StrideB stride_b = make_stride( + int64_t(weight.stride(1)), Int<1>{}, int64_t(0)); + auto stride_c = cutlass::make_cute_packed_stride( + StrideC{}, make_shape(m, n, 1)); + auto stride_d = cutlass::make_cute_packed_stride( + StrideD{}, make_shape(m, n, 1)); + + typename GemmKernel::MainloopArguments mainloop{ + reinterpret_cast(activation.data_ptr()), + stride_a, + reinterpret_cast(weight.data_ptr()), + stride_b, + }; + typename GemmKernel::EpilogueArguments epilogue{ + Definition::prepare_epilogue( + activation_scale.data_ptr(), + weight_scale.data_ptr(), + bias ? reinterpret_cast(bias->data_ptr()) : nullptr), + nullptr, + stride_c, + reinterpret_cast(output.data_ptr()), + stride_d, + }; + typename Gemm::Arguments arguments{ + cutlass::gemm::GemmUniversalMode::kGemm, + {m, n, k, 1}, + mainloop, + epilogue, + }; + arguments.scheduler.max_swizzle_size = swizzle; + + Gemm gemm; + auto status = gemm.can_implement(arguments); + TORCH_CHECK( + status == cutlass::Status::kSuccess, + "CUTLASS cannot implement this shape"); + TORCH_CHECK( + Gemm::get_workspace_size(arguments) == 0, + "Unexpected CUTLASS scheduler workspace"); + auto stream = at::cuda::getCurrentCUDAStream(activation.device().index()); + status = gemm.run(arguments, nullptr, stream); + TORCH_CHECK( + status == cutlass::Status::kSuccess, + "CUTLASS kernel launch failed"); +} + +void validate( + torch::Tensor activation, + torch::Tensor weight, + torch::Tensor activation_scale, + torch::Tensor weight_scale, + c10::optional const& bias, + torch::ScalarType output_dtype) { + TORCH_CHECK( + activation.is_cuda() && weight.is_cuda() && + activation_scale.is_cuda() && weight_scale.is_cuda(), + "inputs and scales must be CUDA tensors"); + TORCH_CHECK( + activation.device() == weight.device() && + activation.device() == activation_scale.device() && + activation.device() == weight_scale.device(), + "inputs and scales must be on the same CUDA device"); + TORCH_CHECK( + activation.scalar_type() == torch::kFloat8_e4m3fn && + weight.scalar_type() == torch::kFloat8_e4m3fn, + "activation and weight must be float8_e4m3fn"); + TORCH_CHECK( + activation.dim() == 2 && weight.dim() == 2, + "activation and weight must be matrices"); + TORCH_CHECK( + activation.stride(1) == 1 && + activation.stride(0) == activation.size(1), + "activation must be contiguous"); + TORCH_CHECK( + weight.stride(0) == 1 && weight.stride(1) == weight.size(0), + "weight must be a transposed contiguous matrix"); + TORCH_CHECK( + activation.size(1) == weight.size(0), + "K dimensions must match"); + TORCH_CHECK( + activation.size(0) > 0 && + activation.size(0) <= std::numeric_limits::max() && + weight.size(1) > 0 && + weight.size(1) <= std::numeric_limits::max() && + activation.size(1) > 0 && + activation.size(1) <= std::numeric_limits::max(), + "M, N and K must be positive int32 values"); + TORCH_CHECK( + activation_scale.scalar_type() == torch::kFloat32 && + weight_scale.scalar_type() == torch::kFloat32, + "scales must be float32"); + TORCH_CHECK( + activation_scale.is_contiguous() && weight_scale.is_contiguous(), + "scales must be contiguous"); + TORCH_CHECK( + activation_scale.numel() == activation.size(0) && + weight_scale.numel() == weight.size(1), + "scale sizes must match the activation rows and weight columns"); + if (bias) { + TORCH_CHECK( + bias->is_cuda() && bias->device() == activation.device(), + "bias must be on the same CUDA device as the inputs"); + TORCH_CHECK( + bias->scalar_type() == output_dtype, + "bias dtype must match the output dtype"); + TORCH_CHECK( + bias->is_contiguous() && bias->dim() == 1 && + bias->size(0) == weight.size(1), + "bias must be a contiguous vector matching the output columns"); + } +} + +template +void launch_config( + torch::Tensor output, + torch::Tensor activation, + torch::Tensor weight, + torch::Tensor activation_scale, + torch::Tensor weight_scale, + c10::optional const& bias, + int64_t config_id) { + KernelConfig const& config = kernel_config(config_id); + if (config.wide_tile) { + launch( + output, + activation, + weight, + activation_scale, + weight_scale, + bias, + config.swizzle); + } else { + launch( + output, + activation, + weight, + activation_scale, + weight_scale, + bias, + config.swizzle); + } +} + +template < + typename NarrowDefinition, + typename NarrowDefinitionWithBias, + typename WideDefinition, + typename WideDefinitionWithBias> +int64_t tune_config( + AutotuneKey const& key, + torch::Tensor output, + torch::Tensor activation, + torch::Tensor weight, + torch::Tensor activation_scale, + torch::Tensor weight_scale, + c10::optional const& bias) { + std::lock_guard measurement_lock(autotune_measurement_mutex()); + if (auto cached = cached_config_id(key)) { + return *cached; + } + + auto launch_candidate = [&](int64_t config_id) { + if (bias) { + launch_config( + output, + activation, + weight, + activation_scale, + weight_scale, + bias, + config_id); + } else { + launch_config( + output, + activation, + weight, + activation_scale, + weight_scale, + bias, + config_id); + } + }; + + constexpr int kWarmups = 2; + constexpr int kTrials = 3; + constexpr int kRepeats = 5; + constexpr int kConfigCount = static_cast(kKernelConfigs.size()); + std::array, kConfigCount> timings{}; + + for (int warmup = 0; warmup < kWarmups; ++warmup) { + for (int index = 0; index < kConfigCount; ++index) { + launch_candidate((index + warmup) % kConfigCount); + } + } + + auto stream = at::cuda::getCurrentCUDAStream(activation.device().index()); + int initial_offset = + static_cast((key.m + key.n + key.k) % kConfigCount); + for (int trial = 0; trial < kTrials; ++trial) { + int offset = (initial_offset + trial * 3) % kConfigCount; + for (int index = 0; index < kConfigCount; ++index) { + int config_id = (index + offset) % kConfigCount; + c10::cuda::CUDAEvent start(cudaEventDefault); + c10::cuda::CUDAEvent end(cudaEventDefault); + start.record(stream); + for (int repeat = 0; repeat < kRepeats; ++repeat) { + launch_candidate(config_id); + } + end.record(stream); + end.synchronize(); + timings[config_id][trial] = + start.elapsed_time(end) / static_cast(kRepeats); + } + } + + int64_t best_config_id = kFallbackConfigId; + float best_time = std::numeric_limits::max(); + for (int config_id = 0; config_id < kConfigCount; ++config_id) { + auto samples = timings[config_id]; + std::sort(samples.begin(), samples.end()); + if (samples[kTrials / 2] < best_time) { + best_time = samples[kTrials / 2]; + best_config_id = config_id; + } + } + cache_config(key, best_config_id); + return best_config_id; +} + +template < + typename NarrowDefinition, + typename NarrowDefinitionWithBias, + typename WideDefinition, + typename WideDefinitionWithBias> +torch::Tensor run( + torch::Tensor activation, + torch::Tensor weight, + torch::Tensor activation_scale, + torch::Tensor weight_scale, + c10::optional const& bias, + torch::ScalarType output_dtype, + c10::optional config_id = c10::nullopt) { + validate( + activation, + weight, + activation_scale, + weight_scale, + bias, + output_dtype); + c10::cuda::CUDAGuard guard(activation.device()); + int32_t m = activation.size(0); + int32_t k = activation.size(1); + int32_t n = weight.size(1); + auto output = torch::empty( + {m, n}, + activation.options().dtype(output_dtype)); + + if (!config_id) { + AutotuneKey key{ + activation.device().index(), + m, + n, + k, + output_dtype, + bias.has_value(), + }; + if (auto cached = cached_config_id(key)) { + config_id = *cached; + } else if (autotune_enabled().load(std::memory_order_relaxed)) { + config_id = tune_config< + NarrowDefinition, + NarrowDefinitionWithBias, + WideDefinition, + WideDefinitionWithBias>( + key, + output, + activation, + weight, + activation_scale, + weight_scale, + bias); + } else { + config_id = kFallbackConfigId; + } + } + + if (bias) { + launch_config( + output, + activation, + weight, + activation_scale, + weight_scale, + bias, + *config_id); + } else { + launch_config( + output, + activation, + weight, + activation_scale, + weight_scale, + bias, + *config_id); + } + return output; +} + +torch::Tensor run_with_dtype( + torch::Tensor activation, + torch::Tensor weight, + torch::Tensor activation_scale, + torch::Tensor weight_scale, + torch::ScalarType out_dtype, + c10::optional const& bias, + c10::optional config_id = c10::nullopt) { + if (out_dtype == torch::kBFloat16) { + return run< + NarrowGemm, + NarrowGemmWithBias, + WideGemm, + WideGemmWithBias>( + activation, + weight, + activation_scale, + weight_scale, + bias, + out_dtype, + config_id); + } + TORCH_CHECK( + out_dtype == torch::kFloat16, + "output dtype must be bfloat16 or float16"); + return run< + NarrowGemmFp16, + NarrowGemmFp16WithBias, + WideGemmFp16, + WideGemmFp16WithBias>( + activation, + weight, + activation_scale, + weight_scale, + bias, + out_dtype, + config_id); +} + +} // namespace + +torch::Tensor cutlass_scaled_fp8_mm_f16_accum_sm120( + torch::Tensor activation, + torch::Tensor weight, + torch::Tensor activation_scale, + torch::Tensor weight_scale, + torch::ScalarType out_dtype, + c10::optional const& bias) { + return run_with_dtype( + activation, + weight, + activation_scale, + weight_scale, + out_dtype, + bias); +} + +torch::Tensor cutlass_scaled_fp8_mm_f16_accum_with_config_sm120( + torch::Tensor activation, + torch::Tensor weight, + torch::Tensor activation_scale, + torch::Tensor weight_scale, + torch::ScalarType out_dtype, + c10::optional const& bias, + int64_t config_id) { + return run_with_dtype( + activation, + weight, + activation_scale, + weight_scale, + out_dtype, + bias, + config_id); +} + +int64_t fp8_f16_accum_autotune_cache_abi_sm120() { + return kAutotuneCacheAbi; +} + +std::vector fp8_f16_accum_autotune_configs_sm120() { + std::vector names; + names.reserve(kKernelConfigs.size()); + for (KernelConfig const& config : kKernelConfigs) { + names.emplace_back(config.name); + } + return names; +} + +void set_fp8_f16_accum_autotune_config_sm120( + int64_t device_index, + int64_t m, + int64_t n, + int64_t k, + torch::ScalarType out_dtype, + bool has_bias, + int64_t config_id) { + TORCH_CHECK(device_index >= 0, "device_index must be non-negative"); + TORCH_CHECK( + m > 0 && m <= std::numeric_limits::max() && + n > 0 && n <= std::numeric_limits::max() && + k > 0 && k <= std::numeric_limits::max(), + "M, N and K must be positive int32 values"); + TORCH_CHECK( + out_dtype == torch::kBFloat16 || out_dtype == torch::kFloat16, + "output dtype must be bfloat16 or float16"); + kernel_config(config_id); + + AutotuneKey key{ + static_cast(device_index), + static_cast(m), + static_cast(n), + static_cast(k), + out_dtype, + has_bias, + }; + cache_config(key, config_id); +} + +void set_fp8_f16_accum_autotune_enabled_sm120(bool enabled) { + autotune_enabled().store(enabled, std::memory_order_relaxed); +} + +torch::Tensor get_fp8_f16_accum_autotune_cache_sm120( + int64_t device_index) { + TORCH_CHECK(device_index >= 0, "device_index must be non-negative"); + std::vector> entries; + { + std::shared_lock lock(autotune_cache_mutex()); + entries.reserve(autotune_cache().size()); + for (auto const& entry : autotune_cache()) { + if (entry.first.device_index == device_index) { + entries.push_back(entry); + } + } + } + std::sort( + entries.begin(), + entries.end(), + [](auto const& left, auto const& right) { + auto const& a = left.first; + auto const& b = right.first; + return std::tie( + a.device_index, + a.m, + a.n, + a.k, + a.output_dtype, + a.has_bias) < + std::tie( + b.device_index, + b.m, + b.n, + b.k, + b.output_dtype, + b.has_bias); + }); + + auto result = torch::empty( + {static_cast(entries.size()), 6}, + torch::TensorOptions().dtype(torch::kInt64).device(torch::kCPU)); + auto rows = result.accessor(); + for (int64_t index = 0; index < static_cast(entries.size()); ++index) { + auto const& [key, config_id] = entries[index]; + rows[index][0] = key.m; + rows[index][1] = key.n; + rows[index][2] = key.k; + rows[index][3] = key.output_dtype == torch::kBFloat16 ? 0 : 1; + rows[index][4] = key.has_bias; + rows[index][5] = config_id; + } + return result; +} + +void clear_fp8_f16_accum_autotune_cache_sm120(int64_t device_index) { + std::unique_lock lock(autotune_cache_mutex()); + if (device_index < 0) { + autotune_cache().clear(); + return; + } + + for (auto entry = autotune_cache().begin(); entry != autotune_cache().end();) { + if (entry->first.device_index == device_index) { + entry = autotune_cache().erase(entry); + } else { + ++entry; + } + } +} diff --git a/lightx2v_kernel/include/lightx2v_kernel_ops.h b/lightx2v_kernel/include/lightx2v_kernel_ops.h index 04b380596..c9bc429ef 100644 --- a/lightx2v_kernel/include/lightx2v_kernel_ops.h +++ b/lightx2v_kernel/include/lightx2v_kernel_ops.h @@ -21,6 +21,7 @@ limitations under the License. #include #include +#include #include #include @@ -42,6 +43,38 @@ limitations under the License. /* * From csrc/gemm */ +torch::Tensor cutlass_scaled_fp8_mm_f16_accum_sm120( + torch::Tensor activation, + torch::Tensor weight, + torch::Tensor activation_scale, + torch::Tensor weight_scale, + torch::ScalarType out_dtype, + c10::optional const& bias = c10::nullopt); + +torch::Tensor cutlass_scaled_fp8_mm_f16_accum_with_config_sm120( + torch::Tensor activation, + torch::Tensor weight, + torch::Tensor activation_scale, + torch::Tensor weight_scale, + torch::ScalarType out_dtype, + c10::optional const& bias, + int64_t config_id); + +int64_t fp8_f16_accum_autotune_cache_abi_sm120(); +std::vector fp8_f16_accum_autotune_configs_sm120(); +void set_fp8_f16_accum_autotune_config_sm120( + int64_t device_index, + int64_t m, + int64_t n, + int64_t k, + torch::ScalarType out_dtype, + bool has_bias, + int64_t config_id); +void set_fp8_f16_accum_autotune_enabled_sm120(bool enabled); +torch::Tensor get_fp8_f16_accum_autotune_cache_sm120( + int64_t device_index); +void clear_fp8_f16_accum_autotune_cache_sm120(int64_t device_index); + void scaled_nvfp4_quant_sm120( torch::Tensor& output, torch::Tensor const& input, torch::Tensor& output_sf, torch::Tensor const& input_sf); diff --git a/lightx2v_kernel/python/lightx2v_kernel/gemm.py b/lightx2v_kernel/python/lightx2v_kernel/gemm.py index 8ae4b956e..9ea4f1deb 100644 --- a/lightx2v_kernel/python/lightx2v_kernel/gemm.py +++ b/lightx2v_kernel/python/lightx2v_kernel/gemm.py @@ -1,6 +1,24 @@ import torch +def _fp8_f16_accum_meta(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None): + del scales_a, scales_b, bias + return torch.empty((mat_a.shape[0], mat_b.shape[1]), dtype=out_dtype, device=mat_a.device) + + +FP8_F16_ACCUM_MM_AVAILABLE = hasattr(torch.ops.lightx2v_kernel, "cutlass_scaled_fp8_mm_f16_accum_sm120") +if FP8_F16_ACCUM_MM_AVAILABLE: + _fp8_f16_accum_op = torch.ops.lightx2v_kernel.cutlass_scaled_fp8_mm_f16_accum_sm120.default + if not _fp8_f16_accum_op.has_kernel_for_dispatch_key("Meta"): + torch.library.register_fake(_fp8_f16_accum_op, _fp8_f16_accum_meta) + + +def cutlass_scaled_fp8_mm_f16_accum(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None): + if not FP8_F16_ACCUM_MM_AVAILABLE: + raise ImportError("lightx2v-kernel was built without the SM120 FP8 GEMM with FP16 accumulation") + return _fp8_f16_accum_op(mat_a, mat_b, scales_a, scales_b, out_dtype, bias) + + def cutlass_scaled_nvfp4_mm(mat_a, mat_b, scales_a, scales_b, alpha, bias=None): m, n = mat_a.shape[0], mat_b.shape[0] out = torch.empty((m, n), dtype=torch.bfloat16, device=mat_a.device) From 0dd5a7a9924cb6726fe6bedeb9f9474ad3f1a1cf Mon Sep 17 00:00:00 2001 From: STwangyingrui Date: Fri, 4 Sep 2026 10:08:05 +0000 Subject: [PATCH 2/6] feat(kernel): persist FP8-F16 GEMM autotuning --- .../lightx2v_kernel/fp8_f16_autotune.py | 224 ++++++++++++++++++ 1 file changed, 224 insertions(+) create mode 100644 lightx2v_kernel/python/lightx2v_kernel/fp8_f16_autotune.py diff --git a/lightx2v_kernel/python/lightx2v_kernel/fp8_f16_autotune.py b/lightx2v_kernel/python/lightx2v_kernel/fp8_f16_autotune.py new file mode 100644 index 000000000..34b5486c8 --- /dev/null +++ b/lightx2v_kernel/python/lightx2v_kernel/fp8_f16_autotune.py @@ -0,0 +1,224 @@ +"""Automatic SM120 FP8 GEMM autotuning and persistent dispatch cache.""" + +from __future__ import annotations + +import fcntl +import json +import os +import re +import tempfile +import warnings +from pathlib import Path + +import torch + +_SCHEMA_VERSION = 1 +_KERNEL_NAME = "sm120_fp8_f16_accum" +_DTYPE_NAMES = { + 0: "bfloat16", + 1: "float16", +} + + +def _require_ops() -> None: + required = ( + "fp8_f16_accum_autotune_cache_abi_sm120", + "fp8_f16_accum_autotune_configs_sm120", + "set_fp8_f16_accum_autotune_config_sm120", + "set_fp8_f16_accum_autotune_enabled_sm120", + "get_fp8_f16_accum_autotune_cache_sm120", + "clear_fp8_f16_accum_autotune_cache_sm120", + ) + missing = [name for name in required if not hasattr(torch.ops.lightx2v_kernel, name)] + if missing: + raise ImportError(f"lightx2v-kernel was built without FP8-F16 autotune ops: {missing}") + + +def _device_index(device: torch.device | str | int | None) -> int: + if isinstance(device, int): + return device + device = torch.device("cuda" if device is None else device) + if device.type != "cuda": + raise ValueError(f"FP8-F16 autotune requires a CUDA device, got {device}") + return torch.cuda.current_device() if device.index is None else device.index + + +def _runtime_identity(device_index: int) -> dict: + properties = torch.cuda.get_device_properties(device_index) + return { + "device_name": properties.name, + "compute_capability": [properties.major, properties.minor], + "torch_version": torch.__version__, + "cuda_version": torch.version.cuda, + "cache_abi": torch.ops.lightx2v_kernel.fp8_f16_accum_autotune_cache_abi_sm120(), + } + + +def _default_cache_path(device_index: int) -> Path: + properties = torch.cuda.get_device_properties(device_index) + device_name = re.sub(r"[^a-z0-9]+", "-", properties.name.lower()).strip("-") + cache_root = Path(os.environ.get("XDG_CACHE_HOME") or Path.home() / ".cache") + return cache_root / "lightx2v" / "autotune" / "fp8_f16_accum" / f"sm{properties.major}{properties.minor}-{device_name}.json" + + +def _entry_key(entry: dict) -> tuple: + return ( + int(entry["m"]), + int(entry["n"]), + int(entry["k"]), + entry["out_dtype"], + bool(entry["has_bias"]), + ) + + +def _validate_cache(cache: dict, device_index: int) -> list[dict]: + if cache.get("schema_version") != _SCHEMA_VERSION: + raise ValueError(f"Unsupported FP8-F16 autotune schema: {cache.get('schema_version')}") + if cache.get("kernel") != _KERNEL_NAME: + raise ValueError(f"Unexpected autotune kernel: {cache.get('kernel')!r}") + + expected = _runtime_identity(device_index) + actual = cache.get("runtime") + if actual != expected: + raise ValueError(f"FP8-F16 autotune cache runtime mismatch: expected {expected}, got {actual}") + + config_names = torch.ops.lightx2v_kernel.fp8_f16_accum_autotune_configs_sm120() + entries = cache.get("entries") + if not isinstance(entries, list): + raise ValueError("FP8-F16 autotune cache entries must be a list") + seen = set() + for entry in entries: + key = _entry_key(entry) + if min(key[:3]) <= 0: + raise ValueError(f"GEMM dimensions must be positive, got {key[:3]}") + if key[3] not in _DTYPE_NAMES.values(): + raise ValueError(f"Unsupported FP8-F16 output dtype: {key[3]!r}") + if not isinstance(entry["has_bias"], bool): + raise ValueError(f"has_bias must be a bool, got {entry['has_bias']!r}") + config_id = int(entry["config_id"]) + if not 0 <= config_id < len(config_names): + raise ValueError(f"Invalid FP8-F16 autotune config_id: {config_id}") + if entry.get("config") != config_names[config_id]: + raise ValueError(f"FP8-F16 autotune config name does not match config_id {config_id}") + if key in seen: + raise ValueError(f"Duplicate FP8-F16 autotune cache entry: {key}") + seen.add(key) + return entries + + +def _load_entries(cache_path: Path, device_index: int) -> list[dict]: + if not cache_path.is_file(): + return [] + return _validate_cache(json.loads(cache_path.read_text()), device_index) + + +def _register_entries(entries: list[dict], device_index: int) -> None: + dtype_by_name = { + "bfloat16": torch.bfloat16, + "float16": torch.float16, + } + for entry in entries: + torch.ops.lightx2v_kernel.set_fp8_f16_accum_autotune_config_sm120( + device_index, + int(entry["m"]), + int(entry["n"]), + int(entry["k"]), + dtype_by_name[entry["out_dtype"]], + bool(entry["has_bias"]), + int(entry["config_id"]), + ) + + +def _current_entries(device_index: int) -> list[dict]: + config_names = torch.ops.lightx2v_kernel.fp8_f16_accum_autotune_configs_sm120() + cache = torch.ops.lightx2v_kernel.get_fp8_f16_accum_autotune_cache_sm120(device_index) + entries = [] + for m, n, k, dtype_code, has_bias, config_id in cache.tolist(): + entries.append( + { + "m": m, + "n": n, + "k": k, + "out_dtype": _DTYPE_NAMES[dtype_code], + "has_bias": bool(has_bias), + "config_id": config_id, + "config": config_names[config_id], + } + ) + return entries + + +class Fp8F16AccumAutotuner: + """Manage automatic first-use tuning and its process-independent cache.""" + + def __init__( + self, + cache_path: str | Path | None = None, + device: torch.device | str | int | None = None, + ): + _require_ops() + self.device_index = _device_index(device) + if torch.cuda.get_device_capability(self.device_index) != (12, 0): + raise ValueError("FP8-F16 autotune requires an SM120 device") + self.cache_path = Path(cache_path).expanduser() if cache_path else _default_cache_path(self.device_index) + self._saved_configs = {} + + def start(self) -> int: + """Load compatible winners and enable exact-shape tuning on cache misses.""" + torch.ops.lightx2v_kernel.clear_fp8_f16_accum_autotune_cache_sm120(self.device_index) + try: + entries = _load_entries(self.cache_path, self.device_index) + except (KeyError, OSError, TypeError, ValueError) as error: + warnings.warn(f"Ignoring FP8-F16 autotune cache {self.cache_path}: {error}", stacklevel=2) + entries = [] + _register_entries(entries, self.device_index) + self._saved_configs = {_entry_key(entry): int(entry["config_id"]) for entry in entries} + torch.ops.lightx2v_kernel.set_fp8_f16_accum_autotune_enabled_sm120(True) + return len(entries) + + def save(self) -> int: + """Merge newly tuned winners and atomically persist the cache.""" + current_entries = _current_entries(self.device_index) + current_configs = {_entry_key(entry): int(entry["config_id"]) for entry in current_entries} + if current_configs == self._saved_configs: + return 0 + + lock_path = self.cache_path.with_suffix(self.cache_path.suffix + ".lock") + try: + self.cache_path.parent.mkdir(parents=True, exist_ok=True) + with lock_path.open("a+") as lock_file: + fcntl.flock(lock_file, fcntl.LOCK_EX) + try: + disk_entries = _load_entries(self.cache_path, self.device_index) + except (KeyError, OSError, TypeError, ValueError): + disk_entries = [] + merged = {_entry_key(entry): entry for entry in disk_entries} + merged.update({_entry_key(entry): entry for entry in current_entries}) + entries = [merged[key] for key in sorted(merged)] + payload = { + "schema_version": _SCHEMA_VERSION, + "kernel": _KERNEL_NAME, + "runtime": _runtime_identity(self.device_index), + "entries": entries, + } + temporary_path = None + try: + with tempfile.NamedTemporaryFile( + "w", + dir=self.cache_path.parent, + prefix=self.cache_path.name + ".", + delete=False, + ) as temporary: + temporary_path = Path(temporary.name) + json.dump(payload, temporary, indent=2) + temporary.write("\n") + os.replace(temporary_path, self.cache_path) + finally: + if temporary_path is not None: + temporary_path.unlink(missing_ok=True) + except OSError as error: + warnings.warn(f"Could not persist FP8-F16 autotune cache {self.cache_path}: {error}", stacklevel=2) + return 0 + + self._saved_configs = current_configs + return len(current_configs) From 390feaf2581ddfeb1173a9353f28d6f1f8441904 Mon Sep 17 00:00:00 2001 From: STwangyingrui Date: Fri, 4 Sep 2026 10:08:22 +0000 Subject: [PATCH 3/6] feat(convert): support H3 FP8-F16 accumulation checkpoints --- .../minimax_h3/fp8_f16_accum_policy.py | 36 ++++++++++ tools/convert/converter.py | 40 +++++++++-- tools/convert/quant/h3_fp8_f16_accum.py | 66 +++++++++++++++++++ tools/convert/readme.md | 2 + tools/convert/readme_zh.md | 2 + 5 files changed, 140 insertions(+), 6 deletions(-) create mode 100644 lightx2v/models/networks/minimax_h3/fp8_f16_accum_policy.py create mode 100644 tools/convert/quant/h3_fp8_f16_accum.py diff --git a/lightx2v/models/networks/minimax_h3/fp8_f16_accum_policy.py b/lightx2v/models/networks/minimax_h3/fp8_f16_accum_policy.py new file mode 100644 index 000000000..9757e601e --- /dev/null +++ b/lightx2v/models/networks/minimax_h3/fp8_f16_accum_policy.py @@ -0,0 +1,36 @@ +from pathlib import Path + +from safetensors import safe_open + +FP8_F16_ACCUM_WEIGHT_QMAX = 14.0 +DIT_FP8_F16_ACCUM_ACTIVATION_QMAX = 7.0 +VIDEO_VAE_FP8_F16_ACCUM_ACTIVATION_QMAX = 14.0 +FP8_F16_ACCUM_QUANTIZATION_PROFILE = "h3-fp8-f16-accum" +FP8_F16_ACCUM_PROJECTION_SUFFIXES = ( + ".attn.to_q", + ".attn.to_k", + ".attn.to_v", + ".attn.to_out.0", + ".ff.net.0.proj", + ".ff.net.2", +) + + +def validate_fp8_f16_accum_checkpoint(checkpoint_path): + checkpoint_path = Path(checkpoint_path) + files = (checkpoint_path,) if checkpoint_path.is_file() else tuple(sorted(checkpoint_path.glob("*.safetensors"))) + if not files: + raise FileNotFoundError(f"No safetensors weights found in FP8 checkpoint: {checkpoint_path}") + + for filename in files: + with safe_open(filename, framework="pt", device="cpu") as checkpoint: + metadata = checkpoint.metadata() or {} + profile = metadata.get("quantization_profile") + if profile != FP8_F16_ACCUM_QUANTIZATION_PROFILE: + raise ValueError(f"{filename} requires quantization profile {FP8_F16_ACCUM_QUANTIZATION_PROFILE!r}, got {profile!r}") + try: + weight_qmax = float(metadata.get("weight_qmax")) + except (TypeError, ValueError): + weight_qmax = None + if weight_qmax != FP8_F16_ACCUM_WEIGHT_QMAX: + raise ValueError(f"{filename} requires weight_qmax={FP8_F16_ACCUM_WEIGHT_QMAX}, got {metadata.get('weight_qmax')!r}") diff --git a/tools/convert/converter.py b/tools/convert/converter.py index 6fae91829..d728060e1 100755 --- a/tools/convert/converter.py +++ b/tools/convert/converter.py @@ -24,6 +24,10 @@ if quant_path not in sys.path: sys.path.insert(0, quant_path) +from h3_fp8_f16_accum import ( # noqa: E402 + FP8_F16_ACCUM_QUANTIZATION_PROFILE, + create_h3_fp8_f16_accum_quantization, +) from quant import * # noqa: E402 from lightx2v.utils.lora_loader import LoRALoader # noqa: E402 @@ -310,9 +314,9 @@ def get_key_mapping_rules(direction, model_type): return [rule["backward"] for rule in unified_rules] else: raise ValueError(f"Invalid direction: {direction}") - elif model_type == "h3": - # MiniMax-H3 checkpoints under the Diffusers ``transformer`` or - # ``transformer_ref`` directory already use LightX2V's runtime keys. + elif model_type in {"h3", "h3_video_vae_decoder"}: + # MiniMax-H3 transformer and Video VAE decoder checkpoints already use + # LightX2V's runtime keys. return [] else: raise ValueError(f"Unsupported model type: {model_type}") @@ -331,6 +335,7 @@ def quantize_model( preserve_non_quant_dtype=False, comfyui_mode=False, comfyui_keys=[], + quantization_policy=None, ): """ Quantize model weights in-place @@ -407,7 +412,10 @@ def quantize_model( # Quantize tensor and store results quantizer = CONVERT_WEIGHT_REGISTER[linear_type](tensor) - w_q, scales, extra = quantizer.weight_quant_func(tensor, comfyui_mode) + if quantization_policy is None: + w_q, scales, extra = quantizer.weight_quant_func(tensor, comfyui_mode) + else: + w_q, scales, extra = quantization_policy.quantize_weight(key, tensor, quantizer.weight_quant_func) weight_global_scale = extra.get("weight_global_scale", None) # For nvfp4 convrot_groupsize = extra.get("convrot_groupsize", None) @@ -448,6 +456,8 @@ def quantize_model( logger.info(f"Total final model size: {total_final_size_mb:.2f} MB") logger.info(f"Size reduction in quantized tensors: {size_reduction_mb:.2f} MB ({size_reduction_mb / original_size_mb * 100:.1f}%)") + if quantization_policy is not None: + quantization_policy.validate() if comfyui_mode: weights["scaled_fp8"] = torch.zeros(2, dtype=torch.float8_e4m3fn) @@ -827,6 +837,7 @@ def convert_key(key): preserve_non_quant_dtype=getattr(args, "preserve_non_quant_dtype", False), comfyui_mode=args.comfyui_mode, comfyui_keys=args.comfyui_keys, + quantization_policy=args.quantization_policy, ) os.makedirs(args.output, exist_ok=True) @@ -852,7 +863,8 @@ def convert_key(key): logger.warning("Consider using --save_by_block or default chunked saving for better memory efficiency.") # Save the entire model as a single file - st.save_file(converted_weights, output_path) + metadata = args.quantization_policy.metadata if args.quantization_policy is not None else None + st.save_file(converted_weights, output_path, metadata=metadata) logger.info(f"Model saved successfully to: {output_path} ({total_size_gb:.2f}GB)") except MemoryError: @@ -975,7 +987,7 @@ def main(): parser.add_argument( "-t", "--model_type", - choices=["wan_dit", "h3", "h3_text_encoder", "hunyuan_dit", "wan_t5", "wan_clip", "wan_animate_dit", "qwen_image_dit", "qwen25vl_llm", "z_image_dit", "self_forcing"], + choices=["wan_dit", "h3", "h3_video_vae_decoder", "h3_text_encoder", "hunyuan_dit", "wan_t5", "wan_clip", "wan_animate_dit", "qwen_image_dit", "qwen25vl_llm", "z_image_dit", "self_forcing"], default="wan_dit", help="Model type", ) @@ -1011,6 +1023,7 @@ def main(): parser.add_argument("--comfyui_mode", action="store_true") parser.add_argument("--full_quantized", action="store_true") parser.add_argument("--quantized", action="store_true") + parser.add_argument("--quantization_profile", choices=[FP8_F16_ACCUM_QUANTIZATION_PROFILE]) parser.add_argument("--bits", type=int, default=8, choices=[8], help="Quantization bit width") parser.add_argument( "--device", @@ -1113,6 +1126,12 @@ def _parse_csv_override(v: str | None) -> list[str] | None: # every tensor outside the quantized block linears. "preserve_non_quant_dtype": True, }, + "h3_video_vae_decoder": { + "key_idx": 1, + "target_keys": ["transformer_blocks", "proj_out"], + "ignore_key": None, + "preserve_non_quant_dtype": True, + }, "self_forcing": { "key_idx": 3, "target_keys": ["self_attn", "cross_attn", "ffn"], @@ -1189,6 +1208,15 @@ def _parse_csv_override(v: str | None) -> list[str] | None: else: args.ignore_quant_keys = None + args.quantization_policy = None + if args.quantization_profile is not None: + if not args.quantized or args.linear_type != "fp8" or not args.single_file or args.output_ext != ".safetensors" or args.comfyui_mode: + parser.error("H3 FP8-F16 accumulation conversion requires --quantized --linear_type fp8 --output_ext .safetensors --single_file without --comfyui_mode") + try: + args.quantization_policy = create_h3_fp8_f16_accum_quantization(args.quantization_profile, args.model_type) + except ValueError as profile_error: + parser.error(str(profile_error)) + if os.path.isfile(args.output): raise ValueError("Output path must be a directory, not a file") diff --git a/tools/convert/quant/h3_fp8_f16_accum.py b/tools/convert/quant/h3_fp8_f16_accum.py new file mode 100644 index 000000000..70e009d05 --- /dev/null +++ b/tools/convert/quant/h3_fp8_f16_accum.py @@ -0,0 +1,66 @@ +"""MiniMax-H3 checkpoint policy for FP8 GEMM with FP16 accumulation.""" + +import torch + +from lightx2v.models.networks.minimax_h3.fp8_f16_accum_policy import ( + FP8_F16_ACCUM_PROJECTION_SUFFIXES, + FP8_F16_ACCUM_QUANTIZATION_PROFILE, + FP8_F16_ACCUM_WEIGHT_QMAX, +) + +# DiT has 50 x (six main projections + AdaLN); VAE has 36 x six projections + proj_out. +_EXPECTED_QUANTIZED_COUNTS = { + "h3": 350, + "h3_video_vae_decoder": 217, +} +_EXPECTED_QMAX14_COUNTS = { + "h3": 300, + "h3_video_vae_decoder": 216, +} + + +class H3FP8F16AccumQuantization: + """Assign qmax14 only to H3 projections using FP16 accumulation.""" + + def __init__(self, model_type): + if model_type not in _EXPECTED_QUANTIZED_COUNTS: + raise ValueError(f"{FP8_F16_ACCUM_QUANTIZATION_PROFILE} does not support model_type={model_type!r}") + self.model_type = model_type + self.quantized_count = 0 + self.qmax14_count = 0 + + @property + def metadata(self): + return { + "format": "pt", + "quantization_profile": FP8_F16_ACCUM_QUANTIZATION_PROFILE, + "weight_qmax": str(FP8_F16_ACCUM_WEIGHT_QMAX), + } + + def quantize_weight(self, name, weight, default_quantize): + projection_name = name.removesuffix(".weight") + uses_reduced_range = projection_name.endswith(FP8_F16_ACCUM_PROJECTION_SUFFIXES) + self.quantized_count += 1 + if not uses_reduced_range: + return default_quantize(weight) + + values = weight.float() + scales = values.abs().amax(dim=1, keepdim=True).clamp_min_(1e-8).div_(FP8_F16_ACCUM_WEIGHT_QMAX) + values.div_(scales).clamp_(-FP8_F16_ACCUM_WEIGHT_QMAX, FP8_F16_ACCUM_WEIGHT_QMAX) + + self.qmax14_count += 1 + return values.to(torch.float8_e4m3fn), scales, {} + + def validate(self): + expected_quantized = _EXPECTED_QUANTIZED_COUNTS[self.model_type] + expected_qmax14 = _EXPECTED_QMAX14_COUNTS[self.model_type] + if self.quantized_count != expected_quantized or self.qmax14_count != expected_qmax14: + raise ValueError(f"Unexpected {self.model_type} FP8 conversion coverage: quantized={self.quantized_count}/{expected_quantized}, qmax14={self.qmax14_count}/{expected_qmax14}") + + +def create_h3_fp8_f16_accum_quantization(profile, model_type): + if profile is None: + return None + if profile != FP8_F16_ACCUM_QUANTIZATION_PROFILE: + raise ValueError(f"Unsupported quantization profile: {profile}") + return H3FP8F16AccumQuantization(model_type) diff --git a/tools/convert/readme.md b/tools/convert/readme.md index 16f03c259..bafb617d7 100755 --- a/tools/convert/readme.md +++ b/tools/convert/readme.md @@ -49,6 +49,8 @@ A powerful model weight conversion tool that supports format conversion, quantiz - `mxfp4`: MXFP4 quantization - `mxfp6`: MXFP6 quantization - `mxfp8`: MXFP8 quantization +- `--quantization_profile`: Optional model-specific policy. `h3-fp8-f16-accum` supports `h3` and + `h3_video_vae_decoder` with `--quantized --linear_type fp8 --single_file`. - `--non_linear_dtype`: Non-linear layer data type - `torch.bfloat16`: BF16 - `torch.float16`: FP16 diff --git a/tools/convert/readme_zh.md b/tools/convert/readme_zh.md index 41f9f0be7..56bdca36f 100755 --- a/tools/convert/readme_zh.md +++ b/tools/convert/readme_zh.md @@ -41,6 +41,8 @@ - `int8`（torch.int8） - `fp8`（torch.float8_e4m3fn） - `nvfp4` / `mxfp4` / `mxfp6` / `mxfp8` +- `--quantization_profile`：可选的模型专用策略。`h3-fp8-f16-accum` 支持 `h3` 和 + `h3_video_vae_decoder`，需配合 `--quantized --linear_type fp8 --single_file`。 - `--non_linear_dtype`：非线性层数据类型（`torch.bfloat16` / `torch.float16` / `torch.float32` 默认） - `--device`：量化设备 `cpu` 或 `cuda`（默认） - `--comfyui_mode`：ComfyUI 兼容模式（仅 int8、fp8） From 7e9696a949011a084bf0e5cbffc50f3c166ffa5c Mon Sep 17 00:00:00 2001 From: STwangyingrui Date: Fri, 4 Sep 2026 10:08:39 +0000 Subject: [PATCH 4/6] feat(minimax_h3): enable FP8 GEMM with FP16 accumulation --- lightx2v/common/ops/mm/fp8_f16_accum.py | 49 +++++++++++++++++++ lightx2v/common/ops/mm/mm_weight.py | 32 ++++++++++++ lightx2v/common/ops/mm/triton_kernels.py | 18 +++++++ lightx2v/models/input_encoders/hf/q_linear.py | 27 ++++++++++ lightx2v/models/networks/base_model.py | 1 + lightx2v/models/networks/minimax_h3/model.py | 20 ++++++++ .../minimax_h3/weights/transformer_weights.py | 13 ++++- lightx2v/models/runners/default_runner.py | 17 +++++++ .../video_encoders/hf/minimax_h3/video_vae.py | 33 ++++++++++++- 9 files changed, 206 insertions(+), 4 deletions(-) create mode 100644 lightx2v/common/ops/mm/fp8_f16_accum.py diff --git a/lightx2v/common/ops/mm/fp8_f16_accum.py b/lightx2v/common/ops/mm/fp8_f16_accum.py new file mode 100644 index 000000000..f04034f3b --- /dev/null +++ b/lightx2v/common/ops/mm/fp8_f16_accum.py @@ -0,0 +1,49 @@ +import math + +import torch + +from lightx2v.common.ops.mm.triton_kernels import fp8_quantize_range_triton + +try: + from lightx2v_kernel.gemm import FP8_F16_ACCUM_MM_AVAILABLE, cutlass_scaled_fp8_mm_f16_accum +except ImportError: + FP8_F16_ACCUM_MM_AVAILABLE = False + cutlass_scaled_fp8_mm_f16_accum = None + + +def fp8_f16_accum_mm_unavailable_reason(): + if not FP8_F16_ACCUM_MM_AVAILABLE: + return "the lightx2v-kernel extension does not provide the FP8-F16 accumulation op" + if not torch.cuda.is_available(): + return "CUDA is unavailable" + capability = torch.cuda.get_device_capability() + if capability != (12, 0): + return f"SM120 is required, but the current CUDA capability is SM{capability[0]}{capability[1]}" + return None + + +def fp8_f16_accum_mm_available(): + return fp8_f16_accum_mm_unavailable_reason() is None + + +def validate_fp8_f16_accum_qmax(activation_qmax): + activation_qmax = float(activation_qmax) + fp8_max = torch.finfo(torch.float8_e4m3fn).max + if not math.isfinite(activation_qmax) or not 0 < activation_qmax <= fp8_max: + raise ValueError(f"FP8 activation qmax must be finite and in (0, {fp8_max}], got {activation_qmax}") + return activation_qmax + + +def fp8_f16_accum_linear(input_tensor, weight, weight_scale, bias, activation_qmax): + input_shape = input_tensor.shape + input_matrix = input_tensor.reshape(-1, input_shape[-1]) + quantized, activation_scale = fp8_quantize_range_triton(input_matrix, activation_qmax) + output = cutlass_scaled_fp8_mm_f16_accum( + quantized, + weight, + activation_scale, + weight_scale.float(), + input_tensor.dtype, + bias, + ) + return output.view(*input_shape[:-1], weight.shape[1]) diff --git a/lightx2v/common/ops/mm/mm_weight.py b/lightx2v/common/ops/mm/mm_weight.py index 61613dd5a..d2a84facb 100755 --- a/lightx2v/common/ops/mm/mm_weight.py +++ b/lightx2v/common/ops/mm/mm_weight.py @@ -11,6 +11,11 @@ except ImportError: magi_register_custom_op = None +from lightx2v.common.ops.mm.fp8_f16_accum import ( + fp8_f16_accum_linear, + fp8_f16_accum_mm_available, + validate_fp8_f16_accum_qmax, +) from lightx2v.common.ops.mm.sgl_kernel import sgl_fp8_scaled_mm, sgl_fp8_scaled_mm_meta from lightx2v.common.ops.mm.triton_kernels import ( fp8_gemm_bias_triton, @@ -1977,6 +1982,33 @@ def apply(self, input_tensor): return output_tensor +@MM_WEIGHT_REGISTER("fp8-f16-accum") +class MMWeightWfp8channelAfp8channelF16Accum(MMWeightWfp8channelAfp8channeldynamicSgl): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.fp8_activation_qmax = None + + def enable_fp8_f16_accum(self, activation_qmax): + activation_qmax = validate_fp8_f16_accum_qmax(activation_qmax) + if fp8_f16_accum_mm_available(): + self.fp8_activation_qmax = activation_qmax + + def apply(self, input_tensor): + if self.fp8_activation_qmax is None: + return super().apply(input_tensor) + + output_tensor = fp8_f16_accum_linear( + input_tensor, + self.weight, + self.weight_scale, + self._get_actual_bias(), + self.fp8_activation_qmax, + ) + if self.has_lora_branch: + return output_tensor + self.apply_lora(input_tensor) + return output_tensor + + @MM_WEIGHT_REGISTER("int8-sgl") class MMWeightWint8channelAint8channeldynamicSglActVllm(MMWeightQuantTemplate): """ diff --git a/lightx2v/common/ops/mm/triton_kernels.py b/lightx2v/common/ops/mm/triton_kernels.py index 7e3fa5f7c..a846e9167 100755 --- a/lightx2v/common/ops/mm/triton_kernels.py +++ b/lightx2v/common/ops/mm/triton_kernels.py @@ -69,6 +69,24 @@ def fp8_quantize_triton(x): return quantized.view(x_shape_orig), scales.view(x_shape_orig[:-1]) +def fp8_quantize_range_triton(x, qmax): + x_shape = x.shape + x = x.reshape(-1, x_shape[-1]).contiguous() + quantized = torch.empty_like(x, dtype=torch.float8_e4m3fn) + scales = torch.empty(x.shape[0], dtype=torch.float32, device=x.device) + block_size = next_power_of_2(x_shape[-1]) + fp8_quantize_kernel[(x.shape[0],)]( + x, + quantized, + scales, + x_shape[-1], + block_size, + FP8_MAX_VAL=qmax, + num_warps=8, + ) + return quantized.view(x_shape), scales.view(x_shape[:-1]) + + def upcast_if_fp8(a): if "fp8" in str(a): return torch.float16 diff --git a/lightx2v/models/input_encoders/hf/q_linear.py b/lightx2v/models/input_encoders/hf/q_linear.py index 2142b392a..ed972c17c 100755 --- a/lightx2v/models/input_encoders/hf/q_linear.py +++ b/lightx2v/models/input_encoders/hf/q_linear.py @@ -31,6 +31,11 @@ except ImportError: fp8_linear = None +from lightx2v.common.ops.mm.fp8_f16_accum import ( + fp8_f16_accum_linear, + fp8_f16_accum_mm_available, + validate_fp8_f16_accum_qmax, +) from lightx2v.common.ops.mm.sgl_kernel import sgl_fp8_scaled_mm from lightx2v.common.ops.mm.triton_kernels import fp8_gemm_bias_triton, fp8_gemm_triton, fp8_quantize_triton, int8_gemm_bias_triton, int8_gemm_triton, int8_quantize_triton from lightx2v_platform.ops.mm.mthreads_musa.fp8_scaled_mm import fp8_linear as musa_fp8_linear @@ -310,6 +315,28 @@ def maybe_cast(t): return self +class F16AccumQuantLinearFp8(SglQuantLinearFp8): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.fp8_activation_qmax = None + + def enable_fp8_f16_accum(self, activation_qmax): + activation_qmax = validate_fp8_f16_accum_qmax(activation_qmax) + if fp8_f16_accum_mm_available(): + self.fp8_activation_qmax = activation_qmax + + def forward(self, input_tensor): + if self.fp8_activation_qmax is None: + return super().forward(input_tensor) + return fp8_f16_accum_linear( + input_tensor, + self.weight.t(), + self.weight_scale, + self.bias, + self.fp8_activation_qmax, + ) + + class MusaQuantLinearFp8(nn.Module): """MUSA W8A8 FP8 linear with per-channel weights and per-token inputs.""" diff --git a/lightx2v/models/networks/base_model.py b/lightx2v/models/networks/base_model.py index 35a8ceaf7..c666842f2 100644 --- a/lightx2v/models/networks/base_model.py +++ b/lightx2v/models/networks/base_model.py @@ -122,6 +122,7 @@ def _check_dit_quantized(self): "int8-q8f", "int8-convrot", "fp8-b128-deepgemm", + "fp8-f16-accum", "fp8-sgl", "int8-sgl", "int8-torchao", diff --git a/lightx2v/models/networks/minimax_h3/model.py b/lightx2v/models/networks/minimax_h3/model.py index 19ca82476..250044b14 100644 --- a/lightx2v/models/networks/minimax_h3/model.py +++ b/lightx2v/models/networks/minimax_h3/model.py @@ -7,7 +7,13 @@ from loguru import logger from safetensors import safe_open +from lightx2v.common.ops.mm.fp8_f16_accum import fp8_f16_accum_mm_unavailable_reason from lightx2v.models.networks.base_model import BaseTransformerModel +from lightx2v.models.networks.minimax_h3.fp8_f16_accum_policy import ( + DIT_FP8_F16_ACCUM_ACTIVATION_QMAX, + FP8_F16_ACCUM_WEIGHT_QMAX, + validate_fp8_f16_accum_checkpoint, +) from lightx2v.models.networks.minimax_h3.infer.module_io import MiniMaxH3SequenceParallelState from lightx2v.models.networks.minimax_h3.infer.offload import MiniMaxH3OffloadTransformerInfer from lightx2v.models.networks.minimax_h3.infer.post_infer import MiniMaxH3PostInfer @@ -23,6 +29,7 @@ H3_CHANNEL_QUANT_SCHEMES = { "fp8-q8f", + "fp8-f16-accum", "fp8-musa", "fp8-sgl", "fp8-torchao", @@ -64,6 +71,19 @@ def __init__(self, model_path, config, device, lora_path=None, lora_strength=1.0 raise NotImplementedError(f"MiniMax-H3 quantized inference requires a per-output-channel FP8/INT8 scheme; got {quant_scheme!r}. Supported schemes: {sorted(H3_CHANNEL_QUANT_SCHEMES)}") if not config.get("dit_quantized_ckpt"): raise ValueError("MiniMax-H3 quantized inference requires dit_quantized_ckpt") + if quant_scheme == "fp8-f16-accum": + validate_fp8_f16_accum_checkpoint(config["dit_quantized_ckpt"]) + fallback_reason = fp8_f16_accum_mm_unavailable_reason() + if config.get("tensor_parallel", False): + logger.info("MiniMax-H3 DiT FP8-F16 accumulation falls back to FP8-SGL under tensor parallel") + elif fallback_reason is not None: + logger.warning("MiniMax-H3 DiT FP8-F16 accumulation requested but {}; falling back to FP8-SGL", fallback_reason) + else: + logger.info( + "MiniMax-H3 DiT FP8-F16 accumulation enabled for Q/K/V, attention output, and FFN projections (weight qmax={}, activation qmax={})", + FP8_F16_ACCUM_WEIGHT_QMAX, + DIT_FP8_F16_ACCUM_ACTIVATION_QMAX, + ) elif config.get("dit_quant_scheme", "Default") != "Default": raise ValueError("MiniMax-H3 dit_quant_scheme requires a dit_quantized_ckpt") if config.get("cpu_offload", False) and config.get("offload_granularity", "model") not in {"model", "block"}: diff --git a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py index 9634eb595..e654cfa6b 100644 --- a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py +++ b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py @@ -2,19 +2,24 @@ import torch.distributed as dist from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList +from lightx2v.models.networks.minimax_h3.fp8_f16_accum_policy import ( + DIT_FP8_F16_ACCUM_ACTIVATION_QMAX, + FP8_F16_ACCUM_PROJECTION_SUFFIXES, +) from lightx2v.models.networks.minimax_h3.infer.triton_ops import MiniMaxH3TritonRope # noqa: F401 from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, ROPE_REGISTER def _linear(config, name, bias=False, create_cuda_buffer=False, tp_split=None): lora_prefix = "transformer_blocks" + quant_scheme = config.get("dit_quant_scheme", "Default") if config.get("tensor_parallel", False) and tp_split is not None: tp_group = config["device_mesh"].get_group(mesh_dim="tensor_p") tp_mm_type = config.get("tp_mm_type", "TensorParallel") return MM_WEIGHT_REGISTER[tp_mm_type]( weight_name=f"{name}.weight", bias_name=f"{name}.bias" if bias else None, - mm_type=config.get("dit_quant_scheme", "Default"), + mm_type=quant_scheme, tp_group=tp_group, tp_rank=dist.get_rank(tp_group), tp_size=dist.get_world_size(tp_group), @@ -23,12 +28,16 @@ def _linear(config, name, bias=False, create_cuda_buffer=False, tp_split=None): create_cuda_buffer=create_cuda_buffer, lora_prefix=lora_prefix, ) - return MM_WEIGHT_REGISTER[config.get("dit_quant_scheme", "Default")]( + + linear = MM_WEIGHT_REGISTER[quant_scheme]( f"{name}.weight", f"{name}.bias" if bias else None, create_cuda_buffer=create_cuda_buffer, lora_prefix=lora_prefix, ) + if quant_scheme == "fp8-f16-accum" and name.endswith(FP8_F16_ACCUM_PROJECTION_SUFFIXES): + linear.enable_fp8_f16_accum(DIT_FP8_F16_ACCUM_ACTIVATION_QMAX) + return linear def _rms(config, name, eps, create_cuda_buffer=False): diff --git a/lightx2v/models/runners/default_runner.py b/lightx2v/models/runners/default_runner.py index 2d5ff4c0f..ca768c2af 100755 --- a/lightx2v/models/runners/default_runner.py +++ b/lightx2v/models/runners/default_runner.py @@ -86,6 +86,7 @@ class DefaultRunner(BaseRunner): def __init__(self, config): super().__init__(config) self.progress_callback = None + self._fp8_f16_accum_autotuner = None self.reuse_cache_path = self.config.get("reuse_cache_path") if self.enable_reuse and not self.reuse_cache_path: raise ValueError("enable_reuse requires reuse_cache_path") @@ -223,6 +224,9 @@ def warmup(self): if dist.is_initialized() and dist.get_world_size() > 1: dist.barrier() + if self._fp8_f16_accum_autotuner: + self._fp8_f16_accum_autotuner.save() + def run_warmup(self): raise NotImplementedError(f"Warmup is not supported for {type(self).__name__}") @@ -252,6 +256,17 @@ def init_modules(self): self.run_input_encoder = self._run_input_encoder_local_i2av elif self.config["task"] == "sr": self.run_input_encoder = self._run_input_encoder_local_sr + + uses_fp8_f16_accum = any(self.config.get(key) == "fp8-f16-accum" for key in ("dit_quant_scheme", "video_vae_quant_scheme")) + if uses_fp8_f16_accum: + try: + from lightx2v_kernel.fp8_f16_autotune import Fp8F16AccumAutotuner + + self._fp8_f16_accum_autotuner = Fp8F16AccumAutotuner(self.config.get("fp8_f16_accum_autotune_cache")) + entry_count = self._fp8_f16_accum_autotuner.start() + logger.info(f"FP8-F16 GEMM autotune loaded {entry_count} entries from {self._fp8_f16_accum_autotuner.cache_path}") + except (ImportError, ValueError) as error: + logger.warning(f"FP8-F16 GEMM autotune is unavailable: {error}") self.config.lock() # lock config to avoid modification def set_init_device(self): @@ -649,6 +664,8 @@ def run_pipeline(self, input_info): if GET_RECORDER_MODE(): monitor_cli.lightx2v_worker_request_success.inc() + if self._fp8_f16_accum_autotuner: + self._fp8_f16_accum_autotuner.save() return gen_video_final def switch_lora(self, lora_path: str, strength: float = 1.0): diff --git a/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py b/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py index 318348b2d..7658e5ca4 100644 --- a/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py +++ b/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py @@ -43,6 +43,12 @@ import torch.nn.functional as F from loguru import logger +from lightx2v.common.ops.mm.fp8_f16_accum import fp8_f16_accum_mm_unavailable_reason +from lightx2v.models.networks.minimax_h3.fp8_f16_accum_policy import ( + FP8_F16_ACCUM_WEIGHT_QMAX, + VIDEO_VAE_FP8_F16_ACCUM_ACTIVATION_QMAX, + validate_fp8_f16_accum_checkpoint, +) from lightx2v.models.video_encoders.hf.minimax_h3.weights import ( SafetensorsSubsetReport, load_safetensors_subset, @@ -552,7 +558,7 @@ def __init__( attn_type: str = "torch_sdpa", ) -> None: super().__init__() - if quant_scheme not in {None, "fp8-musa", "fp8-sgl"}: + if quant_scheme not in {None, "fp8-f16-accum", "fp8-musa", "fp8-sgl"}: raise NotImplementedError(f"Unsupported MiniMax-H3 video VAE quantization scheme: {quant_scheme!r}") if attn_type not in {"torch_sdpa", "sage_attn2"}: raise ValueError(f"Unsupported MiniMax-H3 video VAE attention type: {attn_type!r}; expected torch_sdpa or sage_attn2") @@ -649,8 +655,18 @@ def _pack_decoder_fp8_qkv(self) -> None: for block in self.decoder.transformer_blocks: block.attn._pack_fp8_qkv() + def _configure_fp8_f16_accum_linears(self) -> None: + # Packed QKV, attention output, and FFN use the validated H3 shapes. + for block in self.decoder.transformer_blocks: + block.attn.to_qkv.enable_fp8_f16_accum(VIDEO_VAE_FP8_F16_ACCUM_ACTIVATION_QMAX) + block.attn.to_out[0].enable_fp8_f16_accum(VIDEO_VAE_FP8_F16_ACCUM_ACTIVATION_QMAX) + block.ff.net[0].proj.enable_fp8_f16_accum(VIDEO_VAE_FP8_F16_ACCUM_ACTIVATION_QMAX) + block.ff.net[2].enable_fp8_f16_accum(VIDEO_VAE_FP8_F16_ACCUM_ACTIVATION_QMAX) + def _make_fp8_linear(self, linear: nn.Linear) -> nn.Module: - if self.quant_scheme == "fp8-musa": + if self.quant_scheme == "fp8-f16-accum": + from lightx2v.models.input_encoders.hf.q_linear import F16AccumQuantLinearFp8 as linear_cls + elif self.quant_scheme == "fp8-musa": from lightx2v.models.input_encoders.hf.q_linear import MusaQuantLinearFp8 as linear_cls elif self.quant_scheme == "fp8-sgl": from lightx2v.models.input_encoders.hf.q_linear import SglQuantLinearFp8 as linear_cls @@ -703,6 +719,17 @@ def from_pretrained( if (checkpoint_path is None) != (quant_scheme is None): raise ValueError("MiniMax-H3 video VAE checkpoint_path and quant_scheme must be configured together") weight_path = checkpoint_path if checkpoint_path is not None else vae_dir + if quant_scheme == "fp8-f16-accum": + validate_fp8_f16_accum_checkpoint(weight_path) + fallback_reason = fp8_f16_accum_mm_unavailable_reason() + if fallback_reason is None: + logger.info( + "MiniMax-H3 Video VAE FP8-F16 accumulation enabled for packed QKV, attention output, and FFN projections (weight qmax={}, activation qmax={})", + FP8_F16_ACCUM_WEIGHT_QMAX, + VIDEO_VAE_FP8_F16_ACCUM_ACTIVATION_QMAX, + ) + else: + logger.warning("MiniMax-H3 Video VAE FP8-F16 accumulation requested but {}; falling back to FP8-SGL", fallback_reason) with (vae_dir / "config.json").open("r", encoding="utf-8") as handle: config = json.load(handle) @@ -723,6 +750,8 @@ def from_pretrained( if quant_scheme is not None: # Pack only after loading the checkpoint's original Q/K/V keys. model._pack_decoder_fp8_qkv() + if quant_scheme == "fp8-f16-accum": + model._configure_fp8_f16_accum_linears() model._prepare_inference_dtypes() model.eval().requires_grad_(False) if not cpu_offload: From db5037121974abdf01ece421dbd77aea9cc1b3d0 Mon Sep 17 00:00:00 2001 From: STwangyingrui Date: Fri, 4 Sep 2026 10:08:55 +0000 Subject: [PATCH 5/6] docs(minimax_h3): document FP8-F16 accumulation --- ...ax_h3_fp8_4step_5090_with_fp8_vae_sla.json | 8 +-- .../source/method_tutorials/quantization.md | 72 ++++++++++++++++++- .../source/method_tutorials/quantization.md | 68 +++++++++++++++++- 3 files changed, 142 insertions(+), 6 deletions(-) diff --git a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_with_fp8_vae_sla.json b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_with_fp8_vae_sla.json index 3676c0136..2042b0584 100755 --- a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_with_fp8_vae_sla.json +++ b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_with_fp8_vae_sla.json @@ -36,11 +36,11 @@ "audio_channels": 2, "keep_latents_dtype_in_scheduler": true, "dit_quantized": true, - "dit_quant_scheme": "fp8-sgl", - "dit_quantized_ckpt": "/path/to/models/minimax_h3/h3_quantized/fp8/minimax_h3_fp8.safetensors", + "dit_quant_scheme": "fp8-f16-accum", + "dit_quantized_ckpt": "/path/to/models/minimax_h3/h3_quantized/minimax_h3_dit_fp8_f16_accum.safetensors", "video_vae_quantized": true, - "video_vae_quant_scheme": "fp8-sgl", - "video_vae_quantized_ckpt": "/path/to/models/minimax_h3/h3_quantized/fp8/minimax_h3_video_vae_fp8_sgl_bias_fp16.safetensors", + "video_vae_quant_scheme": "fp8-f16-accum", + "video_vae_quantized_ckpt": "/path/to/models/minimax_h3/h3_quantized/minimax_h3_video_vae_fp8_f16_accum.safetensors", "lora_dynamic_apply": true, "lora_configs": [ { diff --git a/docs/EN/source/method_tutorials/quantization.md b/docs/EN/source/method_tutorials/quantization.md index da355a92a..cde7f1413 100644 --- a/docs/EN/source/method_tutorials/quantization.md +++ b/docs/EN/source/method_tutorials/quantization.md @@ -13,6 +13,7 @@ LightX2V supports quantized inference for DIT, T5, and CLIP models, reducing mem | `fp8-vllm` | FP8 channel symmetric | FP8 channel dynamic symmetric | [VLLM](https://github.com/vllm-project/vllm) | H100/H200/H800, RTX 40 series, etc. | | `int8-vllm` | INT8 channel symmetric | INT8 channel dynamic symmetric | [VLLM](https://github.com/vllm-project/vllm) | A100/A800, RTX 30/40 series, etc. | | `fp8-sgl` | FP8 channel symmetric | FP8 channel dynamic symmetric | [SGL](https://github.com/sgl-project/sglang/tree/main/sgl-kernel) | H100/H200/H800, RTX 40 series, etc. | +| `fp8-f16-accum` | FP8 channel symmetric | FP8 row-wise dynamic symmetric | CUTLASS FP16 accumulation | RTX 5090 (SM120) | | `int8-sgl` | INT8 channel symmetric | INT8 channel dynamic symmetric | [SGL](https://github.com/sgl-project/sglang/tree/main/sgl-kernel) | A100/A800, RTX 30/40 series, etc. | | `fp8-q8f` | FP8 channel symmetric | FP8 channel dynamic symmetric | [Q8-Kernels](https://github.com/KONAKONA666/q8_kernels) | RTX 40 series, L40S, etc. | | `int8-q8f` | INT8 channel symmetric | INT8 channel dynamic symmetric | [Q8-Kernels](https://github.com/KONAKONA666/q8_kernels) | RTX 40 series, L40S, etc. | @@ -67,7 +68,7 @@ For detailed quantization tool usage, refer to: [Model Conversion Documentation] #### Supported Quantization Modes -DIT quantization modes (`dit_quant_scheme`) support: `fp8-vllm`, `int8-vllm`, `fp8-sgl`, `int8-sgl`, `fp8-q8f`, `int8-q8f`, `int8-torchao`, `int4-g128-marlin`, `fp8-b128-deepgemm` +DIT quantization modes (`dit_quant_scheme`) support: `fp8-vllm`, `int8-vllm`, `fp8-sgl`, `fp8-f16-accum`, `int8-sgl`, `fp8-q8f`, `int8-q8f`, `int8-torchao`, `int4-g128-marlin`, `fp8-b128-deepgemm` #### Configuration Example @@ -81,6 +82,75 @@ DIT quantization modes (`dit_quant_scheme`) support: `fp8-vllm`, `int8-vllm`, `f > 💡 **Tip**: When there's only one DIT model in the script's `model_path`, `dit_quantized_ckpt` doesn't need to be specified separately. +#### MiniMax-H3 FP8 with FP16 Accumulation + +On RTX 5090, MiniMax-H3 can use FP8 inputs with FP16 accumulation through `fp8-f16-accum`. Convert +the weights with the `h3-fp8-f16-accum` profile; regular `fp8-sgl` checkpoints are not compatible. +DiT and Video VAE decoder are converted separately. The profile selects the qmax-14 projections, +keeps standard FP8 quantization for the remaining layers, and records the policy in safetensors +metadata. + +```bash +python tools/convert/converter.py \ + --source /path/to/MiniMax-H3/transformer \ + --output /path/to/h3_quantized \ + --output_name minimax_h3_dit_fp8_f16_accum \ + --output_ext .safetensors \ + --model_type h3 \ + --device cuda \ + --quantized \ + --bits 8 \ + --linear_type fp8 \ + --quantization_profile h3-fp8-f16-accum \ + --single_file + +python tools/convert/converter.py \ + --source /path/to/MiniMax-H3/vae \ + --output /path/to/h3_quantized \ + --output_name minimax_h3_video_vae_fp8_f16_accum \ + --output_ext .safetensors \ + --model_type h3_video_vae_decoder \ + --device cuda \ + --quantized \ + --bits 8 \ + --linear_type fp8 \ + --quantization_profile h3-fp8-f16-accum \ + --single_file +``` + +```json +{ + "dit_quantized": true, + "dit_quant_scheme": "fp8-f16-accum", + "dit_quantized_ckpt": "/path/to/minimax_h3_dit_fp8_f16_accum.safetensors", + "video_vae_quantized": true, + "video_vae_quant_scheme": "fp8-f16-accum", + "video_vae_quantized_ckpt": "/path/to/minimax_h3_video_vae_fp8_f16_accum.safetensors" +} +``` + +Activations use dynamic row-wise quantization with `scale = max(abs(x)) / qmax`. Reducing qmax +increases the scale and lowers the raw values accumulated in FP16, at the cost of fewer effective FP8 +levels. In the validated MiniMax-H3 workload, DiT produced non-finite FFN-out values with qmax 14 and +12, while qmax 7 completed every denoising step. Video VAE decoder remained finite and had the lowest +error with qmax 14. The current H3 policy therefore fixes activation qmax to 7 for DiT and 14 for +Video VAE, avoiding mismatches between runtime configuration and checkpoint conversion. + +The kernel is enabled only for DiT Q/K/V, attention output, and FFN projections, and for Video VAE +packed QKV, attention output, and FFN projections. It falls back to `fp8-sgl` when the extension is +unavailable or the device is not SM120. DiT tensor parallel also currently uses the `fp8-sgl` fallback. +Initialization logs report the effective scope or fallback reason. All other pipeline settings are +independent of this quantization mode. + +The kernel automatically tunes its CUTLASS tile and swizzle for each exact GEMM shape. The first use +of an unseen shape benchmarks the built-in candidates in C++; later calls perform only a C++ cache +lookup. After warmup or a request, the runner merges new winners into a persistent, device-specific +cache under `$XDG_CACHE_HOME/lightx2v/autotune/fp8_f16_accum` or +`~/.cache/lightx2v/autotune/fp8_f16_accum`. A later process loads compatible entries automatically. +Enabling `warmup` moves the tuning cost out of the first request when warmup covers the production +shapes. The optional `fp8_f16_accum_autotune_cache` setting overrides the cache file location. Cache +entries are rejected when the device, CUDA, PyTorch, or kernel ABI differs. + ### T5 Model Quantization #### Supported Quantization Modes diff --git a/docs/ZH_CN/source/method_tutorials/quantization.md b/docs/ZH_CN/source/method_tutorials/quantization.md index 311367cc6..11f609c36 100644 --- a/docs/ZH_CN/source/method_tutorials/quantization.md +++ b/docs/ZH_CN/source/method_tutorials/quantization.md @@ -13,6 +13,7 @@ LightX2V 支持对 DIT、T5 和 CLIP 模型进行量化推理，通过降低模 | `fp8-vllm` | FP8 通道对称 | FP8 通道动态对称 | [VLLM](https://github.com/vllm-project/vllm) | H100/H200/H800, RTX 40系等 | | `int8-vllm` | INT8 通道对称 | INT8 通道动态对称 | [VLLM](https://github.com/vllm-project/vllm) | A100/A800, RTX 30/40系等 | | `fp8-sgl` | FP8 通道对称 | FP8 通道动态对称 | [SGL](https://github.com/sgl-project/sglang/tree/main/sgl-kernel) | H100/H200/H800, RTX 40系等 | +| `fp8-f16-accum` | FP8 通道对称 | FP8 行动态对称 | CUTLASS FP16 累加 | RTX 5090（SM120） | | `int8-sgl` | INT8 通道对称 | INT8 通道动态对称 | [SGL](https://github.com/sgl-project/sglang/tree/main/sgl-kernel) | A100/A800, RTX 30/40系等 | | `fp8-q8f` | FP8 通道对称 | FP8 通道动态对称 | [Q8-Kernels](https://github.com/KONAKONA666/q8_kernels) | RTX 40系, L40S等 | | `int8-q8f` | INT8 通道对称 | INT8 通道动态对称 | [Q8-Kernels](https://github.com/KONAKONA666/q8_kernels) | RTX 40系, L40S等 | @@ -67,7 +68,7 @@ huggingface-cli download lightx2v/Encoders-Lightx2v \ #### 支持的量化模式 -DIT 量化模式（`dit_quant_scheme`）支持：`fp8-vllm`、`int8-vllm`、`fp8-sgl`、`int8-sgl`、`fp8-q8f`、`int8-q8f`、`int8-torchao`、`int4-g128-marlin`、`fp8-b128-deepgemm` +DIT 量化模式（`dit_quant_scheme`）支持：`fp8-vllm`、`int8-vllm`、`fp8-sgl`、`fp8-f16-accum`、`int8-sgl`、`fp8-q8f`、`int8-q8f`、`int8-torchao`、`int4-g128-marlin`、`fp8-b128-deepgemm` #### 配置示例 @@ -81,6 +82,71 @@ DIT 量化模式（`dit_quant_scheme`）支持：`fp8-vllm`、`int8-vllm`、`fp8 > 💡 **提示**：当运行脚本的 `model_path` 中只有一个 DIT 模型时，`dit_quantized_ckpt` 可以不用单独指定。 +#### MiniMax-H3 FP8 FP16 累加 + +RTX 5090 上的 MiniMax-H3 可以通过 `fp8-f16-accum` 使用 FP8 输入和 FP16 累加。权重需要用 +`h3-fp8-f16-accum` profile 转换；普通 `fp8-sgl` checkpoint 不兼容。DiT 和 Video VAE decoder +分别转换，profile 会独立选择使用 qmax 14 的投影层、保留其他层的标准 FP8 量化，并把策略写入 +safetensors metadata。 + +```bash +python tools/convert/converter.py \ + --source /path/to/MiniMax-H3/transformer \ + --output /path/to/h3_quantized \ + --output_name minimax_h3_dit_fp8_f16_accum \ + --output_ext .safetensors \ + --model_type h3 \ + --device cuda \ + --quantized \ + --bits 8 \ + --linear_type fp8 \ + --quantization_profile h3-fp8-f16-accum \ + --single_file + +python tools/convert/converter.py \ + --source /path/to/MiniMax-H3/vae \ + --output /path/to/h3_quantized \ + --output_name minimax_h3_video_vae_fp8_f16_accum \ + --output_ext .safetensors \ + --model_type h3_video_vae_decoder \ + --device cuda \ + --quantized \ + --bits 8 \ + --linear_type fp8 \ + --quantization_profile h3-fp8-f16-accum \ + --single_file +``` + +```json +{ + "dit_quantized": true, + "dit_quant_scheme": "fp8-f16-accum", + "dit_quantized_ckpt": "/path/to/minimax_h3_dit_fp8_f16_accum.safetensors", + "video_vae_quantized": true, + "video_vae_quant_scheme": "fp8-f16-accum", + "video_vae_quantized_ckpt": "/path/to/minimax_h3_video_vae_fp8_f16_accum.safetensors" +} +``` + +激活按行动态量化，`scale = max(abs(x)) / qmax`。减小 qmax 会扩大 scale，从而降低 FP16 +累加器中的原始数值范围，但也会减少 FP8 有效量化级数。实测中 DiT 的 qmax 14 和 12 会在 FFN-out +产生非有限值，qmax 7 可完成全部去噪步骤；Video VAE decoder 在 qmax 14 下保持有限且误差最小。因此 +当前 H3 策略固定使用 DiT activation qmax 7 和 Video VAE activation qmax 14，避免运行配置与 +checkpoint 的转换策略错配。 + +DiT 仅对 Q/K/V、attention output 和 FFN projection 启用该内核，Video VAE 仅对 packed QKV、 +attention output 和 FFN projection 启用。扩展不可用或设备不是 SM120 时会回退到 `fp8-sgl`； +DiT tensor parallel 当前也回退到 `fp8-sgl`。初始化日志会打印实际启用范围或回退原因。 +其他 pipeline 配置与该量化模式相互独立。 + +该内核会按精确 GEMM shape 自动调优 CUTLASS tile 和 swizzle。首次遇到新 shape 时在 C++ 内遍历 +内置候选；后续调用只执行 C++ cache 查询。warmup 或请求结束后，runner 会把新增结果合并进与设备 +绑定的持久化 cache。默认路径为 `$XDG_CACHE_HOME/lightx2v/autotune/fp8_f16_accum`，未设置 +`XDG_CACHE_HOME` 时则使用 `~/.cache/lightx2v/autotune/fp8_f16_accum`。后续进程会自动加载兼容 +结果；若 warmup 覆盖正式请求的 shape，首次调优开销也会在请求前完成。高级用户可通过 +`fp8_f16_accum_autotune_cache` 指定其他 cache 文件。设备、CUDA、PyTorch 或 kernel ABI 不一致时， +已有条目不会被复用。 + ### T5 模型量化 #### 支持的量化模式 From a527461aafd7debcd2ad0e48df53db574b8fca55 Mon Sep 17 00:00:00 2001 From: STwangyingrui Date: Tue, 8 Sep 2026 13:28:34 +0000 Subject: [PATCH 6/6] refactor(kernel): keep FP8-F16 autotuning process-local --- .../source/method_tutorials/quantization.md | 11 +- .../source/method_tutorials/quantization.md | 8 +- lightx2v/models/runners/default_runner.py | 16 -- lightx2v_kernel/csrc/common_extension.cc | 27 --- .../fp8_f16_accum_scaled_mm_kernels_sm120.cu | 148 +----------- lightx2v_kernel/include/lightx2v_kernel_ops.h | 14 -- .../lightx2v_kernel/fp8_f16_autotune.py | 224 ------------------ 7 files changed, 16 insertions(+), 432 deletions(-) delete mode 100644 lightx2v_kernel/python/lightx2v_kernel/fp8_f16_autotune.py diff --git a/docs/EN/source/method_tutorials/quantization.md b/docs/EN/source/method_tutorials/quantization.md index cde7f1413..142d3b62d 100644 --- a/docs/EN/source/method_tutorials/quantization.md +++ b/docs/EN/source/method_tutorials/quantization.md @@ -143,13 +143,10 @@ Initialization logs report the effective scope or fallback reason. All other pip independent of this quantization mode. The kernel automatically tunes its CUTLASS tile and swizzle for each exact GEMM shape. The first use -of an unseen shape benchmarks the built-in candidates in C++; later calls perform only a C++ cache -lookup. After warmup or a request, the runner merges new winners into a persistent, device-specific -cache under `$XDG_CACHE_HOME/lightx2v/autotune/fp8_f16_accum` or -`~/.cache/lightx2v/autotune/fp8_f16_accum`. A later process loads compatible entries automatically. -Enabling `warmup` moves the tuning cost out of the first request when warmup covers the production -shapes. The optional `fp8_f16_accum_autotune_cache` setting overrides the cache file location. Cache -entries are rejected when the device, CUDA, PyTorch, or kernel ABI differs. +of an unseen shape benchmarks the built-in candidates in C++ and keeps the winner in a process-local +C++ cache; later calls in the same process perform only a cache lookup. Enabling `warmup` moves the +tuning cost out of the first request when warmup covers the production shapes. A restarted process +tunes its shapes again and does not write to the user's cache directory. ### T5 Model Quantization diff --git a/docs/ZH_CN/source/method_tutorials/quantization.md b/docs/ZH_CN/source/method_tutorials/quantization.md index 11f609c36..63a8b1660 100644 --- a/docs/ZH_CN/source/method_tutorials/quantization.md +++ b/docs/ZH_CN/source/method_tutorials/quantization.md @@ -140,12 +140,8 @@ DiT tensor parallel 当前也回退到 `fp8-sgl`。初始化日志会打印实 其他 pipeline 配置与该量化模式相互独立。 该内核会按精确 GEMM shape 自动调优 CUTLASS tile 和 swizzle。首次遇到新 shape 时在 C++ 内遍历 -内置候选；后续调用只执行 C++ cache 查询。warmup 或请求结束后，runner 会把新增结果合并进与设备 -绑定的持久化 cache。默认路径为 `$XDG_CACHE_HOME/lightx2v/autotune/fp8_f16_accum`，未设置 -`XDG_CACHE_HOME` 时则使用 `~/.cache/lightx2v/autotune/fp8_f16_accum`。后续进程会自动加载兼容 -结果；若 warmup 覆盖正式请求的 shape，首次调优开销也会在请求前完成。高级用户可通过 -`fp8_f16_accum_autotune_cache` 指定其他 cache 文件。设备、CUDA、PyTorch 或 kernel ABI 不一致时， -已有条目不会被复用。 +内置候选，winner 保存在当前进程的 C++ cache 中；同一进程的后续调用只执行 cache 查询。若 warmup +覆盖正式请求的 shape，首次调优开销会在请求前完成。进程重启后会重新调优一次，不写用户目录。 ### T5 模型量化 diff --git a/lightx2v/models/runners/default_runner.py b/lightx2v/models/runners/default_runner.py index ca768c2af..d13b73d40 100755 --- a/lightx2v/models/runners/default_runner.py +++ b/lightx2v/models/runners/default_runner.py @@ -86,7 +86,6 @@ class DefaultRunner(BaseRunner): def __init__(self, config): super().__init__(config) self.progress_callback = None - self._fp8_f16_accum_autotuner = None self.reuse_cache_path = self.config.get("reuse_cache_path") if self.enable_reuse and not self.reuse_cache_path: raise ValueError("enable_reuse requires reuse_cache_path") @@ -224,9 +223,6 @@ def warmup(self): if dist.is_initialized() and dist.get_world_size() > 1: dist.barrier() - if self._fp8_f16_accum_autotuner: - self._fp8_f16_accum_autotuner.save() - def run_warmup(self): raise NotImplementedError(f"Warmup is not supported for {type(self).__name__}") @@ -257,16 +253,6 @@ def init_modules(self): elif self.config["task"] == "sr": self.run_input_encoder = self._run_input_encoder_local_sr - uses_fp8_f16_accum = any(self.config.get(key) == "fp8-f16-accum" for key in ("dit_quant_scheme", "video_vae_quant_scheme")) - if uses_fp8_f16_accum: - try: - from lightx2v_kernel.fp8_f16_autotune import Fp8F16AccumAutotuner - - self._fp8_f16_accum_autotuner = Fp8F16AccumAutotuner(self.config.get("fp8_f16_accum_autotune_cache")) - entry_count = self._fp8_f16_accum_autotuner.start() - logger.info(f"FP8-F16 GEMM autotune loaded {entry_count} entries from {self._fp8_f16_accum_autotuner.cache_path}") - except (ImportError, ValueError) as error: - logger.warning(f"FP8-F16 GEMM autotune is unavailable: {error}") self.config.lock() # lock config to avoid modification def set_init_device(self): @@ -664,8 +650,6 @@ def run_pipeline(self, input_info): if GET_RECORDER_MODE(): monitor_cli.lightx2v_worker_request_success.inc() - if self._fp8_f16_accum_autotuner: - self._fp8_f16_accum_autotuner.save() return gen_video_final def switch_lora(self, lora_path: str, strength: float = 1.0): diff --git a/lightx2v_kernel/csrc/common_extension.cc b/lightx2v_kernel/csrc/common_extension.cc index 2061d95d2..cdded046a 100644 --- a/lightx2v_kernel/csrc/common_extension.cc +++ b/lightx2v_kernel/csrc/common_extension.cc @@ -22,33 +22,6 @@ TORCH_LIBRARY_FRAGMENT(lightx2v_kernel, m) { torch::kCUDA, &cutlass_scaled_fp8_mm_f16_accum_with_config_sm120); - m.def("fp8_f16_accum_autotune_cache_abi_sm120() -> int"); - m.impl( - "fp8_f16_accum_autotune_cache_abi_sm120", - &fp8_f16_accum_autotune_cache_abi_sm120); - m.def("fp8_f16_accum_autotune_configs_sm120() -> str[]"); - m.impl( - "fp8_f16_accum_autotune_configs_sm120", - &fp8_f16_accum_autotune_configs_sm120); - m.def( - "set_fp8_f16_accum_autotune_config_sm120(int device_index, int m, int n, int k, " - "ScalarType out_dtype, bool has_bias, int config_id) -> ()"); - m.impl( - "set_fp8_f16_accum_autotune_config_sm120", - &set_fp8_f16_accum_autotune_config_sm120); - m.def("set_fp8_f16_accum_autotune_enabled_sm120(bool enabled) -> ()"); - m.impl( - "set_fp8_f16_accum_autotune_enabled_sm120", - &set_fp8_f16_accum_autotune_enabled_sm120); - m.def( - "get_fp8_f16_accum_autotune_cache_sm120(int device_index) -> Tensor"); - m.impl( - "get_fp8_f16_accum_autotune_cache_sm120", - &get_fp8_f16_accum_autotune_cache_sm120); - m.def("clear_fp8_f16_accum_autotune_cache_sm120(int device_index=-1) -> ()"); - m.impl( - "clear_fp8_f16_accum_autotune_cache_sm120", - &clear_fp8_f16_accum_autotune_cache_sm120); m.def( "cutlass_scaled_nvfp4_mm_sm120(Tensor! out, Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, Tensor " diff --git a/lightx2v_kernel/csrc/gemm/fp8_f16_accum_scaled_mm_kernels_sm120.cu b/lightx2v_kernel/csrc/gemm/fp8_f16_accum_scaled_mm_kernels_sm120.cu index 115c00814..790f09b74 100644 --- a/lightx2v_kernel/csrc/gemm/fp8_f16_accum_scaled_mm_kernels_sm120.cu +++ b/lightx2v_kernel/csrc/gemm/fp8_f16_accum_scaled_mm_kernels_sm120.cu @@ -7,13 +7,11 @@ #include #include -#include #include #include #include #include #include -#include #include #include #include @@ -183,25 +181,20 @@ using NarrowGemmFp16WithBias = using WideGemmFp16WithBias = GemmDefinition, cutlass::half_t, true>; -// Bump the ABI whenever a config mapping or candidate implementation changes. -constexpr int64_t kAutotuneCacheAbi = 1; -constexpr int64_t kFallbackConfigId = 0; - struct KernelConfig { bool wide_tile; int swizzle; - char const* name; }; constexpr std::array kKernelConfigs = {{ - {false, 1, "tile_128x128x64_swizzle_1"}, - {false, 2, "tile_128x128x64_swizzle_2"}, - {false, 4, "tile_128x128x64_swizzle_4"}, - {false, 8, "tile_128x128x64_swizzle_8"}, - {true, 1, "tile_128x256x64_swizzle_1"}, - {true, 2, "tile_128x256x64_swizzle_2"}, - {true, 4, "tile_128x256x64_swizzle_4"}, - {true, 8, "tile_128x256x64_swizzle_8"}, + {false, 1}, + {false, 2}, + {false, 4}, + {false, 8}, + {true, 1}, + {true, 2}, + {true, 4}, + {true, 8}, }}; KernelConfig const& kernel_config(int64_t config_id) { @@ -259,11 +252,6 @@ std::mutex& autotune_measurement_mutex() { return mutex; } -std::atomic& autotune_enabled() { - static std::atomic enabled{false}; - return enabled; -} - std::optional cached_config_id(AutotuneKey const& key) { std::shared_lock lock(autotune_cache_mutex()); auto entry = autotune_cache().find(key); @@ -518,7 +506,7 @@ int64_t tune_config( } } - int64_t best_config_id = kFallbackConfigId; + int64_t best_config_id = 0; float best_time = std::numeric_limits::max(); for (int config_id = 0; config_id < kConfigCount; ++config_id) { auto samples = timings[config_id]; @@ -571,7 +559,7 @@ torch::Tensor run( }; if (auto cached = cached_config_id(key)) { config_id = *cached; - } else if (autotune_enabled().load(std::memory_order_relaxed)) { + } else { config_id = tune_config< NarrowDefinition, NarrowDefinitionWithBias, @@ -584,8 +572,6 @@ torch::Tensor run( activation_scale, weight_scale, bias); - } else { - config_id = kFallbackConfigId; } } @@ -685,117 +671,3 @@ torch::Tensor cutlass_scaled_fp8_mm_f16_accum_with_config_sm120( bias, config_id); } - -int64_t fp8_f16_accum_autotune_cache_abi_sm120() { - return kAutotuneCacheAbi; -} - -std::vector fp8_f16_accum_autotune_configs_sm120() { - std::vector names; - names.reserve(kKernelConfigs.size()); - for (KernelConfig const& config : kKernelConfigs) { - names.emplace_back(config.name); - } - return names; -} - -void set_fp8_f16_accum_autotune_config_sm120( - int64_t device_index, - int64_t m, - int64_t n, - int64_t k, - torch::ScalarType out_dtype, - bool has_bias, - int64_t config_id) { - TORCH_CHECK(device_index >= 0, "device_index must be non-negative"); - TORCH_CHECK( - m > 0 && m <= std::numeric_limits::max() && - n > 0 && n <= std::numeric_limits::max() && - k > 0 && k <= std::numeric_limits::max(), - "M, N and K must be positive int32 values"); - TORCH_CHECK( - out_dtype == torch::kBFloat16 || out_dtype == torch::kFloat16, - "output dtype must be bfloat16 or float16"); - kernel_config(config_id); - - AutotuneKey key{ - static_cast(device_index), - static_cast(m), - static_cast(n), - static_cast(k), - out_dtype, - has_bias, - }; - cache_config(key, config_id); -} - -void set_fp8_f16_accum_autotune_enabled_sm120(bool enabled) { - autotune_enabled().store(enabled, std::memory_order_relaxed); -} - -torch::Tensor get_fp8_f16_accum_autotune_cache_sm120( - int64_t device_index) { - TORCH_CHECK(device_index >= 0, "device_index must be non-negative"); - std::vector> entries; - { - std::shared_lock lock(autotune_cache_mutex()); - entries.reserve(autotune_cache().size()); - for (auto const& entry : autotune_cache()) { - if (entry.first.device_index == device_index) { - entries.push_back(entry); - } - } - } - std::sort( - entries.begin(), - entries.end(), - [](auto const& left, auto const& right) { - auto const& a = left.first; - auto const& b = right.first; - return std::tie( - a.device_index, - a.m, - a.n, - a.k, - a.output_dtype, - a.has_bias) < - std::tie( - b.device_index, - b.m, - b.n, - b.k, - b.output_dtype, - b.has_bias); - }); - - auto result = torch::empty( - {static_cast(entries.size()), 6}, - torch::TensorOptions().dtype(torch::kInt64).device(torch::kCPU)); - auto rows = result.accessor(); - for (int64_t index = 0; index < static_cast(entries.size()); ++index) { - auto const& [key, config_id] = entries[index]; - rows[index][0] = key.m; - rows[index][1] = key.n; - rows[index][2] = key.k; - rows[index][3] = key.output_dtype == torch::kBFloat16 ? 0 : 1; - rows[index][4] = key.has_bias; - rows[index][5] = config_id; - } - return result; -} - -void clear_fp8_f16_accum_autotune_cache_sm120(int64_t device_index) { - std::unique_lock lock(autotune_cache_mutex()); - if (device_index < 0) { - autotune_cache().clear(); - return; - } - - for (auto entry = autotune_cache().begin(); entry != autotune_cache().end();) { - if (entry->first.device_index == device_index) { - entry = autotune_cache().erase(entry); - } else { - ++entry; - } - } -} diff --git a/lightx2v_kernel/include/lightx2v_kernel_ops.h b/lightx2v_kernel/include/lightx2v_kernel_ops.h index c9bc429ef..a8eaaafdf 100644 --- a/lightx2v_kernel/include/lightx2v_kernel_ops.h +++ b/lightx2v_kernel/include/lightx2v_kernel_ops.h @@ -60,20 +60,6 @@ torch::Tensor cutlass_scaled_fp8_mm_f16_accum_with_config_sm120( c10::optional const& bias, int64_t config_id); -int64_t fp8_f16_accum_autotune_cache_abi_sm120(); -std::vector fp8_f16_accum_autotune_configs_sm120(); -void set_fp8_f16_accum_autotune_config_sm120( - int64_t device_index, - int64_t m, - int64_t n, - int64_t k, - torch::ScalarType out_dtype, - bool has_bias, - int64_t config_id); -void set_fp8_f16_accum_autotune_enabled_sm120(bool enabled); -torch::Tensor get_fp8_f16_accum_autotune_cache_sm120( - int64_t device_index); -void clear_fp8_f16_accum_autotune_cache_sm120(int64_t device_index); void scaled_nvfp4_quant_sm120( torch::Tensor& output, torch::Tensor const& input, torch::Tensor& output_sf, torch::Tensor const& input_sf); diff --git a/lightx2v_kernel/python/lightx2v_kernel/fp8_f16_autotune.py b/lightx2v_kernel/python/lightx2v_kernel/fp8_f16_autotune.py deleted file mode 100644 index 34b5486c8..000000000 --- a/lightx2v_kernel/python/lightx2v_kernel/fp8_f16_autotune.py +++ /dev/null @@ -1,224 +0,0 @@ -"""Automatic SM120 FP8 GEMM autotuning and persistent dispatch cache.""" - -from __future__ import annotations - -import fcntl -import json -import os -import re -import tempfile -import warnings -from pathlib import Path - -import torch - -_SCHEMA_VERSION = 1 -_KERNEL_NAME = "sm120_fp8_f16_accum" -_DTYPE_NAMES = { - 0: "bfloat16", - 1: "float16", -} - - -def _require_ops() -> None: - required = ( - "fp8_f16_accum_autotune_cache_abi_sm120", - "fp8_f16_accum_autotune_configs_sm120", - "set_fp8_f16_accum_autotune_config_sm120", - "set_fp8_f16_accum_autotune_enabled_sm120", - "get_fp8_f16_accum_autotune_cache_sm120", - "clear_fp8_f16_accum_autotune_cache_sm120", - ) - missing = [name for name in required if not hasattr(torch.ops.lightx2v_kernel, name)] - if missing: - raise ImportError(f"lightx2v-kernel was built without FP8-F16 autotune ops: {missing}") - - -def _device_index(device: torch.device | str | int | None) -> int: - if isinstance(device, int): - return device - device = torch.device("cuda" if device is None else device) - if device.type != "cuda": - raise ValueError(f"FP8-F16 autotune requires a CUDA device, got {device}") - return torch.cuda.current_device() if device.index is None else device.index - - -def _runtime_identity(device_index: int) -> dict: - properties = torch.cuda.get_device_properties(device_index) - return { - "device_name": properties.name, - "compute_capability": [properties.major, properties.minor], - "torch_version": torch.__version__, - "cuda_version": torch.version.cuda, - "cache_abi": torch.ops.lightx2v_kernel.fp8_f16_accum_autotune_cache_abi_sm120(), - } - - -def _default_cache_path(device_index: int) -> Path: - properties = torch.cuda.get_device_properties(device_index) - device_name = re.sub(r"[^a-z0-9]+", "-", properties.name.lower()).strip("-") - cache_root = Path(os.environ.get("XDG_CACHE_HOME") or Path.home() / ".cache") - return cache_root / "lightx2v" / "autotune" / "fp8_f16_accum" / f"sm{properties.major}{properties.minor}-{device_name}.json" - - -def _entry_key(entry: dict) -> tuple: - return ( - int(entry["m"]), - int(entry["n"]), - int(entry["k"]), - entry["out_dtype"], - bool(entry["has_bias"]), - ) - - -def _validate_cache(cache: dict, device_index: int) -> list[dict]: - if cache.get("schema_version") != _SCHEMA_VERSION: - raise ValueError(f"Unsupported FP8-F16 autotune schema: {cache.get('schema_version')}") - if cache.get("kernel") != _KERNEL_NAME: - raise ValueError(f"Unexpected autotune kernel: {cache.get('kernel')!r}") - - expected = _runtime_identity(device_index) - actual = cache.get("runtime") - if actual != expected: - raise ValueError(f"FP8-F16 autotune cache runtime mismatch: expected {expected}, got {actual}") - - config_names = torch.ops.lightx2v_kernel.fp8_f16_accum_autotune_configs_sm120() - entries = cache.get("entries") - if not isinstance(entries, list): - raise ValueError("FP8-F16 autotune cache entries must be a list") - seen = set() - for entry in entries: - key = _entry_key(entry) - if min(key[:3]) <= 0: - raise ValueError(f"GEMM dimensions must be positive, got {key[:3]}") - if key[3] not in _DTYPE_NAMES.values(): - raise ValueError(f"Unsupported FP8-F16 output dtype: {key[3]!r}") - if not isinstance(entry["has_bias"], bool): - raise ValueError(f"has_bias must be a bool, got {entry['has_bias']!r}") - config_id = int(entry["config_id"]) - if not 0 <= config_id < len(config_names): - raise ValueError(f"Invalid FP8-F16 autotune config_id: {config_id}") - if entry.get("config") != config_names[config_id]: - raise ValueError(f"FP8-F16 autotune config name does not match config_id {config_id}") - if key in seen: - raise ValueError(f"Duplicate FP8-F16 autotune cache entry: {key}") - seen.add(key) - return entries - - -def _load_entries(cache_path: Path, device_index: int) -> list[dict]: - if not cache_path.is_file(): - return [] - return _validate_cache(json.loads(cache_path.read_text()), device_index) - - -def _register_entries(entries: list[dict], device_index: int) -> None: - dtype_by_name = { - "bfloat16": torch.bfloat16, - "float16": torch.float16, - } - for entry in entries: - torch.ops.lightx2v_kernel.set_fp8_f16_accum_autotune_config_sm120( - device_index, - int(entry["m"]), - int(entry["n"]), - int(entry["k"]), - dtype_by_name[entry["out_dtype"]], - bool(entry["has_bias"]), - int(entry["config_id"]), - ) - - -def _current_entries(device_index: int) -> list[dict]: - config_names = torch.ops.lightx2v_kernel.fp8_f16_accum_autotune_configs_sm120() - cache = torch.ops.lightx2v_kernel.get_fp8_f16_accum_autotune_cache_sm120(device_index) - entries = [] - for m, n, k, dtype_code, has_bias, config_id in cache.tolist(): - entries.append( - { - "m": m, - "n": n, - "k": k, - "out_dtype": _DTYPE_NAMES[dtype_code], - "has_bias": bool(has_bias), - "config_id": config_id, - "config": config_names[config_id], - } - ) - return entries - - -class Fp8F16AccumAutotuner: - """Manage automatic first-use tuning and its process-independent cache.""" - - def __init__( - self, - cache_path: str | Path | None = None, - device: torch.device | str | int | None = None, - ): - _require_ops() - self.device_index = _device_index(device) - if torch.cuda.get_device_capability(self.device_index) != (12, 0): - raise ValueError("FP8-F16 autotune requires an SM120 device") - self.cache_path = Path(cache_path).expanduser() if cache_path else _default_cache_path(self.device_index) - self._saved_configs = {} - - def start(self) -> int: - """Load compatible winners and enable exact-shape tuning on cache misses.""" - torch.ops.lightx2v_kernel.clear_fp8_f16_accum_autotune_cache_sm120(self.device_index) - try: - entries = _load_entries(self.cache_path, self.device_index) - except (KeyError, OSError, TypeError, ValueError) as error: - warnings.warn(f"Ignoring FP8-F16 autotune cache {self.cache_path}: {error}", stacklevel=2) - entries = [] - _register_entries(entries, self.device_index) - self._saved_configs = {_entry_key(entry): int(entry["config_id"]) for entry in entries} - torch.ops.lightx2v_kernel.set_fp8_f16_accum_autotune_enabled_sm120(True) - return len(entries) - - def save(self) -> int: - """Merge newly tuned winners and atomically persist the cache.""" - current_entries = _current_entries(self.device_index) - current_configs = {_entry_key(entry): int(entry["config_id"]) for entry in current_entries} - if current_configs == self._saved_configs: - return 0 - - lock_path = self.cache_path.with_suffix(self.cache_path.suffix + ".lock") - try: - self.cache_path.parent.mkdir(parents=True, exist_ok=True) - with lock_path.open("a+") as lock_file: - fcntl.flock(lock_file, fcntl.LOCK_EX) - try: - disk_entries = _load_entries(self.cache_path, self.device_index) - except (KeyError, OSError, TypeError, ValueError): - disk_entries = [] - merged = {_entry_key(entry): entry for entry in disk_entries} - merged.update({_entry_key(entry): entry for entry in current_entries}) - entries = [merged[key] for key in sorted(merged)] - payload = { - "schema_version": _SCHEMA_VERSION, - "kernel": _KERNEL_NAME, - "runtime": _runtime_identity(self.device_index), - "entries": entries, - } - temporary_path = None - try: - with tempfile.NamedTemporaryFile( - "w", - dir=self.cache_path.parent, - prefix=self.cache_path.name + ".", - delete=False, - ) as temporary: - temporary_path = Path(temporary.name) - json.dump(payload, temporary, indent=2) - temporary.write("\n") - os.replace(temporary_path, self.cache_path) - finally: - if temporary_path is not None: - temporary_path.unlink(missing_ok=True) - except OSError as error: - warnings.warn(f"Could not persist FP8-F16 autotune cache {self.cache_path}: {error}", stacklevel=2) - return 0 - - self._saved_configs = current_configs - return len(current_configs)