[Bugfix][SM70] Fix decode collapse next to a long chunked prefill (#490)

flash-attention-v100/kernel/fused_mha_forward_paged.cu

@@ -2529,7 +2529,8 @@ flash_attention_forward_paged_d256_bm32_phase_body(
     const __half* __restrict__ V_cache, __half* __restrict__ Out,
     float* __restrict__ softmax_lse, const int* __restrict__ block_table,
     const int* __restrict__ seqused_k, int H, int M, int max_num_blocks_per_seq,
-    int num_kv_heads, int64_t k_block_stride, int64_t k_token_stride,
+    int page_block_size, int num_kv_heads, int64_t k_block_stride,
+    int64_t k_token_stride,
     int64_t k_head_stride, int64_t v_block_stride, int64_t v_token_stride,
     int64_t v_head_stride, float softmax_scale,
     float* __restrict__ split_tmp_out, float* __restrict__ split_tmp_row_max,
@@ -2607,14 +2608,12 @@ flash_attention_forward_paged_d256_bm32_phase_body(
       const int token_offset = tid * D256_BM32_PHASE_PAGE_SIZE;
       if (token_offset < valid_k_rows) {
         const int global_token_idx = start_col + token_offset;
-        const int virtual_block_idx =
-            global_token_idx / D256_BM32_PHASE_PAGE_BLOCK_SIZE;
+        const int virtual_block_idx = global_token_idx / page_block_size;
         shared.page_idx[tid] =
             __ldg(&block_table[shared.batch_id * max_num_blocks_per_seq +
                                virtual_block_idx]);
         shared.page_offset[tid] =
-            global_token_idx -
-            virtual_block_idx * D256_BM32_PHASE_PAGE_BLOCK_SIZE;
+            global_token_idx - virtual_block_idx * page_block_size;
         shared.k_tile_ptr[tid] = reinterpret_cast<uint64_t>(
             K_cache + (int64_t)shared.page_idx[tid] * k_block_stride +
             (int64_t)shared.page_offset[tid] * k_token_stride +
@@ -2933,35 +2932,37 @@ __launch_bounds__(D256_BM32_PHASE_THREADS, 2) void flash_attention_forward_paged
     const __half* __restrict__ V_cache, __half* __restrict__ Out,
     float* __restrict__ softmax_lse, const int* __restrict__ block_table,
     const int* __restrict__ seqused_k, int H, int M, int max_num_blocks_per_seq,
-    int num_kv_heads, int64_t k_block_stride, int64_t k_token_stride,
+    int page_block_size, int num_kv_heads, int64_t k_block_stride,
+    int64_t k_token_stride,
     int64_t k_head_stride, int64_t v_block_stride, int64_t v_token_stride,
     int64_t v_head_stride, float softmax_scale) {
   flash_attention_forward_paged_d256_bm32_phase_body<IS_CAUSAL, ALL_P,
                                                      PAIR_SCRATCH, false>(
       Q, K_cache, V_cache, Out, softmax_lse, block_table, seqused_k, H, M,
-      max_num_blocks_per_seq, num_kv_heads, k_block_stride, k_token_stride,
-      k_head_stride, v_block_stride, v_token_stride, v_head_stride,
-      softmax_scale, nullptr, nullptr, nullptr, 0);
+      max_num_blocks_per_seq, page_block_size, num_kv_heads, k_block_stride,
+      k_token_stride, k_head_stride, v_block_stride, v_token_stride,
+      v_head_stride, softmax_scale, nullptr, nullptr, nullptr, 0);
 }
 
 template <bool CHECK_SPLIT_EMPTY>
 __global__
 __launch_bounds__(D256_BM32_PHASE_THREADS, 2) void flash_attention_forward_paged_d256_bm32_splitkv3_partial_kernel(
     const __half* __restrict__ Q, const __half* __restrict__ K_cache,
     const __half* __restrict__ V_cache, const int* __restrict__ block_table,
-    int H, int M, int actual_n, int max_num_blocks_per_seq, int num_kv_heads,
-    int64_t k_block_stride, int64_t k_token_stride, int64_t k_head_stride,
+    int H, int M, int actual_n, int max_num_blocks_per_seq, int page_block_size,
+    int num_kv_heads, int64_t k_block_stride, int64_t k_token_stride,
+    int64_t k_head_stride,
     int64_t v_block_stride, int64_t v_token_stride, int64_t v_head_stride,
     float softmax_scale, float* __restrict__ split_tmp_out,
     float* __restrict__ split_tmp_row_max,
     float* __restrict__ split_tmp_row_sum) {
   flash_attention_forward_paged_d256_bm32_phase_body<true, true, true, true,
                                                      CHECK_SPLIT_EMPTY>(
       Q, K_cache, V_cache, nullptr, nullptr, block_table, nullptr, H, M,
-      max_num_blocks_per_seq, num_kv_heads, k_block_stride, k_token_stride,
-      k_head_stride, v_block_stride, v_token_stride, v_head_stride,
-      softmax_scale, split_tmp_out, split_tmp_row_max, split_tmp_row_sum,
-      actual_n);
+      max_num_blocks_per_seq, page_block_size, num_kv_heads, k_block_stride,
+      k_token_stride, k_head_stride, v_block_stride, v_token_stride,
+      v_head_stride, softmax_scale, split_tmp_out, split_tmp_row_max,
+      split_tmp_row_sum, actual_n);
 }
 
 template <bool IS_CAUSAL, bool ALL_P, bool PAIR_SCRATCH>
@@ -2991,9 +2992,10 @@ void launch_flash_attention_forward_paged_d256_bm32_phase_kernel(
           reinterpret_cast<const __half*>(V_cache.data_ptr()),
           reinterpret_cast<__half*>(Out.data_ptr()),
           softmax_lse.data_ptr<float>(), block_table.data_ptr<int>(),
-          seq_lens.data_ptr<int>(), H, M, max_num_blocks_per_seq, num_kv_heads,
-          k_block_stride, k_token_stride, k_head_stride, v_block_stride,
-          v_token_stride, v_head_stride, softmax_scale);
+          seq_lens.data_ptr<int>(), H, M, max_num_blocks_per_seq,
+          static_cast<int>(K_cache.size(1)), num_kv_heads, k_block_stride,
+          k_token_stride, k_head_stride, v_block_stride, v_token_stride,
+          v_head_stride, softmax_scale);
 }
 
 __global__
