Blockwise Scaling for FP8

[manishucsd]

Summary

As we adopt narrower datatypes, traditional scaling methods struggle to maintain accuracy, particularly with 8-bit floating-point types (e.g., e5m2_t, e4m3_t). The typical GEMM operation uses tensorwise scaling with D = alpha * (A @ B) + beta * C, but narrower datatypes necessitate more finer-grained scaling techniques. This PR adds blockwise scaling strategy to improve accuracy while making an effort to not loose performance. Before we dive deep into blockwise scaling below is a glossary of various scaling methods:

1. Tensorwise Scaling: Uses a single scaling factor per tensor, applied in the epilogue.
2. Rowwise Scaling: Uses a row vector for scaling, with dimensions Mx1 for operand A and 1xN for operand B, avoiding the scaling along the reduction dimension. This can also be handled in the epilogue with EpilogueVisitorTree.
3. Blockwise Scaling (this diff): Introduces a 2D scaling tensor, assigning one scaling value per CTA Block. Since this scaling involves the reduction dimension (M, N, K), it must be applied during the mainloop, impacting performance. This PR implements blockwise scaling for CUTLASS F8 GEMM, staging scaling tensors via shared memory, and preparing for future support of groupwise scaling.
4. Groupwise Scaling: Uses a 2D scaling tensor with multiple scaling values per CTA Block. Scaling granularity is independent of CTA Block configuration, allowing greater flexibility for future implementations.

This enhancement focuses on improving GEMM accuracy for narrow datatypes, balancing the trade-off between performance and precision with the addition of blockwise scaling support.

Blockwise Scaling

The figure below illustrates a blockwise scaled GEMM, with operand tensors A and B shown in grey, block scaling tensors in blue, and output in green. In this implementation, we load operand tensors using UTMALDG and block scaling tensors using LDGSTS, transferring them from global memory to shared memory. Block scaling tensor loads are issued for the same stage as the operand tensor loads. To ensure proper synchronization for LDGSTS, we use cutlass::arch::cpasync_barrier_arrive with noinc modifier. We have modified the PipelineTmaAsync class to accommodate a variable number of producer thread arrival events to support this functionality effectively.

Performance

For the graph below, I used CUDA Toolkit 12.3.2. Please note that with the latest toolkit 12.6.2, I observe LDLs and STLs in the SASS and the performance of block scaling is terrible. Thus, I stick with 12.3.2 for further performance optimizations and look for source of improvements. Please find the SASS attached for the example 54_hopper_fp8_warp_specialized_gemm (F8 with Slow Accumulation, FADDs after QGMMAs inside the mainloop) and the new 64_hopper_fp8_warp_specialized_gemm_with_blockwise_scaling (F8 with Slow Accum, FFMA after QGMMAs inside the mainloop).

NCU Profile (Kernel compiled with CUDA Toolkit 12.3.2)

I see one major stall for both FADD version and FFMA version soon after the QGMMA, waiting for for the accumulator to be ready to apply the promotions and scaling, respectively. I don't see any other difference other that that this stall is larger for FFMA.

FADD Version (Example 54 with slow accumulation and tensorwise scaling. Modified to have same tiling and kernel schedule as 64)

FFMA Version (Example 64 with slow accumulation and blockwise scaling)

Technically, for a large GEMM shape with cooperative schedule, we would expect both version to be running the same performance. Let me know if you have more input on what we are missing here to match the performance. We will eventually need a good implementation of blockwise scaling kernel in CUTLASS for plain F8 GEMMs and also take these learnings to FlashAttention-3 F8 Scaling.

Attn: @IonThruster , @hwu36 , @thakkarV

[thakkarV]

Super cool! Thanks for upstream :) Will do a full review soon. One comment to make to start would to please not extend the existing primary types in CUTLASS. Eg. The new collective builder should just be another specialized dispatch of the existing one. If you need to pass in extra arguments, you can include them in the builder dispatch policy itself. We want to ensure that there is always a single point of entry at each conceptual level.

