[Example] Add HISA: hierarchical sparse attention indexer - #2069

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@examples/dsa_hisa/block_sparse_mqa_fp8.py`:
- Line 210: The unpacking of q_fp8.shape assigns an unused H (M, H, D =
q_fp8.shape) — change it to use a throwaway variable (e.g., M, _, D =
q_fp8.shape or M, _H, D = q_fp8.shape) so linters stop complaining; also replace
any Unicode multiplication signs (e.g., '×') used around line 315 with the ASCII
asterisk '*' (search for occurrences in strings/comments or expressions) to
address the RUF003 warnings.
- Around line 71-101: The code copies IndexK and reads IndexKScale for a full
block using block_s_i and block_N before applying the cu_k_s/cu_k_e mask, which
can read past the end when seq_len_kv % kv_block_size != 0; fix by performing
guarded loads or padding: when computing T.copy(IndexK[block_s_i : block_s_i +
block_N, :], index_k_shared) and the subsequent IndexKScale loads, ensure you
only copy entries where block_s_i + bn_i < seq_len_kv (e.g., copy a shortened
slice and fill the remaining index_k_shared rows with safe defaults or use
conditional copies), and similarly guard the IndexKScale access
(IndexKScale[block_s_i + bn_i]) with block_s_i + bn_i < seq_len_kv; apply the
same guarded-load/padding logic to the other kernel section mentioned (lines
132-160) so no out-of-bounds reads occur.

In `@examples/dsa_hisa/clean_and_maintain_logits.py`:
- Around line 31-37: The loop can produce idx values >= logits width
(seq_len_kv) when block_K > actual width, causing out-of-bounds writes to
Logits; wrap or extend the existing conditions to ensure any write to Logits[bx,
idx] is guarded by a bounds check (e.g., require idx < seq_len_kv) before
assigning +inf or -inf. Locate the pipelined loops using T.Pipelined/T.ceildiv
and the variables idx, block_K, threads, tx, cu_k_s, cu_k_e and add a single idx
< seq_len_kv guard (or combine it into the existing if conditions) so no write
occurs when idx is outside [0, seq_len_kv-1].

In `@examples/dsa_hisa/fp8_block_mean_pooling.py`:
- Around line 48-66: The loop currently does an unconditional T.copy and
unconditional KScale loads which can read past seq_len_k for the last pooling
tile; change the loads to be guarded by cur_tl_block_size (or pad K/KScale to
block_N) so you only copy/assign up to cur_tl_block_size entries and set the
remaining lanes to zero. Specifically, replace the unconditional
T.copy(K[tl_block_s : tl_block_s + block_N, :], index_k) and the for bn_i in
T.Parallel(block_N): scale[bn_i] = KScale[tl_block_s + bn_i] with
guarded/limited loads that use cur_tl_block_size (or conditional bn_i <
cur_tl_block_size) to avoid indexing beyond seq_len_k, and ensure any untouched
index_k/scale lanes are explicitly zeroed so the subsequent index_k scaling and
T.reduce_sum are safe.

In `@examples/dsa_hisa/hisa.py`:
- Line 82: Replace the Unicode multiplication character (×) with the ASCII
letter "x" in the comments inside hisa.py—specifically update the comment
"blocks' raw tokens (block_topk_eff blocks × k_block_size tokens" and the other
comment occurrences that use "×" (also present in the nearby comments around the
token/block size descriptions and later comment at the end of the file). Search
for comments containing the "×" glyph and change them to use "x" so they satisfy
Ruff RUF003 while keeping the original comment wording intact.
- Around line 101-132: The function currently returns topk_indices with shape
[M, topk_tokens_eff] when topk_tokens > block_sparse_logits.shape[-1], but the
docstring promises [M, topk_tokens] with invalid slots set to -1; after
computing and masking topk_indices (symbols: topk_tokens, topk_tokens_eff,
block_sparse_logits, relevant_topk_indices, topk_indices), if topk_tokens_eff <
topk_tokens create a new int32 tensor filled with -1 of shape [M, topk_tokens]
and copy the existing topk_indices into the leftmost topk_tokens_eff columns
(leave remaining columns as -1), then return that padded tensor so the function
always returns shape [M, topk_tokens].

In `@examples/dsa_hisa/pool_mqa_fp8.py`:
- Line 219: Several comments in examples/dsa_hisa/pool_mqa_fp8.py use the
Unicode multiplication sign (×) which triggers Ruff RUF003; replace the Unicode
× with ASCII x in those comments (e.g., change "fp8×fp8 GEMM" to "fp8xfp8 GEMM"
and any other occurrences on the same comment lines). Search for the '×'
character in the file and update the comment text accordingly, keeping only
ASCII characters to silence the RUF003 warnings.
- Around line 79-106: The kernel currently assumes full tiles for queries and K
(block_Q, block_N) causing out-of-bounds reads from
blocked_kv_fp8/blocked_kv_scale and writes to Logits when seq_len_blocked_kv or
the final query tile is partial; fix by adding bounds checks and guarded
loads/stores or by computing padded tile extents before the loops: clamp
cu_k_s_min/cu_k_e_max and the per-tile index offsets using seq_len_blocked_kv
and CuSeqLenBlockedKS/CuSeqLenBlockedKE, use guarded copy/load for IndexBlockedK
and IndexBlockedKScale when cu_k_s_min + nbn_i*block_N + bn_i >=
seq_len_blocked_kv, and only write to Logits when seq_len_i + bq_i and the K
index are within valid lengths (use CuSeqLenBlockedKS/KE and seq_len_blocked_kv
to decide); alternatively ensure inputs are pre-padded to block_Q/block_N
multiples before kernel launch.

In `@examples/dsa_hisa/README.md`:
- Line 87: Update the three fenced code blocks containing the snippets starting
with "block_k_score[m, n] = ..." , the one containing "k_abs = blk *
kv_block_size + i" and the line "(1.1) fp8_native_block_mean_pooling           
K, k_scale → blocked_k, blocked_k_scale" to include a language identifier by
replacing the opening ``` with ```text so the fences are annotated (e.g.,
```text block_k_score..., ```text k_abs = ..., ```text (1.1)
fp8_native_block_mean_pooling ...).