@@ -3086,7 +3088,8 @@ void launch_flash_attention_forward_paged_d256_bm32_splitkv3_kernel(
             reinterpret_cast<const __half*>(K_cache.data_ptr()),
             reinterpret_cast<const __half*>(V_cache.data_ptr()),
             block_table.data_ptr<int>(), H, M, actual_n, max_num_blocks_per_seq,
-            num_kv_heads, k_block_stride, k_token_stride, k_head_stride,
+            static_cast<int>(K_cache.size(1)), num_kv_heads, k_block_stride,
+            k_token_stride, k_head_stride,
             v_block_stride, v_token_stride, v_head_stride, softmax_scale,
             split_tmp_out.data_ptr<float>(),
             split_tmp_row_max.data_ptr<float>(),
@@ -3098,7 +3101,8 @@ void launch_flash_attention_forward_paged_d256_bm32_splitkv3_kernel(
             reinterpret_cast<const __half*>(K_cache.data_ptr()),
             reinterpret_cast<const __half*>(V_cache.data_ptr()),
             block_table.data_ptr<int>(), H, M, actual_n, max_num_blocks_per_seq,
-            num_kv_heads, k_block_stride, k_token_stride, k_head_stride,
+            static_cast<int>(K_cache.size(1)), num_kv_heads, k_block_stride,
+            k_token_stride, k_head_stride,
             v_block_stride, v_token_stride, v_head_stride, softmax_scale,
             split_tmp_out.data_ptr<float>(),
             split_tmp_row_max.data_ptr<float>(),
@@ -3210,7 +3214,10 @@ void launcher_flash_attention_forward_paged(
     if (use_low_smem) {
       const bool use_d256_bm32_phase =
           env_flag_default_enabled("VLLM_FLASH_V100_PREFILL_D256_BM32_PHASE") &&
-          page_block_size == D256_BM32_PHASE_PAGE_BLOCK_SIZE &&
+          (page_block_size == D256_BM32_PHASE_PAGE_BLOCK_SIZE ||
+           (page_block_size % D256_BM32_PHASE_PAGE_SIZE == 0 &&
+            env_flag_enabled(
+                "VLLM_FLASH_V100_PREFILL_D256_BM32_ANY_PAGE"))) &&
           M >= D256_BM32_PHASE_BLOCK_M && bfla_mask_ptr == nullptr &&
           window_size_left < 0 && window_size_right < 0;
       if (use_d256_bm32_phase) {

tests/kernels/attention/test_sm70_flash_v100_paged_prefill_any_page.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""D256 BM32 phase paged prefill at KV page sizes other than 784.
+
+The align-mode attention block is 784 tokens on fp16 KV, but MTP pads the Mamba
+page and makes it 816. With ``VLLM_FLASH_V100_PREFILL_D256_BM32_ANY_PAGE=1``
+the BM32 phase kernel accepts any page size that is a multiple of its 16-token
+page slot; the output must match the page-784 path and a dense reference
+exactly, and with the flag off page 784 must be unaffected.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+FLAG = "VLLM_FLASH_V100_PREFILL_D256_BM32_ANY_PAGE"
+NUM_HEADS = 12
+NUM_KV_HEADS = 2
+HEAD_DIM = 256
+
+
+def _paged_layout(k_lin, v_lin, page: int, seed: int):
+    seq_len = k_lin.shape[0]
+    num_blocks = (seq_len + page - 1) // page
+    k = torch.zeros(
+        num_blocks, page, NUM_KV_HEADS, HEAD_DIM, dtype=k_lin.dtype, device=k_lin.device
+    )
+    v = torch.zeros_like(k)
+    k.view(-1, NUM_KV_HEADS, HEAD_DIM)[:seq_len] = k_lin
+    v.view(-1, NUM_KV_HEADS, HEAD_DIM)[:seq_len] = v_lin
+    gen = torch.Generator(device="cpu").manual_seed(seed)
+    perm = torch.randperm(num_blocks, generator=gen).to(k_lin.device)
+    block_table = torch.argsort(perm).to(torch.int32).unsqueeze(0)
+    seq_lens = torch.tensor([seq_len], dtype=torch.int32, device=k_lin.device)
+    return k[perm].contiguous(), v[perm].contiguous(), block_table, seq_lens
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+@pytest.mark.parametrize("any_page", [False, True])
+@pytest.mark.parametrize("page", [816, 896])
+@pytest.mark.parametrize("query_len", [32, 784])
+@torch.inference_mode()
+def test_bm32_phase_any_page_matches_page784_and_dense(
+    monkeypatch: pytest.MonkeyPatch, any_page: bool, page: int, query_len: int
+) -> None:
+    if torch.cuda.get_device_capability() != (7, 0):
+        pytest.skip("FlashAttention-V100 is SM70/V100 only")
+    fi = pytest.importorskip("flash_attn_v100.flash_attn_interface")
+    if any_page:
+        monkeypatch.setenv(FLAG, "1")
+    else:
+        monkeypatch.delenv(FLAG, raising=False)
+
+    torch.manual_seed(1234)
+    device = "cuda"
+    seq_len = 6 * 784 + 100  # several pages, unaligned tail
+    scale = HEAD_DIM**-0.5
+    k_lin = torch.randn(
+        seq_len, NUM_KV_HEADS, HEAD_DIM, dtype=torch.float16, device=device
+    )
+    v_lin = torch.randn_like(k_lin)
+    query = torch.randn(
+        1, query_len, NUM_HEADS, HEAD_DIM, dtype=torch.float16, device=device
+    )
+
+    dense = fi.flash_attn_func(
+        query, k_lin.unsqueeze(0), v_lin.unsqueeze(0), causal=True, softmax_scale=scale
+    )
+    outs = {}
+    for p in (784, page):
+        k, v, block_table, seq_lens = _paged_layout(k_lin, v_lin, p, seed=p)
+        outs[p] = fi.flash_attn_prefill_paged(
+            query, k, v, block_table, seq_lens, softmax_scale=scale, causal=True
+        )
+    torch.accelerator.synchronize()
+
+    assert torch.isfinite(outs[page]).all()
+    torch.testing.assert_close(outs[784], dense, atol=2e-3, rtol=1e-2)
+    torch.testing.assert_close(outs[page], dense, atol=2e-3, rtol=1e-2)
+    torch.testing.assert_close(outs[page], outs[784], atol=2e-3, rtol=1e-2)

tests/kernels/attention/test_sm70_flash_v100_prefix_decode_rows.py

@@ -0,0 +1,169 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Mixed prefill+decode batches: small-query rows may take the decode route.
+
+Opt-in via ``VLLM_FLASH_V100_PREFILL_PREFIX_DECODE_ROWS``. With the flag off
+the per-sequence paged prefill loop is untouched. With it on, rows with
+``1 <= q <= VLLM_FLASH_V100_SMALLQ_DECODE_MAX_Q`` and prefix context (a
+resident decoder, or an MTP/DFlash verify row) are expanded token-wise into
+one paged-decode call and must match the prefill kernel within fp16
+tolerance, while the chunk row stays bit-identical to the flag-off output
+(same kernel, same inputs).
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+FLAG = "VLLM_FLASH_V100_PREFILL_PREFIX_DECODE_ROWS"
+NUM_HEADS = 12
+NUM_KV_HEADS = 2
+HEAD_DIM = 256
+
+
+def _make_impl(kv_cache_dtype: str):
+    from vllm.v1.attention.backends.flash_attn_v100 import FlashAttnV100Impl
+
+    return FlashAttnV100Impl(
+        num_heads=NUM_HEADS,
+        head_size=HEAD_DIM,
+        scale=HEAD_DIM**-0.5,
+        num_kv_heads=NUM_KV_HEADS,
+        alibi_slopes=None,
+        sliding_window=None,
+        kv_cache_dtype=kv_cache_dtype,
+    )
+
+
+def _make_mixed_batch(
+    *,
+    device: str,
+    block_size: int,
+    kv_cache_dtype: str,
+    chunk_len: int,
+    chunk_context: int,
+    small_rows: list[tuple[int, int]],
+):
+    """Row 0 is a chunked-prefill row; the rest are (q_len, seq_len) rows."""
+    query_lens = [chunk_len] + [q for q, _ in small_rows]
+    seq_lens = [chunk_context + chunk_len] + [s for _, s in small_rows]
+    blocks_per_row = [(s + block_size - 1) // block_size for s in seq_lens]
+    total_blocks = sum(blocks_per_row) + 1  # +1 keeps the K/V axis unambiguous
+    kv_cache = torch.randn(
+        2,
+        total_blocks,
+        block_size,
+        NUM_KV_HEADS,
+        HEAD_DIM,
+        dtype=torch.float16,
+        device=device,
+    )
+    if kv_cache_dtype == "fp8_e5m2":
+        # Real E5M2 values (no NaN/Inf bit patterns), stored as the uint8
+        # bytes the paged cache carries.
+        kv_cache = kv_cache.to(torch.float8_e5m2).view(torch.uint8)
+    max_blocks = max(blocks_per_row)
+    block_table = torch.zeros(len(seq_lens), max_blocks, dtype=torch.int32)
+    next_block = 1
+    for row, n in enumerate(blocks_per_row):
+        block_table[row, :n] = torch.arange(next_block, next_block + n)
+        next_block += n
+    query_start_loc_cpu = torch.tensor(
+        [0] + list(torch.cumsum(torch.tensor(query_lens), 0)), dtype=torch.int32
+    )
+    seq_lens_cpu = torch.tensor(seq_lens, dtype=torch.int32)
+    num_tokens = int(query_start_loc_cpu[-1].item())
+    query = torch.randn(
+        num_tokens, NUM_HEADS, HEAD_DIM, dtype=torch.float16, device=device
+    )
+    attn_metadata = SimpleNamespace(
+        query_start_loc=query_start_loc_cpu.to(device),
+        query_start_loc_cpu=query_start_loc_cpu,
+        seq_lens=seq_lens_cpu.to(device),
+        seq_lens_cpu=seq_lens_cpu,
+        block_table=block_table.to(device),
+        num_actual_tokens=num_tokens,
+        max_query_len=max(query_lens),
+        causal=True,
+        max_model_len=262144,
+    )
+    return query, kv_cache, attn_metadata, query_start_loc_cpu
+
+
+def _run(impl, query, kv_cache, attn_metadata) -> torch.Tensor:
+    layer = SimpleNamespace(_k_scale_float=1.0, _v_scale_float=1.0)
+    output = torch.empty_like(query)
+    impl._flash_v100_prefill_with_prefix(
+        layer, query, None, None, kv_cache, attn_metadata, output
+    )
+    torch.accelerator.synchronize()
+    return output
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_e5m2"])
+@pytest.mark.parametrize("block_size", [16, 784])
+@pytest.mark.parametrize(
+    "small_rows",
+    [
+        [(1, 4097)],  # one resident decoder
+        [(1, 4097), (1, 65537)],  # two resident decoders
+        [(5, 4101), (1, 20000)],  # MTP verify row (K=4) plus a decoder
+        [(16, 8208)],  # largest small-q row
+    ],
+)
+@torch.inference_mode()
+def test_prefix_prefill_small_query_rows_match_prefill_route(
+    monkeypatch: pytest.MonkeyPatch,
+    kv_cache_dtype: str,
+    block_size: int,
+    small_rows: list[tuple[int, int]],
+) -> None:
+    if torch.cuda.get_device_capability() != (7, 0):
+        pytest.skip("FlashAttention-V100 is SM70/V100 only")
+    pytest.importorskip("flash_attn_v100")
+    import vllm.v1.attention.backends.flash_attn_v100 as backend
+
+    torch.manual_seed(1234)
+    device = "cuda"
+    impl = _make_impl(kv_cache_dtype)
+    if not (impl.use_flash_v100_prefill_paged and impl.use_flash_v100_decode):
+        pytest.skip("paged prefill and decode ops are both required")
+    query, kv_cache, attn_metadata, query_start_loc = _make_mixed_batch(
+        device=device,
+        block_size=block_size,
+        kv_cache_dtype=kv_cache_dtype,
+        chunk_len=96,
+        chunk_context=1000,
+        small_rows=small_rows,
+    )
+    chunk_end = int(query_start_loc[1].item())
+    spans = list(zip(query_start_loc[1:-1].tolist(), query_start_loc[2:].tolist()))
+
+    # Observe route selection without enabling the summary env (which would
+    # register an atexit hook in the test process).
+    routes: list[str] = []
+    monkeypatch.setattr(backend, "_record_route", routes.append)
+
+    monkeypatch.delenv(FLAG, raising=False)
+    out_off = _run(impl, query, kv_cache, attn_metadata)
+    assert not any(r.startswith("prefill_prefix_decode_rows") for r in routes)
+
+    monkeypatch.setenv(FLAG, "1")
+    routes.clear()
+    out_on = _run(impl, query, kv_cache, attn_metadata)
+    assert any(r.startswith("prefill_prefix_decode_rows") for r in routes), routes
+
+    # Chunk row: same kernel, same inputs -> bit-identical.
+    assert torch.equal(out_on[:chunk_end], out_off[:chunk_end])
+    # Small-q rows: decode kernel versus paged prefill kernel; both read the
+    # same cache bytes, so only accumulation order differs.
+    atol = 1e-2 if kv_cache_dtype == "fp8_e5m2" else 5e-3
+    for start, end in spans:
+        torch.testing.assert_close(
+            out_on[start:end], out_off[start:end], atol=atol, rtol=1e-2
+        )
+    assert torch.isfinite(out_on).all()