diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/cutlass_kernels/CMakeLists.txt index 7c1d0791c316..1ba3ce6caf75 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/CMakeLists.txt @@ -18,6 +18,9 @@ set(CMAKE_POSITION_INDEPENDENT_CODE ON) set(CMAKE_CUDA_RESOLVE_DEVICE_SYMBOLS ON) +# Common CUDA architectures for FP8/FP16/BF16 GEMM kernels +set(CUTLASS_COMMON_ARCHS 89 90 100f 120f) + # The Python executable will only be defined if building with Torch support. If # not, we need to find it here. if(NOT Python3_EXECUTABLE) @@ -76,15 +79,17 @@ function(process_target target_name enable_hopper enable_blackwell) if(${enable_blackwell} AND ("100" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG OR "103" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG + OR "107" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG OR "120" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG OR "121" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG )) target_compile_options(${target_name} PRIVATE "-DCUTLASS_ENABLE_GDC_FOR_SM100=1") - # Both 100 and 103 support these kernels + # SM100 family (100, 103, 107) support these kernels if("100" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG - OR "103" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG) + OR "103" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG + OR "107" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG) # No kernels should be parsed, unless blackwell is specified. This is a # build time improvement target_compile_definitions(${target_name} @@ -205,12 +210,12 @@ add_cuda_architectures(fpA_intB_gemm_src 89) add_instantiations(fpA_intB_gemm_src ${INSTANTIATION_GENERATION_DIR}/gemm) add_library(fb_gemm_src STATIC ${FBGEMM_SRC_CU} ${FBGEMM_CU_INSTANTIATIONS}) -set_cuda_architectures(fb_gemm_src 89 90 100f 120f) +set_cuda_architectures(fb_gemm_src ${CUTLASS_COMMON_ARCHS}) # add_instantiations(fb_gemm_src # ${INSTANTIATION_GENERATION_DIR}/fp8_rowwise_gemm) add_library(fp8_blockscale_gemm_src STATIC ${FP8_BLOCKSCALE_GEMM_SRC_CU}) -set_cuda_architectures(fp8_blockscale_gemm_src 89 90 100f 120f) +set_cuda_architectures(fp8_blockscale_gemm_src ${CUTLASS_COMMON_ARCHS}) set(GEMM_SWIGLU_SM90_SRC_CU ${CMAKE_CURRENT_SOURCE_DIR}/fused_gated_gemm/gemm_swiglu_e4m3.cu) @@ -264,7 +269,7 @@ if(USING_OSS_CUTLASS_MOE_GEMM) process_target(_moe_gemm_fp4 false true) add_library(_moe_gemm_fp8 OBJECT ${MOE_GEMM_SRC_CU_FP8}) - set_cuda_architectures(_moe_gemm_fp8 89 90 100f 120f) + set_cuda_architectures(_moe_gemm_fp8 ${CUTLASS_COMMON_ARCHS}) process_target(_moe_gemm_fp8 true true) add_instantiations(moe_gemm_src ${INSTANTIATION_GENERATION_DIR}/gemm_grouped) diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp b/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp index bcc0eb1165ac..dec7aa4c3dc5 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp @@ -382,10 +382,19 @@ std::vector get_candidate_configs_sm100_dynamic_cluster_shape std::vector candidate_configs; if ((config & CutlassGemmConfig::FP4_ONLY) != 0) { + // FP4 block-scaled types only support the TMA epilogue schedule on SM107. + // SM107 uses the shared tile set below; the SM100-only tiles are not enabled for it. + if (sm == 107 && schedule != EpilogueScheduleType::TMA) + { + return {}; + } + if (sm == 100) { + // FP4 block-scaled types only support TMA epilogue schedule if (schedule != EpilogueScheduleType::TMA) return {}; + candidate_configs.push_back(CutlassGemmConfig{CutlassTileConfigSM100::CtaShape128x64x128B, MainloopScheduleType::AUTO, schedule, cluster1sm, dynamic_cluster_shape, fallback_cluster_shape, sm}); if (supports_2sm) @@ -499,6 +508,11 @@ std::vector get_candidate_configs_sm100( ClusterShape::Undefined, sm}, }; #else + if (tensorrt_llm::common::isSM100Family(sm) && sm != 103 && sm != 107) + { + TLLM_LOG_INFO("Reassigned sm version to 100 for unknown sm version belonging to SM100 family"); + sm = 100; + } if (config & CutlassGemmConfig::GROUPED_GEMM) { std::vector candidate_configs; diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_preprocessors.cpp b/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_preprocessors.cpp index 6bd24a972fa7..6028acf6e2fb 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_preprocessors.cpp +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_preprocessors.cpp @@ -134,7 +134,7 @@ LayoutDetails getLayoutDetailsForTransform(QuantType quant_type, int arch) { return getLayoutDetailsForArch(quant_type); } - else if (arch == 100) + else if (isSM100Family(arch) && arch != 103) { return getLayoutDetailsForArch(quant_type); } @@ -619,7 +619,7 @@ void preprocess_weights_for_mixed_gemm(int8_t* preprocessed_quantized_weight, in src_buf.swap(dst_buf); } - if (arch != 100 && arch != 103) + if (!isSM100Family(arch)) { TLLM_LOG_INFO("add_bias_and_interleave_quantized_tensor_inplace"); add_bias_and_interleave_quantized_tensor_inplace(src_buf.data(), num_elts, quant_type); diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/fp4_gemm_template.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/fp4_gemm_template.h index e5772b5b39e7..0a269f1db821 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/fp4_gemm_template.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/fp4_gemm_template.h @@ -436,7 +436,7 @@ size_t CutlassFp4GemmRunner::dispatchToArch(T* D, void const* A, { if constexpr (fp4GemmType == FP4GemmType::W4A8_MXFP4_MXFP8) { - if (mSm == 100 || mSm == 103) + if (tk::isSM100Family(mSm)) { return dispatchMXFP8xMXFP4GemmCTAShapeSm100(D, A, B, input_sf, weight_sf, global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, stream, occupancy, bias); @@ -449,7 +449,7 @@ size_t CutlassFp4GemmRunner::dispatchToArch(T* D, void const* A, } else if constexpr (fp4GemmType == FP4GemmType::W8A8_MXFP8_MXFP8) { - if (mSm == 100 || mSm == 103) + if (tk::isSM100Family(mSm)) { return dispatchMXFP8xMXFP8GemmCTAShapeSm100(D, A, B, input_sf, weight_sf, global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, stream, occupancy); @@ -462,21 +462,26 @@ size_t CutlassFp4GemmRunner::dispatchToArch(T* D, void const* A, } else if constexpr (fp4GemmType == FP4GemmType::W4A4_NVFP4_NVFP4) { - if (mSm == 103) + if (tk::isSM100Family(mSm)) { #ifdef COMPILE_BLACKWELL_SM103_TMA_GEMMS - return dispatchNVFP4xNVFP4GemmCTAShapeSm10x(D, A, B, input_sf, weight_sf, - global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, stream, occupancy, bias); + if (mSm == 103) + { + return dispatchNVFP4xNVFP4GemmCTAShapeSm10x(D, A, B, input_sf, weight_sf, + global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, stream, occupancy, bias); + } + else + { + return dispatchNVFP4xNVFP4GemmCTAShapeSm10x(D, A, B, input_sf, weight_sf, + global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, stream, occupancy, bias); + } #else + // SM107, SM100, and other SM100 family members all use the same cutlass::arch::Sm100 kernels (compiled with + // 100f) return dispatchNVFP4xNVFP4GemmCTAShapeSm10x(D, A, B, input_sf, weight_sf, global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, stream, occupancy, bias); #endif } - else if (mSm == 100) - { - return dispatchNVFP4xNVFP4GemmCTAShapeSm10x(D, A, B, input_sf, weight_sf, - global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, stream, occupancy, bias); - } else if (mSm == 120 || mSm == 121) { return dispatchNVFP4xNVFP4GemmCTAShapeSm120(D, A, B, input_sf, weight_sf, global_sf, m, n, k, @@ -514,7 +519,7 @@ std::vector CutlassFp4GemmRunner::getCon std::vector candidateConfigs; - if (mSm == 100 || mSm == 103) + if (tk::isSM100Family(mSm)) { std::vector tilesSm10x = { tkc::CutlassTileConfigSM100::CtaShape128x128x256B, diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_rowwise_gemm/fp8_rowwise_gemm_template.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_rowwise_gemm/fp8_rowwise_gemm_template.h index 0d601060ee2a..b52bf6e57681 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_rowwise_gemm/fp8_rowwise_gemm_template.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_rowwise_gemm/fp8_rowwise_gemm_template.h @@ -687,7 +687,7 @@ size_t CutlassFp8RowwiseGemmRunner::dispatchToArch(void* D, void const* A, vo return dispatchGemmToCutlassSm90(D, A, B, C_bias, quantOption, m, n, k, scale_d0, scale_d1, gemmConfig, workspace, workspaceBytes, stream, occupancy); } - else if (mSm == 100 || mSm == 103) + else if (tk::isSM100Family(mSm)) { return dispatchGemmToCutlassSm100(D, A, B, C_bias, quantOption, m, n, k, scale_d0, scale_d1, gemmConfig, workspace, workspaceBytes, stream, occupancy); @@ -759,7 +759,7 @@ std::vector CutlassFp8RowwiseGemmRunner::getConfigs() } } } - else if (mSm == 100 || mSm == 103) + else if (tk::isSM100Family(mSm)) { std::vector tilesSm100 = { tkc::CutlassTileConfigSM100::CtaShape64x32x128B, diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm_template.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm_template.h index b554ea2c8d06..33043066793c 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm_template.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm_template.h @@ -465,7 +465,7 @@ void CutlassFpAIntBGemmRunner +template void invokeMxFP8Quantization(int b, int m, int n, int padded_n, T const* input, int64_t* output, int32_t* SFOuput, QuantizationSFLayout layout, int multiProcessorCount, cudaStream_t stream) { - // Fixed SF_VEC_SIZE as 32 - static constexpr int SF_VEC_SIZE = 32; + static_assert(SF_VEC_SIZE == 32 || SF_VEC_SIZE == 128, "MXFP8 quantization supports SF vector sizes 32 and 128."); + static_assert(SF_OUTPUT_VEC_SIZE == 32 || SF_OUTPUT_VEC_SIZE == SF_VEC_SIZE, + "MXFP8 output SF vector size must be 32 or match the quantization SF vector size."); // Grid, Block size. // Each thread converts 8 values. @@ -217,8 +218,9 @@ void invokeMxFP8Quantization(int b, int m, int n, int padded_n, T const* input, config.numAttrs = 1; config.attrs = attrs; cudaLaunchKernelEx(&config, - quantize_with_block_size, b, m, n, padded_n, - input, nullptr, reinterpret_cast(output), reinterpret_cast(SFOuput), layout); + quantize_with_block_size, + b, m, n, padded_n, input, nullptr, reinterpret_cast(output), reinterpret_cast(SFOuput), + layout); } //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -402,8 +404,8 @@ template void invokeFP4Quantization(int b, int m, int n, half const* i template void invokeFP4Quantization(int b, int m, int n, half const* input, float const* SFScale, int64_t* output, int32_t* SFOuput, bool useUE8M0, QuantizationSFLayout layout, int multiProcessorCount, cudaStream_t stream); -template void invokeMxFP8Quantization(int b, int m, int n, int padded_n, half const* input, int64_t* output, - int32_t* SFOuput, QuantizationSFLayout layout, int multiProcessorCount, cudaStream_t stream); +template void invokeMxFP8Quantization(int b, int m, int n, int padded_n, half const* input, + int64_t* output, int32_t* SFOuput, QuantizationSFLayout layout, int multiProcessorCount, cudaStream_t stream); template void computePerTokenGlobalScaleForFP4Quantization(int b, int m, int n, half const* input, int const* tokensPerBatch, float* globalScale, int multiProcessorCount, cudaStream_t stream); #ifdef ENABLE_BF16 @@ -413,8 +415,12 @@ template void invokeFP4Quantization<__nv_bfloat16, 16>(int b, int m, int n, __nv template void invokeFP4Quantization<__nv_bfloat16, 32>(int b, int m, int n, __nv_bfloat16 const* input, float const* SFScale, int64_t* output, int32_t* SFOuput, bool useUE8M0, QuantizationSFLayout layout, int multiProcessorCount, cudaStream_t stream); -template void invokeMxFP8Quantization<__nv_bfloat16>(int b, int m, int n, int padded_n, __nv_bfloat16 const* input, - int64_t* output, int32_t* SFOuput, QuantizationSFLayout layout, int multiProcessorCount, cudaStream_t stream); +template void invokeMxFP8Quantization<__nv_bfloat16, 32, 32>(int b, int m, int n, int padded_n, + __nv_bfloat16 const* input, int64_t* output, int32_t* SFOuput, QuantizationSFLayout layout, int multiProcessorCount, + cudaStream_t stream); +template void invokeMxFP8Quantization<__nv_bfloat16, 128, 32>(int b, int m, int n, int padded_n, + __nv_bfloat16 const* input, int64_t* output, int32_t* SFOuput, QuantizationSFLayout layout, int multiProcessorCount, + cudaStream_t stream); template void computePerTokenGlobalScaleForFP4Quantization<__nv_bfloat16>(int b, int m, int n, __nv_bfloat16 const* input, int const* tokensPerBatch, float* globalScale, int multiProcessorCount, cudaStream_t stream); diff --git a/cpp/tensorrt_llm/kernels/quantization.cuh b/cpp/tensorrt_llm/kernels/quantization.cuh index e4a5edab9b1f..4b986cfe3b82 100644 --- a/cpp/tensorrt_llm/kernels/quantization.cuh +++ b/cpp/tensorrt_llm/kernels/quantization.cuh @@ -280,7 +280,7 @@ constexpr int CVT_FP4_THREADS_PER_WARP = 32; constexpr int CVT_FP8_TO_FP4_ELTS_PER_THREAD = 16; // Membermask for the __shfl_xor_sync butterfly among the NUM_THREADS_PER_SF -// lanes that share one scale factor. The xor-1/xor-2 exchange never crosses +// lanes that share one scale factor. The butterfly exchange never crosses // this aligned lane group, so only the group has to converge on the shuffle. // // Do not widen the mask to the full warp. A sync shuffle waits until every @@ -299,9 +299,11 @@ constexpr int CVT_FP8_TO_FP4_ELTS_PER_THREAD = 16; template inline __device__ uint32_t cvt_sf_group_shfl_mask() { - static_assert(NUM_THREADS_PER_SF == 2 || NUM_THREADS_PER_SF == 4, "Unsupported SF group size."); + static_assert( + NUM_THREADS_PER_SF >= 2 && NUM_THREADS_PER_SF <= 32 && (NUM_THREADS_PER_SF & (NUM_THREADS_PER_SF - 1)) == 0, + "The SF group size must be a power of two between 2 and 32."); constexpr uint32_t groupSize = static_cast(NUM_THREADS_PER_SF); - constexpr uint32_t groupMask = (1U << groupSize) - 1U; + constexpr uint32_t groupMask = 0xFFFFFFFFU >> (32U - groupSize); uint32_t laneId = 0; asm("mov.u32 %0, %%laneid;" : "=r"(laneId)); return groupMask << (laneId & ~(groupSize - 1U)); @@ -628,8 +630,30 @@ __device__ uint64_t cvt_warp_fp8_to_fp4(PackedVec& vec, float SFScaleVal, #endif } -// Quantizes the provided PackedVec into the uint64_t output -template +// Stores one scale value into the adjacent consumer slots of a K4 SF atom. +template +inline __device__ void cvt_store_replicated_sf(uint8_t* SFout, uint8_t sfValue) +{ + static_assert( + SF_REPLICATION == 1 || SF_REPLICATION == 2 || SF_REPLICATION == 4, "SF replication must divide a K4 atom."); + if (SFout) + { + if constexpr (SF_REPLICATION == 4) + { + *reinterpret_cast(SFout) = static_cast(sfValue) * 0x01010101U; + } + else if constexpr (SF_REPLICATION == 2) + { + *reinterpret_cast(SFout) = static_cast(sfValue) * 0x0101U; + } + else + { + *SFout = sfValue; + } + } +} + +template __device__ uint64_t cvt_warp_fp16_to_mxfp8(PackedVec& vec, uint8_t* SFout) { #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) @@ -644,16 +668,22 @@ __device__ uint64_t cvt_warp_fp16_to_mxfp8(PackedVec& vec, uint8_t* SFout) } constexpr int CVT_NUM_THREADS_PER_SF = SF_VEC_SIZE / CVT_ELTS_PER_THREAD; - // Get the absolute maximum among all 16 values (two threads for 16, four threads for 32). + // Get the absolute maximum among all values that share one scale factor. uint32_t const sfGroupMask = cvt_sf_group_shfl_mask(); - localMax = cuda_max(__shfl_xor_sync(sfGroupMask, localMax, 1), localMax); - if constexpr (CVT_NUM_THREADS_PER_SF == 4) +#pragma unroll + for (int offset = 1; offset < CVT_NUM_THREADS_PER_SF; offset *= 2) { - localMax = cuda_max(__shfl_xor_sync(sfGroupMask, localMax, 2), localMax); + localMax = cuda_max(__shfl_xor_sync(sfGroupMask, localMax, offset), localMax); } // Get the final absolute maximum values. float vecMax = float(cuda_max(localMax.x, localMax.y)); + // Match fp8_quantize_1x128's handling of an all-zero block. + if constexpr (SF_VEC_SIZE == 128) + { + vecMax = fmaxf(vecMax, 1e-10f); + } + // Get the SF (max value of the vector / max value of mxfp8). float SFValue = vecMax * reciprocal_approximate_ftz(448.0f); // 8 bits representation of the SF. @@ -666,11 +696,9 @@ __device__ uint64_t cvt_warp_fp16_to_mxfp8(PackedVec& vec, uint8_t* SFout) // Get the output scale (reciprocal of the SFValue). float outputScale = vecMax != 0.f ? reciprocal_approximate_ftz(SFValue) : 0.0f; - if (SFout) - { - // Write the SF to global memory (STG.8). - *SFout = fp8SFVal; - } + // Store one byte per consumer scale slot. For example, a K128 quantization + // scale is replicated into four K32 slots for an MXFP8 consumer. + cvt_store_replicated_sf(SFout, fp8SFVal); // Convert the input to float. float2 fp2Vals[CVT_ELTS_PER_THREAD / 2]; @@ -723,8 +751,7 @@ inline __host__ __device__ int64_t get_sf_out_offset_128x4( int32_t kTileIdx = (kIdx / 4); int64_t kTileStride = 32 * outerMStride; // 512 - // SF vector size 16 or 32. We round the "numCols" up to a multiple of 64 or 128. - // It is the same as rounding the "numColVecs" up to a multiple of 4. + // Round the number of scale-factor columns up to a multiple of four. int32_t numKTiles = (numColVecs + 4 - 1) / 4; int32_t mTileIdx = mIdx / (32 * 4); int64_t mTileStride = numKTiles * kTileStride; @@ -740,24 +767,28 @@ inline __host__ __device__ int64_t get_sf_out_offset_128x4( return SFOffset; } -template +template __device__ uint8_t* cvt_quant_get_sf_out_offset(std::optional batchIdx, int rowIdx, int colVecIdx, std::optional numRows, int numColVecs, SFType* SFout, QuantizationSFLayout layout) { #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) // Each thread holds one vector. - static_assert(CVT_NUM_THREADS_PER_SF == 1 || CVT_NUM_THREADS_PER_SF == 2 || CVT_NUM_THREADS_PER_SF == 4); + static_assert(CVT_NUM_THREADS_PER_SF >= 1 && CVT_NUM_THREADS_PER_SF <= 32 + && (CVT_NUM_THREADS_PER_SF & (CVT_NUM_THREADS_PER_SF - 1)) == 0, + "The number of threads per SF must be a power of two no greater than 32."); + static_assert( + SF_REPLICATION == 1 || SF_REPLICATION == 2 || SF_REPLICATION == 4, "SF replication must divide a K4 atom."); - // One pair of threads write one SF to global memory. + // One aligned group of threads writes one SF to global memory. // TODO: stage through smem for packed STG.32 // is it better than STG.8 from 4 threads ? if (threadIdx.x % CVT_NUM_THREADS_PER_SF == 0) { if (layout == QuantizationSFLayout::SWIZZLED) { - // SF vector index (16 elements share one SF in the K dimension). - // numRows and numCols are unpadded. - int32_t kIdx = colVecIdx / CVT_NUM_THREADS_PER_SF; + // Output SF vector index. A quantization scale can be replicated + // into multiple adjacent consumer scale slots. + int32_t kIdx = colVecIdx / CVT_NUM_THREADS_PER_SF * SF_REPLICATION; int32_t mIdx = rowIdx; auto SFOffset = get_sf_out_offset_128x4(batchIdx, mIdx, kIdx, numRows, numColVecs); @@ -766,7 +797,7 @@ __device__ uint8_t* cvt_quant_get_sf_out_offset(std::optional batchIdx, int else if (layout == QuantizationSFLayout::LINEAR) { // Linear row-major layout, no padding required. - int32_t KTileIdx = colVecIdx / CVT_NUM_THREADS_PER_SF; + int32_t KTileIdx = colVecIdx / CVT_NUM_THREADS_PER_SF * SF_REPLICATION; int32_t numKTiles = numColVecs; int64_t mTileStride = numKTiles; @@ -785,7 +816,8 @@ __device__ uint8_t* cvt_quant_get_sf_out_offset(std::optional batchIdx, int return nullptr; } -template +template __global__ void #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) __launch_bounds__(512, 4) quantize_with_block_size( @@ -803,7 +835,12 @@ quantize_with_block_size( : CVT_ELTS_PER_THREAD; using PackedVec = PackedVec; - static constexpr int CVT_NUM_THREADS_PER_SF = SF_VEC_SIZE / ELTS_PER_THREAD; // 2 or 4 + static constexpr int CVT_NUM_THREADS_PER_SF = SF_VEC_SIZE / ELTS_PER_THREAD; + static constexpr int SF_REPLICATION = SF_VEC_SIZE / SF_OUTPUT_VEC_SIZE; + static_assert(SF_VEC_SIZE % ELTS_PER_THREAD == 0, "SF vector size must be divisible by elements per thread."); + static_assert(SF_VEC_SIZE % SF_OUTPUT_VEC_SIZE == 0, "Quantization SF size must be divisible by output SF size."); + static_assert( + SF_REPLICATION == 1 || SF_REPLICATION == 2 || SF_REPLICATION == 4, "SF replication must divide a K4 atom."); static_assert(sizeof(PackedVec) == sizeof(Type) * ELTS_PER_THREAD, "Vec size is not matched."); // Get the global scaling factor, which will be applied to the SF. @@ -816,7 +853,7 @@ quantize_with_block_size( // The number of padded rows considering 128x4 SF layout. int numPaddedRowsForSf = isSfSwizzledLayout ? PadUpFn(numRows, 128) : numRows; - int numColsForSf = isSfSwizzledLayout ? PadUpFn(numPaddedCols, 4 * SF_VEC_SIZE) : numPaddedCols; + int numColsForSf = isSfSwizzledLayout ? PadUpFn(numPaddedCols, 4 * SF_OUTPUT_VEC_SIZE) : numPaddedCols; // The number of threads in the column dimension。 // Note that numCols/numPaddedCols/numColsForSf are guaranteed to be multiples of ELTS_PER_THREAD. @@ -847,14 +884,12 @@ quantize_with_block_size( std::optional optionalNumRows = numRows; // The SF output pointer. - auto sf_out = cvt_quant_get_sf_out_offset( - optionalBatchIdx, rowIdx, colIdx, optionalNumRows, numPaddedCols / SF_VEC_SIZE, SFout, layout); + auto sf_out = cvt_quant_get_sf_out_offset( + optionalBatchIdx, rowIdx, colIdx, optionalNumRows, numPaddedCols / SF_OUTPUT_VEC_SIZE, SFout, + layout); // Set the SF padding to 0. - if (sf_out != nullptr) - { - sf_out[0] = 0x00; - } + cvt_store_replicated_sf(sf_out, 0x00); } } } @@ -869,8 +904,9 @@ quantize_with_block_size( std::optional optionalNumRows = numRows; // The SF output pointer. - auto sf_out = cvt_quant_get_sf_out_offset( - optionalBatchIdx, rowIdx, colIdx, optionalNumRows, numPaddedCols / SF_VEC_SIZE, SFout, layout); + auto sf_out = cvt_quant_get_sf_out_offset( + optionalBatchIdx, rowIdx, colIdx, optionalNumRows, numPaddedCols / SF_OUTPUT_VEC_SIZE, SFout, + layout); // The input tensor offset. int64_t inOffset = static_cast(batchIdx * numRows + rowIdx) * numColThreads + colIdx; @@ -896,10 +932,7 @@ quantize_with_block_size( if (colIdx >= numColThreads) { // Set the SF padding to 0. - if (sf_out != nullptr) - { - sf_out[0] = 0x00; - } + cvt_store_replicated_sf(sf_out, 0x00); } else { @@ -920,7 +953,7 @@ quantize_with_block_size( else if constexpr (quantization_type == BlockScaleQuantizationType::FP16_TO_MXFP8) { reinterpret_cast(out)[outOffset] - = cvt_warp_fp16_to_mxfp8(in_vec, sf_out); + = cvt_warp_fp16_to_mxfp8(in_vec, sf_out); } } } diff --git a/cpp/tensorrt_llm/kernels/quantization.h b/cpp/tensorrt_llm/kernels/quantization.h index e8b0d83abbd3..f2aaacf5ffa8 100644 --- a/cpp/tensorrt_llm/kernels/quantization.h +++ b/cpp/tensorrt_llm/kernels/quantization.h @@ -77,7 +77,7 @@ template void invokeFP4Quantization(int b, int m, int n, T const* input, float const* globalScale, int64_t* output, int32_t* SFOuput, bool useUE8M0, QuantizationSFLayout layout, int multiProcessorCount, cudaStream_t stream = 0); -template +template void invokeMxFP8Quantization(int b, int m, int n, int padded_n, T const* input, int64_t* output, int32_t* SFOuput, QuantizationSFLayout layout, int multiProcessorCount, cudaStream_t stream = 0); diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/KernelRunner.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/KernelRunner.cpp index b98256b9e761..5ddcaa31d88f 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/KernelRunner.cpp +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/KernelRunner.cpp @@ -89,7 +89,12 @@ TrtllmGenGemmRunner::TrtllmGenGemmRunner(TrtllmGenGemmRunnerOptions const& optio } } - TLLM_CHECK_WITH_INFO(mPassingConfigIndices.size() != 0, "No kernel found for the given output type"); + TLLM_CHECK_WITH_INFO(mPassingConfigIndices.size() != 0, + "No kernel found for the given options: eltTypeA: %s, eltTypeB: %s, outputType: %s, deepSeekFp8: %d, " + "transposeMmaOutput: %d, gpuSM: %d", + tg::dtypeToString(mOptions.eltTypeA).c_str(), tg::dtypeToString(mOptions.eltTypeB).c_str(), + tg::dtypeToString(mOptions.outputType).c_str(), mOptions.deepSeekFp8, mOptions.transposeMmaOutput, + gpuNativeSmVersion); } size_t TrtllmGenGemmRunner::getWorkspaceSizeInBytes(int32_t m, int32_t n, int32_t k) diff --git a/cpp/tensorrt_llm/thop/fp8BlockScalingGemm.cpp b/cpp/tensorrt_llm/thop/fp8BlockScalingGemm.cpp index 06f0c1373c35..2735a76d0b1f 100644 --- a/cpp/tensorrt_llm/thop/fp8BlockScalingGemm.cpp +++ b/cpp/tensorrt_llm/thop/fp8BlockScalingGemm.cpp @@ -244,14 +244,25 @@ extern torch::Tensor fp8_block_scaling_gemm(torch::Tensor const& mat1, torch::Te torch::Tensor const& mat1Scale, torch::Tensor const& mat2Scale) { auto const sm = tensorrt_llm::common::getSMVersion(); - switch (sm) + if (tensorrt_llm::common::isSM100Family(sm)) + { + return fp8_block_scale_gemm_blackwell(mat1, mat2, mat1Scale, mat2Scale); + } + else if (sm == 90) + { + return fp8_block_scaling_gemm_hopper(mat1, mat2, mat1Scale, mat2Scale); + } + else if (sm == 89) + { + return fp8_block_scaling_gemm_ada(mat1, mat2, mat1Scale, mat2Scale); + } + else if (sm == 120) + { + return fp8_block_scale_gemm_blackwell_geforce(mat1, mat2, mat1Scale, mat2Scale); + } + else { - case 103: return fp8_block_scale_gemm_blackwell(mat1, mat2, mat1Scale, mat2Scale); - case 100: return fp8_block_scale_gemm_blackwell(mat1, mat2, mat1Scale, mat2Scale); - case 90: return fp8_block_scaling_gemm_hopper(mat1, mat2, mat1Scale, mat2Scale); - case 89: return fp8_block_scaling_gemm_ada(mat1, mat2, mat1Scale, mat2Scale); - case 120: return fp8_block_scale_gemm_blackwell_geforce(mat1, mat2, mat1Scale, mat2Scale); - default: TORCH_CHECK(false, "Unsupported SM version for FP8 block scaling GEMM"); + TORCH_CHECK(false, "Unsupported SM version for FP8 block scaling GEMM"); } } diff --git a/cpp/tensorrt_llm/thop/fp8Quantize.cpp b/cpp/tensorrt_llm/thop/fp8Quantize.cpp index 43eea8cff838..6e0620d082a0 100644 --- a/cpp/tensorrt_llm/thop/fp8Quantize.cpp +++ b/cpp/tensorrt_llm/thop/fp8Quantize.cpp @@ -17,6 +17,7 @@ #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h" #include "tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.h" +#include "tensorrt_llm/kernels/quantization.h" #include "tensorrt_llm/thop/thUtils.h" #include @@ -145,12 +146,11 @@ std::tuple fp8_batched_quantize_1x128_permute102(at::Ten // Fused 1x128 FP8 quantize + UE8M0 packing (SM100 only). // -// Drop-in replacement for the (fp8_quantize_1x128 → get_mn_major_tma_aligned_packed_ue8m0_tensor) -// two-kernel sequence used by deep_gemm fp8 block-scale GEMMs. Returns the -// packed UE8M0 scale tensor (int32, MN-major, TMA-aligned) directly so -// deep_gemm's transform_sf_into_required_layout falls through the -// `(INT, 1, gran_k)` branch and skips its own pack call. -std::tuple fp8_quantize_1x128_packed_ue8m0(at::Tensor const& self) +// Returns the legacy deep_gemm MN-major packed int32 scales (K % 16 == 0) by +// default. useR128c4Layout opts into K32-addressable slots in the standard +// R128c4 layout (K % 128 == 0); quantization stays K128, so each UE8M0 scale is +// replicated into four adjacent K32 slots. +std::tuple fp8_quantize_1x128_packed_ue8m0(at::Tensor const& self, bool useR128c4Layout) { CHECK_TH_CUDA(self); CHECK_CONTIGUOUS(self); @@ -165,46 +165,67 @@ std::tuple fp8_quantize_1x128_packed_ue8m0(at::Tensor co TORCH_CHECK(m <= std::numeric_limits::max(), "M must be within int32"); TORCH_CHECK(n <= std::numeric_limits::max(), "N must be within int32"); + auto const num_n_blocks = (n + 127) / 128; + auto const num_packed_sf_k = (num_n_blocks + 3) / 4; + auto stream = at::cuda::getCurrentCUDAStream(self.get_device()); + + if (useR128c4Layout) + { + constexpr int kQuantSfVecSize = 128; + constexpr int kOutputSfVecSize = 32; + TORCH_CHECK(n % kQuantSfVecSize == 0, "self.sizes()[1] must be a multiple of 128 for R128c4, but got ", n); + + at::Tensor valueE4M3 + = at::detail::empty_cuda({m, n}, at::ScalarType::Float8_e4m3fn, self.device(), /* stride */ std::nullopt); + auto const numOutputSf = n / kOutputSfVecSize; + auto const sfSize + = tensorrt_llm::computeSwizzledLayoutSFSize(static_cast(m), static_cast(numOutputSf)); + at::Tensor scaleFP8SF + = at::detail::empty_cuda({sfSize}, at::ScalarType::Byte, self.device(), /* stride */ std::nullopt); + + if (m > 0 && n > 0) + { +#ifdef ENABLE_BF16 + const thread_local int multiProcessorCount = tensorrt_llm::common::getMultiProcessorCount(); + tensorrt_llm::kernels::invokeMxFP8Quantization<__nv_bfloat16, kQuantSfVecSize, kOutputSfVecSize>(1, + static_cast(m), static_cast(n), static_cast(n), + reinterpret_cast<__nv_bfloat16 const*>(self.data_ptr()), + reinterpret_cast(valueE4M3.data_ptr()), reinterpret_cast(scaleFP8SF.data_ptr()), + tensorrt_llm::QuantizationSFLayout::SWIZZLED, multiProcessorCount, stream); +#else + C10_THROW_ERROR(NotImplementedError, "BFloat16 must be enabled to quantize a BF16 tensor to MXFP8."); +#endif + } + + return {valueE4M3, scaleFP8SF}; + } + TORCH_CHECK(n % 16 == 0, "self.sizes()[1] must be a multiple of 16, but got ", n); - // FP8 output is row-major [m, n] with the same alignment used by the legacy path. - // The legacy path pads M to a multiple of 4; replicate that to avoid layout surprises. + // Legacy deep_gemm path pads M to a multiple of four. auto const m_padded = (m + 4 - 1) / 4 * 4; - at::Tensor valueE4M3 = at::detail::empty_cuda( {m_padded, n}, at::ScalarType::Float8_e4m3fn, self.device(), /* stride */ std::nullopt); - // Packed scale physical layout: [num_packed_sf_k, m_aligned] uint32, MN-contiguous in memory. - // deep_gemm's get_mn_major_tma_aligned_packed_ue8m0_tensor returns a strided VIEW with - // PyTorch shape `[mn, packed_sf_k]` and strides `(1, tma_aligned_mn)`. We build the same - // strided view so deep_gemm's transform_sf_into_required_layout falls into the - // `(INT, 1, gran_k)` branch and skips its own pack call. - auto const num_n_blocks = (n + 127) / 128; - auto const num_packed_sf_k = (num_n_blocks + 3) / 4; constexpr int kTmaAlignedUint32Elems = 4; // 16 bytes / sizeof(uint32_t) - auto const m_aligned = (m_padded + kTmaAlignedUint32Elems - 1) / kTmaAlignedUint32Elems * kTmaAlignedUint32Elems; - - // Allocate physical buffer [num_packed_sf_k, m_aligned] (K-major in memory). - // The kernel writes packed=0 for the [m, m_aligned) tail rows itself, so no - // host-side zero-init is needed. + auto const scaleLeadingDim + = (m_padded + kTmaAlignedUint32Elems - 1) / kTmaAlignedUint32Elems * kTmaAlignedUint32Elems; at::Tensor packedBuf = at::detail::empty_cuda( - {num_packed_sf_k, m_aligned}, at::ScalarType::Int, self.device(), /* stride */ std::nullopt); - - auto stream = at::cuda::getCurrentCUDAStream(self.get_device()); + {num_packed_sf_k, scaleLeadingDim}, at::ScalarType::Int, self.device(), /* stride */ std::nullopt); tensorrt_llm::kernels::fp8_blockscale_gemm::launch_fp8_quantize_1x128_packed_bf16_e4m3( reinterpret_cast<__nv_fp8_e4m3*>(valueE4M3.data_ptr()), reinterpret_cast(packedBuf.data_ptr()), reinterpret_cast<__nv_bfloat16 const*>(self.data_ptr()), static_cast(m), static_cast(n), - static_cast(m_aligned), stream); + static_cast(scaleLeadingDim), stream); - // Wrap the [num_packed_sf_k, m_aligned] memory as a [m, num_packed_sf_k] strided tensor + // Wrap the [num_packed_sf_k, scaleLeadingDim] memory as a [m, num_packed_sf_k] strided tensor // matching deep_gemm's get_mn_major_tma_aligned_packed_ue8m0_tensor return contract: // shape = (m, num_packed_sf_k) - // stride = (1, m_aligned) + // stride = (1, scaleLeadingDim) at::Tensor packedScale = at::from_blob( packedBuf.data_ptr(), /* sizes */ {m, num_packed_sf_k}, - /* strides */ {1, m_aligned}, + /* strides */ {1, scaleLeadingDim}, /* deleter */ [keep = packedBuf](void*) mutable {}, packedBuf.options()); return {valueE4M3.slice(0, 0, m), packedScale}; @@ -249,7 +270,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) { m.def("fp8_quantize_1x128(Tensor input, bool use_ue8m0=False) -> (Tensor, Tensor)"); m.def("fp8_batched_quantize_1x128_permute102(Tensor input) -> (Tensor, Tensor)"); - m.def("fp8_quantize_1x128_packed_ue8m0(Tensor input) -> (Tensor, Tensor)"); + m.def("fp8_quantize_1x128_packed_ue8m0(Tensor input, bool use_r128c4_layout=False) -> (Tensor, Tensor)"); m.def("fp8_quantize_1x128_cutedsl_ue8m0(Tensor input) -> (Tensor, Tensor)"); } diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 9815b2e51406..31cec1dfc264 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -863,15 +863,18 @@ def _(pe: torch.Tensor, nope: torch.Tensor): return packed, scale @torch.library.register_fake("trtllm::fp8_quantize_1x128_packed_ue8m0") - def _(input: torch.Tensor): - # Returns (fp8_e4m3 [m, k], packed_ue8m0_int32 [m, packed_sf_k]) - # matching deep_gemm.get_mn_major_tma_aligned_packed_ue8m0_tensor's return shape. + def _(input: torch.Tensor, + use_r128c4_layout: bool = False) -> tuple[torch.Tensor, torch.Tensor]: m, k = input.shape[0], input.shape[1] num_n_blocks = (k + 127) // 128 num_packed_sf_k = (num_n_blocks + 3) // 4 - return torch.empty_like(input, - dtype=torch.float8_e4m3fn), input.new_empty( - (m, num_packed_sf_k), dtype=torch.int32) + if use_r128c4_layout: + m_padded = (m + 127) // 128 * 128 + sf_k_padded = ((k + 31) // 32 + 3) // 4 * 4 + scale = input.new_empty((m_padded * sf_k_padded), dtype=torch.uint8) + else: + scale = input.new_empty((m, num_packed_sf_k), dtype=torch.int32) + return torch.empty_like(input, dtype=torch.float8_e4m3fn), scale @torch.library.register_fake("trtllm::fp8_quantize_1x128_cutedsl_ue8m0") def _(input: torch.Tensor): diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 713da2bdeab6..b5c97ca55eea 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -1789,12 +1789,10 @@ def _( def _fp8_quantize_1x128_ue8m0(input: torch.Tensor, tactic: int): """Dispatch FP8 1x128 quantization to CUDA or Triton kernel. - When the CUDA path is selected on SM100 and ``TRTLLM_FUSED_FP8_QUANT_PACK=1`` - is set, the fused ``fp8_quantize_1x128_packed_ue8m0`` op is used and the - follow-on ``get_mn_major_tma_aligned_packed_ue8m0_tensor`` call is skipped: - the new op writes packed-UE8M0 (int32) scales directly in the layout - deep_gemm expects, so deep_gemm's internal layout transform falls into the - pre-packed branch and skips its own pack kernel as well. + On SM100 with ``TRTLLM_FUSED_FP8_QUANT_PACK=1``, the fused + ``fp8_quantize_1x128_packed_ue8m0`` op already emits the legacy packed-UE8M0 + (int32) layout deep_gemm expects, so the follow-on + ``get_mn_major_tma_aligned_packed_ue8m0_tensor`` call is skipped. """ TACTIC_TRITON = 1 if tactic == TACTIC_TRITON: @@ -1803,7 +1801,8 @@ def _fp8_quantize_1x128_ue8m0(input: torch.Tensor, tactic: int): a_sf.transpose(0, 1)) return a, a_sf if _USE_FUSED_FP8_QUANT_PACK and get_sm_version() >= 100: - a, a_sf = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0(input) + # Legacy MN-major packed layout, requested explicitly. + a, a_sf = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0(input, False) return a, a_sf a, a_sf = torch.ops.trtllm.fp8_quantize_1x128(input, use_ue8m0=True) a_sf = deep_gemm.get_mn_major_tma_aligned_packed_ue8m0_tensor( diff --git a/tests/unittest/_torch/thop/parallel/test_fp8_block_scale_gemm.py b/tests/unittest/_torch/thop/parallel/test_fp8_block_scale_gemm.py index 41bb017cbd18..69882b8dd285 100644 --- a/tests/unittest/_torch/thop/parallel/test_fp8_block_scale_gemm.py +++ b/tests/unittest/_torch/thop/parallel/test_fp8_block_scale_gemm.py @@ -64,7 +64,7 @@ def test_fp8_block_scale_deep_gemm(dtype, m, k, n): @pytest.mark.skipif( - getSMVersion() != 100 and getSMVersion() != 89 and getSMVersion() != 120, + getSMVersion() not in (100, 107, 89, 120), reason="The test is for Blackwell and Ada only. Current SM is %d." % getSMVersion(), ) @@ -447,7 +447,7 @@ def test_fp8_blockscale_gemm_reference(): @pytest.mark.skipif( - getSMVersion() != 100, + getSMVersion() not in (100, 107), reason="The kernel only supports Blackwell. Current SM is %d." % getSMVersion(), ) diff --git a/tests/unittest/_torch/thop/parallel/test_fp8_quantize.py b/tests/unittest/_torch/thop/parallel/test_fp8_quantize.py index 10a8dae82131..495a2028dce0 100644 --- a/tests/unittest/_torch/thop/parallel/test_fp8_quantize.py +++ b/tests/unittest/_torch/thop/parallel/test_fp8_quantize.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -377,9 +377,19 @@ def decode_ue8m0_int32_to_float(int32_tensor): # Tests for fp8_quantize_1x128_packed_ue8m0 (SM100 fused quant+pack) # --------------------------------------------------------------------------- +# The op guards on the C++ isSM100Family() (SM100-109), which is wider than the +# utils.util helper of the same name. The Triton comparisons further down keep +# using that helper: they are limited to the archs Triton can target, not to the +# archs the op supports. +skip_if_not_sm100_family = pytest.mark.skipif( + not 100 <= getSMVersion() < 110, + reason="fp8_quantize_1x128_packed_ue8m0 is SM100-family only. " + "Current SM is %d." % getSMVersion(), +) + def _decode_packed_int32_ue8m0(packed_int32): - """Decode int32 packed UE8M0 scales to (exponent, value) per byte.""" + """Decode each packed int32 into four UE8M0 exponent bytes.""" b0 = (packed_int32 >> 0) & 0xFF b1 = (packed_int32 >> 8) & 0xFF b2 = (packed_int32 >> 16) & 0xFF @@ -387,121 +397,145 @@ def _decode_packed_int32_ue8m0(packed_int32): return torch.stack([b0, b1, b2, b3], dim=-1) -@pytest.mark.skipif(not isSM100Family(), - reason="fp8_quantize_1x128_packed_ue8m0 is SM100 only.") +def _legacy_packed_to_r128c4(packed_int32, m, num_quant_sf_k, sf_replication=1): + """Convert packed quant scales to replicated R128c4 consumer slots.""" + num_packed_quant_sf_k = (num_quant_sf_k + 3) // 4 + m_padded = (m + 127) // 128 * 128 + num_output_sf_k = num_quant_sf_k * sf_replication + num_packed_output_sf_k = (num_output_sf_k + 3) // 4 + expected = torch.zeros((m_padded * num_packed_output_sf_k * 4), + dtype=torch.uint8, + device=packed_int32.device) + + logical = _decode_packed_int32_ue8m0(packed_int32.contiguous().view( + torch.int32)).reshape(m, num_packed_quant_sf_k * 4)[:, :num_quant_sf_k] + logical = logical.repeat_interleave(sf_replication, dim=1) + + m_idx = torch.arange(m, device=packed_int32.device, + dtype=torch.int64).view(-1, 1) + sf_k_idx = torch.arange(num_output_sf_k, + device=packed_int32.device, + dtype=torch.int64).view(1, -1) + offsets = (((m_idx // 128 * num_packed_output_sf_k + sf_k_idx // 4) * 32 + + m_idx % 32) * 4 + (m_idx % 128) // 32) * 4 + sf_k_idx % 4 + expected[offsets.flatten()] = logical.flatten().to(torch.uint8) + return expected + + +@skip_if_not_sm100_family @pytest.mark.parametrize("m,k", [ (1, 128), (3, 256), (4, 512), (7, 512), + (31, 384), + (32, 512), + (33, 640), (16, 7168), (127, 4096), + (128, 512), + (129, 512), (1024, 7168), pytest.param(262141, 128, id="grid-y-overflow"), ]) -def test_fp8_quantize_1x128_packed_ue8m0_matches_legacy(m, k): - """The fused packed op should produce the same FP8 output and UE8M0 - scales as the legacy (fp8_quantize_1x128 → pack) two-kernel sequence. - The legacy path is the canonical reference for what deep_gemm expects. - """ +def test_fp8_quantize_1x128_packed_ue8m0_r128c4(m, k): + """K128 scales are replicated into K32-addressable R128c4 slots.""" from tensorrt_llm.quantization.utils import fp8_utils torch.manual_seed(0) x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) - fused_fp8, fused_packed = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0( - x) + fused_fp8, fused_scale = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0( + x, True) - # Legacy: unpacked quant + manual pack + # Legacy quantization is the value/scale reference. Convert its MN-major + # packed scales to the standard byte layout for a physical-layout check. ref_fp8, ref_scale = torch.ops.trtllm.fp8_quantize_1x128(x, use_ue8m0=True) - # ref_scale shape from fp8_quantize_1x128: [num_n_blocks, m_padded] float - # Convert to the same packed (m, num_packed_sf_k) int32 layout as fused. ref_packed = fp8_utils.get_col_major_tma_aligned_packed_tensor( ref_scale[:, :m].t().contiguous().to(torch.float32)) + num_quant_sf_k = (k + 127) // 128 + expected_scale = _legacy_packed_to_r128c4(ref_packed, + m, + num_quant_sf_k, + sf_replication=4) - # FP8 bytes must be bit-identical for the valid [0, m) region. assert torch.equal(fused_fp8.view(torch.uint8), ref_fp8.view(torch.uint8)[:m]), \ f"FP8 mismatch for ({m}, {k})" + assert fused_scale.dtype == torch.uint8 + assert fused_scale.shape == expected_scale.shape + assert torch.equal(fused_scale, expected_scale), \ + f"R128c4 UE8M0 mismatch for ({m}, {k})" + + +@skip_if_not_sm100_family +@pytest.mark.parametrize("m,k", [(1, 128), (33, 640)]) +def test_fp8_quantize_1x128_packed_ue8m0_r128c4_zero_blocks(m, k): + """All-zero blocks retain the legacy kernel's UE8M0 scale encoding.""" + from tensorrt_llm.quantization.utils import fp8_utils - # Packed scales: compare element-by-element. - fused_contig = fused_packed.contiguous().view(torch.int32) - ref_contig = ref_packed.contiguous().view(torch.int32) - assert fused_contig.shape == ref_contig.shape, \ - f"shape mismatch fused={fused_contig.shape} ref={ref_contig.shape}" - assert torch.equal( - fused_contig, - ref_contig), (f"Packed UE8M0 mismatch for ({m}, {k}): " - f"fused[0,:]={fused_contig[0]} ref[0,:]={ref_contig[0]}") + x = torch.zeros((m, k), device="cuda", dtype=torch.bfloat16) + fused_fp8, fused_scale = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0( + x, True) + ref_fp8, ref_scale = torch.ops.trtllm.fp8_quantize_1x128(x, use_ue8m0=True) + ref_packed = fp8_utils.get_col_major_tma_aligned_packed_tensor( + ref_scale[:, :m].t().contiguous().to(torch.float32)) + expected_scale = _legacy_packed_to_r128c4(ref_packed, + m, + k // 128, + sf_replication=4) + assert torch.equal(fused_fp8.view(torch.uint8), + ref_fp8.view(torch.uint8)[:m]) + assert torch.equal(fused_scale, expected_scale) -@pytest.mark.skipif(not isSM100Family(), - reason="fp8_quantize_1x128_packed_ue8m0 is SM100 only.") + +@skip_if_not_sm100_family @pytest.mark.parametrize("m,k", [ (1, 128), (3, 256), (7, 512), (13, 7168), (127, 4096), + (129, 640), ]) -def test_fp8_quantize_1x128_packed_ue8m0_padded_rows_are_zero(m, k): - """Padded rows [m, m_aligned) of the physical packed scale buffer must be 0. - The op returns a `(m, num_packed_sf_k)` strided view that hides the padded - rows; allocate a sentinel buffer immediately before the call so any - unwritten ints in the padded region surface as the sentinel. - """ - if m % 4 == 0: - pytest.skip("m is already TMA-aligned; no padded rows to check") - +def test_fp8_quantize_1x128_packed_ue8m0_r128c4_padding_is_zero(m, k): + """R128c4 M padding is initialized and every K32 slot is populated.""" torch.manual_seed(0) - m_padded = ((m + 3) // 4) * 4 - num_packed_sf_k = ((k + 127) // 128 + 3) // 4 - total_int32 = num_packed_sf_k * m_padded - - # Poison the CUDA caching allocator: allocate-fill-free a buffer of the - # exact size, then call the op. The op's empty_cuda is likely to land in - # the same slab, so any unwritten ints surface as the sentinel. - SENTINEL = 0x5A5A5A # 5,921,370 — fits in int32 and is distinctive - poison = torch.empty((num_packed_sf_k, m_padded), - dtype=torch.int32, - device="cuda") - poison.fill_(SENTINEL) - del poison - torch.cuda.synchronize() - x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) - _, packed = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0(x) + _, scale = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0(x, True) + + num_sf_k = k // 32 + num_packed_sf_k = (num_sf_k + 3) // 4 + valid = torch.zeros_like(scale, dtype=torch.bool) + m_idx = torch.arange(m, device="cuda", dtype=torch.int64).view(-1, 1) + sf_k_idx = torch.arange(num_sf_k, device="cuda", + dtype=torch.int64).view(1, -1) + offsets = (( + (m_idx // 128 * num_packed_sf_k + sf_k_idx // 4) * 32 + m_idx % 32) * 4 + + (m_idx % 128) // 32) * 4 + sf_k_idx % 4 + valid[offsets.flatten()] = True + assert torch.count_nonzero(scale[~valid]).item() == 0 + + +@skip_if_not_sm100_family +@pytest.mark.parametrize("m,k", [(3, 256), (7, 512), (127, 4096)]) +def test_fp8_quantize_1x128_packed_ue8m0_legacy_layout(m, k): + """The compatibility mode retains deep_gemm's packed SF contract.""" + from tensorrt_llm.quantization.utils import fp8_utils - # Physical layout: int32[num_packed_sf_k][m_padded] starting at packed.data_ptr(). - # The returned tensor's storage_size() reflects only the logical view, so - # reach into it via the data_ptr + cudaMemcpyDeviceToDevice into a fresh - # full-sized tensor. - physical = torch.empty((num_packed_sf_k, m_padded), - dtype=torch.int32, - device="cuda") - # Use cudart memcpy via torch's torch.cuda.memory functions - src_ptr = packed.data_ptr() - dst_ptr = physical.data_ptr() - nbytes = total_int32 * 4 - # Build a 1-D byte view at src_ptr by creating a fresh tensor of the - # appropriate size; this requires an untyped storage of full extent. The - # storage created by `at::from_blob` is sized to the strided view's reach - # (m, num_packed_sf_k) so we need an out-of-band copy. cudart via torch: - torch.cuda.synchronize() - import ctypes - libcudart = ctypes.CDLL("libcudart.so") - err = libcudart.cudaMemcpy(ctypes.c_void_p(dst_ptr), - ctypes.c_void_p(src_ptr), - ctypes.c_size_t(nbytes), - ctypes.c_int(3)) # cudaMemcpyDeviceToDevice - assert err == 0, f"cudaMemcpy failed with err={err}" - torch.cuda.synchronize() + torch.manual_seed(0) + x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + fused_fp8, fused_packed = \ + torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0(x, False) + ref_fp8, ref_scale = torch.ops.trtllm.fp8_quantize_1x128(x, use_ue8m0=True) + ref_packed = fp8_utils.get_col_major_tma_aligned_packed_tensor( + ref_scale[:, :m].t().contiguous().to(torch.float32)) - padded_tail = physical[:, m:m_padded] - nonzero = int((padded_tail != 0).sum().item()) - assert nonzero == 0, ( - f"Padded rows [{m}, {m_padded}) must be zero; got {nonzero} non-zero " - f"int32 in shape ({num_packed_sf_k}, {m_padded - m})") + assert torch.equal(fused_fp8.view(torch.uint8), + ref_fp8.view(torch.uint8)[:m]) + assert fused_packed.dtype == torch.int32 + assert torch.equal(fused_packed.contiguous(), ref_packed.contiguous()) # ---------------------------------------------------------------------------