Two fixes for the heuristic TopK decode path:

1. Add +1 offset to pre_idx in _gather_prev_topk_for_decode: the saved
   TopK came from a query at position P (last MTP pos of previous step),
   while the current step's first query is at P+1. Shifting by +1
   preserves relative distances under RoPE for a more accurate initial
   threshold in the heuristic kernel.

2. Propagate enable_heuristic_topk through model_config.py: the
   DeepSeekV32 config builder was reconstructing DeepSeekSparseAttention
   Config without forwarding enable_heuristic_topk from the user's
   sparse_attention_config, causing the field to always default to False.

Made-with: claude-4.6-opus-high
Signed-off-by: longcheng-nv <[email protected]>

…uffers

Redesign the heuristic TopK pre_idx management to be fully CUDA Graph
compatible. Replace the old Indexer-side dict/slot approach (which used
.item() and torch.tensor() forbidden during graph capture) with
pre-allocated metadata buffers using a feedback loop pattern.

Key changes:
- Per-layer 3D buffer heuristic_prev_topk on metadata, allocated via
  get_empty(capture_graph=True) for stable addresses across replays.
- Shared staging buffer for +1 RoPE offset, all in-place ops.
- Each graph replay's write becomes the next replay's read.
- Heuristic now works with CUDA Graphs enabled (no need to disable).
- Refine heuristic_topk.cuh kernel for improved convergence.

Made-with: claude-4.6-opus-high
Signed-off-by: longcheng-nv <[email protected]>

…aluation

Extend eval_longbench_v1.py to support DeepSeek DSA sparse attention
alongside existing RocketKV, with heuristic TopK and MTP speculative
decoding options.

Key changes:
- Add --dsa_sparse flag as mutually exclusive with --rocket_sparse
- Add --enable_heuristic_topk for DSA heuristic TopK pre_idx reuse
- Add --mtp for MTP speculative decoding layers
- Update usage examples for both RocketKV and DSA workflows

Made-with: claude-4.6-opus-high
Signed-off-by: longcheng-nv <[email protected]>

…tions

Three kernel-level optimizations (OPT5/OPT6/OPT7) to reduce latency of
the single-CTA heuristic TopK on Blackwell (sm_100):

OPT5 — Skip Phase 3 blockCountGE re-scan when Phase 2 already confirmed
the candidate count is in [TOP_K, MAX_CANDIDATES] (done==1).  Eliminates
one full N-element scan in the common fast-convergence case.

OPT6 — Increase NUM_BINS from 256 to 2048 and replace the serial O(2048)
K-th-bin scan with a two-step parallel search: each warp accumulates its
128-bin slice (NUM_BINS/NUM_WARPS), then thread 0 locates the target
warp in 16 steps and one warp-lane scans 128 bins.  Serial depth drops
from 2048 to 144 steps.  Shared memory grows from ~50 KB to ~59 KB; the
histogram is now 8 KB instead of 1 KB.

OPT7 — Cache per-thread element counts into smem->per_thread_counts[]
inside blockCountGE.  Phase 3 prefix-sum reuses these cached values
instead of repeating the N-scan, saving one full global-memory pass.

Also add distribution-parameterised correctness tests for the heuristic
indexer_topk_decode path in test_indexer_topk.py:

- Four logit families: beta (bounded), logistic (heavy-tailed), lognorm
  (positively skewed), weibull_min (right-skewed extreme-value)
- MTP correlation: consecutive rows within each batch element share tail
  logits (next_n up to 3)
- pre_idx accuracy sweep: success_ratio in {0.5, 0.9}
- Tolerance set to full_range/256 to accept histogram-bin boundary ties

Made-with: claude-4.6-opus-high
Signed-off-by: longcheng-nv <[email protected]>

…flicts

Pass caller-owned scratch buffer through heuristic TopK pipeline to
eliminate cudaMallocAsync/cudaFreeAsync inside the kernel, enabling
CUDA Graph capture. Also resolves stash merge conflicts in test file
and adds seq_lens clamping for valid row lengths with next_n > 1.

Signed-off-by: longcheng-nv <[email protected]>
Made-with: claude-4.6-opus-high

…pK tests

