[GH-ISSUE #10] Is DeepGEMM directly applicable to backward in training? #7

New issue

Closed

opened 2026-10-26 21:10:13 +00:00 by bsakaguchi · 33 comments

bsakaguchi commented 2026-10-26 21:10:13 +00:00

Owner

Copy link

Originally created by @YouJiacheng on GitHub (Feb 26, 2025).
Original GitHub issue: https://github.com/deepseek-ai/DeepGEMM/issues/10

bwd of GEMM is two GEMMs, but I wonder if I need to take some special care of the range of gradients?

Originally created by @YouJiacheng on GitHub (Feb 26, 2025). Original GitHub issue: https://github.com/deepseek-ai/DeepGEMM/issues/10 bwd of GEMM is two GEMMs, but I wonder if I need to take some special care of the range of gradients?

bsakaguchi closed this issue 2026-10-26 21:10:13 +00:00

bsakaguchi commented 2026-10-26 21:10:15 +00:00

Author

Owner

Copy link

@YouJiacheng commented on GitHub (Feb 26, 2025):

Oh, so we must write a quantization kernel that produces the correct lhs_scales and rhs_scales.

<!-- gh-comment-id:2684129960 --> @YouJiacheng commented on GitHub (Feb 26, 2025): Oh, so we must write a quantization kernel that produces the correct `lhs_scales` and `rhs_scales`.

bsakaguchi commented 2026-10-26 21:10:15 +00:00

Author

Owner

Copy link

@zhipeng93 commented on GitHub (Feb 26, 2025):

We need a fp8 gemm with 128x1 LHS scaling and 1x128 RHS scaling?

<!-- gh-comment-id:2684160045 --> @zhipeng93 commented on GitHub (Feb 26, 2025): We need a fp8 gemm with 128x1 LHS scaling and 1x128 RHS scaling?

bsakaguchi commented 2026-10-26 21:10:16 +00:00

Author

Owner

Copy link

@LyricZhao commented on GitHub (Feb 26, 2025):

We provide this library mainly for inference. So this library only supports DGRAD, not WGRAD.

In my understanding, WGRAD support needs more than a GEMM kernel, but also some utility fused kernels (e.g. transposing, fused with casting, fused with SwiGLU, fused with MoE layout). We want this library to be clean, so we didn't open-source them.

We may later release the WGRAD kernel, we will discuss about it internally :)

<!-- gh-comment-id:2684239143 --> @LyricZhao commented on GitHub (Feb 26, 2025): We provide this library mainly for inference. So this library only supports DGRAD, not WGRAD. In my understanding, WGRAD support needs more than a GEMM kernel, but also some utility fused kernels (e.g. transposing, fused with casting, fused with SwiGLU, fused with MoE layout). We want this library to be clean, so we didn't open-source them. We may later release the WGRAD kernel, we will discuss about it internally :)

bsakaguchi commented 2026-10-26 21:10:16 +00:00

Author

Owner

Copy link

@YouJiacheng commented on GitHub (Feb 26, 2025):

Thank you Chenggang!
Yup I forgot that WGRAD needs to transpose matrices.

<!-- gh-comment-id:2684252429 --> @YouJiacheng commented on GitHub (Feb 26, 2025): Thank you Chenggang! Yup I forgot that WGRAD needs to transpose matrices.

bsakaguchi commented 2026-10-26 21:10:16 +00:00

Author

Owner

Copy link

@thefacetakt commented on GitHub (Feb 27, 2025):

We may later release the WGRAD kernel, we will discuss about it internally :)

@LyricZhao, Could you hint at the approach you are taking for now?

Since it seems infeasible to perform (128, 1) x (1, 128) block-quantized gemm (because hopper tensor cores do not support scales), I am guessing you are using requantization of some kind. Do you use (1, 128) x (128,128) blocks for wgrad or (1, 128) x (128, 1)?

<!-- gh-comment-id:2687730306 --> @thefacetakt commented on GitHub (Feb 27, 2025): > We may later release the WGRAD kernel, we will discuss about it internally :) @LyricZhao, Could you hint at the approach you are taking for now? Since it seems infeasible to perform (128, 1) x (1, 128) block-quantized gemm (because hopper tensor cores do not support scales), I am guessing you are using requantization of some kind. Do you use (1, 128) x (128,128) blocks for wgrad or (1, 128) x (128, 1)?

bsakaguchi commented 2026-10-26 21:10:17 +00:00

Author

Owner

Copy link

@soundOfDestiny commented on GitHub (Feb 28, 2025):

We may later release the WGRAD kernel, we will discuss about it internally :)

@LyricZhao, Could you hint at the approach you are taking for now?

Since it seems infeasible to perform (128, 1) x (1, 128) block-quantized gemm (because hopper tensor cores do not support scales), I am guessing you are using requantization of some kind. Do you use (1, 128) x (128,128) blocks for wgrad or (1, 128) x (128, 1)?

(128, 128) block scale is only used in weight matrix and won't be used in activation.

<!-- gh-comment-id:2690128036 --> @soundOfDestiny commented on GitHub (Feb 28, 2025): > > We may later release the WGRAD kernel, we will discuss about it internally :) > > [@LyricZhao](https://github.com/LyricZhao), Could you hint at the approach you are taking for now? > > Since it seems infeasible to perform (128, 1) x (1, 128) block-quantized gemm (because hopper tensor cores do not support scales), I am guessing you are using requantization of some kind. Do you use (1, 128) x (128,128) blocks for wgrad or (1, 128) x (128, 1)? (128, 128) block scale is only used in weight matrix and won't be used in activation.

bsakaguchi commented 2026-10-26 21:10:17 +00:00

Author

Owner

Copy link

@TheTinyTeddy commented on GitHub (Mar 12, 2025):

We may later release the WGRAD kernel, we will discuss about it internally :)

@LyricZhao, Could you hint at the approach you are taking for now?
Since it seems infeasible to perform (128, 1) x (1, 128) block-quantized gemm (because hopper tensor cores do not support scales), I am guessing you are using requantization of some kind. Do you use (1, 128) x (128,128) blocks for wgrad or (1, 128) x (128, 1)?

(128, 128) block scale is only used in weight matrix and won't be used in activation.

So if not using the block scale for activation matrix then what should I use to calculate wgrad=x^T * dy since per token per tile quantizations are used on both A and B matrices (both are activations), which result in very inefficient GEMM (since the accumulation dimension K have different scales for every element), i.e. the GEMM is decomposed to many vector outer products

<!-- gh-comment-id:2717065115 --> @TheTinyTeddy commented on GitHub (Mar 12, 2025): > > > We may later release the WGRAD kernel, we will discuss about it internally :) > > > > > > [@LyricZhao](https://github.com/LyricZhao), Could you hint at the approach you are taking for now? > > Since it seems infeasible to perform (128, 1) x (1, 128) block-quantized gemm (because hopper tensor cores do not support scales), I am guessing you are using requantization of some kind. Do you use (1, 128) x (128,128) blocks for wgrad or (1, 128) x (128, 1)? > > (128, 128) block scale is only used in weight matrix and won't be used in activation. So if not using the block scale for activation matrix then what should I use to calculate wgrad=x^T * dy since per token per tile quantizations are used on both A and B matrices (both are activations), which result in very inefficient GEMM (since the accumulation dimension K have different scales for every element), i.e. the GEMM is decomposed to many vector outer products

bsakaguchi commented 2026-10-26 21:10:17 +00:00

Author

Owner

Copy link

@soundOfDestiny commented on GitHub (Mar 12, 2025):

We may later release the WGRAD kernel, we will discuss about it internally :)

