diff --git a/cpp/tensorrt_llm/kernels/quantization.cu b/cpp/tensorrt_llm/kernels/quantization.cu index f6cf68053094..0ecea32593c6 100644 --- a/cpp/tensorrt_llm/kernels/quantization.cu +++ b/cpp/tensorrt_llm/kernels/quantization.cu @@ -191,6 +191,125 @@ void invokeFP4Quantization(int b, int m, int n, T const* input, float const* SFS } } +//////////////////////////////////////////////////////////////////////////////////////////////////// +// MXFP8 quantization of a few rows with swizzled scale factors + +namespace +{ + +// One CTA converts 1024 columns (eight 128-column scale-factor atoms) of one row. +constexpr int kMxfp8SmallMThreads = 128; +// Rows 0..31 of a swizzled 128x4 atom are the first entries of its 32 16-byte slots. +constexpr int kMxfp8SmallMMaxRows = 32; + +// MXFP8 quantization (K32 UE8M0 scales, swizzled 128x4 layout) of one unpadded [numRows, numCols] matrix with +// numRows <= 32 and numCols % 128 == 0. +// +// quantize_with_block_size gives each CTA one padded row, so a few rows launch 128 CTAs and most of them only write +// zero padding scale factors one byte at a time. Here each thread converts one 8-element vector with +// cvt_warp_fp16_to_mxfp8, the function that quantize_with_block_size uses, so the values and scale factors are bitwise +// the same. The grid is (ceil(numCols / 1024), numRows), and the scale factors are written in 16-byte stores: +// +// A 128-column atom of the swizzled layout has 32 slots of 16 bytes. Slot r holds the four K32 scales of rows r, +// r + 32, r + 64 and r + 96. For r < numRows the slot is [four scales of row r, 12 zero bytes]; the other slots are +// zero. Every byte of the scale-factor buffer is written. +template +__global__ void __launch_bounds__(kMxfp8SmallMThreads) + quantize_mxfp8_small_m_swizzled(int32_t numRows, int32_t numCols, T const* in, uint64_t* out, uint4* SFout) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + // 16 threads convert the 128 columns of one atom. + constexpr int kVecsPerAtom = 4 * 32 / CVT_ELTS_PER_THREAD; + constexpr int kAtomsPerCta = kMxfp8SmallMThreads / kVecsPerAtom; + constexpr int kSlotsPerAtom = 32; + + int const numVecs = numCols / CVT_ELTS_PER_THREAD; + int const rowIdx = blockIdx.y; + int const vecIdx = blockIdx.x * kMxfp8SmallMThreads + threadIdx.x; + // numCols % 128 == 0, so the 16 threads of an atom are all active or all inactive. + bool const active = vecIdx < numVecs; + int64_t const offset = static_cast(rowIdx) * numVecs + vecIdx; + + cudaGridDependencySynchronize(); + + // Inactive threads convert zeros, so that every lane takes part in the shuffles below. + PackedVec vec{}; + if (active) + { + vec = reinterpret_cast const*>(in)[offset]; + } + uint8_t sf = 0; + uint64_t const values = cvt_warp_fp16_to_mxfp8(vec, &sf); + if (active) + { + out[offset] = values; + } + + // The first lane of each 16-lane group stores the slot of this row: its own scale and those of lanes +4, +8, +12. + int const lane = threadIdx.x % 32; + int const groupLane = lane & ~(kVecsPerAtom - 1); + uint32_t const sf1 = __shfl_sync(0xffffffffU, static_cast(sf), groupLane + 4); + uint32_t const sf2 = __shfl_sync(0xffffffffU, static_cast(sf), groupLane + 8); + uint32_t const sf3 = __shfl_sync(0xffffffffU, static_cast(sf), groupLane + 12); + if (active && lane == groupLane) + { + int const atomIdx = vecIdx / kVecsPerAtom; + SFout[atomIdx * kSlotsPerAtom + rowIdx] + = make_uint4(static_cast(sf) | (sf1 << 8) | (sf2 << 16) | (sf3 << 24), 0U, 0U, 0U); + } + + // The CTAs of one column chunk share the zero slots of the rows numRows..31 of its atoms. + int const firstAtom = blockIdx.x * kAtomsPerCta; + int const numAtoms = min(kAtomsPerCta, numCols / 128 - firstAtom); + int const numPadRows = kSlotsPerAtom - numRows; + for (int i = rowIdx * kMxfp8SmallMThreads + threadIdx.x; i < numAtoms * numPadRows; + i += numRows * kMxfp8SmallMThreads) + { + SFout[(firstAtom + i / numPadRows) * kSlotsPerAtom + numRows + i % numPadRows] = make_uint4(0U, 0U, 0U, 0U); + } + + // The PDL contract of quantize_with_block_size: every store is visible before any thread triggers completion. + __threadfence(); + __syncthreads(); + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +// Runs quantize_mxfp8_small_m_swizzled when the problem fits it. Returns false when the caller must run +// quantize_with_block_size. +template +bool launchMxFP8SmallMSwizzled(int b, int m, int n, int padded_n, T const* input, int64_t* output, int32_t* SFOuput, + QuantizationSFLayout layout, cudaStream_t stream) +{ + // Up to 2048 columns both kernels run at the launch latency on B200, and quantize_with_block_size is not slower. + constexpr int kMinCols = 2048; + static int const smVersion = tensorrt_llm::common::getSMVersion(); + bool const fits = (smVersion == 100 || smVersion == 103) && b == 1 && layout == QuantizationSFLayout::SWIZZLED + && m >= 1 && m <= kMxfp8SmallMMaxRows && n > kMinCols && n == padded_n && n % 128 == 0 + && reinterpret_cast(input) % 16 == 0 && reinterpret_cast(SFOuput) % 16 == 0; + if (!fits) + { + return false; + } + + int const numVecs = n / CVT_ELTS_PER_THREAD; + cudaLaunchConfig_t config; + config.gridDim = dim3((numVecs + kMxfp8SmallMThreads - 1) / kMxfp8SmallMThreads, m); + config.blockDim = dim3(kMxfp8SmallMThreads); + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config.numAttrs = 1; + config.attrs = attrs; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, quantize_mxfp8_small_m_swizzled, m, n, input, + reinterpret_cast(output), reinterpret_cast(SFOuput))); + return true; +} + +} // namespace + //////////////////////////////////////////////////////////////////////////////////////////////////// // MXFP8 Quantization @@ -202,6 +321,14 @@ void invokeMxFP8Quantization(int b, int m, int n, int padded_n, T const* input, 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."); + if constexpr (SF_VEC_SIZE == 32 && SF_OUTPUT_VEC_SIZE == 32) + { + if (launchMxFP8SmallMSwizzled(b, m, n, padded_n, input, output, SFOuput, layout, stream)) + { + return; + } + } + if (b == 0 || m == 0 || n == 0) { return; diff --git a/tests/unittest/_torch/kernels/parallel/test_fp8_quantize.py b/tests/unittest/_torch/kernels/parallel/test_fp8_quantize.py index 39cf287fa29f..21495ff72f31 100644 --- a/tests/unittest/_torch/kernels/parallel/test_fp8_quantize.py +++ b/tests/unittest/_torch/kernels/parallel/test_fp8_quantize.py @@ -303,6 +303,89 @@ def test_mxfp8_quantize_tiny_block_is_finite(dtype, is_sf_swizzled_layout): 4] != 0, "a block just above the threshold lost its value" +def _mxfp8_input_with_special_blocks(m, k, dtype, seed): + """Random values, and in every third row some special 32-element blocks: + all zero, one value below the smallest UE8M0 scale, the dtype maximum, + subnormal values and negative zeros.""" + torch.random.manual_seed(seed) + a = torch.randn([m, k], dtype=torch.float) * 16 + blocks = a.view(m, k // 32, 32) + num_blocks = k // 32 + tiny = 448.0 * 2.0**-127 + finfo = torch.finfo(dtype) + for row in range(0, m, 3): + special = [(row + i) % num_blocks for i in range(min(5, num_blocks))] + blocks[row, special[0]] = 0.0 + if len(special) > 1: + blocks[row, special[1]] = 0.0 + blocks[row, special[1], 7] = tiny * 0.5 + if len(special) > 2: + blocks[row, special[2], 3] = finfo.max + if len(special) > 3: + blocks[row, special[3]] = finfo.smallest_normal * 0.25 + if len(special) > 4: + blocks[row, special[4]] = -0.0 + return a.to(dtype).cuda().contiguous() + + +@pytest.mark.parametrize("m", [1, 2, 6, 12, 31, 32, 33, 128]) +@pytest.mark.parametrize("k", [1024, 2048, 2176, 2880, 5120, 7168, 16384]) +@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16]) +@skip_pre_blackwell_unittest +def test_mxfp8_quantize_swizzled_matches_linear_layout(m, k, dtype): + """The swizzled output is the linear output with interleaved scales. + + The check covers every byte of both buffers, including the zero scale + factors of the padding rows. Small m, wide k takes the small-M kernel; + the other shapes take quantize_with_block_size. + """ + a = _mxfp8_input_with_special_blocks(m, k, dtype, seed=m * 131 + k) + + fp8_linear, sf_linear = torch.ops.trtllm.mxfp8_quantize(a, False, 32) + ref_sf = torch.ops.trtllm.block_scale_interleave( + sf_linear.view(torch.uint8).view(m, k // 32).contiguous()) + + # Release 0xFF blocks of the output sizes to the caching allocator, so + # that a byte that the kernel does not write is likely to show as 0xFF. + sf_size = (m + 127) // 128 * 128 * ((k // 32 + 3) // 4 * 4) + poison = [ + torch.full((size, ), 0xFF, dtype=torch.uint8, device="cuda") + for size in (m * k, sf_size) + ] + del poison + + fp8_swizzled, sf_swizzled = torch.ops.trtllm.mxfp8_quantize(a, True, 32) + torch.cuda.synchronize() + + assert sf_swizzled.numel() == sf_size + assert torch.equal(fp8_swizzled.view(torch.uint8), + fp8_linear.view(torch.uint8)) + assert torch.equal(sf_swizzled.view(torch.uint8), ref_sf.view(torch.uint8)) + + +@pytest.mark.parametrize("m,k", [(6, 5120), (32, 7168), (64, 5120)]) +@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16]) +@skip_pre_blackwell_unittest +def test_mxfp8_quantize_swizzled_cuda_graph_replays(m, k, dtype): + """Graph replays with new inputs give the bytes of eager calls.""" + x = _mxfp8_input_with_special_blocks(m, k, dtype, seed=0) + # Run once outside of the capture. + torch.ops.trtllm.mxfp8_quantize(x, True, 32) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + fp8, sf = torch.ops.trtllm.mxfp8_quantize(x, True, 32) + + for seed in range(1, 6): + x.copy_(_mxfp8_input_with_special_blocks(m, k, dtype, seed=seed)) + graph.replay() + ref_fp8, ref_sf = torch.ops.trtllm.mxfp8_quantize(x, True, 32) + torch.cuda.synchronize() + assert torch.equal(fp8.view(torch.uint8), ref_fp8.view(torch.uint8)) + assert torch.equal(sf.view(torch.uint8), ref_sf.view(torch.uint8)) + + def _run_megamoe_prepare(hidden_states, token_selected_experts, token_final_scales): m, k = hidden_states.shape