diff --git a/3rdparty/patches/msa_strided_paged_kv.patch b/3rdparty/patches/msa_strided_paged_kv.patch index a1be1ca5cc82..73e05b4333cb 100644 --- a/3rdparty/patches/msa_strided_paged_kv.patch +++ b/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(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(pack_factor); + params.num_ctas = static_cast(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(maybe_kv_block_indexes.value().size(0) +- * maybe_kv_block_indexes.value().size(1) +- * maybe_kv_block_indexes.value().size(2)) : 0; ++ ? static_cast((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(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(batch_idx_s) * params.load.kv_block_indexes_stride_q ++ + static_cast(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(¶ms.load.kv_block_indexes[kbi_off_s + valid_blks - 1]) < 0) ++ while (valid_blks > 0 ++ && __ldg(¶ms.load.kv_block_indexes[kbi_off_s ++ + static_cast(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(¶ms.load.kv_block_indexes[kbi_off_s + page_idx]); ++ int pos = __ldg(¶ms.load.kv_block_indexes[kbi_off_s ++ + static_cast(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(batch_idx) * params.kv_block_indexes_stride_q ++ + static_cast(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(¶ms.kv_block_indexes[kv_block_offset + logical_page]) ++ ? __ldg(¶ms.kv_block_indexes[kv_block_offset ++ + static_cast(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(¶ms.kv_block_indexes[kv_block_offset + logical_page]) ++ ? __ldg(¶ms.kv_block_indexes[kv_block_offset ++ + static_cast(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") diff --git a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py index 34ad72226453..bad46bcf5a4a 100644 --- a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py @@ -4,8 +4,9 @@ Most of these validate backend selection and decode scratch-buffer sizing without launching kernels; the CUDA-gated tests cover the fused cache-scatter -parity and strided-cache kernel paths. Numerical parity against the Triton -reference is covered by the SM100 integration accuracy test. +parity, strided-cache kernel paths, and sparse block-index stride contract. +Numerical parity against the Triton reference is covered by the SM100 +integration accuracy test. """ from types import SimpleNamespace @@ -1289,3 +1290,86 @@ def run(**resolution): torch.testing.assert_close( split[num_ctx_tokens:], reference[num_ctx_tokens:], rtol=6e-2, atol=6e-2 ) + + +def _run_msa_sparse_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + plan: tuple, + kv_indices: torch.Tensor, + selected: torch.Tensor, +) -> torch.Tensor: + from fmha_sm100.api import fmha_sm100 + + 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 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("query_len", [1, 5]) +@pytest.mark.parametrize("num_kv_splits", [1, 2]) +def test_msa_sparse_attention_honors_noncontiguous_block_indexes( + query_len: int, + num_kv_splits: int, +) -> None: + from fmha_sm100.api import fmha_sm100_plan + + device = torch.device("cuda", 0) + torch.cuda.set_device(device) + 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 selected_strided.stride() == (4, total_selector_rows * 4, 1) + assert not selected_strided.is_contiguous() + assert torch.equal(selected_strided, selected_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, + ) + short_plan = plan[3] + assert short_plan["MM-SA-Nv"] is False + assert short_plan["num_kv_splits"] == num_kv_splits + + expected = _run_msa_sparse_attention(q, k, v, plan, kv_indices, selected_contiguous) + actual = _run_msa_sparse_attention(q, k, v, plan, kv_indices, selected_strided) + torch.testing.assert_close(actual, expected, rtol=0, atol=0)