[IonThruster]

Is it implicit that Element Block Scale will always be same as ElementAccum ? if so, would be good to add some static assert in the collective.

[manishucsd]

For now it is the case that ElementBlockScale = ElementAccumulator. I still prefer to write have an alias type in the code to increase readability and find the scale tensors quickly while reading and searching. Maybe in the future we will need different datatypes for ElementBlockScale and ElementAccumulator. Although, I don't see it happening as long as accumulation is in F32. I have added the static_asserts, in case user tries and set them differently, as we don't have NumericConvertor in the GmmaFP8Accumulation::scale_core.

[IonThruster]

Two suggestions :

• One can do way more than just scaling accums with this approach - so you might optionally want to consider may be generically adding an interface for allowing user to apply a point-wise operation on an accumulator per block.

• This impl places assumptions / restrictions on types of scale and accum. So this can be enhanced to call an appropriate sub-function / impl to ensure optimal scaling - that way in the future if one decides to pass a different scale type - it can just be extended with a new overload / specialization.

[manishucsd]

This is a great suggestion and we can consider it in the future. For now, we are looking for scaling the accumulators blockwise and possibly group-wise. Let us discuss more on this soon.

[manishucsd]

Pushed a few commits to address some of the comments. Please take a look. I kept the commits after initial commit separate so different comments can be reviewed easily. I will squash all commits before the merge.

[manishucsd]

Super cool! Thanks for upstream :) Will do a full review soon. One comment to make to start would to please not extend the existing primary types in CUTLASS. Eg. The new collective builder should just be another specialized dispatch of the existing one. If you need to pass in extra arguments, you can include them in the builder dispatch policy itself. We want to ensure that there is always a single point of entry at each conceptual level.

Thanks @thakkarV for the comment. Please the diff here. This diff makes the CollectiveBuilder entry point another specialization by adding a new dispatch policy. I had something similar in the TODO which also removed now. Please ignore the accidentally CMakeList change, you won't see that change in the full diff.

[zhyncs]

Hi @manishucsd could you help resolve the conflicts with the main branch?
Hi @thakkarV @hwu36 could you please review and merge this PR as soon as possible? SGLang team is eager to use this feature. Thanks!

[soundOfDestiny]

Hi @manishucsd could you help resolve the conflicts with the main branch? Hi @thakkarV @hwu36 could you please review and merge this PR as soon as possible? SGLang team is eager to use this feature. Thanks!

Conflicts can be resolved trivially and I have pushed it to https://github.com/soundOfDestiny/cutlass/tree/f8_blockwise_scaling_pr_branch_resolve_conflict_blockwise

