[Heisenberg-Yin]

@DarkSharpness could you take another look? This revision is rebuilt on dev 7c4a95f. After deepseek-ai#462 / deepseek-ai#471, SM100 paged MQA logits are BF16-only, so the previous design (FP32 histogram, token Q tile, a separate BF16 producer) has nothing left to attach to. The histogram is now an optional output of dev's own kernel. It is the producer side of LiteTopK decode (sgl-project/sglang#40764).

Motivation

In DSA and DeepSeek-V4.1 decode, the top-k after the paged MQA logits starts with a pass that only builds a coarse histogram of each row. The logits kernel already holds every score in registers, so it can count that histogram while it writes the logits, and the SGLang selector skips the pass.

What this adds

• API: fp8_fp4_paged_mqa_logits(..., indices=None, histogram=None). histogram is an optional int32 [rows, 1024] tensor (contiguous, 8-byte aligned).

  • The kernel adds to it: for every stored logit (column < context_len, not NaN) it increments one bin. Callers zero the tensor first, for example with a memset node in the same CUDA graph.
  • Bins run in descending order: FP16-RN exponent + 4 mantissa bits for |x| < 16, unit bins up to 223, a saturating bin at each end, and -0 counted as +0.
  • Logits are bit-identical with and without the histogram.

• Coverage: SM100, MXFP8 and MXFP4, 32 or 64 heads, pages of 32 / 64 / 128, any number of rows per request. SM90 and other head counts are rejected with a clear error.

• Without a histogram nothing changes. The generated source and SASS of all 49 compared kernels equal dev's. The only difference is that the parameter block grows by 8 bytes.

• epilogue/coarse_histogram.cuh (new):

  • The bin encoding.
  • RowWindowHistogram: 16-bit counters in shared memory for a window of two Q blocks. They are flushed with packed 64-bit global adds when a task leaves the window, after 170 splits, and at exit.
  • Scores are binned two per register on a warp-uniform path (all fine, all unit, or mixed). Anything else takes an exact scalar path.

• Kernel, scheduler and host:

  • sm100_mqa_logits.cuh gains if constexpr hooks in the shared core.
  • The paged scheduler caps a task at 170 splits, only when the histogram is on, so the 16-bit counters cannot wrap.
  • MXFP4 at head dim 128 keeps 7 KV stages instead of 8 to make room for the window. MXFP8 is unchanged.

• Schema:

  fp8_fp4_paged_mqa_logits(Tensor q, Tensor? q_sf, Tensor fused_kv_cache, Tensor weights, Tensor context_lens,
      Tensor block_table, Tensor schedule_meta, int max_context_len, Tensor? indices, Tensor(a!)? histogram) -> Tensor
  

One commit: 12 files, +617 / -34. The previous revision had 19 files, +2150 / -44.

Changes since the previous revision (b8648e3 on dev b296ce3)

• Rebuilt on dev 7c4a95f. Removed:
  • the FP32 histogram and the token Q tile (fp8_paged_mqa_logits and FP32 logits no longer exist);
  • the separate BF16 producer and tokens_per_request (dev's kernel is that kernel now);
  • the 16384-row assert.
• Review comments:
  • The flush is now a single flush in coarse_histogram.cuh. It uses plain comparisons and adds bin pairs through pack_bin_pair.
  • There is now one histogram type, RowWindowHistogram, plus the disabled default NoHistogram.
• Not included: the previous BF16 producer's 6-token Q tile. It will be a separate PR; its cost is shown under "Producer cost" below.

Validation (B200, CUDA 13.0, torch 2.14.1)

• test_paged_mqa_logits_histogram (new, also in run_tests.sh's release sanitizer list): 128 / 128 cases, edge lengths 0-131,071 with 1-6 rows per request (PDL off and on), forced SM counts 1 / 7 / 100, and CUDA graph replay.
• test_attention.py, release profile: prefill 4 / 240, paged 4 / 900, histogram 4 / 128, sparse 8 / 161.
• compute-sanitizer memcheck and synccheck on the histogram test: 0 errors.
• SASS against dev:
  • 49 / 49 kernels without a histogram are identical (38 paged, 8 contiguous, 3 metadata);
  • all 30 histogram variants use 128 registers with no spills.
• Bin math, checked exhaustively: for all 65,536 BF16 values and all 2^32 packed pairs, the bins equal the FP32 definition.
• [Perf] [NVIDIA] attention: Add opt-in LiteTopK decode top-k for DSA and DeepSeek-V4.1 on SM100 sglang#40764: its kernel test passes 23 / 23 with this build.

Producer cost (real data)

Method: B200, cold (1 GiB flush, then a single-call CUDA graph replay). The same package runs with and without the histogram, interleaved. The A/A spread is 0.98-1.01.

GLM-5.2, index layers 0 and 38: MXFP8 with 32 heads, converted the way sgl-project/sglang#44350 does. 383 cells: B 1-128 × L 8K-1M × next_n 1-4.

• Geomean +2.28 %; median extra +0.42 µs. The maximum is +9.84 % (B 4, 8K), and no cell is above +10 %.
• The cost falls with batch size: +4.0-4.3 % at B 1, +0.5-0.6 % at B 128. At ~1M it is at most +0.6 %.

DeepSeek-V4.1, layer 20: MXFP4 with 32 heads. 142 cells, averaged over two runs.

tokens per request	1	2	3	4	5	6
geomean	+1.25 %	+4.18 %	+5.46 %	+5.38 %	+13.85 %	+14.54 %
median extra (µs)	+0.14	+0.64	+0.98	+0.69	+1.87	+2.66

• 5-6 tokens per request: the cost is highest here, up to +21 % at B 64, ~1M (that cell was power-capped).
• Against the previous revision's BF16 producer, both with the histogram:
  • at 1-4 tokens per request the two are on par (1.002-1.008×);
  • at 5 / 6 tokens, this revision is 1.135× / 1.114× slower, because it runs dev's 4-row Q tile where the old producer ran its 6-token tile;
  • the follow-up 6-token tile PR targets this;
  • the logits of the two producers are bit-identical in all 108 compared cells.

[Heisenberg-Yin]

@DarkSharpness could you take a look? This is the producer side for LiteTopK decode. The paged MQA logits kernel emits the coarse histogram that the SGLang top-k v2 streaming and cluster paths currently build in their first pass over the logits.

[DarkSharpness]

why not use normal count[i].x == 0 || count[i].y == 0 (or first pack into 1 unsigned long long)? this gives very poor readability.

[DarkSharpness]

it seems that there are 2 class with name CoarseHistogram? why?