[GH-ISSUE #102] [Help] Starter script for nn.Linear for fwd,bwd

bsakaguchi commented

2026-10-26 21:11:23 +00:00

Owner

Originally created by @heuristicoder on GitHub (May 17, 2025).
Original GitHub issue: https://github.com/deepseek-ai/DeepGEMM/issues/102

It would be great to show an example that implements simple nn.Linear with both forward and backward pass. A small starter script would be helpful.

Originally created by @heuristicoder on GitHub (May 17, 2025). Original GitHub issue: https://github.com/deepseek-ai/DeepGEMM/issues/102 It would be great to show an example that implements simple nn.Linear with both forward and backward pass. A small starter script would be helpful.

bsakaguchi commented

2026-10-26 21:11:24 +00:00

Author

Owner

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

We currectly don't encourage users to this, as we are not open-sourcing some fast util kernels (e.g. casting/scaling factor transposing), so if you don't have such things in this Linear, the performance will be bad.

We will open-source them later, with example Linear module as well.

If you don't care about performance, here is a pseduo code (this is somehow basic, you could also search on the Internet for some general tutorials):

def forward(x_bf16):
    x_fp8 = cast(x_bf16)
    out = create_tensor()
    deep_gemm.gemm_fp8_fp8_bf16_nt(x_fp8, weight_fp8, out)
    return out

def dgrad(out_grad_bf16):
    out_grad_fp8 = cast(out_grad_bf16)
    grad = create_tensor()
    deep_gemm.gemm_fp8_fp8_bf16_nt(out_grad_fp8, transpose(weight_fp8), grad)
    return grad

def wgrad(out_grad_bf16):
    out_grad_fp8 = cast(out_grad_bf16)
    grad = create_tensor()
    deep_gemm.wgrad_gemm_fp8_fp8_fp32_nt(transpose(out_grad_fp8), transpose(x_fp8), grad)
    return grad

@LyricZhao commented on GitHub (May 22, 2025): We currectly don't encourage users to this, as we are not open-sourcing some fast util kernels (e.g. casting/scaling factor transposing), so if you don't have such things in this Linear, the performance will be bad. We will open-source them later, with example Linear module as well. If you don't care about performance, here is a pseduo code (this is somehow basic, you could also search on the Internet for some general tutorials): ```python def forward(x_bf16): x_fp8 = cast(x_bf16) out = create_tensor() deep_gemm.gemm_fp8_fp8_bf16_nt(x_fp8, weight_fp8, out) return out def dgrad(out_grad_bf16): out_grad_fp8 = cast(out_grad_bf16) grad = create_tensor() deep_gemm.gemm_fp8_fp8_bf16_nt(out_grad_fp8, transpose(weight_fp8), grad) return grad def wgrad(out_grad_bf16): out_grad_fp8 = cast(out_grad_bf16) grad = create_tensor() deep_gemm.wgrad_gemm_fp8_fp8_fp32_nt(transpose(out_grad_fp8), transpose(x_fp8), grad) return grad ```

bsakaguchi commented

2026-10-26 21:11:25 +00:00

Author

Owner

@heuristicoder commented on GitHub (May 22, 2025):

@LyricZhao Thanks,

Also another question would be deep_gemm.gemm_fp8_fp8_bf16_nt, transposes matrix RHS internally. Hence the shape of RHS is set as:

       rhs: the first element is an FP8 tensor (typed `torch.float8_e4m3fn`) of shape `[n, k]`,
       the second element is an FP32 128x128 scaling tensor for RHS of shape `[⌈n / 128⌉, ⌈k / 128⌉]`.

However I observed that scales of RHS are not transposed internally when computing C scales. We are computing it as LHS_scale * RHS_scale. But intuitively, shouldn't it be LHS_scales * RHS_scales.T since we wanna do x @ W.T ?

   // Load B scales with math warp-groups
   if (threadIdx.x >= 32) {
       auto num_previous_lines = scheduler.get_global_idx<false>(ceil_div(SHAPE_N, BLOCK_K), 0, 0, m_block_idx);
       auto local_scales_b = scales_b + (num_previous_lines + ((n_block_idx * BLOCK_N) / BLOCK_K)) * SHAPE_K_SCALES;
       for (uint32_t i = threadIdx.x - 32; i < num_scales_b; i += kNumMathThreads - 32)
           st_shared(smem_scales_b + i, __ldg(local_scales_b + i));
   }

Am I missing something here?

@heuristicoder commented on GitHub (May 22, 2025): @LyricZhao Thanks, Also another question would be deep_gemm.gemm_fp8_fp8_bf16_nt, transposes matrix RHS internally. Hence the shape of RHS is set as: ``` rhs: the first element is an FP8 tensor (typed `torch.float8_e4m3fn`) of shape `[n, k]`, the second element is an FP32 128x128 scaling tensor for RHS of shape `[⌈n / 128⌉, ⌈k / 128⌉]`. ``` However I observed that scales of RHS are not transposed internally when computing C scales. We are computing it as LHS_scale * RHS_scale. But intuitively, shouldn't it be LHS_scales * RHS_scales.T since we wanna do x @ W.T ? ``` // Load B scales with math warp-groups if (threadIdx.x >= 32) { auto num_previous_lines = scheduler.get_global_idx<false>(ceil_div(SHAPE_N, BLOCK_K), 0, 0, m_block_idx); auto local_scales_b = scales_b + (num_previous_lines + ((n_block_idx * BLOCK_N) / BLOCK_K)) * SHAPE_K_SCALES; for (uint32_t i = threadIdx.x - 32; i < num_scales_b; i += kNumMathThreads - 32) st_shared(smem_scales_b + i, __ldg(local_scales_b + i)); } ``` Am I missing something here?

bsakaguchi commented

2026-10-26 21:11:25 +00:00

Author

Owner

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

auto local_scales_b = scales_b + (num_previous_lines + ((n_block_idx * BLOCK_N) / BLOCK_K)) * SHAPE_K_SCALES; reads a whole row of (ceil_div(n, 128), ceil_div(k, 128)), the row shape is (1, ceil_div(k, 128)), which is used for later all K blocks.

@LyricZhao commented on GitHub (May 23, 2025): `auto local_scales_b = scales_b + (num_previous_lines + ((n_block_idx * BLOCK_N) / BLOCK_K)) * SHAPE_K_SCALES;` reads a whole row of `(ceil_div(n, 128), ceil_div(k, 128))`, the row shape is `(1, ceil_div(k, 128))`, which is used for later all K blocks.

bsakaguchi commented

2026-10-26 21:11:26 +00:00

Author

Owner

@heuristicoder commented on GitHub (May 28, 2025):

Got it, thanks. Is there any update on when all the kernels for a full Linear layer will be released?

@heuristicoder commented on GitHub (May 28, 2025): Got it, thanks. Is there any update on when all the kernels for a full Linear layer will be released?

bsakaguchi commented

2026-10-26 21:11:26 +00:00

Author

Owner

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

Maybe by the end of the year :)

