[Performance][MLA] Use FP16 logits for sparse indexer - #52696
WoosukKwon wants to merge 11 commits into
Conversation
Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Add forced backend selection to the CUDA-graph benchmark so dispatch boundaries can be measured directly on captured model logits. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Process packed FP16 score bits directly, refine only the exact remaining key bits, and benchmark fixed-stride CUDA graph captures separately from live KV length. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
|
I see the duplicate-work check didn't catch this (i closed it a while ago), so I just wanted to reference the PR here just in case there’s anything useful you can get from it: #40811 |
|
@LopezCastroRoberto Oh thanks. I didn't know you tried the same idea. Here, I used FP16 instead of BF16 since I though precision matters more than range. Could you please let me know why you closed the previous PR? |
I used BF16 because DSv4 QAT apparently quantized the indexer scores to BF16, so I went with it for consistency: "We further quantize the index scores 𝐼:,: from FP32 to BF16 during this QAT process. This optimization achieves a 2× speedup for the top-k selector, while preserving a 99.7% recall rate of KV entries." Source: https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/blob/main/DeepSeek_V4.pdf
When I created the cooperative topK version (#43008), I wasn't getting significant speedups from the index scores quantization as I did in #40811 with topk v1 (persistent topK). So I just decided to close it and maybe revisit in the future. Great to see there was a way to do it :) I always trusted this was a nice optimization to apply, or at least to have as a configurable parameter |
Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
|
@LopezCastroRoberto Got it. Thanks for the explanation! Using FP16 gives me higher recall (99.96% median) than BF16. Given the low probability of having more than 2,048 logits exceed |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude Code Review
This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.
Tip: disable this comment in your organization's Code Review settings.
|
/ci run |
|
✅ Triggered Buildkite CI #84371 for commit |
|
@LopezCastroRoberto Can you please review this PR? |
Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
|
Hi @WoosukKwon, the pre-commit checks have failed. Please run: uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-filesThen, commit the changes and push to your branch. For future commits, |
|
/ci run |
|
✅ Triggered Buildkite CI #84421 for commit |
|
This pull request has merge conflicts that must be resolved before it can be |
Resolve the sparse indexer test import conflict by retaining both FP16 dtype coverage and upstream MQA-only coverage. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
|
/ci run |
|
@LopezCastroRoberto Can you please review? |
|
✅ Triggered Buildkite CI #84751 for commit |
There was a problem hiding this comment.
@WoosukKwon I feel like we should verify at least:
(1) FP32 performance-regression coverage for every modified path, i.e., top_k_per_row_prefill, top_k_per_row_prefill, cooperative_topk, persistent_topk, and filtered_topk. Just to make sure the code changes do not introduce any unexpected overhead in the FP32 existing path.
(2) FP16-versus-FP32 prefill benchmarks. The reported performance tables cover decode topk, but it is unclear to me if the FP16 prefill code has been benchmarked and what's the speedup obtained.
(3) Since the PR makes FP16 logits the default, I think we should make sure this doesn't affect real accuracy numbers. This is probably correct for DSv4, since they claimed to have quantized the logits during QAT to BF16. But it is unclear to me if this is safe for other models using DSA, such as DSv3.2. MRCR accy evals for different context lengths can be good for this this, imo. See EleutherAI/lm-evaluation-harness#3754.
(4, maybe optional) Matched end-to-end FP32-versus-FP16 serving benchmarks. The current e2e results validate retrieval accuracy, but do not demonstrate an e2e performance benefit. I think it is good to check both TTFT and TPOT look healthy.
| def _benchmark_graph( | ||
| launch: Callable[[], None], | ||
| graph_calls: int, | ||
| replays: int, | ||
| samples: int, | ||
| ) -> float: |
There was a problem hiding this comment.
nit: I generally think we should use triton.testing.do_bench_cudagraph instead of CUDAGraph+events to measure kernel-wise performance, e.g., this. Not only because it is better for time measurement, but it also reduces the LoC significantly.
| } else { | ||
| launch_cooperative_cluster<TopK, 4>(params, ct::kSmemSize4, stream); | ||
| launch_cooperative_cluster<InputType, TopK, 4>( | ||
| params, ct::kSmemSize4<InputType>, stream); |
There was a problem hiding this comment.
I wonder if this heuristic holds for BF16 input scores, i.e., CS=16 if num_rows <=4, CS=8 if num_rows<=8, otherwise CS=4.
Have you checked this? Maybe leave a TODO if you think this is out of the scope of this PR.
| const bool is_half = | ||
| logits.scalar_type() == torch::headeronly::ScalarType::Half; | ||
| if (is_half) { | ||
| STD_TORCH_CHECK(max_smem_per_block >= 128 * 1024, |
There was a problem hiding this comment.
I think this is too strict for the half path? Seems like half of this workspace is needed, right?
constexpr int kBuffers =
FilteredTopKTraits<DType>::NUM_REFINE_ROUNDS == 1 ? 1 : 2;
return sizeof(int) * kBuffers * 16 * 1024;
i.e., 4 bytes × 1 buffer × 16,384 entries = 65,536 bytes = 64 KiB
| "128KB of shared memory per block"); | ||
| } | ||
|
|
||
| if ((num_rows > 32 || is_half) && max_smem_per_block >= 128 * 1024) { |
There was a problem hiding this comment.
Why is this changing the logic? This always skips persistent_topk. Just making sure this is on purpose.
I feel that the logic to select cooperative/persistent/filtered/top_k_per_row_decode is a bit difficult to follow now.
| if logits_dtype != torch.float16: | ||
| return False | ||
| if max_seq_len <= 32768: | ||
| return False | ||
| if max_seq_len <= 65536: | ||
| return num_rows >= 768 | ||
| if max_seq_len <= 100000: | ||
| return num_rows >= 512 | ||
| return max_seq_len <= 131072 and num_rows >= 1024 |
There was a problem hiding this comment.
Should we re-run this on Hopper to refine this heuristic? I think we should at least refer to this PR in a comment to justify where this numbers come from.
| and topk_tokens in (512, 1024, 2048) | ||
| and num_rows <= 32 | ||
| and logits.stride(0) % 4 == 0 # TMA 16-byte alignment | ||
| and logits.stride(0) % (16 // logits.element_size()) == 0 |
| parser.add_argument("--samples", type=int, default=7) | ||
| parser.add_argument( | ||
| "--backend", | ||
| choices=("auto", "cooperative", "persistent", "decode"), |
There was a problem hiding this comment.
Should we add top_k_per_row_prefill to this script? It also supports FP16 now.
Co-authored-by: OpenAI Codex <codex@openai.com>
|
This pull request has merge conflicts that must be resolved before it can be |
2 participants
Summary
with an explicit
float16/float32engine overridewithout widening the materialized logits tensor
native FP16 refinement, and retune persistent/decode crossover tiers
correctness, engine-policy, DeepGEMM, DCP, and odd-length coverage
warm-up and inference do not depend on a thread-local config context
paged MTP decode with exact packed-cache-layout correctness tests
Why
GLM-5.2's sparse indexer materializes a large logits tensor before selecting
top-k indices. Keeping the tensor and selection native FP16 reduces memory
traffic.
--attention-config.indexer_logits_dtypeacceptsauto,float16,or
float32;autochooses FP16 on CUDA and FP32 elsewhere.The final kernel transforms both packed halves with one 32-bit operation and
reuses the exact keys in histogram and scatter. This reduced the unified
FP16 kernel from 57 to 32 registers/thread at 1,024 rows and raised measured
occupancy from 49.7% to 94.9%.
The dtype is passed as a constant custom-op argument. Resolving it in the
operator body originally failed during real-model warm-up because the worker
callback does not run inside
set_current_vllm_config; resolving it when theSparseAttnIndexermodule is constructed fixes eager warm-up and graph replay.DeepSeek-V4 Flash's optional MXFP4 indexer cache uses the same FP16 logits path.
The model-specific tests exercise 64 indexer heads, the packed 68-byte cache
layout, and
next_n=3(MTP=2), comparing both prefill and paged decode againstdequantized MXFP4 references.
Performance
All selector measurements use captured real-model GLM logits, top-k 2,048,
warmed CUDA graphs, and the repository
.venv. CUDA events bracket graphreplays and are not captured in the graph. Both the custom top-k and reference
torch.topkrun on CUDA; CPU transfers occur only in separate pytest failurediagnostics. Every result was checked against
torch.topkbefore timing.Production-style fixed 256K decode graph
The logits row stride and captured
max_seq_lenremain 256K while only thelive KV length changes, matching reuse of a maximum-length graph. Entries are
FP16 us / FP32 us (FP32 / FP16 speedup):FP16 is faster in all 25 production-style shapes: 1.042x-1.733x.
Exact-stride matrix
FP16 speedup over FP32 after retuning the selector:
FP16 is faster in all 28 exact-stride shapes: 1.048x-1.839x. The packed-key
change invalidated the old decode crossover: persistent now wins through 32K,
and measured row thresholds are 768 for 32K-64K, 512 for 64K-100K, and 1,024
for 100K-131K.
Accuracy
End-to-end long-context recall
The final PR code was evaluated with the real
nvidia/GLM-5.2-NVFP4checkpoint and the standard Paul Graham NIAH corpus. The paired runs used TP4,
greedy decoding, FP8 E4M3 KV cache, 2,048-token chunked prefill, one sequence
at a time, eager execution, disabled prefix caching, disabled FlashInfer
autotuning, and no GVR. The only changed engine setting was
--attention-config.indexer_logits_dtype=float32versusfloat16.Each context used needle depths 10%, 50%, and 90%, with two independently
generated seven-digit keys per cell:
Server-reported prompt lengths were 50,000-50,002, 100,001-100,002,
200,002, and 250,001-250,002 tokens. This is deterministic single-needle
retrieval coverage, not a claim that every downstream task is unchanged.
DeepSeek-V4 Flash FP4 indexer-cache NIAH
The real
deepseek-ai/DeepSeek-V4-Flashcheckpoint was evaluated with TP4,MTP=2, FP8 main KV cache, FP16 indexer logits, eager execution, disabled
prefix caching, and disabled FlashInfer autotuning. Both arms used the same
ordinary fused-MoE backend because the local mega-MoE extension has a stale
ABI. The only paired change was an FP8 indexer cache versus
use_fp4_indexer_cache=true, which startup resolved tomxfp4andconfirmed with
Using MXFP4 indexer cache for Lightning Indexer.The standard Paul Graham NIAH corpus used one deterministic seven-digit key
per cell and greedy decoding:
Needle depths were 10%, 50%, and 90%. Server prompt-token counts matched
between paired arms for all 12 cells and ranged from 49,983 to 249,984, so no
prompt was truncated. This is deterministic single-needle retrieval coverage,
not a claim that every downstream task is unchanged. Request latency is not
reported because the control arm absorbed first-use JIT compilation.
Real-input top-k set recall
To isolate selector-set changes, FP32 logits captured from real GLM decode
for BEAM documents were rounded to FP16 and exact top-2,048 was recomputed.
Recall is
|topk(FP16) intersect topk(FP32)| / 2048.Two layer-74 snapshots with 1,024 real rows each and 199,401-199,481 live KV
tokens produced:
A separate context sweep used two selector positions and 64 real rows per
length:
This recall measurement isolates FP16 score rounding on real distributions;
the end-to-end NIAH table above measures the actual native FP16 DeepGEMM and
top-k execution path.
Validation
cmake --build build/gvr_native_fp16 --target _C_stable_libtorch -j16settings after moving dtype resolution to module construction
torch.topkat 10K, 50K, 100K,200K, and 256K with a fixed 256K stride
correctness before CUDA-graph timing
all 65,536 FP16 encodings
with
use_fp4_indexer_cache=true; all 12 paired responses were identicalMXFP4-to-FP16 prefill and paged
next_n=3decode against dequantizedreferences; 2 full DCP FP16/FP32 decode-parity cases passed earlier
SWEEP_MAX_CONCURRENCY=1,8,32,256 SWEEP_NUM_PROMPTS_FACTOR=3 bash ../glm_bench.shand a matched concurrency-1,024 factor-1 run completed1,915 requests per dtype with no failures
uv pip install -e . --torch-backend=autowas attempted. The full editablebuild is currently blocked on unmodified current-main QUTLASS code including
ATen extension headers while
TORCH_TARGET_VERSIONis defined; both changedCUDA targets build successfully.
Duplicate-work check
This does not duplicate #52149, which fixes persistent top-k candidate
overflow, or #48726, which adds an opt-in fused DSA indexer/top-k path. This PR
changes dtype production and native selection for the existing materialized-
logits path, including DeepGEMM and DCP.
AI assistance
AI assistance was used to implement and benchmark this change. Before this
draft is marked ready, the submitting human must review every changed line,
reproduce the relevant validation, and be prepared to explain and defend the
change end-to-end.