deepseek-ai/DeepGEMM

Fork 0

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

[PR #435] [CLOSED] Improve SM100 MX MQA logits clean-output performance #443

New issue

Closed

opened 2026-10-26 21:17:56 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:17:56 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/435
Author: @OutWhite
Created: 9/10/2026
Status: ❌ Closed

Base: main ← Head: codex/sm100-mx-mqa-clean-performance

📝 Commits (2)

0185f6b Improve SM100 MX MQA logits clean-output performance
197189e Document SM100 MQA cleanup scheduling choices

📊 Changes

1 file changed (+34 additions, -16 deletions)

View changed files

📝 deep_gemm/include/deep_gemm/impls/sm100_mqa_logits.cuh (+34 -16)

📄 Description

On GB200 with NVCC 13.3.33, placing MXFP8/MXFP4 MQA-logits cleanup after the CTA's main computation improves clean_logits=True performance by about 1.6–2.3x for the H64/D128 cases below, with identical complete outputs. This PR proposes that scheduling change for the SM100 MX paths.

Move the MX cleaner after the CTA's producer/consumer work and TMEM release. Reuse the existing cleaner through a local lambda, and synchronize with a separate named barrier (ID 1) so it cannot overlap the math-only barrier (ID 0). The grid-stride cleanup ranges, in-coverage masks, metadata layout, and public API are unchanged. FP8 retains concurrent cleanup; moving FP8 cleanup as well caused a 5–8% slowdown in some exploratory cases. The lambda's kCleanLogits guard also avoids instantiating contiguous-only cleaner methods for paged schedulers.

Performance

Baseline: 66081d4c9c7d7c44f13fea402e5b622aa0f409c2. Same GB200 GPU (152 SMs), NVCC 13.3.33, default DeepGEMM JIT flags targeting sm_100f, driver 580.167.08, PyTorch 2.13.0+cu130; PDL disabled.

H=64, D=128, FP32 weights/output, uncompressed logits. Q/K/weights are all zero, MX scales are one, and each row uses [0, KV-Q+row). These are CUDA Graph GPU timings of the complete clean_logits=True call, including all cleaning. Metadata is precomputed outside timing. Five calls per graph with distinct retained output buffers; 40 samples, rotating/reversing variant order, immediate warmup before each timed replay; median per-call latency. The timing runs use normal GPU clocks. Separate Nsight captures below use controlled base clocks.

Q / KV tokens	Format	Metadata	Before (us)	After (us)	Speedup
2048 / 32768	MXFP8	no	645.9	412.0	1.57x
2048 / 32768	MXFP8	yes	817.7	406.3	2.01x
2048 / 32768	MXFP4	no	806.4	407.9	1.98x
2048 / 32768	MXFP4	yes	780.2	403.0	1.94x
8192 / 65536	MXFP8	no	4994.6	3022.2	1.65x
8192 / 65536	MXFP8	yes	7093.2	3083.6	2.30x
8192 / 65536	MXFP4	no	5990.3	2968.9	2.02x
8192 / 65536	MXFP4	yes	6008.1	3054.5	1.97x

The clean_logits=False control stays within 0.2% between builds for these eight points. A further 18-shape screen (both metadata modes) includes H=8/32/64, Q=KV=8192, CP windows, and BF16 weights/output: MX latencies improve in every screened case; FP8 is within 0.2% of baseline. The improvement also appears with random inputs. The table uses an all-zero workload to reduce input-dependent power effects.

Timing data, including p10/p90 (microseconds)

Q,KV,format,metadata,before_median_us,before_p10_us,before_p90_us,after_median_us,after_p10_us,after_p90_us
2048,32768,mxfp8,0,645.859,636.832,675.693,411.974,411.77,412.186
2048,32768,mxfp8,1,817.683,812.563,825.261,406.333,406.227,406.65
2048,32768,mxfp4,0,806.432,806.042,806.842,407.888,407.469,408.301
2048,32768,mxfp4,1,780.218,779.808,780.218,402.966,402.566,403.59
8192,65536,mxfp8,0,4994.589,4950.765,5020.403,3022.163,3021.344,3022.982
8192,65536,mxfp8,1,7093.178,7083.347,7106.08,3083.603,3082.784,3084.013
8192,65536,mxfp4,0,5990.317,5989.088,5991.52,2968.915,2966.861,2970.547
8192,65536,mxfp4,1,6008.093,6007.13,6008.755,3054.506,3052.454,3056.557

Investigation notes

A few local ablations helped distinguish cleanup bandwidth costs from sensitivity to code placement. These use H64/D128, Q=2048, KV=32768:

Compile out only the cleaner warp's work, keeping the main-loop invalid-column mask: MXFP8 drops from 641.3 to 397.2 us; MXFP4 from 806.0 to 416.0 us. This is a diagnostic variant, not a complete-output implementation.
Keep the exact same compiled kernel and disable the cleaner via a runtime gate: MXFP8 is 633.1 us with cleanup versus 632.8 us without; MXFP4 is 805.1 versus 805.5 us. The gate toggles the first row's end by one element; split counts and metadata remain unchanged (metadata was checked bitwise). This suggests that the stores themselves do not explain the timing difference.
Removing only the main-loop output mask gives similar no-metadata timings to baseline: 635.1 us for MXFP8 and 800.3 us for MXFP4.
Changing the ordering of the specialized-warp branches also changes performance, including for clean_logits=False. The cleaner-loop unrolling and inline-PTX store variants showed similar timings to baseline.

With the final patch, a controlled-clock Nsight Compute 2026.2 comparison for MXFP8, Q=2048, KV=32768, no metadata shows:

Metric	Before	After
SM clock (GHz)	1.132	1.128
Registers / thread	168.000	168.000
Dynamic shared memory (decimal KB / CTA)	222.208	222.208
DRAM reads (decimal MB)	22.218	22.249
DRAM writes (decimal MB)	211.368	209.558
Long-scoreboard stalls (`per_issue_active` ratio)	3.519	1.529
Tensor-pipe active (% of elapsed cycles)	34.927	53.430

Exploratory PC samples place much of the additional wait around the math-side TMEM-result barrier check. Together, these results suggest that cleaner placement affects main-loop execution; memory traffic and resource allocation are largely unchanged, and neither binary has register spills. The precise low-level dependency remains to be investigated. The proposed change avoids this sensitivity on the tested configuration while retaining the existing output semantics.

Validation

Existing upstream test_mqa_logits: 32 sampled configurations pass (including its 20-repeat bitwise checks, scheduled-path equality, full mask checks, and reference comparisons).
Existing upstream test_paged_mqa_logits: 4 sampled configurations pass, exercising the shared core with paged schedulers.
One-off comparison: all 48 complete outputs are bitwise identical to baseline. Covers FP8/MXFP8/MXFP4, FP32/BF16 weights and output, both metadata modes, random inputs, full/causal and irregular windows, nonzero starts, empty rows, non-aligned lengths (Q=129/KV=4100), and all-zero inputs; also checked against a quantized PyTorch reference.
Compute Sanitizer on 32 compact MX output configurations (Q=17/KV=260, both metadata modes, FP32/BF16, varied masks): synccheck PASS; memcheck PASS.

No permanent test-matrix or benchmark-file additions. Existing upstream tests can be run from a built checkout with:

CUDA_VISIBLE_DEVICES=0 DG_MQA_NUM_CASES=32 python - <<'PY'
import os, random, sys, torch, deep_gemm
sys.path.insert(0, 'tests')
import test_attention
torch.manual_seed(0)
random.seed(0)
deep_gemm.set_pdl(False)
test_attention.test_mqa_logits()
os.environ['DG_MQA_NUM_CASES'] = '4'
test_attention.test_paged_mqa_logits()
PY

Standalone performance reproducer (writes every timing sample to JSON)

Save as pr-reproduce.py and run it once with each built revision, using the same GPU and NVCC. Both builds and all JIT compilations must finish before timed sampling; the script warms every specialization before capture. It verifies the zero output and complete negative-infinity mask before timing.

"""Run with the baseline and PR DeepGEMM builds on the same GPU/toolchain.

CUDA_VISIBLE_DEVICES=0 python pr-reproduce.py --output before.json
CUDA_VISIBLE_DEVICES=0 python pr-reproduce.py --output after.json
"""
import argparse
import gc
import json
import statistics
from functools import partial

import torch
import deep_gemm as dg


def measure(functions, samples=40, calls=5):
    graphs = {}
    for name, fn in functions.items():
        for _ in range(3):
            fn()
        torch.cuda.synchronize()
        graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(graph):
            outputs = [fn() for _ in range(calls)]
        graphs[name] = (graph, outputs)
        for _ in range(5):
            graph.replay()
    torch.cuda.synchronize()
    begin, end = [torch.cuda.Event(enable_timing=True) for _ in range(2)]
    times = {name: [] for name in functions}
    names = list(functions)
    for sample in range(samples):
        order = names[sample % len(names):] + names[:sample % len(names)]
        if sample % 2:
            order.reverse()
        for name in order:
            graphs[name][0].replay()
            begin.record()
            graphs[name][0].replay()
            end.record()
            end.synchronize()
            times[name].append(begin.elapsed_time(end) * 1000 / calls)
    return {
        name: {
            'median_us': statistics.median(values),
            'p10_us': sorted(values)[int(.1 * (samples - 1))],
            'p90_us': sorted(values)[int(.9 * (samples - 1))],
            'samples_us': values,
        }
        for name, values in times.items()
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--output', required=True)
    args = parser.parse_args()
    dg.set_pdl(False)
    result = []
    for m, n in ((2048, 32768), (8192, 65536)):
        for fmt in ('mxfp8', 'mxfp4'):
            fp4 = fmt == 'mxfp4'
            q = torch.zeros((m, 64, 64 if fp4 else 128), device='cuda',
                            dtype=torch.int8 if fp4 else torch.float8_e4m3fn)
            k = torch.zeros((n, q.shape[-1]), device='cuda', dtype=q.dtype)
            qs = torch.full((m, 64), 0x7f7f7f7f, device='cuda', dtype=torch.int32)
            ks = torch.full((n,), 0x7f7f7f7f, device='cuda', dtype=torch.int32)
            w = torch.zeros((m, 64), device='cuda', dtype=torch.float32)
            starts = torch.zeros(m, device='cuda', dtype=torch.int32)
            ends = torch.arange(m, device='cuda', dtype=torch.int32) + n - m
            metadata = dg.get_mqa_logits_metadata(starts, ends, n, 64)
            functions = {
                f'clean{int(clean)}_meta{int(use_meta)}': partial(
                    dg.fp8_fp4_mqa_logits, (q, qs), (k, ks), w, starts, ends,
                    clean_logits=clean, logits_dtype=torch.float32,
                    schedule_meta=metadata if use_meta else None,
                )
                for clean in (False, True) for use_meta in (False, True)
            }
            for use_meta in (False, True):
                out = functions[f'clean1_meta{int(use_meta)}']()
                cols = torch.arange(n, device='cuda')[None, :]
                for row in range(0, m, 128):
                    valid = cols < ends[row:row + 128, None]
                    chunk = out[row:row + 128]
                    assert torch.equal(chunk.isneginf(), ~valid)
                    assert (chunk.masked_fill(~valid, 0) == 0).all()
                del out
            timings = measure(functions)
            result.append(dict(shape=[m, n, 64, 128], format=fmt, timings=timings))
            with open(args.output, 'w') as f:
                json.dump(result, f, indent=2)
            print(m, n, fmt, {k: round(v['median_us'], 3) for k, v in timings.items()}, flush=True)
            del functions, q, k, qs, ks, w, starts, ends, metadata, chunk, valid, cols
            gc.collect()
            torch.cuda.empty_cache()


if __name__ == '__main__':
    main()

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepGEMM/pull/435 **Author:** [@OutWhite](https://github.com/OutWhite) **Created:** 9/10/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `codex/sm100-mx-mqa-clean-performance` --- ### 📝 Commits (2) - [`0185f6b`](https://github.com/deepseek-ai/DeepGEMM/commit/0185f6beb7319f7beae31d03fd3ac3f34e8b3b4b) Improve SM100 MX MQA logits clean-output performance - [`197189e`](https://github.com/deepseek-ai/DeepGEMM/commit/197189e994e6212d541cc79a270c115916c459ed) Document SM100 MQA cleanup scheduling choices ### 📊 Changes **1 file changed** (+34 additions, -16 deletions) <details> <summary>View changed files</summary> 📝 `deep_gemm/include/deep_gemm/impls/sm100_mqa_logits.cuh` (+34 -16) </details> ### 📄 Description On GB200 with NVCC 13.3.33, placing MXFP8/MXFP4 MQA-logits cleanup after the CTA's main computation improves `clean_logits=True` performance by about 1.6–2.3x for the H64/D128 cases below, with identical complete outputs. This PR proposes that scheduling change for the SM100 MX paths. Move the MX cleaner after the CTA's producer/consumer work and TMEM release. Reuse the existing cleaner through a local lambda, and synchronize with a separate named barrier (ID 1) so it cannot overlap the math-only barrier (ID 0). The grid-stride cleanup ranges, in-coverage masks, metadata layout, and public API are unchanged. FP8 retains concurrent cleanup; moving FP8 cleanup as well caused a 5–8% slowdown in some exploratory cases. The lambda's `kCleanLogits` guard also avoids instantiating contiguous-only cleaner methods for paged schedulers. ## Performance Baseline: `66081d4c9c7d7c44f13fea402e5b622aa0f409c2`. Same GB200 GPU (152 SMs), NVCC 13.3.33, default DeepGEMM JIT flags targeting `sm_100f`, driver 580.167.08, PyTorch 2.13.0+cu130; PDL disabled. H=64, D=128, FP32 weights/output, uncompressed logits. Q/K/weights are all zero, MX scales are one, and each row uses `[0, KV-Q+row)`. These are CUDA Graph GPU timings of the **complete `clean_logits=True` call**, including all cleaning. Metadata is precomputed outside timing. Five calls per graph with distinct retained output buffers; 40 samples, rotating/reversing variant order, immediate warmup before each timed replay; median per-call latency. The timing runs use normal GPU clocks. Separate Nsight captures below use controlled base clocks. | Q / KV tokens | Format | Metadata | Before (us) | After (us) | Speedup | |---|---|---|---:|---:|---:| | 2048 / 32768 | MXFP8 | no | 645.9 | 412.0 | 1.57x | | 2048 / 32768 | MXFP8 | yes | 817.7 | 406.3 | 2.01x | | 2048 / 32768 | MXFP4 | no | 806.4 | 407.9 | 1.98x | | 2048 / 32768 | MXFP4 | yes | 780.2 | 403.0 | 1.94x | | 8192 / 65536 | MXFP8 | no | 4994.6 | 3022.2 | 1.65x | | 8192 / 65536 | MXFP8 | yes | 7093.2 | 3083.6 | 2.30x | | 8192 / 65536 | MXFP4 | no | 5990.3 | 2968.9 | 2.02x | | 8192 / 65536 | MXFP4 | yes | 6008.1 | 3054.5 | 1.97x | The `clean_logits=False` control stays within 0.2% between builds for these eight points. A further 18-shape screen (both metadata modes) includes H=8/32/64, Q=KV=8192, CP windows, and BF16 weights/output: MX latencies improve in every screened case; FP8 is within 0.2% of baseline. The improvement also appears with random inputs. The table uses an all-zero workload to reduce input-dependent power effects. <details> <summary>Timing data, including p10/p90 (microseconds)</summary> ```csv Q,KV,format,metadata,before_median_us,before_p10_us,before_p90_us,after_median_us,after_p10_us,after_p90_us 2048,32768,mxfp8,0,645.859,636.832,675.693,411.974,411.77,412.186 2048,32768,mxfp8,1,817.683,812.563,825.261,406.333,406.227,406.65 2048,32768,mxfp4,0,806.432,806.042,806.842,407.888,407.469,408.301 2048,32768,mxfp4,1,780.218,779.808,780.218,402.966,402.566,403.59 8192,65536,mxfp8,0,4994.589,4950.765,5020.403,3022.163,3021.344,3022.982 8192,65536,mxfp8,1,7093.178,7083.347,7106.08,3083.603,3082.784,3084.013 8192,65536,mxfp4,0,5990.317,5989.088,5991.52,2968.915,2966.861,2970.547 8192,65536,mxfp4,1,6008.093,6007.13,6008.755,3054.506,3052.454,3056.557 ``` </details> ## Investigation notes A few local ablations helped distinguish cleanup bandwidth costs from sensitivity to code placement. These use H64/D128, Q=2048, KV=32768: - Compile out only the cleaner warp's work, keeping the main-loop invalid-column mask: MXFP8 drops from 641.3 to 397.2 us; MXFP4 from 806.0 to 416.0 us. This is a diagnostic variant, not a complete-output implementation. - Keep the exact same compiled kernel and disable the cleaner via a runtime gate: MXFP8 is 633.1 us with cleanup versus 632.8 us without; MXFP4 is 805.1 versus 805.5 us. The gate toggles the first row's end by one element; split counts and metadata remain unchanged (metadata was checked bitwise). This suggests that the stores themselves do not explain the timing difference. - Removing only the main-loop output mask gives similar no-metadata timings to baseline: 635.1 us for MXFP8 and 800.3 us for MXFP4. - Changing the ordering of the specialized-warp branches also changes performance, including for `clean_logits=False`. The cleaner-loop unrolling and inline-PTX store variants showed similar timings to baseline. With the final patch, a controlled-clock Nsight Compute 2026.2 comparison for MXFP8, Q=2048, KV=32768, no metadata shows: | Metric | Before | After | |---|---:|---:| | SM clock (GHz) | 1.132 | 1.128 | | Registers / thread | 168.000 | 168.000 | | Dynamic shared memory (decimal KB / CTA) | 222.208 | 222.208 | | DRAM reads (decimal MB) | 22.218 | 22.249 | | DRAM writes (decimal MB) | 211.368 | 209.558 | | Long-scoreboard stalls (`per_issue_active` ratio) | 3.519 | 1.529 | | Tensor-pipe active (% of elapsed cycles) | 34.927 | 53.430 | Exploratory PC samples place much of the additional wait around the math-side TMEM-result barrier check. Together, these results suggest that cleaner placement affects main-loop execution; memory traffic and resource allocation are largely unchanged, and neither binary has register spills. The precise low-level dependency remains to be investigated. The proposed change avoids this sensitivity on the tested configuration while retaining the existing output semantics. ## Validation - Existing upstream `test_mqa_logits`: 32 sampled configurations pass (including its 20-repeat bitwise checks, scheduled-path equality, full mask checks, and reference comparisons). - Existing upstream `test_paged_mqa_logits`: 4 sampled configurations pass, exercising the shared core with paged schedulers. - One-off comparison: all 48 complete outputs are bitwise identical to baseline. Covers FP8/MXFP8/MXFP4, FP32/BF16 weights and output, both metadata modes, random inputs, full/causal and irregular windows, nonzero starts, empty rows, non-aligned lengths (Q=129/KV=4100), and all-zero inputs; also checked against a quantized PyTorch reference. - Compute Sanitizer on 32 compact MX output configurations (Q=17/KV=260, both metadata modes, FP32/BF16, varied masks): synccheck PASS; memcheck PASS. No permanent test-matrix or benchmark-file additions. Existing upstream tests can be run from a built checkout with: ```bash CUDA_VISIBLE_DEVICES=0 DG_MQA_NUM_CASES=32 python - <<'PY' import os, random, sys, torch, deep_gemm sys.path.insert(0, 'tests') import test_attention torch.manual_seed(0) random.seed(0) deep_gemm.set_pdl(False) test_attention.test_mqa_logits() os.environ['DG_MQA_NUM_CASES'] = '4' test_attention.test_paged_mqa_logits() PY ``` <details> <summary>Standalone performance reproducer (writes every timing sample to JSON)</summary> Save as `pr-reproduce.py` and run it once with each built revision, using the same GPU and NVCC. Both builds and all JIT compilations must finish before timed sampling; the script warms every specialization before capture. It verifies the zero output and complete negative-infinity mask before timing. ```python """Run with the baseline and PR DeepGEMM builds on the same GPU/toolchain. CUDA_VISIBLE_DEVICES=0 python pr-reproduce.py --output before.json CUDA_VISIBLE_DEVICES=0 python pr-reproduce.py --output after.json """ import argparse import gc import json import statistics from functools import partial import torch import deep_gemm as dg def measure(functions, samples=40, calls=5): graphs = {} for name, fn in functions.items(): for _ in range(3): fn() torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): outputs = [fn() for _ in range(calls)] graphs[name] = (graph, outputs) for _ in range(5): graph.replay() torch.cuda.synchronize() begin, end = [torch.cuda.Event(enable_timing=True) for _ in range(2)] times = {name: [] for name in functions} names = list(functions) for sample in range(samples): order = names[sample % len(names):] + names[:sample % len(names)] if sample % 2: order.reverse() for name in order: graphs[name][0].replay() begin.record() graphs[name][0].replay() end.record() end.synchronize() times[name].append(begin.elapsed_time(end) * 1000 / calls) return { name: { 'median_us': statistics.median(values), 'p10_us': sorted(values)[int(.1 * (samples - 1))], 'p90_us': sorted(values)[int(.9 * (samples - 1))], 'samples_us': values, } for name, values in times.items() } def main(): parser = argparse.ArgumentParser() parser.add_argument('--output', required=True) args = parser.parse_args() dg.set_pdl(False) result = [] for m, n in ((2048, 32768), (8192, 65536)): for fmt in ('mxfp8', 'mxfp4'): fp4 = fmt == 'mxfp4' q = torch.zeros((m, 64, 64 if fp4 else 128), device='cuda', dtype=torch.int8 if fp4 else torch.float8_e4m3fn) k = torch.zeros((n, q.shape[-1]), device='cuda', dtype=q.dtype) qs = torch.full((m, 64), 0x7f7f7f7f, device='cuda', dtype=torch.int32) ks = torch.full((n,), 0x7f7f7f7f, device='cuda', dtype=torch.int32) w = torch.zeros((m, 64), device='cuda', dtype=torch.float32) starts = torch.zeros(m, device='cuda', dtype=torch.int32) ends = torch.arange(m, device='cuda', dtype=torch.int32) + n - m metadata = dg.get_mqa_logits_metadata(starts, ends, n, 64) functions = { f'clean{int(clean)}_meta{int(use_meta)}': partial( dg.fp8_fp4_mqa_logits, (q, qs), (k, ks), w, starts, ends, clean_logits=clean, logits_dtype=torch.float32, schedule_meta=metadata if use_meta else None, ) for clean in (False, True) for use_meta in (False, True) } for use_meta in (False, True): out = functions[f'clean1_meta{int(use_meta)}']() cols = torch.arange(n, device='cuda')[None, :] for row in range(0, m, 128): valid = cols < ends[row:row + 128, None] chunk = out[row:row + 128] assert torch.equal(chunk.isneginf(), ~valid) assert (chunk.masked_fill(~valid, 0) == 0).all() del out timings = measure(functions) result.append(dict(shape=[m, n, 64, 128], format=fmt, timings=timings)) with open(args.output, 'w') as f: json.dump(result, f, indent=2) print(m, n, fmt, {k: round(v['median_us'], 3) for k, v in timings.items()}, flush=True) del functions, q, k, qs, ks, w, starts, ends, metadata, chunk, valid, cols gc.collect() torch.cuda.empty_cache() if __name__ == '__main__': main() ``` </details> --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-26 21:17:56 +00:00

closed this issue
added the
pull-request
label

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

No description provided.

Rows
Columns