@LyricZhao commented on GitHub (May 28, 2025): Maybe by the end of the year :)

bsakaguchi commented

2026-10-26 21:11:27 +00:00

Author

Owner

@goldhuang commented on GitHub (Jul 26, 2025):

We currectly don't encourage users to this, as we are not open-sourcing some fast util kernels (e.g. casting/scaling factor transposing), so if you don't have such things in this Linear, the performance will be bad.

We will open-source them later, with example Linear module as well.

If you don't care about performance, here is a pseduo code (this is somehow basic, you could also search on the Internet for some general tutorials):

def forward(x_bf16):
x_fp8 = cast(x_bf16)
out = create_tensor()
deep_gemm.gemm_fp8_fp8_bf16_nt(x_fp8, weight_fp8, out)
return out

def dgrad(out_grad_bf16):
out_grad_fp8 = cast(out_grad_bf16)
grad = create_tensor()
deep_gemm.gemm_fp8_fp8_bf16_nt(out_grad_fp8, transpose(weight_fp8), grad)
return grad

def wgrad(out_grad_bf16):
out_grad_fp8 = cast(out_grad_bf16)
grad = create_tensor()
deep_gemm.wgrad_gemm_fp8_fp8_fp32_nt(transpose(out_grad_fp8), transpose(x_fp8), grad)
return grad

