[GH-ISSUE #20] How to skip certain expert in grouped gemm masked layout

bsakaguchi commented

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

Owner

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

Hi, thanks for the great implementation.
I'm trying to use this grouped gemm kernel into FusedMoE. I tried to set the mask_m as 0 for the experts that I want to skip but it got me nan for the output. Wondering is there a best practice for skipping certain experts? Which should be a very common situation under low batch size decoding stage.

Originally created by @DiegoD94 on GitHub (Feb 27, 2025). Original GitHub issue: https://github.com/deepseek-ai/DeepGEMM/issues/20 Hi, thanks for the great implementation. I'm trying to use this grouped gemm kernel into FusedMoE. I tried to set the mask_m as 0 for the experts that I want to skip but it got me nan for the output. Wondering is there a best practice for skipping certain experts? Which should be a very common situation under low batch size decoding stage.

bsakaguchi closed this issue

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

bsakaguchi commented

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

Author

Owner

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

The current impl already supports that, try this test (notice masked_m[1] = 0 # !!!! Skipping the expert 1 in the test):

def test_m_grouped_gemm_masked() -> None:
    print('Testing grouped masked GEMM:')

    for num_groups, m in ((2, 512), (4, 256)):
        for k, n in ((7168, 4096), (2048, 7168), ):
            # Test correctness
            masked_m_candidates = list(filter(lambda candidate: candidate <= m, (64, 128, 192, 256, 320, 384)))
            for i in range(10):
                x_fp8, y_fp8, out, ref_out = construct_grouped(num_groups, m, k, n, is_masked=True)
                masked_m = torch.empty((num_groups, ), device='cuda', dtype=torch.int)
                for j in range(num_groups):
                    masked_m[j] = random.choice(masked_m_candidates)
                masked_m[1] = 0  # !!!! Skipping the expert 1
                expected_m = min(int(masked_m.float().mean()) + 1, m)
                deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_masked(x_fp8, y_fp8, out, masked_m, expected_m)
                for j in range(num_groups):
                    diff = calc_diff(out[j, :masked_m[j].item()], ref_out[j, :masked_m[j].item()])
                    # Skipping expert 1 checks
                    assert j == 1 or diff < 0.001, f'{m=}, {k=}, {n=}, {j=}, masked_m={masked_m[j]}, {num_groups=}, {diff:.5f}'

            # noinspection PyShadowingNames
            def test_func():
                # Construct new tensors every time to avoid L2 cache acceleration
                x_fp8, y_fp8, out, ref_out = construct_grouped(num_groups, m, k, n, is_masked=True)
                masked_m = torch.ones((num_groups, ), device='cuda', dtype=torch.int) * m
                deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_masked(x_fp8, y_fp8, out, masked_m, m)

            # Test performance with fixed shapes
            t = bench_kineto(test_func, 'fp8_gemm', suppress_kineto_output=True)
            print(f' > Performance ({num_groups=}, m_per_group={m:4}, n={n:4}, k={k:4}): {t * 1e6:4.0f} us | '
                  f'throughput: {2 * num_groups * m * n * k / t / 1e12:4.0f} TFLOPS, '
                  f'{(num_groups * (m * k + k * n + m * n * 2)) / 1e9 / t:4.0f} GB/s')
    print()

For the skipped experts, the NaN output is in expectation as no computation for that expert and no memory writes (no cleaning them into zero).

@LyricZhao commented on GitHub (Feb 27, 2025): The current impl already supports that, try this test (notice `masked_m[1] = 0 # !!!! Skipping the expert 1` in the test): ```python def test_m_grouped_gemm_masked() -> None: print('Testing grouped masked GEMM:') for num_groups, m in ((2, 512), (4, 256)): for k, n in ((7168, 4096), (2048, 7168), ): # Test correctness masked_m_candidates = list(filter(lambda candidate: candidate <= m, (64, 128, 192, 256, 320, 384))) for i in range(10): x_fp8, y_fp8, out, ref_out = construct_grouped(num_groups, m, k, n, is_masked=True) masked_m = torch.empty((num_groups, ), device='cuda', dtype=torch.int) for j in range(num_groups): masked_m[j] = random.choice(masked_m_candidates) masked_m[1] = 0 # !!!! Skipping the expert 1 expected_m = min(int(masked_m.float().mean()) + 1, m) deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_masked(x_fp8, y_fp8, out, masked_m, expected_m) for j in range(num_groups): diff = calc_diff(out[j, :masked_m[j].item()], ref_out[j, :masked_m[j].item()]) # Skipping expert 1 checks assert j == 1 or diff < 0.001, f'{m=}, {k=}, {n=}, {j=}, masked_m={masked_m[j]}, {num_groups=}, {diff:.5f}' # noinspection PyShadowingNames def test_func(): # Construct new tensors every time to avoid L2 cache acceleration x_fp8, y_fp8, out, ref_out = construct_grouped(num_groups, m, k, n, is_masked=True) masked_m = torch.ones((num_groups, ), device='cuda', dtype=torch.int) * m deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_masked(x_fp8, y_fp8, out, masked_m, m) # Test performance with fixed shapes t = bench_kineto(test_func, 'fp8_gemm', suppress_kineto_output=True) print(f' > Performance ({num_groups=}, m_per_group={m:4}, n={n:4}, k={k:4}): {t * 1e6:4.0f} us | ' f'throughput: {2 * num_groups * m * n * k / t / 1e12:4.0f} TFLOPS, ' f'{(num_groups * (m * k + k * n + m * n * 2)) / 1e9 / t:4.0f} GB/s') print() ``` For the skipped experts, the NaN output is in expectation as no computation for that expert and no memory writes (no cleaning them into zero).

bsakaguchi commented

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

Author

Owner

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

Ahh thanks for your reply! That makes sense, got asserted for the nan diff at the beginning, but yeah that should be the expected behavior, closing this for now :)

@DiegoD94 commented on GitHub (Feb 27, 2025): Ahh thanks for your reply! That makes sense, got asserted for the nan diff at the beginning, but yeah that should be the expected behavior, closing this for now :)