@LyricZhao, Could you hint at the approach you are taking for now?
Since it seems infeasible to perform (128, 1) x (1, 128) block-quantized gemm (because hopper tensor cores do not support scales), I am guessing you are using requantization of some kind. Do you use (1, 128) x (128,128) blocks for wgrad or (1, 128) x (128, 1)?

(128, 128) block scale is only used in weight matrix and won't be used in activation.

So if not using the block scale for activation matrix then what should I use to calculate wgrad=x^T * dy since per token per tile quantizations are used on both A and B matrices (both are activations), which result in very inefficient GEMM (since the accumulation dimension K have different scales for every element), i.e. the GEMM is decomposed to many vector outer products

Requantization is inevitable in your case.

<!-- gh-comment-id:2717073904 --> @soundOfDestiny commented on GitHub (Mar 12, 2025): > > > > We may later release the WGRAD kernel, we will discuss about it internally :) > > > > > > > > > [@LyricZhao](https://github.com/LyricZhao), Could you hint at the approach you are taking for now? > > > Since it seems infeasible to perform (128, 1) x (1, 128) block-quantized gemm (because hopper tensor cores do not support scales), I am guessing you are using requantization of some kind. Do you use (1, 128) x (128,128) blocks for wgrad or (1, 128) x (128, 1)? > > > > > > (128, 128) block scale is only used in weight matrix and won't be used in activation. > > So if not using the block scale for activation matrix then what should I use to calculate wgrad=x^T * dy since per token per tile quantizations are used on both A and B matrices (both are activations), which result in very inefficient GEMM (since the accumulation dimension K have different scales for every element), i.e. the GEMM is decomposed to many vector outer products Requantization is inevitable in your case.

bsakaguchi commented 2026-10-26 21:10:18 +00:00

Author

Owner

Copy link

@TheTinyTeddy commented on GitHub (Mar 12, 2025):

We may later release the WGRAD kernel, we will discuss about it internally :)

@LyricZhao, Could you hint at the approach you are taking for now?
Since it seems infeasible to perform (128, 1) x (1, 128) block-quantized gemm (because hopper tensor cores do not support scales), I am guessing you are using requantization of some kind. Do you use (1, 128) x (128,128) blocks for wgrad or (1, 128) x (128, 1)?

(128, 128) block scale is only used in weight matrix and won't be used in activation.

So if not using the block scale for activation matrix then what should I use to calculate wgrad=x^T * dy since per token per tile quantizations are used on both A and B matrices (both are activations), which result in very inefficient GEMM (since the accumulation dimension K have different scales for every element), i.e. the GEMM is decomposed to many vector outer products

Requantization is inevitable in your case.

So do you mean rather than per token per 128 tile quantization for x during inference, I should be using per channel per 128 tile quantization for x^T and the same per channel per 128 tile for dy?

However, this means I need to quantize input x twice (1 for per token 1 for per channel) or save x in bf16 format during inference? both approaches increase memory for training (deepseek v3 paper explicitly says "the FP8 Wgrad GEMM allows activations to be stored in FP8 for use in the backward pass")

<!-- gh-comment-id:2717141484 --> @TheTinyTeddy commented on GitHub (Mar 12, 2025): > > > > > We may later release the WGRAD kernel, we will discuss about it internally :) > > > > > > > > > > > > [@LyricZhao](https://github.com/LyricZhao), Could you hint at the approach you are taking for now? > > > > Since it seems infeasible to perform (128, 1) x (1, 128) block-quantized gemm (because hopper tensor cores do not support scales), I am guessing you are using requantization of some kind. Do you use (1, 128) x (128,128) blocks for wgrad or (1, 128) x (128, 1)? > > > > > > > > > (128, 128) block scale is only used in weight matrix and won't be used in activation. > > > > > > So if not using the block scale for activation matrix then what should I use to calculate wgrad=x^T * dy since per token per tile quantizations are used on both A and B matrices (both are activations), which result in very inefficient GEMM (since the accumulation dimension K have different scales for every element), i.e. the GEMM is decomposed to many vector outer products > > Requantization is inevitable in your case. So do you mean rather than per token per 128 tile quantization for x during inference, I should be using per channel per 128 tile quantization for x^T and the same per channel per 128 tile for dy? However, this means I need to quantize input x twice (1 for per token 1 for per channel) or save x in bf16 format during inference? both approaches increase memory for training (deepseek v3 paper explicitly says "the FP8 Wgrad GEMM allows activations to be stored in FP8 for use in the backward pass")

bsakaguchi commented 2026-10-26 21:10:18 +00:00

Author

Owner

Copy link

@soundOfDestiny commented on GitHub (Mar 12, 2025):

We may later release the WGRAD kernel, we will discuss about it internally :)

@LyricZhao, Could you hint at the approach you are taking for now?
Since it seems infeasible to perform (128, 1) x (1, 128) block-quantized gemm (because hopper tensor cores do not support scales), I am guessing you are using requantization of some kind. Do you use (1, 128) x (128,128) blocks for wgrad or (1, 128) x (128, 1)?

(128, 128) block scale is only used in weight matrix and won't be used in activation.

So if not using the block scale for activation matrix then what should I use to calculate wgrad=x^T * dy since per token per tile quantizations are used on both A and B matrices (both are activations), which result in very inefficient GEMM (since the accumulation dimension K have different scales for every element), i.e. the GEMM is decomposed to many vector outer products

Requantization is inevitable in your case.

So do you mean rather than per token per 128 tile quantization for x during inference, I should be using per channel per 128 tile quantization for x^T and the same per channel per 128 tile for dy?

However, this means I need to quantize input x twice (1 for per token 1 for per channel) or save x in bf16 format during inference? both approaches increase memory for training (deepseek v3 paper explicitly says "the FP8 Wgrad GEMM allows activations to be stored in FP8 for use in the backward pass")

P.S. Usually we don't call the forward pass in training 'inference'.

<!-- gh-comment-id:2717152910 --> @soundOfDestiny commented on GitHub (Mar 12, 2025): > > > > > > We may later release the WGRAD kernel, we will discuss about it internally :) > > > > > > > > > > > > > > > [@LyricZhao](https://github.com/LyricZhao), Could you hint at the approach you are taking for now? > > > > > Since it seems infeasible to perform (128, 1) x (1, 128) block-quantized gemm (because hopper tensor cores do not support scales), I am guessing you are using requantization of some kind. Do you use (1, 128) x (128,128) blocks for wgrad or (1, 128) x (128, 1)? > > > > > > > > > > > > (128, 128) block scale is only used in weight matrix and won't be used in activation. > > > > > > > > > So if not using the block scale for activation matrix then what should I use to calculate wgrad=x^T * dy since per token per tile quantizations are used on both A and B matrices (both are activations), which result in very inefficient GEMM (since the accumulation dimension K have different scales for every element), i.e. the GEMM is decomposed to many vector outer products > > > > > > Requantization is inevitable in your case. > > So do you mean rather than per token per 128 tile quantization for x during inference, I should be using per channel per 128 tile quantization for x^T and the same per channel per 128 tile for dy? > > However, this means I need to quantize input x twice (1 for per token 1 for per channel) or save x in bf16 format during inference? both approaches increase memory for training (deepseek v3 paper explicitly says "the FP8 Wgrad GEMM allows activations to be stored in FP8 for use in the backward pass") P.S. Usually we don't call the forward pass in training 'inference'.

bsakaguchi commented 2026-10-26 21:10:19 +00:00

Author

Owner

Copy link

@TheTinyTeddy commented on GitHub (Mar 12, 2025):

We may later release the WGRAD kernel, we will discuss about it internally :)

@LyricZhao, Could you hint at the approach you are taking for now?
Since it seems infeasible to perform (128, 1) x (1, 128) block-quantized gemm (because hopper tensor cores do not support scales), I am guessing you are using requantization of some kind. Do you use (1, 128) x (128,128) blocks for wgrad or (1, 128) x (128, 1)?

