[SM100] Add exact BF16 paged MQA histograms for next-n 1–6 by Heisenberg-Yin · Pull Request #96 · sgl-project/DeepGEMM

Heisenberg-Yin · 2026-10-03T05:31:41Z

@Fridge003 @DarkSharpness could you take a look when you are free?

Thanks to @TKONIY for the collaboration on this work.

Motivation

The companion SGLang#42351 uses a score histogram to select the exact top-512 index-K slots for DeepSeek-V4.1 decode and target verify. This PR ports the upstream BF16 paged MQA producer into sgl-project dev and adds an optional exact histogram of its stored scores, including next-n 5/6 KV reuse.

The companion SGLang PR is an incremental change after SGLang #40764. The existing FP32 module stays in #40764; BF16 producer kernels, bindings and producer tests stay here. SGLang adds the BF16 selector and serving integration and requires this PR's BF16 APIs.

Code sources and diff scope

The BF16 scoring arithmetic and device core are adapted from deepseek-ai/DeepGEMM#462, merge commit 057ca5964aae0879ff2e0eb71ee05a3cb0ba3df7. deepseek-ai#462 is both an implementation source and the independent numerical/performance reference. BF16 output and BF16 head reduction originate upstream.

Local file under `mqa_bf16/`	Upstream source at `057ca5964aae`	Adaptation
`kernel.cuh`	`impls/sm100_mqa_logits.cuh`	Isolated core, histogram hooks, Q4/Q6 geometry
`scheduler.cuh`	`scheduler/sm100_paged_mqa_logits.cuh`	Cost-balanced schedule and request grouping
`layout.cuh`	`layout/mqa_logits.cuh`	Isolated namespace; separate math-group/TMEM-stage barrier geometry
`packing.cuh`	`common/packing.cuh`	Namespace change
`ring_pipeline.cuh`	`common/ring_pipeline.cuh`	Namespace change and phase handling for the two-stage ring
`helpers.cuh`	Relevant upstream math/PTX helpers	Local compatibility helpers for the dev baseline

The coarse-bin encoding follows the earlier sgl-project/DeepGEMM#94 histogram design. BF16 counter packing, histogram integration, Q6 routing, swizzle, scheduling changes and dev bindings are the work added here. The existing generic MQA APIs keep their implementation and default FP32 behavior.

This PR targets dev at ec226bcf400c11d62755137b916fae9ee359c684. Its 17 files, +1,827 / −1 include 1,372 device-header lines in a separate deep_gemm/mqa_bf16 directory, so adapted upstream code is visible as newly added files. It is a complete isolated port plus histogram extension, rather than a small diff that only adds histogram atomics to the generic dev kernel.

Modifications

Add get_paged_mqa_logits_bf16_metadata and fp4_paged_mqa_logits_bf16 through both pybind and TVM-FFI wrappers (deep_gemm / sgl_deep_gemm).
Count every live non-NaN stored BF16 score once into an additive int32 [R,1024] histogram. Histogram-enabled and histogram-disabled calls return bit-identical live scores. Counts are per row and per score, with no pooling or block-max substitution.
CPU tokens_per_request: 1–4 selects Q4 / 3 TMEM stages; 5–6 selects Q6 / 2 stages, sharing KV loads among adjacent query rows of a request; hint 6 enables shared-memory histogram XOR swizzle.
Use an isolated split-384 schedule balanced by estimated work. Schedule generation runs with live request IDs/lengths during CUDA graph replay.
Add the API contract in docs/mqa_bf16.md and GPU tests in sgl_deep_gemm/tests/test_mqa_bf16.py.

schedule = deep_gemm.get_paged_mqa_logits_bf16_metadata(
    context_lens, 128, deep_gemm.get_num_sms(),
    indices=request_ids, tokens_per_request=n,
)
histogram = torch.zeros((rows, 1024), dtype=torch.int32, device="cuda")
scores = deep_gemm.fp4_paged_mqa_logits_bf16(
    (q_fp4, q_sf), kv_cache, weights, context_lens, block_table,
    schedule, max_context_len,
    indices=request_ids, histogram=histogram, tokens_per_request=n,
)

Scope

SM100, H32/D128, MXFP4 Q [R,1,32,64] with packed int32 UE8M0 scales [R,1,32], paged uint8 KV [pages,128,1,68], BF16 weights [R,32], and BF16 output [R,max_context_len] with row stride aligned to 1536 elements. Context lengths are int32 [R,1]; page tables have one row per score row. CPU hint is 1–6, with at most 16,384 rows.

IDs identify adjacent request runs; page tables match and causal lengths are nondecreasing within each run. Omitting IDs treats rows independently. Schedule and producer use the same hint/IDs/lengths; this metadata is incompatible with the generic MQA schedule. Only live score columns are defined. Callers zero counts before an independent call, or use the companion selector which consumes and clears them.

BF16 scoring follows upstream deepseek-ai#462 semantics and can yield a different top-k from the default FP32 producer. Histogram counts ignore NaNs; the companion selector's nonfinite behavior is documented in the SGLang PR.

Accuracy Tests

