[Aravind-11]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

On the non-compact (pre-SM100) path, _indexer_topk_core materializes the full fp32 score matrix (total_q x max_seqlen_kv) even when callers only need Top-K results. Add scores_output="topk" | "none", which scores query rows in chunks under a 1 GiB budget and keeps only the selected scores. Each chunk reuses the THD kernel path with per-segment q_causal_offsets, so BSHD and THD share one implementation and indices are unchanged. Chunk index, length, and score outputs are allocated once. Sparse indexer loss uses the compact Top-K softmax or the dense-fallback Top-K scores, and raises if neither was produced.

Use it for inference Top-K and for the sparse indexer loss in FusedCSAIndexerSparseAttnFunc; the dense loss keeps the full matrix.

Linked issue: #4468

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• If this PR adds or changes a GPU kernel (Triton, jit_fuser/torch.compile, CUDA extension, TE or external-library dispatch, or a scatter/index accumulation), I have added or updated its bit-exact determinism test and registered it in tests/unit_tests/determinism/kernels/manifest.py (guide)
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Testing

Latest run is at 9d652ae89 on a Colab A100-SXM4-40GB (SM 8.0, Python 3.13.15, torch 2.11.0+cu130, pytest 8.4.2). cuDNN Frontend and Transformer Engine were not installed, so real DSA kernel tests skip.

python3 -m pytest --noconftest -p no:cacheprovider --override-ini='addopts=' -ra --tb=line \
  tests/unit_tests/transformer/experimental_attention_variant/test_csa_indexer_row_chunking.py \
  tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_sparse_attention.py \
  tests/unit_tests/transformer/experimental_attention_variant/test_csa_indexer_loss.py

Suite	Result
test_csa_indexer_row_chunking.py	29 passed
test_csa_fused_sparse_attention.py	108 passed, 27 skipped
test_csa_indexer_loss.py	21 passed

158 passed, 27 skipped in 55.25s. The row-chunking file includes test_chunked_topk_prealloc_peak_is_one_output, which passed on this A100. The fused file includes test_sparse_loss_raises_when_compact_softmax_is_missing, which passed. All 27 skips are _skip_if_real_kernels_unavailable (requires SM9+, found SM8). That includes test_real_indexer_topk_row_chunked_matches_single_pass for chunk rows 48 and 100, and test_real_indexer_topk_thd_row_chunked_matches_single_pass. Those real cuDNN kernels did not execute here.

An earlier A100 run at 763ecf3 was 28 passed, 107 passed / 24 skipped, and 21 passed. At 963bd1c, test_csa_fused_sparse_attention.py was 105 passed, 24 skipped, and 2 failed: test_build_flat_topk_idxs_thd_compact_cpu_cuda_parity (no cuDNN compactify wrapper) and test_indexer_topk_mxfp8_compact_dispatch (MXFP8 flags were not stubbed). Both failed the same way on parent ac100f7. 763ecf3 fixes them.

@buptzyb @hxbai lmk if im headed in the right direction

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yaox12]

/ok to test 2d03243

[copy-pr-bot]

/ok to test 2d03243

@yaox12, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

[yaox12]

/ok to test 2d03243

[yaox12]

/review model=claude mode=strict

[svcnvidia-nemo-ci]

Strict review: bounded CSA indexer score memory on the dense fallback path

I found no blocking defect, so this review is a comment rather than a request for changes. There is one IMPORTANT memory issue and two SUGGESTION items. Details, patches and acceptance checks are in the inline comments.

Totals: CRITICAL: 0, IMPORTANT: 1, SUGGESTION: 2

Findings

1. [IMPORTANT Performance] The chunked loop in _indexer_topk_row_chunked keeps every per-chunk index and score tensor in lists, then calls torch.cat on them. For a moment, the full Top-K indices and scores exist twice, about 2 x total_q x topk x 8 bytes. That is bigger than the 1 GiB score budget once total_q x topk x 8 > 1 GiB, which is the long-context case this path is for. Fix: allocate the outputs before the loop and copy each chunk into its rows.
2. [SUGGESTION Simplification] With sparse loss, indexer_scores is now always None. The prepare_sparse_loss branch at L3026-L3037 can no longer run; the only way to reach it is to fail the assert. Replace it with an explicit error.
3. [SUGGESTION Testing] No test runs the multi-chunk path on real cuDNN kernels. In that path BSHD inputs go through the THD form of the kernel, with made-up cu_seqlens, segments that have no query rows, and causal offsets for segments cut mid-sequence. The new tests check this only against the PyTorch stand-in, yet the new test docstring says the GPU suite covers it. Add a real-kernel test that monkeypatches the budget to force chunking.

Checked and found correct

These are source-level checks; I could not run anything.