bsakaguchi commented

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

Author

Owner

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

Sry for re-opening this issue, but I observed an unexpected benchmark result.
I was trying to benchmark low batch size decoding situation, where assume batch size = 8, then it should be at most 64 experts get routed, and the rest should be skipped.
I'm setting the group number as 256 for TP situation and 32 for EP situation, then I tried to skip certain experts by setting the masking, however I found that the latency won't reduce much even with a lot mask m = 0, could i be that I have some mis understanding of using the kernel?

Below is my latency benchmark and also my test code. All tested on H200 GPU though.:

Testing grouped masked GEMM:
 > Performance (num_groups=256, m_per_group=  64, real_m=   0, n=7168, k=4096): 1847 us | throughput:  521 TFLOPS, 4233 GB/s
 > Performance (num_groups=256, m_per_group=  64, real_m=  16, n=7168, k=4096): 1847 us | throughput:  521 TFLOPS, 4233 GB/s
 > Performance (num_groups=256, m_per_group=  64, real_m=  32, n=7168, k=4096): 1847 us | throughput:  521 TFLOPS, 4233 GB/s
 > Performance (num_groups=256, m_per_group=  64, real_m=  64, n=7168, k=4096): 1846 us | throughput:  521 TFLOPS, 4235 GB/s

def test_m_grouped_gemm_masked() -> None:
    print('Testing grouped masked GEMM:')

    for num_groups in [256]:
        m = 64
        for k, n in ((4096, 7168), (7168, 2048)):
            for real_m in [0, 16,32,64,128,256]:
            # Test correctness
            
                for i in range(10):
                    x_fp8, y_fp8, out, ref_out = construct_grouped(num_groups, m, k, n, is_masked=True)
                    masked_m = torch.empty((num_groups, ), device='cuda', dtype=torch.int)
                    for j in range(num_groups):
                        masked_m[j] = real_m
                    if real_m == 0:
                        masked_m[-1] = 1 ## Set at least 1 expert has 1 token to prevent assertion.
                    expected_m = max(real_m, 1) ## expected_m has to be > 0 to prevent assertion.
                    deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_masked(x_fp8, y_fp8, out, masked_m, expected_m)

                # noinspection PyShadowingNames
                def test_func():
                    # Construct new tensors every time to avoid L2 cache acceleration
                    x_fp8, y_fp8, out, ref_out = construct_grouped(num_groups, m, k, n, is_masked=True)
                    masked_m = torch.ones((num_groups, ), device='cuda', dtype=torch.int) * m
                    deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_masked(x_fp8, y_fp8, out, masked_m, m)

                # Test performance with fixed shapes
                t = bench_kineto(test_func, 'fp8_gemm', suppress_kineto_output=True)
                print(f' > Performance ({num_groups=}, m_per_group={m:4}, real_m={real_m:4}, n={n:4}, k={k:4}): {t * 1e6:4.0f} us | '
                    f'throughput: {2 * num_groups * m * n * k / t / 1e12:4.0f} TFLOPS, '
                    f'{(num_groups * (m * k + k * n + m * n * 2)) / 1e9 / t:4.0f} GB/s')
    print()