Fresh focused test run on the complete SGLang PR head: 55 passed — 15 dispatch/storage/metadata unit tests, 12 producer tests, and 28 BF16 selector/serving GPU tests. The GPU suite covers n=1–6, PDL off/on, exact counts, repeated accumulation, zero/page/split/top-k boundaries, shared and changing request IDs, dynamic CUDA graphs, FullTopK/SparseTable, storage lifetime, fallback and BF16 candidate maxima. The producer test is an unchanged copy of the repository test run outside the package directory; its source hash is recorded.

Every one of the 90 cold benchmark cells verifies live BF16 bits against the independent upstream deepseek-ai#462 producer, exact histograms against a per-score oracle, selection values against torch.topk, valid unique physical slots and reusable scratch. LiteTopK additionally matches (BF16 ordered score key descending, physical slot ascending) exactly. The default FP32 chain is checked against top-k of its own FP32 scores. Outputs are checked before and after timing.

The separate fresh qualification has 69 passing records: 4 large-batch cases up to 768 rows, 51 graph replays across page/split/top-k boundaries, 8 selector edge cases (including ties, signed zero, nonfinite scores and overflow), graph storage growth, and 5 actual FullTopK hook cases. Both live producer bindings come from the dev-targeted #96 build; the existing sparse APIs are used directly. The 90-cell cold matrix and 69-record qualification apply the live SGLang diff to local main snapshot a8b54f138e; the complete PR head additionally passes the same 55-test suite. All timed selector/top-k files match the complete head byte for byte. Three unrelated Hopper-path base changes are recorded in the source audit. Earlier measurements are retained separately and are not substituted for these reruns.

Speed Tests and Profiling

These fresh measurements use the dev-targeted TVM-FFI build from this PR, with all device headers verified against its live diff. Upstream deepseek-ai#462 is the BF16 reference. The following numbers are CUPTI producer kernel durations from the randomized measurements, not isolated histogram execution costs. Reference BF16 runs in a DeepSelect chain, final no-hist runs alone, and final histogram runs in the LiteTopK chain. Mapping, selection and metadata are not separate summands in these producer durations, but their scheduling and resource effects can affect the durations.

Each timed cold replay contains one call, with cache eviction before the call. The producer-duration change is an observed cold measurement under these chain contexts. It is not an isolated instruction-cost estimate; a matched producer-only histogram/no-hist ablation with reset costs controlled would be needed for that. PDL is enabled in these measurements; its on/off performance contribution was not measured.

Cold-cache measurements use NVIDIA B200 (148 SMs), Torch 2.13.0+cu130, CUDA 13.0, and PDL enabled. Comparisons within each cell run on the same GPU; n=1–3 use GPU4 and n=4–6 GPU6. The matrix is n=1–6 × B={1,16,64} × L={4096,16384,65536,262144,1048342}: 90 cold cells. B counts requests, R=B×n score rows, and L compressed KV positions.

Inputs are captured DeepSeek-V4.1 layer-20 MXFP4 Q/KV and BF16-rounded head weights, extended to each shape with randomly permuted physical pages and each query row's causal length. The default FP32 producer receives one independent page-table/length row per query, as the serving path does; the paired BF16 producer can reuse KV among adjacent equal request IDs. L is the live context, while score allocation retains the captured ~1M-column maximum window; only live columns are defined.

Each cell uses 80 randomized paired blocks. A 1 GiB buffer is zeroed before every single-call CUDA graph replay to evict the preceding cache contents. The flush is outside the timed span; CUPTI measures each chain from its first kernel start to its last kernel end. Outputs and scratch are reused. Physical-slot mapping, histogram consumption and scratch cleanup are included. Q projection/packing, schedule/plan construction, serving weight conversion and CPU launch overhead are outside this matrix.

Clocks and power settings are unchanged. Idle gates and per-process monitoring reject affected cells with a 10 s / 6 s margin, including their preparation and profiling window. Raw samples, kernel launch signatures, clocks, rejected attempts and source/binary hashes accompany the results.

Geometric means over 15 cold shapes per n. Ratios compare observed producer kernel durations in their respective chains; the last column is the histogram-chain producer duration divided by standalone no-hist duration, minus one. Use the complete cold chain measurements for the latency/speedup conclusion.

Mode	n	Upstream / no-hist duration	Upstream / hist-chain producer duration	Producer duration change (different chains)
cold	1	0.974×	0.958×	+1.8%
cold	2	0.976×	0.950×	+2.7%
cold	3	0.977×	0.936×	+4.3%
cold	4	0.982×	0.942×	+4.3%
cold	5	1.060×	0.971×	+9.1%
cold	6	1.008×	0.960×	+5.0%

cold: all 90 producer cells