• Chunk metadata: chunk_cu_q is cut to the chunk's rows and rebased. chunk_offsets = base + (row_start - seg_start).clamp(min=0). This gives the same per-row lengths as build_seq_lens, both the Triton kernel and the fallback. It follows the same pattern as the CP layout in cp_utils._build_cp_indexer_layout: clip the query rows, keep the full K, offsets restore the true positions.
• Padding rows: CUDA-graph padding rows past cu_seqlens_q[-1] still get length 0 and -1 indices.
• Sparse-loss probabilities: _sparse_indexer_predict matches _prepare_sparse_loss_fallback. Invalid slots use the fp32 minimum, and padding rows are masked to -1 before the softmax.
• Other loss settings: the dense loss still gets the full score matrix. When loss_coeff <= 0, indexer_scores was never used, so returning nothing there is safe.
• CUDA graphs: the chunk loop has no host syncs, and the chunk count depends only on shapes and Python ints.
• Recursion: the inner calls use use_compact=False and scores_output="full", so they never loop back into the chunked path.
• Compact path and public API: the compact (SM100) path is unchanged, and indexer_topk's signature is unchanged.
• Second commit: the fallback in _compact_flat_topk_idxs produces the same layout as cuDNN compactify; the existing parity test covers this. The MXFP8 dispatch test now stubs HAVE_TE_MXFP8 and HAVE_TRITON, which is consistent with the quantization and kernel it already mocks.
• Mandatory unused-identifier check: every new constant, argument, helper and local has a real use.

Notes

• The style-guide.md the review skill points to does not exist in the base tree. The style check was therefore limited to what the configured linters enforce; no style findings.
• The determinism manifest and coverage tool named in the PR checklist (tests/unit_tests/determinism/kernels/, tools/check_kernel_determinism_coverage.py) are not in this tree either, so I did not raise that requirement.