In addition, I have implemented groupwise scaling granularity along M in A tensor with slight modification on this PR (https://github.com/manishucsd/cutlass/pull/1/files). Is it possible for reviewers to take a look at this PR manishucsd#1 for PR #1932? Thanks!

I also resolved conflicts on my groupwise branch and pushed it to
https://github.com/soundOfDestiny/cutlass/tree/f8_blockwise_scaling_pr_branch_resolve_conflict_groupwise

[manishucsd]

@zhyncs and @soundOfDestiny, I have rebased it on top of the latest commit (CUTLASS 3.6). Apologies for the delay, I was AFK. Looking forward to your use of this feature. We believe this PR will be a good starting point for groupwise, as your commit shows.

Do you only plan to have groups in M-dim? Are there uses case for groups in N-dim? Do you already have a non-CUTLASS groupwise kernel to show how much groupwise help numerically?

Thanks @IonThruster on review and discussion on handling synchronization with multiple producers. Thanks in advance @hwu36 to help merge this.

[soundOfDestiny]

@zhyncs and @soundOfDestiny, I have rebased it on top of the latest commit (CUTLASS 3.6). Apologies for the delay, I was AFK. Looking forward to your use of this feature. We believe this PR will be a good starting point for groupwise, as your commit shows.

Do you only plan to have groups in M-dim? Are there uses case for groups in N-dim? Do you already have a non-CUTLASS groupwise kernel to show how much groupwise help numerically?

Thanks @IonThruster on review and discussion on handling synchronization with multiple producers. Thanks in advance @hwu36 to help merge this.

Thank you very much! I have also updated my branch in manishucsd#1

Do you only plan to have groups in M-dim?

Having a smaller granularity in K-dim than this PR would hurt performance a lot due to the internal structure of WGMMA. We can have a larger granularity in K-dim by repeating the scaling tensor as input.

Are there uses case for groups in N-dim?

I think it's very useful to have groups in N-dim and we can implement groups in N-dim in the future by loading group scale of B into tCrScaleBViewAsC in the same way that we implement groups in M-dim by loading group scale of A into tCrScaleAViewAsC.
However, there still remains a slight difference that when we have groups in both M-dim and N-dim, we cannot pre-calculate the product of scale_a and scale_b because of various combinations in different rows and columns, so we have to calculate the ternary product of accumulator(this_row & this_column), scale_a(this_row) and scale_b(this_column) on-the-fly instead of calculating the binary product of accumulator and scale == scale_a * scale_b.

A special use case for groupwise is when group granularity along M-dim equals to 1, i.e. per-row (in M) x per-block (in K) scale for matrix A. This is used in https://github.com/deepseek-ai/DeepSeek-V3, in which there is also a triton kernel demo for GEMM with the per-row granularity scale (https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/kernel.py):

a_s_ptrs = a_s_ptr + offs_m * k
...
a_s = tl.load(a_s_ptrs)
...
accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]

among

def fp8_gemm_kernel(a_ptr, b_ptr, c_ptr,
                    a_s_ptr, b_s_ptr,
                    M, N: tl.constexpr, K: tl.constexpr,
                    BLOCK_SIZE_M: tl.constexpr,
                    BLOCK_SIZE_N: tl.constexpr,
                    BLOCK_SIZE_K: tl.constexpr):
    pid_m = tl.program_id(axis=0)
    pid_n = tl.program_id(axis=1)
    k = tl.cdiv(K, BLOCK_SIZE_K)
    offs_m = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
    offs_n = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    a_ptrs = a_ptr + offs_m[:, None] * K + offs_k[None, :]
    b_ptrs = b_ptr + offs_n[None, :] * K + offs_k[:, None]
    a_s_ptrs = a_s_ptr + offs_m * k
    b_s_ptrs = b_s_ptr + (offs_n // BLOCK_SIZE_K) * k

    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for i in range(k):
        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - i * BLOCK_SIZE_K, other=0.0)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - i * BLOCK_SIZE_K, other=0.0)
        a_s = tl.load(a_s_ptrs)
        b_s = tl.load(b_s_ptrs)
        accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
        a_ptrs += BLOCK_SIZE_K
        b_ptrs += BLOCK_SIZE_K
        a_s_ptrs += 1
        b_s_ptrs += 1
    c = accumulator.to(c_ptr.dtype.element_ty)
    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = c_ptr + offs_m[:, None] * N + offs_n[None, :]
    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, c, mask=mask)

[zhyncs]

Great to see it merged!!

[zhyncs]

Hi all @manishucsd @hwu36 @soundOfDestiny

I think we also need this https://github.com/manishucsd/cutlass/pull/1/files

May we consider raising another PR to support that?

[hwu36]

Yes, please file a pr to us directly.

[soundOfDestiny]

Hi all @manishucsd @hwu36 @soundOfDestiny

I think we also need this https://github.com/manishucsd/cutlass/pull/1/files

May we consider raising another PR to support that?

I got it. I've just moved https://github.com/manishucsd/cutlass/pull/1/files to #2037
cc @hwu36 Could you please take a look? Thanks.