n	B	L	Upstream BF16	Final no-hist alone	Final hist in selector chain	Producer duration change
1	1	4,096	4.80	5.12	5.38	+5.0%
1	1	16,384	5.15	5.49	5.73	+4.4%
1	1	65,536	6.27	6.54	6.66	+1.7%
1	1	262,144	9.26	9.66	9.74	+0.8%
1	1	1,048,342	21.31	21.76	21.82	+0.3%
1	16	4,096	6.50	6.85	7.26	+6.1%
1	16	16,384	9.71	9.95	10.22	+2.7%
1	16	65,536	21.63	21.86	22.19	+1.5%
1	16	262,144	62.93	63.49	63.28	-0.3%
1	16	1,048,342	204.03	204.66	202.72	-0.9%
1	64	4,096	9.65	10.03	10.34	+3.0%
1	64	16,384	21.65	21.95	22.22	+1.2%
1	64	65,536	63.49	63.78	64.00	+0.4%
1	64	262,144	203.59	203.44	204.72	+0.6%
1	64	1,048,342	752.93	753.09	753.99	+0.1%
2	1	4,096	5.15	5.46	5.73	+5.0%
2	1	16,384	5.46	5.73	5.89	+2.8%
2	1	65,536	6.62	6.98	7.07	+1.4%
2	1	262,144	9.63	9.98	10.21	+2.2%
2	1	1,048,342	21.76	22.08	22.59	+2.3%
2	16	4,096	6.91	7.30	7.71	+5.7%
2	16	16,384	9.89	10.30	10.78	+4.7%
2	16	65,536	22.11	22.37	22.94	+2.6%
2	16	262,144	64.37	64.62	65.90	+2.0%
2	16	1,048,342	205.81	204.78	209.12	+2.1%
2	64	4,096	10.08	10.38	10.75	+3.5%
2	64	16,384	22.42	22.69	23.41	+3.2%
2	64	65,536	64.34	64.86	66.11	+1.9%
2	64	262,144	210.98	211.82	210.96	-0.4%
2	64	1,048,342	779.97	779.01	788.95	+1.3%
3	1	4,096	5.30	5.60	5.95	+6.3%
3	1	16,384	5.70	5.89	6.27	+6.5%
3	1	65,536	6.91	7.26	7.55	+4.0%
3	1	262,144	10.05	10.34	10.78	+4.3%
3	1	1,048,342	22.77	23.38	23.95	+2.5%
3	16	4,096	7.39	7.68	8.42	+9.6%
3	16	16,384	10.40	11.11	11.73	+5.6%
3	16	65,536	23.10	23.62	24.42	+3.4%
3	16	262,144	66.08	66.35	68.10	+2.6%
3	16	1,048,342	220.35	220.80	224.61	+1.7%
3	64	4,096	10.62	10.83	11.92	+10.0%
3	64	16,384	23.17	23.57	24.51	+4.0%
3	64	65,536	67.34	67.17	68.94	+2.6%
3	64	262,144	224.53	223.81	224.43	+0.3%
3	64	1,048,342	817.27	815.11	828.59	+1.7%
4	1	4,096	5.50	5.79	6.10	+5.2%
4	1	16,384	5.73	6.10	6.34	+3.9%
4	1	65,536	7.23	7.49	7.90	+5.6%
4	1	262,144	10.74	10.94	11.52	+5.3%
4	1	1,048,342	23.60	23.70	24.86	+4.9%
4	16	4,096	7.39	7.74	8.54	+10.3%
4	16	16,384	11.14	11.33	12.25	+8.2%
4	16	65,536	24.37	24.22	25.55	+5.5%
4	16	262,144	69.74	69.07	71.46	+3.5%
4	16	1,048,342	232.67	232.94	224.27	-3.7%
4	64	4,096	11.30	11.30	12.70	+12.5%
4	64	16,384	24.26	24.54	25.70	+4.7%
4	64	65,536	69.25	69.74	71.34	+2.3%
4	64	262,144	231.20	229.62	222.11	-3.3%
4	64	1,048,342	816.37	850.20	853.91	+0.4%
5	1	4,096	6.08	6.05	7.33	+21.2%
5	1	16,384	6.34	6.46	7.52	+16.3%
5	1	65,536	8.16	8.19	8.70	+6.2%
5	1	262,144	12.26	12.26	12.67	+3.4%
5	1	1,048,342	28.51	26.03	26.72	+2.6%
5	16	4,096	9.36	8.34	10.08	+20.9%
5	16	16,384	13.73	13.07	14.82	+13.3%
5	16	65,536	29.63	27.06	29.46	+8.9%
5	16	262,144	80.70	72.27	76.10	+5.3%
5	16	1,048,342	274.53	253.83	262.27	+3.3%
5	64	4,096	13.47	12.91	14.94	+15.7%
5	64	16,384	29.12	27.10	29.94	+10.4%
5	64	65,536	78.61	72.38	75.82	+4.8%
5	64	262,144	276.15	255.33	263.73	+3.3%
5	64	1,048,342	1136.20	1046.84	1087.25	+3.9%
6	1	4,096	6.24	6.37	7.50	+17.8%
6	1	16,384	6.49	6.75	7.65	+13.3%
6	1	65,536	8.45	8.43	8.93	+5.9%
6	1	262,144	12.80	13.36	13.47	+0.8%
6	1	1,048,342	29.73	28.48	27.65	-2.9%
6	16	4,096	9.79	8.83	10.72	+21.4%
6	16	16,384	14.14	13.81	15.74	+14.0%
6	16	65,536	31.23	30.24	32.13	+6.2%
6	16	262,144	86.98	85.12	84.42	-0.8%
6	16	1,048,342	305.51	304.39	295.04	-3.1%
6	64	4,096	14.34	14.06	15.65	+11.3%
6	64	16,384	30.46	30.67	32.00	+4.3%
6	64	65,536	86.35	85.73	84.56	-1.4%
6	64	262,144	304.00	305.14	295.31	-3.2%
6	64	1,048,342	1270.28	1306.63	1254.07	-4.0%