In `@examples/dsa_hisa/tilelang_utils.py`:
- Around line 264-270: The current slice uses last_seq_id =
torch.where(cu_seqlens.cumsum(0) >= total_seqlen)[0][0] which raises if no
cumulative sum reaches total_seqlen and prevents the fallback; change the logic
around cu_seqlens/cumsum to first compute the indices (e.g., idxs =
torch.where(cu_seqlens.cumsum(0) >= total_seqlen)[0]) and check idxs.numel() (or
.nelement()) before indexing: if idxs is empty, leave cu_seqlens unchanged so
the subsequent if cu_seqlens.sum() < total_seqlen fallback can run, otherwise
set last_seq_id = idxs[0].item() and slice cu_seqlens =
cu_seqlens[:last_seq_id]; keep references to cu_seqlens, total_seqlen,
last_seq_id and the cumsum check when applying the fix.

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@examples/dsa_hisa/block_sparse_mqa_fp8.py`:
- Around line 229-232: The current test masks out NaNs via `finite` so a kernel
that produces NaN where `ref` is finite can pass; add an explicit NaN mask
assertion before using `finite`: assert that torch.isnan(got) equals
torch.isnan(ref) (e.g., raise "NaN mask differs") so any NaN/finite mismatch
fails, then keep the existing infinity-mask checks and the finite-element
comparison using torch.testing.assert_close on got[finite] vs ref[finite].

---

Duplicate comments:
In `@examples/dsa_hisa/block_sparse_mqa_fp8.py`:
- Line 144: The tuple unpacking M, H, D = q_fp8.shape introduces an unused
variable H and the FLOPs comment contains Unicode multiplication signs; change
the unpacking to M, _, D = q_fp8.shape (or M, H_unused, D) to mark H as
intentionally unused and update any FLOPs comments that use “×” to use the ASCII
asterisk "*" instead so Ruff no longer flags them; search for q_fp8.shape and
the FLOPs comment strings in this file (and the other occurrence around line
249) and apply these two edits.
- Around line 77-101: The code reads a full block_N in IndexK and related loops
even for the final ragged KV block; compute valid_N = max(0, min(block_N,
seq_len_kv - block_s_i)) and use that to limit T.copy into index_k_shared and to
bound all T.Parallel loops (the scale_shared fill from IndexKScale, the
s_reshaped write loop, the per-item cu_k range check that writes logits, and the
final Logits write) so no reads or writes occur past seq_len_kv; keep the
existing cu_k_s_min/cu_k_e_max check but ensure iterations only go 0..valid_N-1
when referencing index_k_shared, scale_shared, s_reshaped, logits, and Logits to
avoid out-of-range K reads.

In `@examples/dsa_hisa/clean_and_maintain_logits.py`:
- Around line 31-37: The loop can compute idx beyond seq_len_kv-1, causing OOB
writes to Logits; update the write conditions in the nested loops (where idx =
n_i * block_K + k_i * threads + tx is computed) to only assign Logits[bx, idx]
when idx < seq_len_kv (e.g., add idx < seq_len_kv to both the infinity and
-infinity branches or combine checks so writes occur only when cu_k_s/cu_k_e
logic AND idx < seq_len_kv are true); ensure the check uses the same
dtype/variables (seq_len_kv, block_K, threads, tx, cu_k_s, cu_k_e, Logits, bx)
so the tail pipelined iteration cannot write out of bounds.

In `@examples/dsa_hisa/fp8_block_mean_pooling.py`:
- Around line 48-63: The code reads past seq_len_k in the last pooling block;
compute cur_tl_block_size = tl_block_e - tl_block_s before doing any loads and
guard the copies and KScale loads so you only read lanes with bn_i <
cur_tl_block_size (or alternatively pad K/KScale to block_N); specifically,
change the T.copy(K[...], index_k) and KScale[tl_block_s + bn_i] accesses to be
conditional on bn_i < cur_tl_block_size (or use a safe masked copy), ensuring
index_k and scale for out-of-range bn_i are set to zero before the subsequent
multiply and T.reduce_sum (referencing tl_block_s, tl_block_e, block_N,
cur_tl_block_size, T.copy, K, KScale, scale, index_k).

In `@examples/dsa_hisa/pool_mqa_fp8.py`:
- Around line 219-246: Summary: Replace the Unicode multiplication sign (×,
U+00D7) used in comments with ASCII 'x' to satisfy Ruff RUF003. Fix: edit the
comment strings and any inline comments around the code using
pool_mqa_attn_return_logits_fp8_interface, do_bench, and the module-level guard
so that occurrences like "M × k_block_size^-1" and any prints or comment lines
use "x" instead of "×"; do not change logic or variable names, only replace the
character in comments and string literals. Ensure all three occurrences noted
near the benchmarking block and the module doc/guard are updated.
- Around line 79-106: The loops and copies (using symbols CuSeqLenBlockedKS,
CuSeqLenBlockedKE, IndexQ, Weights, IndexBlockedK, Logits) can read/write past
the valid sequence/KV ranges when seq_len % block_Q != 0 or seq_len_blocked_kv <
block_N; fix by adding explicit bounds guards or masked/padded copies: before
any T.copy or array read use min/conditional logic against seq_len and
seq_len_blocked_kv (e.g., clamp seq_len_i + bq_i and cu_k_s_min+nbn_i*block_N +
bn_i), and for writes to Logits only emit stores when the target index <
seq_len_blocked_kv and seq_len; alternatively pad IndexQ/Weights/IndexBlockedK
inputs to full tile sizes before launch. Ensure guards are applied for the
T.copy(IndexQ..., Weights...), T.copy(IndexBlockedK...), uses of
CuSeqLenBlockedKS/KE, and the final Logits[...] write so no OOB accesses occur.

---

Nitpick comments:
In `@examples/dsa_hisa/clean_and_maintain_logits.py`:
- Around line 100-109: The benchmark currently allocates a fresh tensor inside
fn() which mixes allocation/RNG overhead into the timing; move the
torch.randn(M, N, device="cuda", dtype=torch.float32) allocation outside of fn()
(e.g., create a reusable buffer before calling do_bench) and inside fn() either
clone/reset the pre-allocated buffer in-place or simply call
clean_and_maintain_logits_interface(logits, cu_ks, cu_ke) on that pre-allocated
tensor so do_bench measures only the kernel work performed by
clean_and_maintain_logits_interface for the [M, N] buffer.