Improvements for: Groupwise scaling along M for FP8 gemm #2095

LucasWilkinson · 2025-02-10T05:30:35Z

Various improvements to "Groupwise scaling along M" (#2037) namely to address: #2087, context vllm-project/vllm#11868 (comment)

Improvements:

Multiple threads now participating in copy A scales
Predication when copying A scale loads, this means if there is partial M tile (due to the problem shape not being evenly divided by the M tile shape)
More commonly used scale layouts, currently CUTLASS uses a layout like:

(M_TILES, ScaleMsPerTile, K_TILES, L), ordered: (2, 0, 1, 3)

this PR moves to a layout of (i.e. standard M-major):

(M / ScaleGranularityM, K_TILES, L), ordered: (1, 0, 2)

making it much easier to integrate into inference libraries

These improvements were part of vLLMs adoption of this kernel https://github.com/vllm-project/vllm/blob/v0.7.1/csrc/cutlass_extensions/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp (PR: vllm-project/vllm#11868) and is in current wide scale use. Our goal is to rely on the CUTLASS implementation but that currently not possible given the issues above.

hwu36 · 2025-02-21T03:01:21Z

@LucasWilkinson , we upstreamed our change to groupwise scaling kernels. there are some conflicts in this PR that needs to be solved.

Our change is mainly:

Extend groupwise scaling gemm to support both M dimension and N dimension groupwise scaling in FP8 GEMM.
In examples/67_hopper_fp8_warp_specialized_gemm_with_blockwise_scaling/67_hopper_fp8_warp_specialized_gemm_with_groupwise_scaling.cu, two parameters ScaleGranularityM and ScaleGranularityNcontrol the scaling mode:


ScaleGranularityM == 128 and ScaleGranularityN == 128 --> 2Dx2D scaling (block-wise scaling, same as 67_hopper_fp8_warp_specialized_gemm_with_blockwise_scaling.cu , 2Dx2D refers to the shape of the scaling factor)

ScaleGranularityM == 1 and ScaleGranularityN == 128 --> 1Dx2D scaling

ScaleGranularityM == 128 and ScaleGranularityN == 1 --> 2Dx1D scaling

ScaleGranularityM == 1 and ScaleGranularityN == 1 --> 1Dx1D scaling

Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>

LucasWilkinson · 2025-02-25T07:01:33Z

@LucasWilkinson , we upstreamed our change to groupwise scaling kernels. there are some conflicts in this PR that needs to be solved.

Our change is mainly:

Extend groupwise scaling gemm to support both M dimension and N dimension groupwise scaling in FP8 GEMM.
In examples/67_hopper_fp8_warp_specialized_gemm_with_blockwise_scaling/67_hopper_fp8_warp_specialized_gemm_with_groupwise_scaling.cu, two parameters ScaleGranularityM and ScaleGranularityNcontrol the scaling mode:


ScaleGranularityM == 128 and ScaleGranularityN == 128 --> 2Dx2D scaling (block-wise scaling, same as 67_hopper_fp8_warp_specialized_gemm_with_blockwise_scaling.cu , 2Dx2D refers to the shape of the scaling factor)

ScaleGranularityM == 1 and ScaleGranularityN == 128 --> 1Dx2D scaling

ScaleGranularityM == 128 and ScaleGranularityN == 1 --> 2Dx1D scaling

ScaleGranularityM == 1 and ScaleGranularityN == 1 --> 1Dx1D scaling

apologies for the delay the PR has been updated, currently I am still vectorizing the loads of B scales along N (like main) but it might actually makes sense to not do this to enable transposing A and B (since we currently have partial tiles along M this would mean partial tiles along N)

ProphetPeng · 2025-02-25T09:41:44Z

include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp

@@ -280,8 +278,11 @@ struct CollectiveMma<
    implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_A>(cute::make_shape(M,K,L), StrideA{});
    constexpr int min_tma_aligned_elements_B = tma_alignment_bits / cutlass::sizeof_bits<ElementB>::value;
    implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_B>(cute::make_shape(N,K,L), StrideB{});
-    /* MMA promotion interval should be a multiple of 4, since each mainloop iteration would issue 4 MMA instructions. */
-    implementable = implementable && (args.mma_promotion_interval % 4 == 0);
+


Is there any promblems when transpose A and transpose B?

currently this assumes full tiles in N and K so if using this for inference where activations may have partial tiles if you transpose it to Y^T = WX^T it may report not implementable, I think im going to update this since ideally in vLLM we'd like to transpose it to use smaller tensor core instructions, we do lose vectorization on the loads then though

hwu36 · 2025-02-25T18:51:19Z

@manishucsd

include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp

...p8_warp_specialized_gemm_with_blockwise_scaling/reference/host/gemm_with_groupwise_scaling.h

hwu36 · 2025-02-26T04:49:47Z

include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp

@@ -575,7 +607,7 @@ struct CollectiveMma<

    tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;

-    GmmaFP8Accumulation accumulation(accum, mainloop_params.mma_promotion_interval, size<2>(tCrA));
+    GmmaFP8Accumulation accumulation(accum, size<2>(TileShape{}) / size<2>(typename TiledMma::AtomShape_MNK{}), size<2>(tCrA));


Maybe still using ScalePromotionInterval here, and move size<2>(TileShape{}) / size<2>(typename TiledMma::AtomShape_MNK{} to can_implement check?

Hmm im not sure I see ScalePromotionInterval, what would be the motivation to not have this determined at compile time? it seems a bit unnecessarily burdensome on the user to have them set mma_promotion_interval manually

In anycase moving this as constexpr somewhere on the top will better for readability.

static constexpr int ScalePromotionInterval = size<2>(TileShape{}) / size<2>(typename TiledMma::AtomShape_MNK{}) and using that here?

@hwu36, So this will be 4 for TileShapeK = 128 and InstructionShape = 32, which is the original case, for TileShape = 64 this will be 2. Will that be not supported?

Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>

manishucsd · 2025-02-26T05:54:02Z

include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp

@@ -147,9 +147,10 @@ struct CollectiveMma<
      cute::conditional_t< ::cutlass::gemm::detail::is_major<0,StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));

  // Block scaling gmem-to-smem copy atom 