The full cold SGLang chain has geometric means of 1.549× against default FP32 SGLang and 1.249× against the faster full/sliced DeepSelect chain. The first comparison includes producer precision/arithmetic and KV-sharing changes. The companion SGLang PR contains every cold chain result and the slower cases.

Cold A/A controls duplicate both the default SGLang and full-width DeepSelect chains with separate buffers/graphs and identical flush treatment. Default SGLang A/A: geometric mean 1.0024, range 0.9502–1.0491; DeepSelect A/A: 0.9988, range 0.9705–1.0348. Near-equal timing ratios are numerical results, not significance claims.

dev already includes the sparse MQA APIs used by SGLang's SparseTable path; this PR adds the BF16 dense producer. Producer and chain timings do not establish full-model serving throughput or model quality.

Add a separate SM100 API for DeepSeek-V4.1 LiteTopK decode: MXFP4 paged MQA logits with BF16 head weights and BF16 output, optionally counting every live non-NaN output score once into a 1024-bin coarse histogram. Adding the histogram does not change the logits. Existing MQA APIs and their FP32 scores are unchanged. - fp4_paged_mqa_logits_bf16(q, kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, indices=None, histogram=None, tokens_per_request=1): Q packed MXFP4 [rows, 1, 32, 64] + UE8M0 scales, KV pages [pages, 128, 1, 68], BF16 weights [rows, 32]. Adjacent equal request IDs share KV loads. - get_paged_mqa_logits_bf16_metadata(context_lens, block_kv=128, num_sms, indices, tokens_per_request): split-384 schedule balanced by estimated split cost, up to 16384 rows; not interchangeable with the existing split-256 metadata. - tokens_per_request is a CPU hint in 1..6: 1..4 select Q4 with three TMEM stages, 5/6 select Q6 with two; 6 also swizzles the histogram's shared memory. - histogram: contiguous int32 [rows, 1024]; counts are added, bins run in descending score order and match the paired SGLang BF16 selector. - Both APIs are exposed through pybind and tvm-ffi (deep_gemm and sgl_deep_gemm). Kernels live in deep_gemm/include/deep_gemm/mqa_bf16/, the contract in docs/mqa_bf16.md, and the test in sgl_deep_gemm/tests/test_mqa_bf16.py (added to run_tests.sh). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Port sgl-project#96 with split384 metadata and Q4/Q6 routing for one through six query tokens. Signed-off-by: Heisenberg-Yin <ziqi003@e.ntu.edu.sg>

…stograms, next-n 1..6 Port the BF16 paged MQA producer of deepseek-ai#462 (merge 057ca59: MXFP4 Q/KV, BF16 head weights, BF16 head reduction and BF16 logits) as an isolated `mqa_bf16` kernel with its own split-384 cost-balanced schedule, and let it emit the same coarse histogram as the FP32 path. - APIs: `get_paged_mqa_logits_bf16_metadata(context_lens, 128, num_sms, indices=, tokens_per_request=)` and `fp4_paged_mqa_logits_bf16(..., indices=, histogram=, tokens_per_request=)` (pybind, TVM-FFI, deep_gemm / sgl_deep_gemm). - Adjacent equal request IDs share KV loads; the CPU hint selects Q4 / 3 TMEM stages (1..4) or Q6 / 2 stages (5..6; 6 also XOR-swizzles the histogram). Q4 keeps deepseek-ai#462's static TMEM stage mapping (its no-histogram kernel is SASS-identical to deepseek-ai#462); Q6 rotates a two-stage TMEM ring. - Exact histogram: every live non-NaN stored BF16 score counted once per row (16-bit shared counters for a window of rows, packed fast path for 2^-14 <= |x| < 16, published with packed 64-bit global atomics). The histogram kernels' scheduler caps a task's splits so the 16-bit overflow guard runs once per task; Q6 counts a 12-row window with 6 KV stages, so a CTA moving to the next request does not publish mid-kernel. Logits are bit-identical with and without the histogram. - Shared with the FP32 path: the coarse-bin encoding (epilogue/ coarse_histogram.cuh, plus a BF16 overload) and NoHistogram; dev's packing, ring pipeline (new parity-phase option for steps longer than a 2-stage ring), math / PTX / UMMA descriptor helpers. Only the kernel core, its shared-memory layout, its scheduler and the histogram accumulator are BF16-specific. - Test: sgl_deep_gemm/tests/test_mqa_bf16.py (n = 1..6, PDL off / on, exact counts, accumulation, CUDA-graph replay with changing lengths and request IDs). Replaces sgl-project#96.