Patches (top-level, per the skill's rule for multi-line changes)

Finding 1: replace the list and torch.cat handling in _indexer_topk_row_chunked:

    topk_indices = torch.empty((total_q, topk), dtype=torch.int32, device=device)
    topk_length = torch.empty((total_q,), dtype=torch.int32, device=device)
    topk_scores = (
        torch.empty((total_q, topk), dtype=torch.float32, device=device)
        if return_topk_scores
        else None
    )
    for row_start in range(0, total_q, chunk_rows):
        row_end = min(row_start + chunk_rows, total_q)
        chunk_cu_q = (cu_q.clamp(min=row_start, max=row_end) - row_start).to(cu_q.dtype)
        chunk_offsets = (base_offsets + (row_start - segment_starts).clamp(min=0)).to(torch.int32)
        chunk_indices, chunk_length, chunk_scores, _ = _indexer_topk_core(...)  # call unchanged
        topk_indices[row_start:row_end].copy_(chunk_indices)
        topk_length[row_start:row_end].copy_(chunk_length)
        if topk_scores is not None:
            topk_scores[row_start:row_end].copy_(_gather_topk_scores(chunk_scores, chunk_indices))
        del chunk_indices, chunk_length, chunk_scores
    # the `if not is_thd:` view back to (b, sq, ...) stays as is

Finding 2: replace the unreachable else branch in FusedCSAIndexerSparseAttnFunc.forward:

                else:
                    raise RuntimeError(
                        "Sparse indexer loss needs the compact Top-K softmax or the "
                        "dense-fallback Top-K scores; neither was produced."
                    )

Overall risk: low

The chunking logic is correct. The remaining work is a small memory fix in the long-context case, removing a dead branch, and adding real-kernel test coverage for the new multi-chunk path.

Review generated by claude

[svcnvidia-nemo-ci]

[IMPORTANT Performance] torch.cat briefly holds two full copies of the Top-K outputs, which works against this path's memory bound.

Failure. Each chunk's chunk_indices (int32) and gathered scores (fp32) are appended to index_chunks / score_chunks. Those lists stay alive until the function returns. L2260-L2262 then allocate a second full-size copy of each.

With scores_output="topk" the peak is about 2 x total_q x topk x 8 bytes. The expected peak is total_q x topk x 8 plus one score chunk of at most 1 GiB. The extra memory appears once total_q x topk x 8 > 1 GiB, which is the long-context case this path exists for.

Example: a 1M-row call (the 1M-token case used in test_chunk_rows_respects_budget_and_alignment) with topk=512 has 2 GiB of indices and 2 GiB of scores. The concat peak is about 8 GiB; with preallocation it would be about 5 GiB.

Correction. Allocate the outputs before the loop and copy each chunk into its row slice. Drop the three lists and the torch.cat calls.

Invariant to restore: at most one copy of the final outputs plus one score chunk are live at any time.

topk_indices = torch.empty((total_q, topk), dtype=torch.int32, device=device)
topk_length = torch.empty((total_q,), dtype=torch.int32, device=device)
topk_scores = (torch.empty((total_q, topk), dtype=torch.float32, device=device)
               if return_topk_scores else None)
for row_start in range(0, total_q, chunk_rows):
    ...  # chunk_cu_q / chunk_offsets / _indexer_topk_core call unchanged
    topk_indices[row_start:row_end].copy_(chunk_indices)
    topk_length[row_start:row_end].copy_(chunk_length)
    if topk_scores is not None:
        topk_scores[row_start:row_end].copy_(_gather_topk_scores(chunk_scores, chunk_indices))
    del chunk_indices, chunk_length, chunk_scores

The dtypes already match the inner THD return:

• indices come back as topk_indices.int();
• sanitize_topk returns int32 lengths;
• the dense scores are fp32.

Preservation.

• The single-pass branch (L2191-L2212) is unchanged.
• Output shapes and dtypes are unchanged: indices (total_q, topk) for THD or (b, sq, topk) for BSHD, int32 lengths, fp32 scores with -inf at -1 slots.
• scores_output="none" still allocates no score output.
• The chunk count still depends only on shapes, so CUDA-graph behaviour is unchanged.

Acceptance. These existing tests must pass unchanged:

• test_bshd_chunked_matches_single_pass
• test_thd_chunked_matches_single_pass, including the padding-row -1 / 0 asserts
• test_thd_chunked_matches_independent_reference
• test_peak_score_block_is_bounded (forward_calls == 16)

Static invariant: no list accumulation or torch.cat of full-size outputs remains in this function.

Optional CUDA test (proposed): use _force_chunk_rows and assert that the increase in torch.cuda.max_memory_allocated() is at most total_q*topk*8 + chunk_rows*sk*4 plus slack.

[svcnvidia-nemo-ci]

[SUGGESTION Simplification] This branch can no longer run.

Failure. When loss_coeff > 0 and sparse_loss, L2850-L2851 always set indexer_scores = None. Only two cases can produce a prediction:

• the compact path supplies compact_predict, because return_softmax=True;
• the dense fallback supplies indexer_topk_scores, because _indexer_topk_row_chunked(return_topk_scores=True) returns a value in both of its branches.

So this else runs only when both are None, and then it stops at assert indexer_scores is not None. The prepare_sparse_loss call and the comment about the fused row kernel describe code that never runs. This is also the only call of csa_indexer_loss_kernels.prepare_sparse_loss in this module.

Readers will assume the fused Triton path is still in use. The bare AssertionError is also unhelpful, and python -O strips it.

Correction. Replace the else body (L3026-L3037) with an explicit error:

                else:
                    raise RuntimeError(
                        "Sparse indexer loss needs the compact Top-K softmax or the "
                        "dense-fallback Top-K scores; neither was produced."
                    )

Optionally, update the step-2 comment at L2825-L2829 ("Top-K dispatch is independent of the configured loss"). The fallback's score output now depends on loss_coeff and sparse_loss.

Preservation.

• The compact branch (L3009-L3016) and the new indexer_topk_scores branch (L3017-L3025) are unchanged.
• The dense-loss branch (L3068 onward) still takes indexer_scores from scores_output="full".
• Keep prepare_sparse_loss and its direct kernel tests (test_sparse_loss_preparation_cpu_fallback, test_sparse_loss_preparation_triton_matches_eager, test_prepare_sparse_loss_compiles_once_across_score_widths). They test the kernel module, not this branch.

Acceptance.

• The existing mocked sparse-loss fused_csa_indexer_sparse_attn tests and test_sparse_predict_matches_full_score_loss_preparation still pass.
• Proposed test: mock a compact dispatch whose wrapper returns "softmax": None, with sparse_loss=True and loss_coeff > 0. It should raise RuntimeError matching compact Top-K softmax; today it raises a bare AssertionError.
• Static invariant: every sparse-loss path that can run sets predict from compact_predict or from indexer_topk_scores.

[svcnvidia-nemo-ci]

[SUGGESTION Testing] No test runs the multi-chunk path on real cuDNN kernels.

Failure. This docstring says the GPU tests in test_csa_fused_sparse_attention.py cover kernel numerics, but none of them reaches the multi-chunk branch:

• The real fallback test TestRealKernelIndexerTopk.test_real_indexer_topk_set_matches_reference uses b=2, sq=128, sk=128, which is 128 KiB of scores. That is far below the 1 GiB _CSA_INDEXER_SCORE_CHUNK_MAX_BYTES, so _indexer_topk_row_chunked always takes its single-pass branch.
• On SM100 that test takes the compact path anyway.

The multi-chunk branch is the only code that:

• (a) sends SBHD/BSHD inputs through the THD form of _DSA.indexer_forward_wrapper (made-up cu_seqlens_q/k plus q_causal_offsets) instead of the BSHD form;
• (b) passes segments that have no query rows but do have K rows;
• (c) passes causal offsets for segments cut mid-sequence.

All three are checked only against ReferenceDSA, which assumes position = offset + local_row. The _indexer_topk_core docstring goes further and says chunked indices and lengths are "identical to 'full'". If cuDNN behaves differently, Top-K is silently wrong only at the long-context sizes where chunking turns on. The legacy CP layout uses a similar clipping pattern, which lowers the risk but does not cover the BSHD-to-THD switch.

Correction. Add a real-kernel test that forces chunking by monkeypatching the budget; no production change is needed. Proposed, in TestRealKernelIndexerTopk:

@pytest.mark.parametrize("chunk_rows", [48, 100])
def test_real_indexer_topk_row_chunked_matches_single_pass(
    self, chunk_rows, monkeypatch, reset_lazy_kernel_state
):
    _skip_if_real_kernels_unavailable()
    # same q_indexer / k_indexer / weights as test_real_indexer_topk_set_matches_reference
    kw = dict(topk=8, ratio=4, indexer_softmax_scale=128**-0.5,
              use_compact=False, deterministic=True)
    full_idx, full_len = indexer_topk(q_indexer, k_indexer, weights, **kw)
    monkeypatch.setattr(dk, "_CSA_INDEXER_SCORE_CHUNK_MAX_BYTES", chunk_rows * 128 * 4)
    monkeypatch.setattr(dk, "_CSA_INDEXER_SCORE_CHUNK_ROW_ALIGNMENT", 1)
    idx, length = indexer_topk(q_indexer, k_indexer, weights, **kw)
    assert torch.equal(length, full_len)
    # per-row set overlap with the reference >= max(1, n - 1), as in the existing test

With chunk_rows=48, one chunk covers rows 96-143 and so crosses the batch boundary at row 128. That exercises a segment with no query rows and a non-zero offset. use_compact=False forces the fallback on SM100 as well.

Also add a THD variant: two segments with q_causal_offsets, for example [5, 0] as in test_indexer_topk_thd_compact_dispatch, where one segment spans two chunks.

To keep the docstring's "identical" claim, also assert torch.equal(idx, full_idx) with deterministic=True. Otherwise, soften that docstring to "same lengths; same Top-K sets up to kernel rounding". Update this module docstring either way.

Preservation. Keep the CPU stand-in tests in this file unchanged. They remain the logic-level checks for row slicing, offsets and padding.

Acceptance.

• The proposed test passes on SM90 and SM100.
• topk_length matches the single-pass run exactly.
• Per-row index-set overlap with the reference is at least n-1.
• A failure would point directly at a THD-versus-BSHD semantics mismatch in cuDNN.

[Aravind-11]

Hi @yaox12, thank you so much for the review!

The three items from the strict review on 2d0324378 are in 9d652ae89.

1. Preallocated Top-K outputs. _indexer_topk_row_chunked allocates the index, length, and score buffers once and copies each chunk into its row slice. test_chunked_topk_prealloc_peak_is_one_output checks that the CUDA peak stays within one output copy plus one score chunk.

2. Sparse-loss error. When the compact Top-K softmax and the dense-fallback Top-K scores are both missing, sparse indexer loss raises RuntimeError. Covered by test_sparse_loss_raises_when_compact_softmax_is_missing. The compact path and the indexer_topk_scores path are unchanged. The step-2 comment now says the dense fallback returns nothing when the loss is off, the selected scores for the sparse loss, and the full matrix for the dense KL loss.

3. Real multi-chunk kernel tests. test_real_indexer_topk_row_chunked_matches_single_pass (chunk rows 48 and 100) and test_real_indexer_topk_thd_row_chunked_matches_single_pass patch the score budget, force use_compact=False and deterministic=True, and require the chunked indices and lengths to match one single pass. chunk_rows=48 crosses the batch boundary at row 128. The THD case uses query lengths [80, 64] and causal offsets [5, 0].

On a Colab A100-SXM4-40GB (SM 8.0, torch 2.11.0+cu130) at 9d652ae89:

• test_csa_indexer_row_chunking.py: 29 passed
• test_csa_fused_sparse_attention.py: 108 passed, 27 skipped
• test_csa_indexer_loss.py: 21 passed

158 passed, 27 skipped in 55.25s. The 27 skips are _skip_if_real_kernels_unavailable (requires SM9+, found SM8), including the two new real-kernel tests. Those run on an SM90 or SM100 runner.

Could you take another look when you have a chance? Thank you.

[yaox12]

/ok to test d19ec33