(128, 128) block scale is only used in weight matrix and won't be used in activation.

So if not using the block scale for activation matrix then what should I use to calculate wgrad=x^T * dy since per token per tile quantizations are used on both A and B matrices (both are activations), which result in very inefficient GEMM (since the accumulation dimension K have different scales for every element), i.e. the GEMM is decomposed to many vector outer products

Requantization is inevitable in your case.

So do you mean rather than per token per 128 tile quantization for x during inference, I should be using per channel per 128 tile quantization for x^T and the same per channel per 128 tile for dy?
However, this means I need to quantize input x twice (1 for per token 1 for per channel) or save x in bf16 format during inference? both approaches increase memory for training (deepseek v3 paper explicitly says "the FP8 Wgrad GEMM allows activations to be stored in FP8 for use in the backward pass")

P.S. Usually we don't call the forward pass in training 'inference'.

My mistake, the terminology should indeed be the forward pass for training. So would you be able to let me know if my previous message is correct regarding the per channel quantization for activations in the Wgrad computation?

<!-- gh-comment-id:2717177732 --> @TheTinyTeddy commented on GitHub (Mar 12, 2025): > > > > > > > We may later release the WGRAD kernel, we will discuss about it internally :) > > > > > > > > > > > > > > > > > > [@LyricZhao](https://github.com/LyricZhao), Could you hint at the approach you are taking for now? > > > > > > Since it seems infeasible to perform (128, 1) x (1, 128) block-quantized gemm (because hopper tensor cores do not support scales), I am guessing you are using requantization of some kind. Do you use (1, 128) x (128,128) blocks for wgrad or (1, 128) x (128, 1)? > > > > > > > > > > > > > > > (128, 128) block scale is only used in weight matrix and won't be used in activation. > > > > > > > > > > > > So if not using the block scale for activation matrix then what should I use to calculate wgrad=x^T * dy since per token per tile quantizations are used on both A and B matrices (both are activations), which result in very inefficient GEMM (since the accumulation dimension K have different scales for every element), i.e. the GEMM is decomposed to many vector outer products > > > > > > > > > Requantization is inevitable in your case. > > > > > > So do you mean rather than per token per 128 tile quantization for x during inference, I should be using per channel per 128 tile quantization for x^T and the same per channel per 128 tile for dy? > > However, this means I need to quantize input x twice (1 for per token 1 for per channel) or save x in bf16 format during inference? both approaches increase memory for training (deepseek v3 paper explicitly says "the FP8 Wgrad GEMM allows activations to be stored in FP8 for use in the backward pass") > > P.S. Usually we don't call the forward pass in training 'inference'. My mistake, the terminology should indeed be the forward pass for training. So would you be able to let me know if my previous message is correct regarding the per channel quantization for activations in the Wgrad computation?

bsakaguchi commented 2026-10-26 21:10:19 +00:00

Author

Owner

Copy link

@soundOfDestiny commented on GitHub (Mar 12, 2025):

We may later release the WGRAD kernel, we will discuss about it internally :)

@LyricZhao, Could you hint at the approach you are taking for now?
Since it seems infeasible to perform (128, 1) x (1, 128) block-quantized gemm (because hopper tensor cores do not support scales), I am guessing you are using requantization of some kind. Do you use (1, 128) x (128,128) blocks for wgrad or (1, 128) x (128, 1)?

(128, 128) block scale is only used in weight matrix and won't be used in activation.

So if not using the block scale for activation matrix then what should I use to calculate wgrad=x^T * dy since per token per tile quantizations are used on both A and B matrices (both are activations), which result in very inefficient GEMM (since the accumulation dimension K have different scales for every element), i.e. the GEMM is decomposed to many vector outer products

Requantization is inevitable in your case.

So do you mean rather than per token per 128 tile quantization for x during inference, I should be using per channel per 128 tile quantization for x^T and the same per channel per 128 tile for dy?
However, this means I need to quantize input x twice (1 for per token 1 for per channel) or save x in bf16 format during inference? both approaches increase memory for training (deepseek v3 paper explicitly says "the FP8 Wgrad GEMM allows activations to be stored in FP8 for use in the backward pass")

P.S. Usually we don't call the forward pass in training 'inference'.

My mistake, the terminology should indeed be the forward pass for training. So would you be able to let me know if my previous message is correct regarding the per channel quantization for activations in the Wgrad computation?

As you mentioned, we don't save x in bf16 format. (deepseek v3 paper explicitly says "the FP8 Wgrad GEMM allows activations to be stored in FP8 for use in the backward pass")