Heisenberg-Yin · 2026-10-09T09:41:54Z

@Fridge003 @DarkSharpness could you take a look when you are free?

Thanks to @TKONIY for the collaboration on this work.

Motivation

The companion SGLang#42351 uses a score histogram to select the exact top-512 index-K slots for DeepSeek-V4.1 decode and target verify. This PR ports the upstream BF16 paged MQA producer into sgl-project dev and adds an optional exact histogram of its stored scores, including next-n 5/6 KV reuse.

The companion SGLang PR is an incremental change after SGLang #40764. The existing FP32 module stays in #40764; BF16 producer kernels, bindings and producer tests stay here. SGLang adds the BF16 selector and serving integration and requires this PR's BF16 APIs.

Code sources and diff scope

The BF16 scoring arithmetic and device core are adapted from deepseek-ai/DeepGEMM#462, merge commit 057ca5964aae0879ff2e0eb71ee05a3cb0ba3df7. deepseek-ai#462 is both an implementation source and the independent numerical/performance reference. BF16 output and BF16 head reduction originate upstream.

Local file under mqa_bf16/ Upstream source at 057ca5964aae Adaptation
kernel.cuh impls/sm100_mqa_logits.cuh Isolated core, histogram hooks, Q4/Q6 geometry
scheduler.cuh scheduler/sm100_paged_mqa_logits.cuh Cost-balanced schedule and request grouping
layout.cuh layout/mqa_logits.cuh Isolated namespace; separate math-group/TMEM-stage barrier geometry
packing.cuh common/packing.cuh Namespace change
ring_pipeline.cuh common/ring_pipeline.cuh Namespace change and phase handling for the two-stage ring
helpers.cuh Relevant upstream math/PTX helpers Local compatibility helpers for the dev baseline
The coarse-bin encoding follows the earlier sgl-project/DeepGEMM#94 histogram design. BF16 counter packing, histogram integration, Q6 routing, swizzle, scheduling changes and dev bindings are the work added here. The existing generic MQA APIs keep their implementation and default FP32 behavior.

This PR targets dev at ec226bcf400c11d62755137b916fae9ee359c684. Its 17 files, +1,827 / −1 include 1,372 device-header lines in a separate deep_gemm/mqa_bf16 directory, so adapted upstream code is visible as newly added files. It is a complete isolated port plus histogram extension, rather than a small diff that only adds histogram atomics to the generic dev kernel.

Modifications

Add get_paged_mqa_logits_bf16_metadata and fp4_paged_mqa_logits_bf16 through both pybind and TVM-FFI wrappers (deep_gemm / sgl_deep_gemm).

Count every live non-NaN stored BF16 score once into an additive int32 [R,1024] histogram. Histogram-enabled and histogram-disabled calls return bit-identical live scores. Counts are per row and per score, with no pooling or block-max substitution.

CPU tokens_per_request: 1–4 selects Q4 / 3 TMEM stages; 5–6 selects Q6 / 2 stages, sharing KV loads among adjacent query rows of a request; hint 6 enables shared-memory histogram XOR swizzle.

Use an isolated split-384 schedule balanced by estimated work. Schedule generation runs with live request IDs/lengths during CUDA graph replay.

Add the API contract in docs/mqa_bf16.md and GPU tests in sgl_deep_gemm/tests/test_mqa_bf16.py.
schedule = deep_gemm.get_paged_mqa_logits_bf16_metadata(
    context_lens, 128, deep_gemm.get_num_sms(),
    indices=request_ids, tokens_per_request=n,
)
histogram = torch.zeros((rows, 1024), dtype=torch.int32, device="cuda")
scores = deep_gemm.fp4_paged_mqa_logits_bf16(
    (q_fp4, q_sf), kv_cache, weights, context_lens, block_table,
    schedule, max_context_len,
    indices=request_ids, histogram=histogram, tokens_per_request=n,
)
Scope

SM100, H32/D128, MXFP4 Q [R,1,32,64] with packed int32 UE8M0 scales [R,1,32], paged uint8 KV [pages,128,1,68], BF16 weights [R,32], and BF16 output [R,max_context_len] with row stride aligned to 1536 elements. Context lengths are int32 [R,1]; page tables have one row per score row. CPU hint is 1–6, with at most 16,384 rows.

IDs identify adjacent request runs; page tables match and causal lengths are nondecreasing within each run. Omitting IDs treats rows independently. Schedule and producer use the same hint/IDs/lengths; this metadata is incompatible with the generic MQA schedule. Only live score columns are defined. Callers zero counts before an independent call, or use the companion selector which consumes and clears them.

BF16 scoring follows upstream deepseek-ai#462 semantics and can yield a different top-k from the default FP32 producer. Histogram counts ignore NaNs; the companion selector's nonfinite behavior is documented in the SGLang PR.

Accuracy Tests