@LyricZhao Is it a good idea to fuse casting and scaling to the gemm kernel? (per token for a and per block for b). Seems feasible. I'm looking into it and find that it will increase the pressure either on shared memory or register usage, as we need extra array for b16 input a and b, beside the fp8 a and b.
I see you mentioned fast util kernels. Did you keep these kernels separate or did you try to fuse them together?
Thanks a lot in advance if you could share some experiences on your side so that I would not waste my time to fuse the kernels but later find it becomes slower!

@goldhuang commented on GitHub (Jul 26, 2025): > We currectly don't encourage users to this, as we are not open-sourcing some fast util kernels (e.g. casting/scaling factor transposing), so if you don't have such things in this Linear, the performance will be bad. > > We will open-source them later, with example Linear module as well. > > If you don't care about performance, here is a pseduo code (this is somehow basic, you could also search on the Internet for some general tutorials): > > def forward(x_bf16): > x_fp8 = cast(x_bf16) > out = create_tensor() > deep_gemm.gemm_fp8_fp8_bf16_nt(x_fp8, weight_fp8, out) > return out > > def dgrad(out_grad_bf16): > out_grad_fp8 = cast(out_grad_bf16) > grad = create_tensor() > deep_gemm.gemm_fp8_fp8_bf16_nt(out_grad_fp8, transpose(weight_fp8), grad) > return grad > > def wgrad(out_grad_bf16): > out_grad_fp8 = cast(out_grad_bf16) > grad = create_tensor() > deep_gemm.wgrad_gemm_fp8_fp8_fp32_nt(transpose(out_grad_fp8), transpose(x_fp8), grad) > return grad @LyricZhao Is it a good idea to fuse casting and scaling to the gemm kernel? (per token for a and per block for b). Seems feasible. I'm looking into it and find that it will increase the pressure either on shared memory or register usage, as we need extra array for b16 input a and b, beside the fp8 a and b. I see you mentioned fast util kernels. Did you keep these kernels separate or did you try to fuse them together? Thanks a lot in advance if you could share some experiences on your side so that I would not waste my time to fuse the kernels but later find it becomes slower!

bsakaguchi commented

2026-10-26 21:11:27 +00:00

Author

Owner

@heuristicoder commented on GitHub (Jul 27, 2025):

That would be a great addition, and makes the library much more useful for a lot of use cases then just inference

@heuristicoder commented on GitHub (Jul 27, 2025): That would be a great addition, and makes the library much more useful for a lot of use cases then just inference

bsakaguchi commented

2026-10-26 21:11:27 +00:00

Author

Owner

@LyricZhao commented on GitHub (Jul 30, 2025):

Is it a good idea to fuse casting and scaling to the gemm kernel?

No, we should fuse cast/scaling into the previous kernels' output, HBM bandwidth could also be saved.

Did you keep these kernels separate or did you try to fuse them together?

Fuse as much as you can.

@LyricZhao commented on GitHub (Jul 30, 2025): > Is it a good idea to fuse casting and scaling to the gemm kernel? No, we should fuse cast/scaling into the previous kernels' output, HBM bandwidth could also be saved. > Did you keep these kernels separate or did you try to fuse them together? Fuse as much as you can.

Rows
Columns

[GH-ISSUE #102] [Help] Starter script for nn.Linear for fwd,bwd #56