[GH-ISSUE #105] How to amortize Python-side workload for WGRAD grouped gemm. #59

New issue

Closed

opened 2026-10-26 21:11:30 +00:00 by bsakaguchi · 1 comment

bsakaguchi commented

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

Owner

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

Hello, I have been trying to adapt your WGRAD GEMM to be used as a grouped WGRAD GEMM.
I see that the CPU is doing a lot of work before launching the kernels, which are very small:

This is roughly the loop I am doing:

            tokens_per_expert = torch.bincount(m_indices.to(torch.int64), minlength=ne)
            offsets = torch.cat([
                torch.zeros(1, dtype=tokens_per_expert.dtype, device=tokens_per_expert.device),
                torch.cumsum(tokens_per_expert, 0),
            ])  # len = ne + 1

            # convert offsets to cpu once to avoid per-iteration gpu-cpu sync when casting to python ints
            # TODO: this is still pretty bad.
            offsets_cpu = offsets.to(torch.int64).cpu().tolist()

            hidden_dim = x_fp8_all.size(0)
            inter_dim = dgate_fp8_all.size(0)

            for expert_idx in range(ne):
                start = offsets_cpu[expert_idx]
                end = offsets_cpu[expert_idx + 1]

                # block indices for 1x128 scaling tensors
                block_start = start // 128
                block_end = end // 128

                # grab fp8 + scale slices (all views are contiguous along last dim)
                lhs1_fp8, lhs1_sc = dgate_fp8_all[:, start:end], dgate_sc_all[:, block_start:block_end]
                lhs2_fp8, lhs2_sc = dup_fp8_all[:, start:end], dup_sc_all[:, block_start:block_end]
                rhs_fp8, rhs_sc = x_fp8_all[:, start:end], x_sc_all[:, block_start:block_end]

                gemm_wgrad_fp8_fp8_fp32((lhs1_fp8, lhs1_sc), (rhs_fp8, rhs_sc), dw1_param_view[expert_idx])
                gemm_wgrad_fp8_fp8_fp32((lhs2_fp8, lhs2_sc), (rhs_fp8, rhs_sc), dv1_param_view[expert_idx])

                lhs3_fp8, lhs3_sc = dy_fp8_all[:, start:end], dy_sc_all[:, block_start:block_end]
                rhs2_fp8, rhs2_sc = hid_fp8_all[:, start:end], hid_sc_all[:, block_start:block_end]

                dw2_scratch = torch.zeros(
                    (hidden_dim, inter_dim), dtype=torch.float32, device=x_fp8_all.device
                )
                gemm_wgrad_fp8_fp8_fp32((lhs3_fp8, lhs3_sc), (rhs2_fp8, rhs2_sc), dw2_scratch)
                dw2_param_view[expert_idx].copy_(dw2_scratch.t())

I am wondering how I can amortize the CPU cost here, since the kernels are quite small.

Originally created by @cassanof on GitHub (May 21, 2025). Original GitHub issue: https://github.com/deepseek-ai/DeepGEMM/issues/105 Hello, I have been trying to adapt your WGRAD GEMM to be used as a grouped WGRAD GEMM. I see that the CPU is doing a lot of work before launching the kernels, which are very small: ![Launch](https://github.com/user-attachments/assets/8631afa9-528f-4d22-bf8d-3644a0a3b3ab) ![Kernel](https://github.com/user-attachments/assets/d350eb6b-2913-4bef-9a1f-3b09acd9b66d) This is roughly the loop I am doing: ```py tokens_per_expert = torch.bincount(m_indices.to(torch.int64), minlength=ne) offsets = torch.cat([ torch.zeros(1, dtype=tokens_per_expert.dtype, device=tokens_per_expert.device), torch.cumsum(tokens_per_expert, 0), ]) # len = ne + 1 # convert offsets to cpu once to avoid per-iteration gpu-cpu sync when casting to python ints # TODO: this is still pretty bad. offsets_cpu = offsets.to(torch.int64).cpu().tolist() hidden_dim = x_fp8_all.size(0) inter_dim = dgate_fp8_all.size(0) for expert_idx in range(ne): start = offsets_cpu[expert_idx] end = offsets_cpu[expert_idx + 1] # block indices for 1x128 scaling tensors block_start = start // 128 block_end = end // 128 # grab fp8 + scale slices (all views are contiguous along last dim) lhs1_fp8, lhs1_sc = dgate_fp8_all[:, start:end], dgate_sc_all[:, block_start:block_end] lhs2_fp8, lhs2_sc = dup_fp8_all[:, start:end], dup_sc_all[:, block_start:block_end] rhs_fp8, rhs_sc = x_fp8_all[:, start:end], x_sc_all[:, block_start:block_end] gemm_wgrad_fp8_fp8_fp32((lhs1_fp8, lhs1_sc), (rhs_fp8, rhs_sc), dw1_param_view[expert_idx]) gemm_wgrad_fp8_fp8_fp32((lhs2_fp8, lhs2_sc), (rhs_fp8, rhs_sc), dv1_param_view[expert_idx]) lhs3_fp8, lhs3_sc = dy_fp8_all[:, start:end], dy_sc_all[:, block_start:block_end] rhs2_fp8, rhs2_sc = hid_fp8_all[:, start:end], hid_sc_all[:, block_start:block_end] dw2_scratch = torch.zeros( (hidden_dim, inter_dim), dtype=torch.float32, device=x_fp8_all.device ) gemm_wgrad_fp8_fp8_fp32((lhs3_fp8, lhs3_sc), (rhs2_fp8, rhs2_sc), dw2_scratch) dw2_param_view[expert_idx].copy_(dw2_scratch.t()) ``` I am wondering how I can amortize the CPU cost here, since the kernels are quite small.

bsakaguchi closed this issue

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

bsakaguchi commented

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

Author

Owner

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

You can try to fuse these kernels with CUDA/triton code or use @torch.compile.

@soundOfDestiny commented on GitHub (May 22, 2025): You can try to fuse these kernels with CUDA/triton code or use `@torch.compile`.

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

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#59

No description provided.

Rows
Columns