Fresh focused test run on the complete SGLang PR head: 55 passed — 15 dispatch/storage/metadata unit tests, 12 producer tests, and 28 BF16 selector/serving GPU tests. The GPU suite covers n=1–6, PDL off/on, exact counts, repeated accumulation, zero/page/split/top-k boundaries, shared and changing request IDs, dynamic CUDA graphs, FullTopK/SparseTable, storage lifetime, fallback and BF16 candidate maxima. The producer test is an unchanged copy of the repository test run outside the package directory; its source hash is recorded.

Every one of the 90 cold benchmark cells verifies live BF16 bits against the independent upstream deepseek-ai#462 producer, exact histograms against a per-score oracle, selection values against torch.topk, valid unique physical slots and reusable scratch. LiteTopK additionally matches (BF16 ordered score key descending, physical slot ascending) exactly. The default FP32 chain is checked against top-k of its own FP32 scores. Outputs are checked before and after timing.

The separate fresh qualification has 69 passing records: 4 large-batch cases up to 768 rows, 51 graph replays across page/split/top-k boundaries, 8 selector edge cases (including ties, signed zero, nonfinite scores and overflow), graph storage growth, and 5 actual FullTopK hook cases. Both live producer bindings come from the dev-targeted #96 build; the existing sparse APIs are used directly. The 90-cell cold matrix and 69-record qualification apply the live SGLang diff to local main snapshot a8b54f138e; the complete PR head additionally passes the same 55-test suite. All timed selector/top-k files match the complete head byte for byte. Three unrelated Hopper-path base changes are recorded in the source audit. Earlier measurements are retained separately and are not substituted for these reruns.

Speed Tests and Profiling

These fresh measurements use the dev-targeted TVM-FFI build from this PR, with all device headers verified against its live diff. Upstream deepseek-ai#462 is the BF16 reference. The following numbers are CUPTI producer kernel durations from the randomized measurements, not isolated histogram execution costs. Reference BF16 runs in a DeepSelect chain, final no-hist runs alone, and final histogram runs in the LiteTopK chain. Mapping, selection and metadata are not separate summands in these producer durations, but their scheduling and resource effects can affect the durations.

Each timed cold replay contains one call, with cache eviction before the call. The producer-duration change is an observed cold measurement under these chain contexts. It is not an isolated instruction-cost estimate; a matched producer-only histogram/no-hist ablation with reset costs controlled would be needed for that. PDL is enabled in these measurements; its on/off performance contribution was not measured.

Cold-cache measurements use NVIDIA B200 (148 SMs), Torch 2.13.0+cu130, CUDA 13.0, and PDL enabled. Comparisons within each cell run on the same GPU; n=1–3 use GPU4 and n=4–6 GPU6. The matrix is n=1–6 × B={1,16,64} × L={4096,16384,65536,262144,1048342}: 90 cold cells. B counts requests, R=B×n score rows, and L compressed KV positions.

Inputs are captured DeepSeek-V4.1 layer-20 MXFP4 Q/KV and BF16-rounded head weights, extended to each shape with randomly permuted physical pages and each query row's causal length. The default FP32 producer receives one independent page-table/length row per query, as the serving path does; the paired BF16 producer can reuse KV among adjacent equal request IDs. L is the live context, while score allocation retains the captured ~1M-column maximum window; only live columns are defined.

Each cell uses 80 randomized paired blocks. A 1 GiB buffer is zeroed before every single-call CUDA graph replay to evict the preceding cache contents. The flush is outside the timed span; CUPTI measures each chain from its first kernel start to its last kernel end. Outputs and scratch are reused. Physical-slot mapping, histogram consumption and scratch cleanup are included. Q projection/packing, schedule/plan construction, serving weight conversion and CPU launch overhead are outside this matrix.

Clocks and power settings are unchanged. Idle gates and per-process monitoring reject affected cells with a 10 s / 6 s margin, including their preparation and profiling window. Raw samples, kernel launch signatures, clocks, rejected attempts and source/binary hashes accompany the results.

Geometric means over 15 cold shapes per n. Ratios compare observed producer kernel durations in their respective chains; the last column is the histogram-chain producer duration divided by standalone no-hist duration, minus one. Use the complete cold chain measurements for the latency/speedup conclusion.

