[None][fix] Honor MSA block-index strides for q_len <= 32

3rdparty/patches/msa_strided_paged_kv.patch

@@ -677,3 +677,314 @@ index 306b416..a564d47 100644
          out.copy_(result)
          return out, None
      return result, None
+diff --git a/python/fmha_sm100/csrc/fmha_sm100_inst.jinja b/python/fmha_sm100/csrc/fmha_sm100_inst.jinja
+index 6949035..e59ecfb 100644
+--- a/python/fmha_sm100/csrc/fmha_sm100_inst.jinja
++++ b/python/fmha_sm100/csrc/fmha_sm100_inst.jinja
+@@ -58,6 +58,8 @@ cudaError_t {{ func_name }}(const FMHACutlassSM100Params& p) {
+       p.kv_indices_ptr, p.kv_page_indptr_ptr, p.total_page_num,
+       p.max_score_ptr, p.max_k_tiles,
+       p.kv_block_indexes_ptr, p.kv_block_num,
++      p.kv_block_indexes_stride_q, p.kv_block_indexes_stride_h,
++      p.kv_block_indexes_stride_k,
+       p.pack_factor, p.q_stride_n_original, p.q_stride_h_original, p.h_r_original,
+       flashinfer::PackGQAUnpackParams{p.max_score_direct_ptr, p.total_qo_len_orig,
+                                       p.o_direct_ptr, p.num_qo_heads_orig,
+diff --git a/python/fmha_sm100/csrc/fmha_sm100_params.h b/python/fmha_sm100/csrc/fmha_sm100_params.h
+index 878a90c..f621511 100644
+--- a/python/fmha_sm100/csrc/fmha_sm100_params.h
++++ b/python/fmha_sm100/csrc/fmha_sm100_params.h
+@@ -54,6 +54,9 @@ struct FMHACutlassSM100Params {
+   int max_k_tiles;
+   int* kv_block_indexes_ptr;
+   int kv_block_num;
++  int64_t kv_block_indexes_stride_q;
++  int64_t kv_block_indexes_stride_h;
++  int64_t kv_block_indexes_stride_k;
+   int pack_factor = 1;
+   int h_r_original = 0;
+   int q_stride_n_original = 0;
+diff --git a/python/fmha_sm100/csrc/fmha_sm100_variant_run.cu.jinja b/python/fmha_sm100/csrc/fmha_sm100_variant_run.cu.jinja
+index 7e21585..03e6c7a 100644
+--- a/python/fmha_sm100/csrc/fmha_sm100_variant_run.cu.jinja
++++ b/python/fmha_sm100/csrc/fmha_sm100_variant_run.cu.jinja
+@@ -124,6 +124,12 @@ void FMHAVariantRun_{{ variant_name }}(ffi::TensorView workspace_buffer, ffi::Te
+   params.kv_block_num = maybe_kv_block_indexes.has_value()
+                             ? static_cast<int>(maybe_kv_block_indexes.value().size(2))
+                             : 0;
++  params.kv_block_indexes_stride_q = maybe_kv_block_indexes.has_value()
++                                         ? maybe_kv_block_indexes.value().stride(0) : 0;
++  params.kv_block_indexes_stride_h = maybe_kv_block_indexes.has_value()
++                                         ? maybe_kv_block_indexes.value().stride(1) : 0;
++  params.kv_block_indexes_stride_k = maybe_kv_block_indexes.has_value()
++                                         ? maybe_kv_block_indexes.value().stride(2) : 0;
+
+   params.pack_factor = static_cast<int>(pack_factor);
+   params.num_ctas = static_cast<int>(packed_work_range.size(0));
+@@ -188,9 +194,12 @@ void FMHAVariantRun_{{ variant_name }}(ffi::TensorView workspace_buffer, ffi::Te
+                          * maybe_max_score.value().size(1)
+                          * maybe_max_score.value().size(2)) : 0;
+   params.gmem_bounds.kv_block_indexes_numel = maybe_kv_block_indexes.has_value()
+-      ? static_cast<int>(maybe_kv_block_indexes.value().size(0)
+-                         * maybe_kv_block_indexes.value().size(1)
+-                         * maybe_kv_block_indexes.value().size(2)) : 0;
++      ? static_cast<int>((maybe_kv_block_indexes.value().size(0) - 1)
++                             * maybe_kv_block_indexes.value().stride(0)
++                         + (maybe_kv_block_indexes.value().size(1) - 1)
++                             * maybe_kv_block_indexes.value().stride(1)
++                         + (maybe_kv_block_indexes.value().size(2) - 1)
++                             * maybe_kv_block_indexes.value().stride(2) + 1) : 0;
+   params.gmem_bounds.split_kv_size = maybe_kv_tile_begin_indices.has_value()
+       ? static_cast<int>(maybe_kv_tile_begin_indices.value().size(0)) : 0;
+   params.gmem_bounds.qo_offsets_size = maybe_qo_offsets.has_value()
+diff --git a/python/fmha_sm100/csrc/include/fmha_cutlass_sm100.cuh b/python/fmha_sm100/csrc/include/fmha_cutlass_sm100.cuh
+index b11abe6..f6f95f7 100644
+--- a/python/fmha_sm100/csrc/include/fmha_cutlass_sm100.cuh
++++ b/python/fmha_sm100/csrc/include/fmha_cutlass_sm100.cuh
+@@ -126,6 +126,9 @@ struct FwdRunner {
+                          int max_k_tiles = 0,
+                          int* kv_block_indexes = nullptr,
+                          int kv_block_num = 0,
++                         int64_t kv_block_indexes_stride_q = 0,
++                         int64_t kv_block_indexes_stride_h = 0,
++                         int64_t kv_block_indexes_stride_k = 0,
+                          int pack_factor = 1,
+                          int q_stride_n_original = 0,
+                          int q_stride_h_original = 0,
+@@ -224,6 +227,7 @@ struct FwdRunner {
+         arguments = {
+           {problem_shape,
+            {{q, layout_Q, k, layout_K, v, layout_V, kv_indices, kv_page_indptr, kv_block_indexes, kv_block_num,
++             kv_block_indexes_stride_q, kv_block_indexes_stride_h, kv_block_indexes_stride_k,
+ #ifdef FMHA_GMEM_BOUNDS_CHECK
+              pack_gqa.gmem_bounds.kv_page_indptr_size, pack_gqa.gmem_bounds.kv_indices_size, pack_gqa.gmem_bounds.kv_block_indexes_numel,
+ #endif
+@@ -245,6 +249,7 @@ struct FwdRunner {
+         arguments = {
+           problem_shape,
+           {{q, layout_Q, k, layout_K, v, layout_V, kv_indices, kv_page_indptr, kv_block_indexes, kv_block_num,
++            kv_block_indexes_stride_q, kv_block_indexes_stride_h, kv_block_indexes_stride_k,
+ #ifdef FMHA_GMEM_BOUNDS_CHECK
+             pack_gqa.gmem_bounds.kv_page_indptr_size, pack_gqa.gmem_bounds.kv_indices_size, pack_gqa.gmem_bounds.kv_block_indexes_numel,
+ #endif
+@@ -436,6 +441,9 @@ cudaError_t run_fmha_fwd(void* workspace_buffer, DTypeIn* q, DTypeIn* k, DTypeIn
+                          int max_k_tiles = 0,
+                          int* kv_block_indexes = nullptr,
+                          int kv_block_num = 0,
++                         int64_t kv_block_indexes_stride_q = 0,
++                         int64_t kv_block_indexes_stride_h = 0,
++                         int64_t kv_block_indexes_stride_k = 0,
+                          int pack_factor = 1,
+                          int q_stride_n_original = 0,
+                          int q_stride_h_original = 0,
+@@ -455,6 +463,7 @@ cudaError_t run_fmha_fwd(void* workspace_buffer, DTypeIn* q, DTypeIn* k, DTypeIn
+       kv_indices, kv_page_indptr, total_page_num,
+       maybe_max_score, max_k_tiles,
+       kv_block_indexes, kv_block_num,
++      kv_block_indexes_stride_q, kv_block_indexes_stride_h, kv_block_indexes_stride_k,
+       pack_factor, q_stride_n_original, q_stride_h_original, h_r_original,
+       pack_gqa, num_ctas);
+ }
+diff --git a/python/fmha_sm100/csrc/include/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp b/python/fmha_sm100/csrc/include/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp
+index 5a5d0f3..843aae8 100644
+--- a/python/fmha_sm100/csrc/include/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp
++++ b/python/fmha_sm100/csrc/include/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp
+@@ -1131,14 +1131,13 @@ struct Sm100FmhaFwdMainloopTmaWarpspecialized {
+     int skip_count, real_masked_count, unmasked_count;
+     int effective_end;
+
+-    int kbi_off_s = 0;
++    int64_t kbi_off_s = 0;
+     if constexpr (kNeedSparse) {
+       int h_r_s = get<3, 0, 0>(problem_shape);
+-      int num_kv_heads_s = get<3, 0, 1>(problem_shape);
+       int kv_head_idx_s = get<2, 0>(blk_coord) / h_r_s;
+       int batch_idx_s = get<2, 1>(blk_coord);
+-      kbi_off_s = batch_idx_s * num_kv_heads_s * params.load.kv_block_num
+-                + kv_head_idx_s * params.load.kv_block_num;
++      kbi_off_s = static_cast<int64_t>(batch_idx_s) * params.load.kv_block_indexes_stride_q
++                + static_cast<int64_t>(kv_head_idx_s) * params.load.kv_block_indexes_stride_h;
+     }
+
+     if constexpr (kNeedSparse) {
+@@ -1154,7 +1153,10 @@ struct Sm100FmhaFwdMainloopTmaWarpspecialized {
+
+       // Count padding blocks (backward scan for -1)
+       int valid_blks = params.load.kv_block_num;
+-      while (valid_blks > 0 && __ldg(&params.load.kv_block_indexes[kbi_off_s + valid_blks - 1]) < 0)
++      while (valid_blks > 0
++             && __ldg(&params.load.kv_block_indexes[kbi_off_s
++                    + static_cast<int64_t>(valid_blks - 1)
++                        * params.load.kv_block_indexes_stride_k]) < 0)
+         valid_blks--;
+
+       int valid_tiles = (valid_blks * KVPageSize + full_tile_kv - 1) / full_tile_kv;
+@@ -1223,7 +1225,8 @@ struct Sm100FmhaFwdMainloopTmaWarpspecialized {
+           cS = domain_offset(make_coord(q_s, INT_MAX / 2), cS_base);
+           return false;
+         }
+-        int pos = __ldg(&params.load.kv_block_indexes[kbi_off_s + page_idx]);
++        int pos = __ldg(&params.load.kv_block_indexes[kbi_off_s
++            + static_cast<int64_t>(page_idx) * params.load.kv_block_indexes_stride_k]);
+         if (pos < 0) {
+           cS = domain_offset(make_coord(q_s, INT_MAX / 2), cS_base);
+           return false;
+diff --git a/python/fmha_sm100/csrc/include/sm100_fmha_load_tma_warpspecialized.hpp b/python/fmha_sm100/csrc/include/sm100_fmha_load_tma_warpspecialized.hpp
+index 0fdcbfa..c577f34 100644
+--- a/python/fmha_sm100/csrc/include/sm100_fmha_load_tma_warpspecialized.hpp
++++ b/python/fmha_sm100/csrc/include/sm100_fmha_load_tma_warpspecialized.hpp
+@@ -112,6 +112,9 @@ struct Sm100FmhaLoadTmaWarpspecialized {
+     int* kv_page_indptr = nullptr;
+     int* kv_block_indexes = nullptr;
+     int kv_block_num = 0;
++    int64_t kv_block_indexes_stride_q = 0;
++    int64_t kv_block_indexes_stride_h = 0;
++    int64_t kv_block_indexes_stride_k = 0;
+ #ifdef FMHA_GMEM_BOUNDS_CHECK
+     int kv_page_indptr_size = 0;
+     int kv_indices_size = 0;
+@@ -148,6 +151,9 @@ struct Sm100FmhaLoadTmaWarpspecialized {
+     int* kv_page_indptr = nullptr;
+     int* kv_block_indexes = nullptr;
+     int kv_block_num = 0;
++    int64_t kv_block_indexes_stride_q = 0;
++    int64_t kv_block_indexes_stride_h = 0;
++    int64_t kv_block_indexes_stride_k = 0;
+ #ifdef FMHA_GMEM_BOUNDS_CHECK
+     int kv_page_indptr_size = 0;
+     int kv_indices_size = 0;
+@@ -218,6 +224,8 @@ struct Sm100FmhaLoadTmaWarpspecialized {
+
+     Params p{tma_load_Q, layout_Q, tma_load_K, layout_K, tma_load_V, layout_V,
+                   args.kv_indices, args.kv_page_indptr, args.kv_block_indexes, args.kv_block_num,
++                  args.kv_block_indexes_stride_q, args.kv_block_indexes_stride_h,
++                  args.kv_block_indexes_stride_k,
+ #ifdef FMHA_GMEM_BOUNDS_CHECK
+                   args.kv_page_indptr_size, args.kv_indices_size, args.kv_block_indexes_numel,
+ #endif
+@@ -525,11 +533,10 @@ struct Sm100FmhaLoadTmaWarpspecialized {
+       constexpr int effective_tile_kv = get<1>(TileShapeQK{});
+       constexpr int tiles_per_page = KVPageSize / effective_tile_kv;
+
+-      int kv_block_offset = 0;
++      int64_t kv_block_offset = 0;
+       if constexpr (kSparseAttnMode == SparseAttnMode::Sparse) {
+-        int num_kv_heads_val = get<3, 0, 1>(params_problem_shape);
+-        kv_block_offset = batch_idx * num_kv_heads_val * params.kv_block_num
+-                        + kv_head_idx * params.kv_block_num;
++        kv_block_offset = static_cast<int64_t>(batch_idx) * params.kv_block_indexes_stride_q
++                        + static_cast<int64_t>(kv_head_idx) * params.kv_block_indexes_stride_h;
+       }
+
+       // Keep full 4D tensor (page_size, D, H_kv, P) — select head at copy time
+@@ -549,7 +556,8 @@ struct Sm100FmhaLoadTmaWarpspecialized {
+         int page_for_lookup;
+         if constexpr (kSparseAttnMode == SparseAttnMode::Sparse) {
+           int sparse_idx = (logical_page < params.kv_block_num)
+-              ? __ldg(&params.kv_block_indexes[kv_block_offset + logical_page])
++              ? __ldg(&params.kv_block_indexes[kv_block_offset
++                    + static_cast<int64_t>(logical_page) * params.kv_block_indexes_stride_k])
+               : -1;
+           page_for_lookup = (sparse_idx >= 0) ? sparse_idx : 0;
+         } else {
+@@ -573,7 +581,8 @@ struct Sm100FmhaLoadTmaWarpspecialized {
+         int page_for_lookup;
+         if constexpr (kSparseAttnMode == SparseAttnMode::Sparse) {
+           int sparse_idx = (logical_page < params.kv_block_num)
+-              ? __ldg(&params.kv_block_indexes[kv_block_offset + logical_page])
++              ? __ldg(&params.kv_block_indexes[kv_block_offset
++                    + static_cast<int64_t>(logical_page) * params.kv_block_indexes_stride_k])
+               : -1;
+           page_for_lookup = (sparse_idx >= 0) ? sparse_idx : 0;
+         } else {
+diff --git a/tests/regression/test_noncontiguous_sparse_block_indexes.py b/tests/regression/test_noncontiguous_sparse_block_indexes.py
+new file mode 100644
+index 0000000..6a162b0
+--- /dev/null
++++ b/tests/regression/test_noncontiguous_sparse_block_indexes.py
+@@ -0,0 +1,85 @@
++#!/usr/bin/env python3
++# SPDX-FileCopyrightText: Copyright (c) 2026 MiniMax
++# SPDX-License-Identifier: MIT
++
++"""Regression for sparse attention consuming a head-major selector view."""
++
++import torch
++
++from fmha_sm100.api import fmha_sm100, fmha_sm100_plan
++
++
++def _run_attention(q, k, v, plan, kv_indices, selected):
++    out = torch.empty_like(q)
++    returned, _ = fmha_sm100(
++        q,
++        k,
++        v,
++        plan,
++        kv_indices=kv_indices,
++        kv_block_indexes=selected,
++        out=out,
++        sm_scale=128.0**-0.5,
++        output_maxscore=False,
++    )
++    torch.cuda.synchronize()
++    assert returned.data_ptr() == out.data_ptr()
++    assert torch.isfinite(out).all()
++    return out
++
++
++def _check_query_length(query_len, num_kv_splits):
++    device = torch.device("cuda", 0)
++    total_selector_rows = query_len + 2
++    q = torch.zeros((query_len, 32, 128), dtype=torch.bfloat16, device=device)
++    k = torch.zeros((8, 2, 128, 128), dtype=torch.bfloat16, device=device)
++    v = torch.empty_like(k)
++    for page in range(8):
++        v[page].fill_(page + 1)
++    kv_indices = torch.arange(8, dtype=torch.int32, device=device)
++
++    intended = torch.tensor([0, 1, 2, 3], dtype=torch.int32, device=device)
++    poison = torch.tensor([4, 5, 6, 7], dtype=torch.int32, device=device)
++    logical = intended.expand(total_selector_rows, 2, 4).clone()
++    # These rows are outside the logical slice but occupy addresses reached by
++    # a pointer-only contiguous read for later query/head pairs.
++    logical[:2, 1, :] = poison
++    backing = logical.permute(1, 0, 2).contiguous()
++    selected_strided = backing.permute(1, 0, 2)[-query_len:]
++    selected_contiguous = logical[-query_len:].contiguous()
++
++    assert torch.equal(selected_strided, selected_contiguous)
++    assert not selected_strided.is_contiguous()
++
++    plan = fmha_sm100_plan(
++        torch.tensor([query_len], dtype=torch.int32),
++        torch.tensor([1024], dtype=torch.int32),
++        32,
++        num_kv_heads=2,
++        qo_offset=torch.tensor([1024 - query_len], dtype=torch.int32),
++        num_kv_splits=num_kv_splits,
++        page_size=128,
++        output_maxscore=False,
++        kv_block_num=4,
++        causal=True,
++        device=device,
++    )
++    assert plan[3]["MM-SA-Nv"] is False
++
++    expected = _run_attention(
++        q, k, v, plan, kv_indices, selected_contiguous
++    )
++    actual = _run_attention(q, k, v, plan, kv_indices, selected_strided)
++    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
++
++
++def test_noncontiguous_sparse_block_indexes():
++    torch.cuda.set_device(0)
++    for query_len in (1, 5):
++        for num_kv_splits in (1, 2):
++            _check_query_length(query_len, num_kv_splits)
++
++
++if __name__ == "__main__":
++    test_noncontiguous_sparse_block_indexes()
++    print("MSA_NONCONTIGUOUS_SPARSE_BLOCK_INDEX_REGRESSION=PASS")