-  using BlockScaleCopyTypeA = cute::uint_byte_t<cute::min(static_cast<int>(sizeof(ElementBlockScale)) * ScaleMsPerTile, 16)>;
+  //  we can have partial tiles in M, so don't vectorize those loads


Why is this restriction only for M and not for N? dim-M usually maps to batch count while dim-N will be model_dimension, a nice multiple of 2? correct?

If this is A_row * B_col groupwise GEMM, it is sometimes required that we do transposed and swap creating an underlying GEMM to be B_row * A_row, swapping M <-> N. This is typically helpful for (a.) mixed-input BF16*F8 which doesn't apply here (b.) M is small say 64, we can swap and transpose to run a bette. I have seen that to give more performance for small M.

Does vectorizing scale_copy_b vs not-vectorizing give any performance improvements? If not, I would suggest that we be symmetric for this kernel in M and N to allow user to apply swap and transpose trick to this kernel.

I was mostly just trying to keep it as close to the original as possible to minimize the chances of perf regressions, but I agree this is much less confusing. And I think we will want to transpose in vLLM in order to use smaller instructions for smaller batch sizes.

pushed an update that enables partial tiles in N

manishucsd · 2025-02-26T06:07:27Z

include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp

+      copy_if(scale_copy_a, tApA_ScaleA, tAgA_ScaleA(_,_,*k_tile_iter), tAsA_ScaleA(_,_,write_stage));
+      copy_if(scale_copy_b, tBpB_ScaleB, tBgB_ScaleB(_,_,*k_tile_iter), tBsB_ScaleB(_,_,write_stage));


can you make sure that this copy_if is issued by only 32 threads? The thread layout of shape 32 (created above) won't be tiled over entire tile by make_tiled_copy, just confirm please using simple printf

Ran

if ((!blockIdx.x && !blockIdx.y && !blockIdx.z)) printf("%d ", threadIdx.x); if (thread0()) printf("\n"); pipeline.producer_commit(smem_pipe_write, cutlass::arch::cpasync_barrier_arrive_noinc);

and got:

... 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 ...

I think we should be good 👍

manishucsd · 2025-02-26T06:08:40Z

include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp

+    Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{});        // (BLK_M,BLK_K,PIPE)
+    Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), SmemLayoutB{});        // (BLK_N,BLK_K,PIPE)


Should TMA related tensor constructions be in lane_predicate as before, no need for all the threads to construct this even in this implementation?

Im not sure, I didn't think this was a big deal since if you look at the 3.6.0 diff with improved the mixed input GEMM (we were told 3.6 had perf improvements for mixed input) in include/cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp you can see that it was updated to have all the threads compute the TMA tensors, not sure what the recommended approach is, or if this particular change had any impact. Would some guidance!

manishucsd · 2025-02-26T06:12:52Z

include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp

@@ -575,7 +607,7 @@ struct CollectiveMma<

    tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;

-    GmmaFP8Accumulation accumulation(accum, mainloop_params.mma_promotion_interval, size<2>(tCrA));
+    GmmaFP8Accumulation accumulation(accum, size<2>(TileShape{}) / size<2>(typename TiledMma::AtomShape_MNK{}), size<2>(tCrA));