Mode n Upstream / no-hist duration Upstream / hist-chain producer duration Producer duration change (different chains)
cold 1 0.974× 0.958× +1.8%
cold 2 0.976× 0.950× +2.7%
cold 3 0.977× 0.936× +4.3%
cold 4 0.982× 0.942× +4.3%
cold 5 1.060× 0.971× +9.1%
cold 6 1.008× 0.960× +5.0%
cold: all 90 producer cells
n B L Upstream BF16 Final no-hist alone Final hist in selector chain Producer duration change
1 1 4,096 4.80 5.12 5.38 +5.0%
1 1 16,384 5.15 5.49 5.73 +4.4%
1 1 65,536 6.27 6.54 6.66 +1.7%
1 1 262,144 9.26 9.66 9.74 +0.8%
1 1 1,048,342 21.31 21.76 21.82 +0.3%
1 16 4,096 6.50 6.85 7.26 +6.1%
1 16 16,384 9.71 9.95 10.22 +2.7%
1 16 65,536 21.63 21.86 22.19 +1.5%
1 16 262,144 62.93 63.49 63.28 -0.3%
1 16 1,048,342 204.03 204.66 202.72 -0.9%
1 64 4,096 9.65 10.03 10.34 +3.0%
1 64 16,384 21.65 21.95 22.22 +1.2%
1 64 65,536 63.49 63.78 64.00 +0.4%
1 64 262,144 203.59 203.44 204.72 +0.6%
1 64 1,048,342 752.93 753.09 753.99 +0.1%
2 1 4,096 5.15 5.46 5.73 +5.0%
2 1 16,384 5.46 5.73 5.89 +2.8%
2 1 65,536 6.62 6.98 7.07 +1.4%
2 1 262,144 9.63 9.98 10.21 +2.2%
2 1 1,048,342 21.76 22.08 22.59 +2.3%
2 16 4,096 6.91 7.30 7.71 +5.7%
2 16 16,384 9.89 10.30 10.78 +4.7%
2 16 65,536 22.11 22.37 22.94 +2.6%
2 16 262,144 64.37 64.62 65.90 +2.0%
2 16 1,048,342 205.81 204.78 209.12 +2.1%
2 64 4,096 10.08 10.38 10.75 +3.5%
2 64 16,384 22.42 22.69 23.41 +3.2%
2 64 65,536 64.34 64.86 66.11 +1.9%
2 64 262,144 210.98 211.82 210.96 -0.4%
2 64 1,048,342 779.97 779.01 788.95 +1.3%
3 1 4,096 5.30 5.60 5.95 +6.3%
3 1 16,384 5.70 5.89 6.27 +6.5%
3 1 65,536 6.91 7.26 7.55 +4.0%
3 1 262,144 10.05 10.34 10.78 +4.3%
3 1 1,048,342 22.77 23.38 23.95 +2.5%
3 16 4,096 7.39 7.68 8.42 +9.6%
3 16 16,384 10.40 11.11 11.73 +5.6%
3 16 65,536 23.10 23.62 24.42 +3.4%
3 16 262,144 66.08 66.35 68.10 +2.6%
3 16 1,048,342 220.35 220.80 224.61 +1.7%
3 64 4,096 10.62 10.83 11.92 +10.0%
3 64 16,384 23.17 23.57 24.51 +4.0%
3 64 65,536 67.34 67.17 68.94 +2.6%
3 64 262,144 224.53 223.81 224.43 +0.3%
3 64 1,048,342 817.27 815.11 828.59 +1.7%
4 1 4,096 5.50 5.79 6.10 +5.2%
4 1 16,384 5.73 6.10 6.34 +3.9%
4 1 65,536 7.23 7.49 7.90 +5.6%
4 1 262,144 10.74 10.94 11.52 +5.3%
4 1 1,048,342 23.60 23.70 24.86 +4.9%
4 16 4,096 7.39 7.74 8.54 +10.3%
4 16 16,384 11.14 11.33 12.25 +8.2%
4 16 65,536 24.37 24.22 25.55 +5.5%
4 16 262,144 69.74 69.07 71.46 +3.5%
4 16 1,048,342 232.67 232.94 224.27 -3.7%
4 64 4,096 11.30 11.30 12.70 +12.5%
4 64 16,384 24.26 24.54 25.70 +4.7%
4 64 65,536 69.25 69.74 71.34 +2.3%
4 64 262,144 231.20 229.62 222.11 -3.3%
4 64 1,048,342 816.37 850.20 853.91 +0.4%
5 1 4,096 6.08 6.05 7.33 +21.2%
5 1 16,384 6.34 6.46 7.52 +16.3%
5 1 65,536 8.16 8.19 8.70 +6.2%
5 1 262,144 12.26 12.26 12.67 +3.4%
5 1 1,048,342 28.51 26.03 26.72 +2.6%
5 16 4,096 9.36 8.34 10.08 +20.9%
5 16 16,384 13.73 13.07 14.82 +13.3%
5 16 65,536 29.63 27.06 29.46 +8.9%
5 16 262,144 80.70 72.27 76.10 +5.3%
5 16 1,048,342 274.53 253.83 262.27 +3.3%
5 64 4,096 13.47 12.91 14.94 +15.7%
5 64 16,384 29.12 27.10 29.94 +10.4%
5 64 65,536 78.61 72.38 75.82 +4.8%
5 64 262,144 276.15 255.33 263.73 +3.3%
5 64 1,048,342 1136.20 1046.84 1087.25 +3.9%
6 1 4,096 6.24 6.37 7.50 +17.8%
6 1 16,384 6.49 6.75 7.65 +13.3%
6 1 65,536 8.45 8.43 8.93 +5.9%
6 1 262,144 12.80 13.36 13.47 +0.8%
6 1 1,048,342 29.73 28.48 27.65 -2.9%
6 16 4,096 9.79 8.83 10.72 +21.4%
6 16 16,384 14.14 13.81 15.74 +14.0%
6 16 65,536 31.23 30.24 32.13 +6.2%
6 16 262,144 86.98 85.12 84.42 -0.8%
6 16 1,048,342 305.51 304.39 295.04 -3.1%
6 64 4,096 14.34 14.06 15.65 +11.3%
6 64 16,384 30.46 30.67 32.00 +4.3%
6 64 65,536 86.35 85.73 84.56 -1.4%
6 64 262,144 304.00 305.14 295.31 -3.2%
6 64 1,048,342 1270.28 1306.63 1254.07 -4.0%
The full cold SGLang chain has geometric means of 1.549× against default FP32 SGLang and 1.249× against the faster full/sliced DeepSelect chain. The first comparison includes producer precision/arithmetic and KV-sharing changes. The companion SGLang PR contains every cold chain result and the slower cases.

