deepseek-ai/DeepGEMM

Fork 0

mirror of https://github.com/deepseek-ai/DeepGEMM.git synced 2026-11-05 00:47:26 +00:00

[GH-ISSUE #121] find out why wgrad mma will cause illegal memory access in specific shape mnk #71

New issue

Closed

opened 2026-10-26 21:11:51 +00:00 by bsakaguchi · 6 comments

bsakaguchi commented

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

Owner

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

in shape mnk like [1536, 256, 2048]will case the error when using deep_gemm.wgrad_gemm_fp8_fp8_fp32_nt(a_fp8, b_fp8, out)
RuntimeError: CUDA error: an illegal memory access was encountered

this is because in fp8_wgrad_gemm_kernel there is a setting

    // TMA Barrier for both divisible and non-divisible cases
    Barrier* full_barriers[kNumStages + 1];
    Barrier* empty_barriers[kNumStages + 1];

and latter it will be used:

    #pragma unroll
    for (int i = 0; i < kNumStages + 1; ++ i) {
        full_barriers[i] = barrier_start_ptr + i;
        empty_barriers[i] = barrier_start_ptr + kNumStages + 1 + i;
    }

    // Initialize barriers
    DG_STATIC_ASSERT(kNumTMAMulticast <= 32, "To many TMA multicast");
    if (threadIdx.x == kNumMathThreads) {
        // NOTES: we always use `lane_idx` to arrive for the `lane_idx`-th CTA in the cluster,
        // even with TMA multicast disabled, we want to make the behavior aligned
        #pragma unroll
        for (int i = 0; i < kNumStages; ++ i) {
            full_barriers[i]->init(1);
            empty_barriers[i]->init(kNumTMAMulticast * kNumMathThreads / 32);
        }
        full_barriers[kNumStages]->init(1);
        empty_barriers[kNumStages]->init(1);

        // Make initialized barrier visible in async proxy
        cutlass::arch::fence_view_async_shared();
        (kNumTMAMulticast > 1) ? cutlass::arch::fence_barrier_init() : void();
    }

but in gemm.py, the method get_smem_config only set kNumStages share memory size:

def get_smem_config(num_stages: int, k: int, block_m: int, block_n: int, block_k: int = 128,        
                    is_fp32_out: bool = False, is_wgrad: bool = False) -> Tuple[int, int, int]:
    assert block_k == 128

    # Try swizzle first, as it does not waste shared memory
    swizzle_mode = get_swizzle_mode(block_n)
    block_n_padding = get_block_n_padding_for_smem_d(
        block_n) if swizzle_mode == 0 else 0

    # NOTES: `scales_b` in a total manner or per-stage manner
    smem_d = block_m * (block_n + block_n_padding) * (4 if is_fp32_out else 2)
    smem_a_per_stage = block_m * block_k
    smem_scales_a_per_stage = block_m * 4
    smem_b_per_stage = block_n * block_k
    smem_scales_b_per_stage = ceil_div(block_n * 4, block_k) * block_k if is_wgrad else 0
    smem_scales_b = ceil_div(k, block_k) * 4 if not is_wgrad else 0
    smem_barrier = num_stages * 8 * 2
    ....

in caculator, the memory caculate in python is 173184 and in cuda it is 173200.
after i set smem_barrier = (num_stages + 1)* 8 * 2, the problem goes well.
so i think maybe this is the cause of the error

Originally created by @fy1214 on GitHub (Jun 26, 2025). Original GitHub issue: https://github.com/deepseek-ai/DeepGEMM/issues/121 in shape mnk like [1536, 256, 2048]will case the error when using `deep_gemm.wgrad_gemm_fp8_fp8_fp32_nt(a_fp8, b_fp8, out)` `RuntimeError: CUDA error: an illegal memory access was encountered` this is because in fp8_wgrad_gemm_kernel there is a setting ``` // TMA Barrier for both divisible and non-divisible cases Barrier* full_barriers[kNumStages + 1]; Barrier* empty_barriers[kNumStages + 1]; ``` and latter it will be used: ``` #pragma unroll for (int i = 0; i < kNumStages + 1; ++ i) { full_barriers[i] = barrier_start_ptr + i; empty_barriers[i] = barrier_start_ptr + kNumStages + 1 + i; } // Initialize barriers DG_STATIC_ASSERT(kNumTMAMulticast <= 32, "To many TMA multicast"); if (threadIdx.x == kNumMathThreads) { // NOTES: we always use `lane_idx` to arrive for the `lane_idx`-th CTA in the cluster, // even with TMA multicast disabled, we want to make the behavior aligned #pragma unroll for (int i = 0; i < kNumStages; ++ i) { full_barriers[i]->init(1); empty_barriers[i]->init(kNumTMAMulticast * kNumMathThreads / 32); } full_barriers[kNumStages]->init(1); empty_barriers[kNumStages]->init(1); // Make initialized barrier visible in async proxy cutlass::arch::fence_view_async_shared(); (kNumTMAMulticast > 1) ? cutlass::arch::fence_barrier_init() : void(); } ``` but in gemm.py, the method get_smem_config only set kNumStages share memory size: ``` def get_smem_config(num_stages: int, k: int, block_m: int, block_n: int, block_k: int = 128, is_fp32_out: bool = False, is_wgrad: bool = False) -> Tuple[int, int, int]: assert block_k == 128 # Try swizzle first, as it does not waste shared memory swizzle_mode = get_swizzle_mode(block_n) block_n_padding = get_block_n_padding_for_smem_d( block_n) if swizzle_mode == 0 else 0 # NOTES: `scales_b` in a total manner or per-stage manner smem_d = block_m * (block_n + block_n_padding) * (4 if is_fp32_out else 2) smem_a_per_stage = block_m * block_k smem_scales_a_per_stage = block_m * 4 smem_b_per_stage = block_n * block_k smem_scales_b_per_stage = ceil_div(block_n * 4, block_k) * block_k if is_wgrad else 0 smem_scales_b = ceil_div(k, block_k) * 4 if not is_wgrad else 0 smem_barrier = num_stages * 8 * 2 .... ``` in caculator, the memory caculate in python is 173184 and in cuda it is 173200. after i set smem_barrier = (num_stages + 1)* 8 * 2, the problem goes well. so i think maybe this is the cause of the error

bsakaguchi closed this issue

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

bsakaguchi commented

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

Author

Owner

@fy1214 commented on GitHub (Jun 27, 2025):

still don't know why some shape can work without (num_stages + 1)

@fy1214 commented on GitHub (Jun 27, 2025): still don't know why some shape can work without (num_stages + 1)

bsakaguchi commented

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

Author

Owner

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

@zheanxu Could you help fixing this?

@LyricZhao commented on GitHub (Jul 1, 2025): @zheanxu Could you help fixing this?

bsakaguchi commented

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

Author

Owner

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

Thanks and merged!

@LyricZhao commented on GitHub (Jul 2, 2025): Thanks and merged!

bsakaguchi commented

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

Author

Owner

@fy1214 commented on GitHub (Jul 2, 2025):

Thanks and merged!

but i still wandering why in some shape can work without fix, like (1536, 1024, 2048) didn't trigger this bug

@fy1214 commented on GitHub (Jul 2, 2025): > Thanks and merged! but i still wandering why in some shape can work without fix, like (1536, 1024, 2048) didn't trigger this bug

bsakaguchi commented

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

Author

Owner

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

The hardware may allocate more shared memory space sometimes (i.e. you can use more bytes than your allocation without illegal memory access; just a guess, undocumented).

@LyricZhao commented on GitHub (Jul 2, 2025): The hardware may allocate more shared memory space sometimes (i.e. you can use more bytes than your allocation without illegal memory access; just a guess, undocumented).

bsakaguchi commented

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

Author

Owner

@fy1214 commented on GitHub (Jul 2, 2025):

The hardware may allocate more shared memory space sometimes (i.e. you can use more bytes than your allocation without illegal memory access; just a guess, undocumented).

yep，i print the pointer in smem find out in the correct case, it allocate a little memory more than we assign, i guess maybe it's because nvidia gpu will aligned the smem adress so that there will generate more smem than we assign.

@fy1214 commented on GitHub (Jul 2, 2025): > The hardware may allocate more shared memory space sometimes (i.e. you can use more bytes than your allocation without illegal memory access; just a guess, undocumented). yep，i print the pointer in smem find out in the correct case, it allocate a little memory more than we assign, i guess maybe it's because nvidia gpu will aligned the smem adress so that there will generate more smem than we assign.

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

No description provided.

Rows
Columns