In anycase moving this as constexpr somewhere on the top will better for readability.

static constexpr int ScalePromotionInterval = size<2>(TileShape{}) / size<2>(typename TiledMma::AtomShape_MNK{}) and using that here?

@hwu36, So this will be 4 for TileShapeK = 128 and InstructionShape = 32, which is the original case, for TileShape = 64 this will be 2. Will that be not supported?

Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>

LucasWilkinson changed the title ~~[WIP][Bugfix] Bug fixes for: Groupwise scaling along M for FP8 gemm~~ Improvements: Groupwise scaling along M for FP8 gemm Feb 10, 2025

LucasWilkinson changed the title ~~Improvements: Groupwise scaling along M for FP8 gemm~~ Improvements for: Groupwise scaling along M for FP8 gemm Feb 10, 2025

LucasWilkinson marked this pull request as ready for review February 10, 2025 21:07

LucasWilkinson mentioned this pull request Feb 10, 2025

[QST] how to use groupwise scaling along M for FP8 gemm to impelement per-token-per-128-channel and blockwise? #2087

Open

yizhang2077 mentioned this pull request Feb 11, 2025

support blockwise fp8 matmul kernel sgl-project/sglang#3267

Merged

4 tasks

LucasWilkinson added 3 commits February 25, 2025 06:13

fix blockwise fp8 kernels

3986bc1

Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>

wip, < 128 not working

b42a7b4

Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>

fix < 128

7f541db

Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>

LucasWilkinson force-pushed the lwilkinson/fix-fp8-blockwise branch from db87722 to 7f541db Compare February 25, 2025 06:50

reduce diff

60ddaa3

Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>

ProphetPeng reviewed Feb 25, 2025

View reviewed changes

LucasWilkinson mentioned this pull request Feb 25, 2025

[Kernel] Update cutlass_scaled_mm to support 2d group (blockwise) scaling vllm-project/vllm#11868

Merged

4 tasks

hwu36 reviewed Feb 25, 2025

View reviewed changes

include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp Outdated Show resolved Hide resolved

hwu36 reviewed Feb 26, 2025

View reviewed changes

...p8_warp_specialized_gemm_with_blockwise_scaling/reference/host/gemm_with_groupwise_scaling.h Outdated Show resolved Hide resolved

hwu36 reviewed Feb 26, 2025

View reviewed changes

review comments

2a9256f

Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>

LucasWilkinson force-pushed the lwilkinson/fix-fp8-blockwise branch from edd90be to 2a9256f Compare February 26, 2025 05:36

manishucsd reviewed Feb 26, 2025

View reviewed changes

support partial n blocks

1dc4ebd

Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Improvements for: Groupwise scaling along M for FP8 gemm #2095

Improvements for: Groupwise scaling along M for FP8 gemm #2095

LucasWilkinson commented Feb 10, 2025 •

edited

Loading

hwu36 commented Feb 21, 2025

LucasWilkinson commented Feb 25, 2025

ProphetPeng Feb 25, 2025

LucasWilkinson Feb 25, 2025

hwu36 commented Feb 25, 2025

hwu36 Feb 26, 2025

LucasWilkinson Feb 26, 2025

manishucsd Feb 26, 2025

manishucsd Feb 26, 2025

LucasWilkinson Feb 26, 2025

LucasWilkinson Feb 26, 2025 •

edited

Loading

manishucsd Feb 26, 2025 •

edited

Loading

LucasWilkinson Feb 26, 2025

manishucsd Feb 26, 2025

LucasWilkinson Feb 26, 2025

manishucsd Feb 26, 2025

		copy_if(scale_copy_a, tApA_ScaleA, tAgA_ScaleA(_,_,*k_tile_iter), tAsA_ScaleA(_,_,write_stage));
		copy_if(scale_copy_b, tBpB_ScaleB, tBgB_ScaleB(_,_,*k_tile_iter), tBsB_ScaleB(_,_,write_stage));

		Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
		Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)

Improvements for: Groupwise scaling along M for FP8 gemm #2095

Are you sure you want to change the base?

Improvements for: Groupwise scaling along M for FP8 gemm #2095

Conversation

LucasWilkinson commented Feb 10, 2025 • edited Loading

hwu36 commented Feb 21, 2025

LucasWilkinson commented Feb 25, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

hwu36 commented Feb 25, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

LucasWilkinson Feb 26, 2025 • edited Loading

Choose a reason for hiding this comment

manishucsd Feb 26, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

LucasWilkinson commented Feb 10, 2025 •

edited

Loading

LucasWilkinson Feb 26, 2025 •

edited

Loading

manishucsd Feb 26, 2025 •

edited

Loading