Cold A/A controls duplicate both the default SGLang and full-width DeepSelect chains with separate buffers/graphs and identical flush treatment. Default SGLang A/A: geometric mean 1.0024, range 0.9502–1.0491; DeepSelect A/A: 0.9988, range 0.9705–1.0348. Near-equal timing ratios are numerical results, not significance claims.

dev already includes the sparse MQA APIs used by SGLang's SparseTable path; this PR adds the BF16 dense producer. Producer and chain timings do not establish full-model serving throughput or model quality.

Closing in favor of #94, which now includes this BF16 DeepSeek-V4.1 producer alongside the FP32 histogram, with the shared bin encoding and helpers deduplicated and a faster histogram path. Please review there.

…stograms, next-n 1..6 Port the BF16 paged MQA producer of deepseek-ai#462 (merge 057ca59: MXFP4 Q/KV, BF16 head weights, BF16 head reduction and BF16 logits) as an isolated `mqa_bf16` kernel with its own split-384 cost-balanced schedule, and let it emit the same coarse histogram as the FP32 path. - APIs: `get_paged_mqa_logits_bf16_metadata(context_lens, 128, num_sms, indices=, tokens_per_request=)` and `fp4_paged_mqa_logits_bf16(..., indices=, histogram=, tokens_per_request=)`: TORCH_LIBRARY ops in csrc/torch_library_api.cpp (`Tensor(a!)? histogram`, no schema defaults), deep_gemm / sgl_deep_gemm wrappers that own the defaults, and pybind. - Adjacent equal request IDs share KV loads; the CPU hint selects Q4 / 3 TMEM stages (1..4) or Q6 / 2 stages (5..6; 6 also XOR-swizzles the histogram). Q4 keeps deepseek-ai#462's static TMEM stage mapping (its no-histogram kernel is SASS-identical to deepseek-ai#462); Q6 rotates a two-stage TMEM ring. Each stage count is the deepest ring that fits (TMEM columns, SM100 shared memory with the histogram window); split 384 and the 3 x 128 math threads are named `Config` constants. The schedule (one CTA, two int32 per row in shared memory) checks its documented limit of 16384 rows. - Exact histogram: every live non-NaN stored BF16 score counted once per row (16-bit shared counters for a window of rows, packed fast path for 2^-14 <= |x| < 16, published with packed 64-bit global atomics). The histogram kernels' scheduler caps a task's splits so the 16-bit overflow guard runs once per task; Q6 counts a 12-row window with 6 KV stages, so a CTA moving to the next request does not publish mid-kernel. Logits are bit-identical with and without the histogram. - Shared with the FP32 path: the coarse-bin encoding (epilogue/ coarse_histogram.cuh, plus a BF16 overload) and NoHistogram; dev's packing, ring pipeline (new parity-phase option for steps longer than a 2-stage ring), math / PTX / UMMA descriptor helpers. Only the kernel core, its shared-memory layout, its scheduler and the histogram accumulator are BF16-specific. - Test: sgl_deep_gemm/tests/test_mqa_bf16.py, a plain script that run_tests.sh runs on SM100 (n = 1..6, PDL off / on, exact counts, accumulation, CUDA-graph replay with changing lengths and request IDs; the schedule's row limit). Replaces sgl-project#96. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Heisenberg-Yin mentioned this pull request Oct 9, 2026

[SM100] Paged MQA logits: optional exact coarse score histogram for LiteTopK decode #94

Open

Heisenberg-Yin closed this Oct 9, 2026

Navigation Menu

Repository navigation

[SM100] Add exact BF16 paged MQA histograms for next-n 1–6 - #96

[SM100] Add exact BF16 paged MQA histograms for next-n 1–6#96
Heisenberg-Yin wants to merge 1 commit into
sgl-project:devfrom
Heisenberg-Yin:dsv41-litetopk-mqa-bf16-dev

Heisenberg-Yin commented Oct 3, 2026 •

edited

Loading

Uh oh!

Heisenberg-Yin commented Oct 9, 2026

Motivation

Code sources and diff scope

Modifications

Scope

Accuracy Tests

Speed Tests and Profiling

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

Heisenberg-Yin commented Oct 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Code sources and diff scope

Modifications

Scope

Accuracy Tests

Speed Tests and Profiling

Uh oh!

Heisenberg-Yin commented Oct 9, 2026

Motivation

Code sources and diff scope

Modifications

Scope

Accuracy Tests

Speed Tests and Profiling

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Heisenberg-Yin commented Oct 3, 2026 •

edited

Loading