@DiegoD94 commented on GitHub (Feb 27, 2025): Sry for re-opening this issue, but I observed an unexpected benchmark result. I was trying to benchmark low batch size decoding situation, where assume batch size = 8, then it should be at most 64 experts get routed, and the rest should be skipped. I'm setting the group number as 256 for TP situation and 32 for EP situation, then I tried to skip certain experts by setting the masking, however I found that the latency won't reduce much even with a lot mask m = 0, could i be that I have some mis understanding of using the kernel? Below is my latency benchmark and also my test code. All tested on H200 GPU though.: ``` Testing grouped masked GEMM: > Performance (num_groups=256, m_per_group= 64, real_m= 0, n=7168, k=4096): 1847 us | throughput: 521 TFLOPS, 4233 GB/s > Performance (num_groups=256, m_per_group= 64, real_m= 16, n=7168, k=4096): 1847 us | throughput: 521 TFLOPS, 4233 GB/s > Performance (num_groups=256, m_per_group= 64, real_m= 32, n=7168, k=4096): 1847 us | throughput: 521 TFLOPS, 4233 GB/s > Performance (num_groups=256, m_per_group= 64, real_m= 64, n=7168, k=4096): 1846 us | throughput: 521 TFLOPS, 4235 GB/s ``` ``` def test_m_grouped_gemm_masked() -> None: print('Testing grouped masked GEMM:') for num_groups in [256]: m = 64 for k, n in ((4096, 7168), (7168, 2048)): for real_m in [0, 16,32,64,128,256]: # Test correctness for i in range(10): x_fp8, y_fp8, out, ref_out = construct_grouped(num_groups, m, k, n, is_masked=True) masked_m = torch.empty((num_groups, ), device='cuda', dtype=torch.int) for j in range(num_groups): masked_m[j] = real_m if real_m == 0: masked_m[-1] = 1 ## Set at least 1 expert has 1 token to prevent assertion. expected_m = max(real_m, 1) ## expected_m has to be > 0 to prevent assertion. deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_masked(x_fp8, y_fp8, out, masked_m, expected_m) # noinspection PyShadowingNames def test_func(): # Construct new tensors every time to avoid L2 cache acceleration x_fp8, y_fp8, out, ref_out = construct_grouped(num_groups, m, k, n, is_masked=True) masked_m = torch.ones((num_groups, ), device='cuda', dtype=torch.int) * m deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_masked(x_fp8, y_fp8, out, masked_m, m) # Test performance with fixed shapes t = bench_kineto(test_func, 'fp8_gemm', suppress_kineto_output=True) print(f' > Performance ({num_groups=}, m_per_group={m:4}, real_m={real_m:4}, n={n:4}, k={k:4}): {t * 1e6:4.0f} us | ' f'throughput: {2 * num_groups * m * n * k / t / 1e12:4.0f} TFLOPS, ' f'{(num_groups * (m * k + k * n + m * n * 2)) / 1e9 / t:4.0f} GB/s') print() ```

bsakaguchi commented

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

Author

Owner

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

@DiegoD94 It seems that real_m is not set correctly in test_func.

            def test_func():
                # Construct new tensors every time to avoid L2 cache acceleration
                x_fp8, y_fp8, out, ref_out = construct_grouped(num_groups, m, k, n, is_masked=True)
                masked_m = torch.ones((num_groups, ), device='cuda', dtype=torch.int) * real_m
                deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_masked(x_fp8, y_fp8, out, masked_m, real_m)

@zheanxu commented on GitHub (Feb 27, 2025): @DiegoD94 It seems that `real_m` is not set correctly in `test_func`. ```c++ def test_func(): # Construct new tensors every time to avoid L2 cache acceleration x_fp8, y_fp8, out, ref_out = construct_grouped(num_groups, m, k, n, is_masked=True) masked_m = torch.ones((num_groups, ), device='cuda', dtype=torch.int) * real_m deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_masked(x_fp8, y_fp8, out, masked_m, real_m) ```

bsakaguchi referenced this issue

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

[PR #16] [MERGED] fix typo #205

bsakaguchi referenced this issue

2026-10-26 21:15:46 +00:00

[PR #188] [CLOSED] disable persistent kernel #283

Rows
Columns

[GH-ISSUE #20] How to skip certain expert in grouped gemm masked layout #16