Add docstrings to all test functions and helpers to improve docstring
coverage. Remove broken duplicate test_indexer_topk_decode_dist that
referenced undefined names (dtype, run_fn). Suppress F811 on intentional
CuTE DSL test redefinitions with different parametrize configurations.

Made-with: claude-4.6-opus-high
Signed-off-by: longcheng-nv <[email protected]>

…n test_dsa_indexer

The DSA backend accesses sparse_attention_config.enable_heuristic_topk
but the mock SparseAttentionConfig in test_dsa_indexer.py was missing
this field, causing all 76 DSA indexer tests to fail with AttributeError.

Made-with: claude-4.6-opus-high
Signed-off-by: longcheng-nv <[email protected]>

Add one-line docstrings to all public functions, classes, and methods
across the 5 Python files modified by this PR to satisfy the CI
docstring coverage gate (was 41%, now ~100%).

Signed-off-by: longcheng-nv <[email protected]>
Made-with: claude-4.6-opus-high

Without heuristic_scratch the call silently falls back to the radix kernel.

Made-with: claude-4.6-opus-high
Signed-off-by: longcheng-nv <[email protected]>

The heuristic TopK kernel requires stride0 divisible by 4 for float4
loads. In production fp8_paged_mqa_logits always returns tensors with
stride aligned to 256, but the test helpers created logits with
arbitrary max_len from random seq_lens, causing assertion failures
when max_len was not a multiple of 4.

Made-with: claude-4.6-opus-high
Signed-off-by: longcheng-nv <[email protected]>

…duplicate allocation

- Remove duplicate `heuristic_scratch_values` allocation from
  `create_expanded_buffers()` (already created in `__post_init__`);
  move resize logic into `update_spec_dec_param()` where shape changes.
- Eliminate `heuristic_pre_idx_staging` buffer entirely: pass
  `heuristic_prev_topk` directly as `pre_idx` to the kernel.
- Move +1 temporal offset into the C++ kernel (`preIdxOffset =
  (rowIdx % next_n) + 1`), removing two Python tensor ops (copy_ and
  += 1) from the CUDA Graph captured region per layer per decode step.

Addresses review comments from lfr-0531 on PR NVIDIA#12385.

Signed-off-by: longcheng-nv <[email protected]>

…on pre-Blackwell

- Add `get_sm_version() >= 100` guard to `enable_heuristic_topk` in both
  `DSAtrtllmAttentionMetadata` and `Indexer`, so the heuristic path
  silently falls back to radix sort on Hopper and older architectures.
- Add `@skip_pre_blackwell` decorator to `test_indexer_topk_decode_dist`
  in `test_indexer_topk.py` and `test_indexer_topk_dist.py`, since the
  heuristic kernel is only validated for Blackwell GPUs.

Fixes CI failure on DGX_H100: test_indexer_topk_decode_dist[8192-2048-2-128-logistic_m0.47_s1.46-0.5].

Made-with: claude-4.6-opus-high
Signed-off-by: longcheng-nv <[email protected]>

…eedback

Revert docstring-only changes in cpp_custom_ops.py, model_config.py, and
llm_args.py that were added in eda0351 but are outside the scope of this
PR. Docstrings in files directly modified by this feature (dsa.py and
test_dsa_indexer.py) are retained.

Addresses feedback from lfr-0531 and kaiyux on PR NVIDIA#12385.

Made-with: claude-4.6-opus-high
Signed-off-by: longcheng-nv <[email protected]>

… and multi-GPU

Add runtime topK == TOP_K validation, compile-time static_assert blocking
non-32-bit index types, remove process-wide static flag in favor of
unconditional cudaFuncSetAttribute, and drop unsafe reinterpret_casts.

Signed-off-by: Long Cheng <[email protected]>
Signed-off-by: longcheng-nv <[email protected]>

Add a parametrized variant to TestDeepSeekV32.test_fp8_blockscale that
exercises the end-to-end heuristic TopK decode path with MTP (next_n=1).
Gated to Blackwell (SM >= 100). Validates accuracy on MMLU and GSM8K
with enable_heuristic_topk=True through the LLM API.

Made-with: claude-4.6-opus-high
Signed-off-by: longcheng-nv <[email protected]>