<!-- gh-comment-id:2717184376 --> @soundOfDestiny commented on GitHub (Mar 12, 2025): > > > > > > > > We may later release the WGRAD kernel, we will discuss about it internally :) > > > > > > > > > > > > > > > > > > > > > [@LyricZhao](https://github.com/LyricZhao), Could you hint at the approach you are taking for now? > > > > > > > Since it seems infeasible to perform (128, 1) x (1, 128) block-quantized gemm (because hopper tensor cores do not support scales), I am guessing you are using requantization of some kind. Do you use (1, 128) x (128,128) blocks for wgrad or (1, 128) x (128, 1)? > > > > > > > > > > > > > > > > > > (128, 128) block scale is only used in weight matrix and won't be used in activation. > > > > > > > > > > > > > > > So if not using the block scale for activation matrix then what should I use to calculate wgrad=x^T * dy since per token per tile quantizations are used on both A and B matrices (both are activations), which result in very inefficient GEMM (since the accumulation dimension K have different scales for every element), i.e. the GEMM is decomposed to many vector outer products > > > > > > > > > > > > Requantization is inevitable in your case. > > > > > > > > > So do you mean rather than per token per 128 tile quantization for x during inference, I should be using per channel per 128 tile quantization for x^T and the same per channel per 128 tile for dy? > > > However, this means I need to quantize input x twice (1 for per token 1 for per channel) or save x in bf16 format during inference? both approaches increase memory for training (deepseek v3 paper explicitly says "the FP8 Wgrad GEMM allows activations to be stored in FP8 for use in the backward pass") > > > > > > P.S. Usually we don't call the forward pass in training 'inference'. > > My mistake, the terminology should indeed be the forward pass for training. So would you be able to let me know if my previous message is correct regarding the per channel quantization for activations in the Wgrad computation? As you mentioned, we don't save x in bf16 format. (deepseek v3 paper explicitly says "the FP8 Wgrad GEMM allows activations to be stored in FP8 for use in the backward pass")

bsakaguchi commented 2026-10-26 21:10:19 +00:00

Author

Owner

Copy link

@NiuMa-1234 commented on GitHub (Mar 18, 2025):

I found the fp8 gemm for Wgrad in BWD requires the shape of input activation to be the multiple of 128. But in moe layer it couldn't always be this way. Is there any necessity to restrain the number of tokens dispatched to each expert?

<!-- gh-comment-id:2733038774 --> @NiuMa-1234 commented on GitHub (Mar 18, 2025): I found the fp8 gemm for Wgrad in BWD requires the shape of input activation to be the multiple of 128. But in moe layer it couldn't always be this way. Is there any necessity to restrain the number of tokens dispatched to each expert?

bsakaguchi commented 2026-10-26 21:10:20 +00:00

Author

Owner

Copy link

@qijiaxing commented on GitHub (Mar 20, 2025):

I have implemented a variant that supports 1d1d scaling, can be used for WGRAD computation.
https://github.com/qijiaxing/DeepGEMM/tree/1d1d
any feedback or suggestions are welcomed!

<!-- gh-comment-id:2739395789 --> @qijiaxing commented on GitHub (Mar 20, 2025): I have implemented a variant that supports 1d1d scaling, can be used for WGRAD computation. https://github.com/qijiaxing/DeepGEMM/tree/1d1d any feedback or suggestions are welcomed!

bsakaguchi commented 2026-10-26 21:10:20 +00:00

Author

Owner

Copy link

@thefacetakt commented on GitHub (Mar 20, 2025):

I have implemented a variant that supports 1d1d scaling, can be used for WGRAD computation. https://github.com/qijiaxing/DeepGEMM/tree/1d1d any feedback or suggestions are welcomed!

Hi @qijiaxing, thanks for publishing 1d1d kernel!

I tried to run test_1d1d.py from your branch, but got an error.

When compiled with -lineinfo and run with cuda-gdb, I am getting the followind error:

Kernel config: sm: 124, block_m: 128, block_n: 112, stages: 6, multicast: 1
Config: N: 2112, BLOCK_N: 112, K: 7168, Stages: 6, K_per_iter: 768, block dim: 384

CUDA Exception: Warp Misaligned Address
The exception was triggered at PC 0x7f4b01bb6ce0  void deep_gemm::fp8_gemm_kernel<2112u, 7168u, 128u, 112u, 128u, 1u, 6u, 128u, 128u, 1u, (deep_gemm::GemmType)0>(__nv_bfloat16*, int*, unsigned int, CUtensorMap_st, CUtensorMap_st, CUtensorMap_st, CUtensorMap_st, CUtensorMap_st)  (copy_sm90_tma.hpp:110 in _ZN4cute16SM90_TMA_LOAD_2D4copyEPKvPmmPvRKiS6_ inlined from tma_utils.cuh:90)

Thread 1 "python" received signal CUDA_EXCEPTION_6, Warp Misaligned Address.
[Switching focus to CUDA kernel 0, grid 16, cluster (4,0,0), block (4,0,0), thread (256,0,0), device 0, sm 0, warp 8, lane 0]
0x00007f4b01bb6fd0 in deep_gemm::fp8_gemm_kernel<2112u, 7168u, 128u, 112u, 128u, 1u, 6u, 128u, 128u, 1u, (deep_gemm::GemmType)0><<<(124,1,1),(384,1,1)>>> ()
    at /home/thefacetak/src/1d1d/DeepGEMM/deep_gemm/jit/../include/cute/arch/copy_sm90_tma.hpp:110 in _ZN4cute16SM90_TMA_LOAD_2D4copyEPKvPmmPvRKiS6_ inlined from tma_utils.cuh:90
110             "r"(crd0), "r"(crd1), "l"(cache_hint)

Any idea what went wrong?

<!-- gh-comment-id:2739501690 --> @thefacetakt commented on GitHub (Mar 20, 2025): > I have implemented a variant that supports 1d1d scaling, can be used for WGRAD computation. https://github.com/qijiaxing/DeepGEMM/tree/1d1d any feedback or suggestions are welcomed! Hi @qijiaxing, thanks for publishing 1d1d kernel! I tried to run test_1d1d.py from your branch, but got an error. When compiled with `-lineinfo` and run with cuda-gdb, I am getting the followind error: ``` Kernel config: sm: 124, block_m: 128, block_n: 112, stages: 6, multicast: 1 Config: N: 2112, BLOCK_N: 112, K: 7168, Stages: 6, K_per_iter: 768, block dim: 384 CUDA Exception: Warp Misaligned Address The exception was triggered at PC 0x7f4b01bb6ce0 void deep_gemm::fp8_gemm_kernel<2112u, 7168u, 128u, 112u, 128u, 1u, 6u, 128u, 128u, 1u, (deep_gemm::GemmType)0>(__nv_bfloat16*, int*, unsigned int, CUtensorMap_st, CUtensorMap_st, CUtensorMap_st, CUtensorMap_st, CUtensorMap_st) (copy_sm90_tma.hpp:110 in _ZN4cute16SM90_TMA_LOAD_2D4copyEPKvPmmPvRKiS6_ inlined from tma_utils.cuh:90) Thread 1 "python" received signal CUDA_EXCEPTION_6, Warp Misaligned Address. [Switching focus to CUDA kernel 0, grid 16, cluster (4,0,0), block (4,0,0), thread (256,0,0), device 0, sm 0, warp 8, lane 0] 0x00007f4b01bb6fd0 in deep_gemm::fp8_gemm_kernel<2112u, 7168u, 128u, 112u, 128u, 1u, 6u, 128u, 128u, 1u, (deep_gemm::GemmType)0><<<(124,1,1),(384,1,1)>>> () at /home/thefacetak/src/1d1d/DeepGEMM/deep_gemm/jit/../include/cute/arch/copy_sm90_tma.hpp:110 in _ZN4cute16SM90_TMA_LOAD_2D4copyEPKvPmmPvRKiS6_ inlined from tma_utils.cuh:90 110 "r"(crd0), "r"(crd1), "l"(cache_hint) ``` Any idea what went wrong?

bsakaguchi commented 2026-10-26 21:10:20 +00:00

Author

Owner

Copy link

@qijiaxing commented on GitHub (Mar 20, 2025):

I have implemented a variant that supports 1d1d scaling, can be used for WGRAD computation. https://github.com/qijiaxing/DeepGEMM/tree/1d1d any feedback or suggestions are welcomed!

Hi @qijiaxing, thanks for publishing 1d1d kernel!

I tried to run test_1d1d.py from your branch, but got an error.

When compiled with -lineinfo and run with cuda-gdb, I am getting the followind error:

Kernel config: sm: 124, block_m: 128, block_n: 112, stages: 6, multicast: 1
Config: N: 2112, BLOCK_N: 112, K: 7168, Stages: 6, K_per_iter: 768, block dim: 384

CUDA Exception: Warp Misaligned Address
The exception was triggered at PC 0x7f4b01bb6ce0  void deep_gemm::fp8_gemm_kernel<2112u, 7168u, 128u, 112u, 128u, 1u, 6u, 128u, 128u, 1u, (deep_gemm::GemmType)0>(__nv_bfloat16*, int*, unsigned int, CUtensorMap_st, CUtensorMap_st, CUtensorMap_st, CUtensorMap_st, CUtensorMap_st)  (copy_sm90_tma.hpp:110 in _ZN4cute16SM90_TMA_LOAD_2D4copyEPKvPmmPvRKiS6_ inlined from tma_utils.cuh:90)

Thread 1 "python" received signal CUDA_EXCEPTION_6, Warp Misaligned Address.
[Switching focus to CUDA kernel 0, grid 16, cluster (4,0,0), block (4,0,0), thread (256,0,0), device 0, sm 0, warp 8, lane 0]
0x00007f4b01bb6fd0 in deep_gemm::fp8_gemm_kernel<2112u, 7168u, 128u, 112u, 128u, 1u, 6u, 128u, 128u, 1u, (deep_gemm::GemmType)0><<<(124,1,1),(384,1,1)>>> ()
    at /home/thefacetak/src/1d1d/DeepGEMM/deep_gemm/jit/../include/cute/arch/copy_sm90_tma.hpp:110 in _ZN4cute16SM90_TMA_LOAD_2D4copyEPKvPmmPvRKiS6_ inlined from tma_utils.cuh:90
110             "r"(crd0), "r"(crd1), "l"(cache_hint)

Any idea what went wrong?

HI @thefacetakt , thanks for your message, I also got a CUDA error when block_n is set to 112. I need to check further how to fix it.
As a temporary workaround, can you try change this line in file gemm_1d1d.py

# block_ns = tuple(range(16, 129, 8))
block_ns = (64, 128)

these two values works for me.

<!-- gh-comment-id:2739810896 --> @qijiaxing commented on GitHub (Mar 20, 2025): > > I have implemented a variant that supports 1d1d scaling, can be used for WGRAD computation. https://github.com/qijiaxing/DeepGEMM/tree/1d1d any feedback or suggestions are welcomed! > > Hi [@qijiaxing](https://github.com/qijiaxing), thanks for publishing 1d1d kernel! > > I tried to run test_1d1d.py from your branch, but got an error. > > When compiled with `-lineinfo` and run with cuda-gdb, I am getting the followind error: > > ``` > Kernel config: sm: 124, block_m: 128, block_n: 112, stages: 6, multicast: 1 > Config: N: 2112, BLOCK_N: 112, K: 7168, Stages: 6, K_per_iter: 768, block dim: 384 > > CUDA Exception: Warp Misaligned Address > The exception was triggered at PC 0x7f4b01bb6ce0 void deep_gemm::fp8_gemm_kernel<2112u, 7168u, 128u, 112u, 128u, 1u, 6u, 128u, 128u, 1u, (deep_gemm::GemmType)0>(__nv_bfloat16*, int*, unsigned int, CUtensorMap_st, CUtensorMap_st, CUtensorMap_st, CUtensorMap_st, CUtensorMap_st) (copy_sm90_tma.hpp:110 in _ZN4cute16SM90_TMA_LOAD_2D4copyEPKvPmmPvRKiS6_ inlined from tma_utils.cuh:90) > > Thread 1 "python" received signal CUDA_EXCEPTION_6, Warp Misaligned Address. > [Switching focus to CUDA kernel 0, grid 16, cluster (4,0,0), block (4,0,0), thread (256,0,0), device 0, sm 0, warp 8, lane 0] > 0x00007f4b01bb6fd0 in deep_gemm::fp8_gemm_kernel<2112u, 7168u, 128u, 112u, 128u, 1u, 6u, 128u, 128u, 1u, (deep_gemm::GemmType)0><<<(124,1,1),(384,1,1)>>> () > at /home/thefacetak/src/1d1d/DeepGEMM/deep_gemm/jit/../include/cute/arch/copy_sm90_tma.hpp:110 in _ZN4cute16SM90_TMA_LOAD_2D4copyEPKvPmmPvRKiS6_ inlined from tma_utils.cuh:90 > 110 "r"(crd0), "r"(crd1), "l"(cache_hint) > ``` > > Any idea what went wrong? HI @thefacetakt , thanks for your message, I also got a CUDA error when block_n is set to 112. I need to check further how to fix it. As a temporary workaround, can you try change this [line](https://github.com/qijiaxing/DeepGEMM/blob/e6f2074a2522d2ac7ef2ac89026405da18d5bef1/deep_gemm/jit_kernels/gemm_1d1d.py#L76) in file `gemm_1d1d.py` ``` # block_ns = tuple(range(16, 129, 8)) block_ns = (64, 128) ``` these two values works for me.

bsakaguchi commented 2026-10-26 21:10:21 +00:00

Author

Owner

Copy link

@thefacetakt commented on GitHub (Mar 20, 2025):

@qijiaxing thanks, that did the trick.

TWIMC, right now I am getting the followind 1d1d perf:

> Performance (m= 4096, n= 2112, k= 7168):  198 us | throughput:  627 TFLOPS,  313 GB/s
 > Performance (m= 4096, n=24576, k= 1536):  505 us | throughput:  612 TFLOPS,  486 GB/s
 > Performance (m= 4096, n=32768, k=  512):  325 us | throughput:  423 TFLOPS,  884 GB/s
 > Performance (m= 4096, n= 7168, k=16384): 1200 us | throughput:  802 TFLOPS,  203 GB/s
 > Performance (m= 4096, n= 4096, k= 7168):  317 us | throughput:  760 TFLOPS,  292 GB/s
 > Performance (m= 4096, n= 7168, k= 2048):  195 us | throughput:  617 TFLOPS,  420 GB/s

<!-- gh-comment-id:2739862714 --> @thefacetakt commented on GitHub (Mar 20, 2025): @qijiaxing thanks, that did the trick. TWIMC, right now I am getting the followind 1d1d perf: ``` > Performance (m= 4096, n= 2112, k= 7168): 198 us | throughput: 627 TFLOPS, 313 GB/s > Performance (m= 4096, n=24576, k= 1536): 505 us | throughput: 612 TFLOPS, 486 GB/s > Performance (m= 4096, n=32768, k= 512): 325 us | throughput: 423 TFLOPS, 884 GB/s > Performance (m= 4096, n= 7168, k=16384): 1200 us | throughput: 802 TFLOPS, 203 GB/s > Performance (m= 4096, n= 4096, k= 7168): 317 us | throughput: 760 TFLOPS, 292 GB/s > Performance (m= 4096, n= 7168, k= 2048): 195 us | throughput: 617 TFLOPS, 420 GB/s ```

bsakaguchi commented 2026-10-26 21:10:21 +00:00

Author

Owner

Copy link

@qijiaxing commented on GitHub (Mar 24, 2025):

Hi @thefacetakt , I fixed the bug you mentioned in this commit. It is due to a TMA 128 bytes alignment required for scale B shared memory storage. Now it should have no limitation on the block_n value any more. You may pull my branch and test again.

<!-- gh-comment-id:2747428274 --> @qijiaxing commented on GitHub (Mar 24, 2025): Hi @thefacetakt , I fixed the bug you mentioned in this [commit](https://github.com/qijiaxing/DeepGEMM/commit/9b1cda1dd91f17e1db97d92c53a3e00c80c9189a). It is due to a TMA 128 bytes alignment required for scale B shared memory storage. Now it should have no limitation on the block_n value any more. You may pull my branch and test again.

bsakaguchi commented 2026-10-26 21:10:21 +00:00

Author

Owner

Copy link

@RuiWang1998 commented on GitHub (Mar 24, 2025):

Hi @qijiaxing,

Does it work with grouped gemm?

<!-- gh-comment-id:2748386779 --> @RuiWang1998 commented on GitHub (Mar 24, 2025): Hi @qijiaxing, Does it work with grouped gemm?

bsakaguchi commented 2026-10-26 21:10:22 +00:00

Author

Owner

Copy link

@young-mu commented on GitHub (Mar 25, 2025):

I have implemented a variant that supports 1d1d scaling, can be used for WGRAD computation. https://github.com/qijiaxing/DeepGEMM/tree/1d1d any feedback or suggestions are welcomed!

@qijiaxing
Perhaps this can NOT be used in WGRAD directly, since WGRAD needs 128x1 LHS scales and 128x1 RHS scales, not 1x128 LHS scales and 1x128 RHS scales.

WGRAD formula: grad_weight = grad_output ^ T * input, in which

• grad_output: (M, N)
• input: (M, K)
• grad_weight: (N, K)

The scaling should be applied along with N/K axis (hidden_states), not M axis

<!-- gh-comment-id:2751382112 --> @young-mu commented on GitHub (Mar 25, 2025): > I have implemented a variant that supports 1d1d scaling, can be used for WGRAD computation. https://github.com/qijiaxing/DeepGEMM/tree/1d1d any feedback or suggestions are welcomed! @qijiaxing Perhaps this can NOT be used in WGRAD directly, since WGRAD needs **128x1 LHS scales and 128x1 RHS scales**, not 1x128 LHS scales and 1x128 RHS scales. WGRAD formula: `grad_weight = grad_output ^ T * input`, in which - grad_output: (M, N) - input: (M, K) - grad_weight: (N, K) The scaling should be applied along with N/K axis (hidden_states), not M axis

bsakaguchi commented 2026-10-26 21:10:22 +00:00

Author

Owner

Copy link

@RuiWang1998 commented on GitHub (Mar 25, 2025):

I have implemented a variant that supports 1d1d scaling, can be used for WGRAD computation. qijiaxing/DeepGEMM@1d1d any feedback or suggestions are welcomed!

@qijiaxing Perhaps this can NOT be used in WGRAD directly, since WGRAD needs 128x1 LHS scales and 128x1 RHS scales, not 1x128 LHS scales and 1x128 RHS scales.

WGRAD formula: grad_weight = grad_output ^ T * input, in which

• grad_output: (M, N)
• input: (N, K)
• grad_weight: (N, K)

The scaling OP should be executed along with N/K axis (hidden_states), not M axis

I don't think that's possible under Hopper. P.S., I think input should be (M, K)?

<!-- gh-comment-id:2751522800 --> @RuiWang1998 commented on GitHub (Mar 25, 2025): > > I have implemented a variant that supports 1d1d scaling, can be used for WGRAD computation. [qijiaxing/DeepGEMM@`1d1d`](https://github.com/qijiaxing/DeepGEMM/tree/1d1d?rgh-link-date=2025-03-25T14%3A05%3A46.000Z) any feedback or suggestions are welcomed! > > [@qijiaxing](https://github.com/qijiaxing?rgh-link-date=2025-03-25T14%3A05%3A46.000Z) Perhaps this can NOT be used in WGRAD directly, since WGRAD needs **128x1 LHS scales and 128x1 RHS scales**, not 1x128 LHS scales and 1x128 RHS scales. > > WGRAD formula: `grad_weight = grad_output ^ T * input`, in which > > * grad_output: (M, N) > * input: (N, K) > * grad_weight: (N, K) > > The scaling OP should be executed along with N/K axis (hidden_states), not M axis I don't think that's possible under Hopper. P.S., I think input should be (M, K)?

bsakaguchi commented 2026-10-26 21:10:23 +00:00

Author

Owner

Copy link

@young-mu commented on GitHub (Mar 25, 2025):

I have implemented a variant that supports 1d1d scaling, can be used for WGRAD computation. qijiaxing/DeepGEMM@1d1d any feedback or suggestions are welcomed!

@qijiaxing Perhaps this can NOT be used in WGRAD directly, since WGRAD needs 128x1 LHS scales and 128x1 RHS scales, not 1x128 LHS scales and 1x128 RHS scales.
WGRAD formula: grad_weight = grad_output ^ T * input, in which

• grad_output: (M, N)
• input: (N, K)
• grad_weight: (N, K)

The scaling OP should be executed along with N/K axis (hidden_states), not M axis

I don't think that's possible under Hopper. P.S., I think input should be (M, K)?

it's a typo, corrected.

<!-- gh-comment-id:2751548784 --> @young-mu commented on GitHub (Mar 25, 2025): > > > I have implemented a variant that supports 1d1d scaling, can be used for WGRAD computation. [qijiaxing/DeepGEMM@`1d1d`](https://github.com/qijiaxing/DeepGEMM/tree/1d1d?rgh-link-date=2025-03-25T14%3A05%3A46.000Z) any feedback or suggestions are welcomed! > > > > > > [@qijiaxing](https://github.com/qijiaxing?rgh-link-date=2025-03-25T14%3A05%3A46.000Z) Perhaps this can NOT be used in WGRAD directly, since WGRAD needs **128x1 LHS scales and 128x1 RHS scales**, not 1x128 LHS scales and 1x128 RHS scales. > > WGRAD formula: `grad_weight = grad_output ^ T * input`, in which > > > > * grad_output: (M, N) > > * input: (N, K) > > * grad_weight: (N, K) > > > > The scaling OP should be executed along with N/K axis (hidden_states), not M axis > > I don't think that's possible under Hopper. P.S., I think input should be (M, K)? it's a typo, corrected.

bsakaguchi commented 2026-10-26 21:10:23 +00:00

Author

Owner

Copy link

@qijiaxing commented on GitHub (Mar 26, 2025):

HI @RuiWang1998 , I didn't check if it works for grouped gemm yet, will let you later on.

@RuiWang1998 @young-mu , concerning WGRAD, my understanding is like this

x shape is [tokens, h_in]
w and dw shape is [h_out, h_in]
y and dy shape is [tokens, h_out]

in math, dw = dy^t @ x
because [h_out, tokens] @ [tokens, h_in] -> [h_out, h_in]
if use GEMM shape M,N,K, it would be 
M = h_out
K = tokens
N = h_in
But GEMM kernel requires tensor B to be transposed, this means we need to use 
dy^t as Tensor A, shape is [M, K] = [h_out, tokens]
x^t as Tensor B, shape is [N, K] = [h_in, tokens]
k is the dim for accumulation, thus every 128 elements in k dim need to share the same scale factor, which is applied after accumulation. 

<!-- gh-comment-id:2753100349 --> @qijiaxing commented on GitHub (Mar 26, 2025): HI @RuiWang1998 , I didn't check if it works for grouped gemm yet, will let you later on. @RuiWang1998 @young-mu , concerning WGRAD, my understanding is like this ``` x shape is [tokens, h_in] w and dw shape is [h_out, h_in] y and dy shape is [tokens, h_out] in math, dw = dy^t @ x because [h_out, tokens] @ [tokens, h_in] -> [h_out, h_in] if use GEMM shape M,N,K, it would be M = h_out K = tokens N = h_in But GEMM kernel requires tensor B to be transposed, this means we need to use dy^t as Tensor A, shape is [M, K] = [h_out, tokens] x^t as Tensor B, shape is [N, K] = [h_in, tokens] k is the dim for accumulation, thus every 128 elements in k dim need to share the same scale factor, which is applied after accumulation. ```

bsakaguchi commented 2026-10-26 21:10:24 +00:00

Author

Owner

Copy link

@RuiWang1998 commented on GitHub (Mar 26, 2025):

Hi @qijiaxing ,

Got it! Thanks for the great work!

<!-- gh-comment-id:2753146577 --> @RuiWang1998 commented on GitHub (Mar 26, 2025): Hi @qijiaxing , Got it! Thanks for the great work!

bsakaguchi commented 2026-10-26 21:10:24 +00:00

Author

Owner

Copy link

@young-mu commented on GitHub (Mar 26, 2025):

@qijiaxing

Got it, the process itself seems fine.
The key question is that if scaling can be applied along with M-axis (tokens in your case). From my understanding, it should be applied along with the N/K axis (hidden states), not the M axis (tokens). However, I don’t have solid proof for this yet.

BTW, the technical report mentions 128x1 tiles in the backward pass, but I couldn't find it in your 1d1d test case. Did I miss something? What's your take on this?

---

I have integrated FP8 DeepGEMM into FPROP, DGRAD and WGRAD processes (with WGRAD using your 1d1d implementation) and tested it on a GPT-2 model (about 410M). After hundreds of iterations, the final loss shows almost no difference compared with the baseline (BF16 + torch.matmul).

It seems that my initial understand was incorrect.
Thanks for your great work.

<!-- gh-comment-id:2753453108 --> @young-mu commented on GitHub (Mar 26, 2025): @qijiaxing Got it, the process itself seems fine. The key question is that if scaling can be applied along with M-axis (tokens in your case). From my understanding, it should be applied along with the N/K axis (hidden states), not the M axis (tokens). However, I don’t have solid proof for this yet. BTW, the [technical report](https://arxiv.org/abs/2412.19437) mentions **128x1 tiles** in the backward pass, but I couldn't find it in your [1d1d test case](https://github.com/qijiaxing/DeepGEMM/blob/1d1d/tests/test_1d1d.py). Did I miss something? What's your take on this? ----- I have integrated FP8 DeepGEMM into FPROP, DGRAD and WGRAD processes (with WGRAD using your 1d1d implementation) and tested it on a GPT-2 model (about 410M). After hundreds of iterations, the final loss shows almost no difference compared with the baseline (BF16 + torch.matmul). It seems that my initial understand was incorrect. Thanks for your great work.

bsakaguchi commented 2026-10-26 21:10:24 +00:00

Author

Owner

Copy link

@MayDomine commented on GitHub (Apr 2, 2025):

HI @RuiWang1998 , I didn't check if it works for grouped gemm yet, will let you later on.

@RuiWang1998 @young-mu , concerning WGRAD, my understanding is like this

x shape is [tokens, h_in]
w and dw shape is [h_out, h_in]
y and dy shape is [tokens, h_out]

in math, dw = dy^t @ x
because [h_out, tokens] @ [tokens, h_in] -> [h_out, h_in]
if use GEMM shape M,N,K, it would be 
M = h_out
K = tokens
N = h_in
But GEMM kernel requires tensor B to be transposed, this means we need to use 
dy^t as Tensor A, shape is [M, K] = [h_out, tokens]
x^t as Tensor B, shape is [N, K] = [h_in, tokens]
k is the dim for accumulation, thus every 128 elements in k dim need to share the same scale factor, which is applied after accumulation. 

@RuiWang1998 In this way, it seems that we still need to do transpose on x and y since we only have a gemm_nt_1d1d kernel. Is there any chance that we can have a gemm_tn_1d1d kernel?

<!-- gh-comment-id:2771786422 --> @MayDomine commented on GitHub (Apr 2, 2025): > HI [@RuiWang1998](https://github.com/RuiWang1998) , I didn't check if it works for grouped gemm yet, will let you later on. > > [@RuiWang1998](https://github.com/RuiWang1998) [@young-mu](https://github.com/young-mu) , concerning WGRAD, my understanding is like this > > ``` > x shape is [tokens, h_in] > w and dw shape is [h_out, h_in] > y and dy shape is [tokens, h_out] > > in math, dw = dy^t @ x > because [h_out, tokens] @ [tokens, h_in] -> [h_out, h_in] > if use GEMM shape M,N,K, it would be > M = h_out > K = tokens > N = h_in > But GEMM kernel requires tensor B to be transposed, this means we need to use > dy^t as Tensor A, shape is [M, K] = [h_out, tokens] > x^t as Tensor B, shape is [N, K] = [h_in, tokens] > k is the dim for accumulation, thus every 128 elements in k dim need to share the same scale factor, which is applied after accumulation. > ``` @RuiWang1998 In this way, it seems that we still need to do transpose on x and y since we only have a `gemm_nt_1d1d` kernel. Is there any chance that we can have a `gemm_tn_1d1d` kernel?

bsakaguchi commented 2026-10-26 21:10:25 +00:00

Author

Owner

Copy link

@qijiaxing commented on GitHub (Apr 3, 2025):

HI @RuiWang1998 , I didn't check if it works for grouped gemm yet, will let you later on.
@RuiWang1998 @young-mu , concerning WGRAD, my understanding is like this

x shape is [tokens, h_in]
w and dw shape is [h_out, h_in]
y and dy shape is [tokens, h_out]

in math, dw = dy^t @ x
because [h_out, tokens] @ [tokens, h_in] -> [h_out, h_in]
if use GEMM shape M,N,K, it would be 
M = h_out
K = tokens
N = h_in
But GEMM kernel requires tensor B to be transposed, this means we need to use 
dy^t as Tensor A, shape is [M, K] = [h_out, tokens]
x^t as Tensor B, shape is [N, K] = [h_in, tokens]
k is the dim for accumulation, thus every 128 elements in k dim need to share the same scale factor, which is applied after accumulation. 

@RuiWang1998 In this way, it seems that we still need to do transpose on x and y since we only have a gemm_nt_1d1d kernel. Is there any chance that we can have a gemm_tn_1d1d kernel?

the reason for NT is about perf. Because each C value is computed from a row of A and a column of B, both are a list of K values, thus make K dim the continuous dim would make memory loading most efficient. This leads to NT for A and B. otherwise perf would not be optimal.

<!-- gh-comment-id:2774805791 --> @qijiaxing commented on GitHub (Apr 3, 2025): > > HI [@RuiWang1998](https://github.com/RuiWang1998) , I didn't check if it works for grouped gemm yet, will let you later on. > > [@RuiWang1998](https://github.com/RuiWang1998) [@young-mu](https://github.com/young-mu) , concerning WGRAD, my understanding is like this > > ``` > > x shape is [tokens, h_in] > > w and dw shape is [h_out, h_in] > > y and dy shape is [tokens, h_out] > > > > in math, dw = dy^t @ x > > because [h_out, tokens] @ [tokens, h_in] -> [h_out, h_in] > > if use GEMM shape M,N,K, it would be > > M = h_out > > K = tokens > > N = h_in > > But GEMM kernel requires tensor B to be transposed, this means we need to use > > dy^t as Tensor A, shape is [M, K] = [h_out, tokens] > > x^t as Tensor B, shape is [N, K] = [h_in, tokens] > > k is the dim for accumulation, thus every 128 elements in k dim need to share the same scale factor, which is applied after accumulation. > > ``` > > [@RuiWang1998](https://github.com/RuiWang1998) In this way, it seems that we still need to do transpose on x and y since we only have a `gemm_nt_1d1d` kernel. Is there any chance that we can have a `gemm_tn_1d1d` kernel? the reason for NT is about perf. Because each C value is computed from a row of A and a column of B, both are a list of K values, thus make K dim the continuous dim would make memory loading most efficient. This leads to NT for A and B. otherwise perf would not be optimal.

bsakaguchi commented 2026-10-26 21:10:25 +00:00

Author

Owner

Copy link

@hxdtest commented on GitHub (May 4, 2025):

HI @RuiWang1998 , I didn't check if it works for grouped gemm yet, will let you later on.
@RuiWang1998 @young-mu , concerning WGRAD, my understanding is like this

x shape is [tokens, h_in]
w and dw shape is [h_out, h_in]
y and dy shape is [tokens, h_out]

in math, dw = dy^t @ x
because [h_out, tokens] @ [tokens, h_in] -> [h_out, h_in]
if use GEMM shape M,N,K, it would be 
M = h_out
K = tokens
N = h_in
But GEMM kernel requires tensor B to be transposed, this means we need to use 
dy^t as Tensor A, shape is [M, K] = [h_out, tokens]
x^t as Tensor B, shape is [N, K] = [h_in, tokens]
k is the dim for accumulation, thus every 128 elements in k dim need to share the same scale factor, which is applied after accumulation. 

@RuiWang1998 In this way, it seems that we still need to do transpose on x and y since we only have a gemm_nt_1d1d kernel. Is there any chance that we can have a gemm_tn_1d1d kernel?

the reason for NT is about perf. Because each C value is computed from a row of A and a column of B, both are a list of K values, thus make K dim the continuous dim would make memory loading most efficient. This leads to NT for A and B. otherwise perf would not be optimal.

DeepSeek v3 report: In our workflow, activations during the forward pass are quantized into 1x128 FP8 tiles and stored. During the backward pass, the matrix needs to be read out, dequantized, transposed, re-quantized into 128x1 tiles, and stored in HBM.

https://github.com/qijiaxing/DeepGEMM/blob/1d1d/deep_gemm/jit_kernels/gemm_1d1d.py#L134 Do a normal GEMM with FP8 inputs and BF16 output, with 1x128 LHS scaling and 1x128 RHS scaling. The method mentioned in DeepSeek V3 technical report and in 1d1d are not the same?

<!-- gh-comment-id:2848996293 --> @hxdtest commented on GitHub (May 4, 2025): > > > HI [@RuiWang1998](https://github.com/RuiWang1998) , I didn't check if it works for grouped gemm yet, will let you later on. > > > [@RuiWang1998](https://github.com/RuiWang1998) [@young-mu](https://github.com/young-mu) , concerning WGRAD, my understanding is like this > > > ``` > > > x shape is [tokens, h_in] > > > w and dw shape is [h_out, h_in] > > > y and dy shape is [tokens, h_out] > > > > > > in math, dw = dy^t @ x > > > because [h_out, tokens] @ [tokens, h_in] -> [h_out, h_in] > > > if use GEMM shape M,N,K, it would be > > > M = h_out > > > K = tokens > > > N = h_in > > > But GEMM kernel requires tensor B to be transposed, this means we need to use > > > dy^t as Tensor A, shape is [M, K] = [h_out, tokens] > > > x^t as Tensor B, shape is [N, K] = [h_in, tokens] > > > k is the dim for accumulation, thus every 128 elements in k dim need to share the same scale factor, which is applied after accumulation. > > > ``` > > > > > > [@RuiWang1998](https://github.com/RuiWang1998) In this way, it seems that we still need to do transpose on x and y since we only have a `gemm_nt_1d1d` kernel. Is there any chance that we can have a `gemm_tn_1d1d` kernel? > > the reason for NT is about perf. Because each C value is computed from a row of A and a column of B, both are a list of K values, thus make K dim the continuous dim would make memory loading most efficient. This leads to NT for A and B. otherwise perf would not be optimal. [DeepSeek v3 report](https://arxiv.org/pdf/2412.19437): In our workflow, activations during the forward pass are quantized into 1x128 FP8 tiles and stored. During the backward pass, the matrix needs to be read out, dequantized, transposed, re-quantized into 128x1 tiles, and stored in HBM. https://github.com/qijiaxing/DeepGEMM/blob/1d1d/deep_gemm/jit_kernels/gemm_1d1d.py#L134 Do a normal GEMM with FP8 inputs and BF16 output, with 1x128 LHS scaling and 1x128 RHS scaling. The method mentioned in DeepSeek V3 technical report and in `1d1d` are not the same?

bsakaguchi commented 2026-10-26 21:10:25 +00:00

Author

Owner

Copy link

@qijiaxing commented on GitHub (May 7, 2025):

@hxdtest the method is the same as DS v3 report.
in BWD, in math, dw = dy^t @ x, x is the RHS tensor, thus its scaling is 128x1. But the actual GEMM kernel, it requires RHS tensor to be transposed, thus we need x^t, whose scaling become 1x128.

<!-- gh-comment-id:2859133391 --> @qijiaxing commented on GitHub (May 7, 2025): @hxdtest the method is the same as DS v3 report. in BWD, in math, `dw = dy^t @ x`, x is the RHS tensor, thus its scaling is 128x1. But the actual GEMM kernel, it requires RHS tensor to be transposed, thus we need `x^t`, whose scaling become 1x128.

bsakaguchi commented 2026-10-26 21:10:26 +00:00

Author

Owner

Copy link

@hxdtest commented on GitHub (May 8, 2025):

@hxdtest the method is the same as DS v3 report. in BWD, in math, dw = dy^t @ x, x is the RHS tensor, thus its scaling is 128x1. But the actual GEMM kernel, it requires RHS tensor to be transposed, thus we need x^t, whose scaling become 1x128.

Many thanks for your reply.

<!-- gh-comment-id:2862355491 --> @hxdtest commented on GitHub (May 8, 2025): > [@hxdtest](https://github.com/hxdtest) the method is the same as DS v3 report. in BWD, in math, `dw = dy^t @ x`, x is the RHS tensor, thus its scaling is 128x1. But the actual GEMM kernel, it requires RHS tensor to be transposed, thus we need `x^t`, whose scaling become 1x128. Many thanks for your reply.

bsakaguchi commented 2026-10-26 21:10:26 +00:00

Author

Owner

Copy link

@hxdtest commented on GitHub (May 9, 2025):

@hxdtest the method is the same as DS v3 report. in BWD, in math, dw = dy^t @ x, x is the RHS tensor, thus its scaling is 128x1. But the actual GEMM kernel, it requires RHS tensor to be transposed, thus we need x^t, whose scaling become 1x128.

@qijiaxing
in BWD, in math, dw = dy^t @ x, x is the RHS tensor, thus its scaling is 128x1. I am still confused about rule of tiling and scaling. When x is the activation, in forward y = x @ w^t , x is the LHS tensor, x's scaling is 1x128, it's tiled in hidden_dim . If the tiling rule is scaling at hidden_dim, it's stilled need 1x128 in RHS.

--- update ---
Oh the rule of tiling and scaling is k is the dim for accumulation, thus every 128 elements in k dim need to share the same scale factor, which is applied after accumulation. Is that right ?

<!-- gh-comment-id:2865463680 --> @hxdtest commented on GitHub (May 9, 2025): > [@hxdtest](https://github.com/hxdtest) the method is the same as DS v3 report. in BWD, in math, `dw = dy^t @ x`, x is the RHS tensor, thus its scaling is 128x1. But the actual GEMM kernel, it requires RHS tensor to be transposed, thus we need `x^t`, whose scaling become 1x128. @qijiaxing ` in BWD, in math, dw = dy^t @ x, x is the RHS tensor, thus its scaling is 128x1. ` I am still confused about rule of tiling and scaling. When x is the activation, in forward y = x @ w^t , x is the LHS tensor, x's scaling is 1x128, it's tiled in hidden_dim . If the tiling rule is scaling at hidden_dim, it's stilled need 1x128 in RHS. --- update --- Oh the rule of tiling and scaling is `k is the dim for accumulation, thus every 128 elements in k dim need to share the same scale factor, which is applied after accumulation. ` Is that right ?

bsakaguchi commented 2026-10-26 21:10:26 +00:00

Author

Owner

Copy link

@qijiaxing commented on GitHub (May 12, 2025):

@hxdtest yes, the scaling tile is determined by the GEMM algorithm. Even though x is used in both the FWD GEMM and WGRAD GEMM, these two GEMMs are independent. It's not necessary for x to have the same scaling.

<!-- gh-comment-id:2873055783 --> @qijiaxing commented on GitHub (May 12, 2025): @hxdtest yes, the scaling tile is determined by the GEMM algorithm. Even though x is used in both the FWD GEMM and WGRAD GEMM, these two GEMMs are independent. It's not necessary for x to have the same scaling.

bsakaguchi commented 2026-10-26 21:10:27 +00:00

Author

Owner

Copy link

@LyricZhao commented on GitHub (May 14, 2025):

Official weight gradient GEMM support is in #95, will merge today.

<!-- gh-comment-id:2878854656 --> @LyricZhao commented on GitHub (May 14, 2025): Official weight gradient GEMM support is in #95, will merge today.

bsakaguchi referenced this issue 2026-10-26 21:18:19 +00:00

[PR #470] [CLOSED] Sync/upstream dev e1f418c2 #477

Sign in to join this conversation.

No Branch/Tag specified

Branches Tags

main

public-sync-20261014

0930-release

nv_dev

fix-mega-moe-fence

public-release-2609

public-release-2607

public-release-2606

mega-update

public-release-260416

ko3n1g/ci/build-wheels

nv_dev_f8e8fb5

nv_dev_950d31a

nv_dev_67fc648

nv_dev_c491439

nv_dev_4ff3f54

v2.1.1.post3

v2.1.1.post2

v2.1.1.post1

v2.1.1

v2.1.0

Labels

Clear labels   

pull-request

Mirrored from GitHub Pull Request

No labels

pull-request

Milestone

Clear milestone

No items

No milestone

Projects

Clear projects

No items

No project

Assignees

Clear assignees

No assignees

1 participant

Notifications

Subscribe

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

deepseek-ai/DeepGEMM#7

No description provided.