[None][perf] Add a small-M MXFP8 quantization kernel for swizzled scale factors

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 <typename T>
+__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<int64_t>(rowIdx) * numVecs + vecIdx;
+
+    cudaGridDependencySynchronize();
+
+    // Inactive threads convert zeros, so that every lane takes part in the shuffles below.
+    PackedVec<T> vec{};
+    if (active)
+    {
+        vec = reinterpret_cast<PackedVec<T> const*>(in)[offset];
+    }
+    uint8_t sf = 0;
+    uint64_t const values = cvt_warp_fp16_to_mxfp8<T, 32>(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<uint32_t>(sf), groupLane + 4);
+    uint32_t const sf2 = __shfl_sync(0xffffffffU, static_cast<uint32_t>(sf), groupLane + 8);
+    uint32_t const sf3 = __shfl_sync(0xffffffffU, static_cast<uint32_t>(sf), groupLane + 12);
+    if (active && lane == groupLane)
+    {
+        int const atomIdx = vecIdx / kVecsPerAtom;
+        SFout[atomIdx * kSlotsPerAtom + rowIdx]
+            = make_uint4(static_cast<uint32_t>(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 <typename T>
+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<uintptr_t>(input) % 16 == 0 && reinterpret_cast<uintptr_t>(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<T>, m, n, input,
+        reinterpret_cast<uint64_t*>(output), reinterpret_cast<uint4*>(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;

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