diff --git a/flash-attention-v100/kernel/fused_mha_forward_paged.cu b/flash-attention-v100/kernel/fused_mha_forward_paged.cu index 0a5a55f406..20de8125ab 100644 --- a/flash-attention-v100/kernel/fused_mha_forward_paged.cu +++ b/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( K_cache + (int64_t)shared.page_idx[tid] * k_block_stride + (int64_t)shared.page_offset[tid] * k_token_stride + @@ -2933,15 +2932,16 @@ __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( 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 @@ -2949,8 +2949,9 @@ __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, @@ -2958,10 +2959,10 @@ __launch_bounds__(D256_BM32_PHASE_THREADS, 2) void flash_attention_forward_paged flash_attention_forward_paged_d256_bm32_phase_body( 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 @@ -2991,9 +2992,10 @@ void launch_flash_attention_forward_paged_d256_bm32_phase_kernel( reinterpret_cast(V_cache.data_ptr()), reinterpret_cast<__half*>(Out.data_ptr()), softmax_lse.data_ptr(), block_table.data_ptr(), - seq_lens.data_ptr(), 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(), H, M, max_num_blocks_per_seq, + static_cast(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(K_cache.data_ptr()), reinterpret_cast(V_cache.data_ptr()), block_table.data_ptr(), H, M, actual_n, max_num_blocks_per_seq, - num_kv_heads, k_block_stride, k_token_stride, k_head_stride, + static_cast(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(), split_tmp_row_max.data_ptr(), @@ -3098,7 +3101,8 @@ void launch_flash_attention_forward_paged_d256_bm32_splitkv3_kernel( reinterpret_cast(K_cache.data_ptr()), reinterpret_cast(V_cache.data_ptr()), block_table.data_ptr(), H, M, actual_n, max_num_blocks_per_seq, - num_kv_heads, k_block_stride, k_token_stride, k_head_stride, + static_cast(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(), split_tmp_row_max.data_ptr(), @@ -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) { diff --git a/tests/kernels/attention/test_sm70_flash_v100_paged_prefill_any_page.py b/tests/kernels/attention/test_sm70_flash_v100_paged_prefill_any_page.py new file mode 100644 index 0000000000..98b00a0fe6 --- /dev/null +++ b/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) diff --git a/tests/kernels/attention/test_sm70_flash_v100_prefix_decode_rows.py b/tests/kernels/attention/test_sm70_flash_v100_prefix_decode_rows.py new file mode 100644 index 0000000000..57b59108f1 --- /dev/null +++ b/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() diff --git a/vllm/envs.py b/vllm/envs.py index 402623b4d5..a65e0d75ab 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -108,6 +108,7 @@ VLLM_FORCE_AOT_LOAD: bool = False VLLM_USE_MEGA_AOT_ARTIFACT: bool = False VLLM_USE_TRITON_AWQ: bool = False + VLLM_1CAT_PREFILL_PACE_STEPS: int = 0 VLLM_1CAT_ENABLE_SM70_MTP_DEFAULTS: bool = False VLLM_1CAT_ENABLE_QWEN35_MTP_DEFAULTS: bool = False VLLM_1CAT_DISABLE_SM70_MTP_DEFAULTS: bool = False @@ -401,6 +402,7 @@ VLLM_SM70_GDN_MIXED_QKV_CONTIGUOUS: bool = False VLLM_SM70_DECODE_TILE_PROFILE: bool = False VLLM_FLASH_V100_ROUTE_SUMMARY: bool = False + VLLM_FLASH_V100_PREFILL_PREFIX_DECODE_ROWS: bool = False VLLM_FLASH_V100_FP8_PREFILL_BRIDGE: bool = True VLLM_FLASH_V100_DECODE_FP8_XQA_MIN_SEQ_LEN: int = 16384 VLLM_FLASH_V100_KERNEL_BLOCK_SIZE16: bool = False @@ -410,6 +412,7 @@ VLLM_FLASH_V100_PREFILL_D256_SCALAR_QK: bool = False VLLM_FLASH_V100_PREFILL_D256_BM32: bool = False VLLM_FLASH_V100_PREFILL_D256_BM32_PHASE: bool = True + VLLM_FLASH_V100_PREFILL_D256_BM32_ANY_PAGE: bool = False VLLM_FLASH_V100_PREFILL_D256_BM32_ALL_P: bool = True VLLM_FLASH_V100_PREFILL_D256_BM32_PAIR_SCRATCH: bool = True VLLM_FLASH_V100_PREFILL_D256_OUTPUT_STRIDE_268: bool = True @@ -1641,6 +1644,12 @@ def _resolve_rust_frontend_path() -> str | None: # 1Cat SM70 public-profile MTP opt-ins/opt-outs. These are consumed while # building EngineArgs and must be registered so environment validation does # not warn users that our own documented knobs are unknown. + # Experimental (1CatAI/1Cat-vLLM#490): while another running request is + # decoding, a request that is still prefilling is scheduled a chunk only + # every N engine steps; the other N-1 steps are decode-only. 0 = off. + "VLLM_1CAT_PREFILL_PACE_STEPS": lambda: int( + os.getenv("VLLM_1CAT_PREFILL_PACE_STEPS", "0") + ), "VLLM_1CAT_ENABLE_SM70_MTP_DEFAULTS": lambda: bool( int(os.getenv("VLLM_1CAT_ENABLE_SM70_MTP_DEFAULTS", "0")) ), @@ -2889,6 +2898,12 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_FLASH_V100_ROUTE_SUMMARY": lambda: bool( int(os.getenv("VLLM_FLASH_V100_ROUTE_SUMMARY", "0")) ), + # Experimental (1CatAI/1Cat-vLLM#490): inside a mixed prefill+decode batch, + # run the q=1 rows through the paged decode kernels instead of the + # per-sequence paged prefill kernel. Default off; default path unchanged. + "VLLM_FLASH_V100_PREFILL_PREFIX_DECODE_ROWS": lambda: bool( + int(os.getenv("VLLM_FLASH_V100_PREFILL_PREFIX_DECODE_ROWS", "0")) + ), "VLLM_FLASH_V100_FP8_PREFILL_BRIDGE": lambda: bool( int(os.getenv("VLLM_FLASH_V100_FP8_PREFILL_BRIDGE", "1")) ), @@ -2916,6 +2931,12 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_FLASH_V100_PREFILL_D256_BM32_PHASE": lambda: bool( int(os.getenv("VLLM_FLASH_V100_PREFILL_D256_BM32_PHASE", "1")) ), + # Experimental (1CatAI/1Cat-vLLM#490): let the D256 BM32 phase paged-prefill + # kernel run at any KV page size that is a multiple of 16 (MTP makes the + # align-mode block 816, which otherwise falls off the page-784 fast path). + "VLLM_FLASH_V100_PREFILL_D256_BM32_ANY_PAGE": lambda: bool( + int(os.getenv("VLLM_FLASH_V100_PREFILL_D256_BM32_ANY_PAGE", "0")) + ), "VLLM_FLASH_V100_PREFILL_D256_BM32_ALL_P": lambda: bool( int(os.getenv("VLLM_FLASH_V100_PREFILL_D256_BM32_ALL_P", "1")) ), diff --git a/vllm/v1/attention/backends/flash_attn_v100.py b/vllm/v1/attention/backends/flash_attn_v100.py index a9b98d7c9a..dab19c5b6a 100644 --- a/vllm/v1/attention/backends/flash_attn_v100.py +++ b/vllm/v1/attention/backends/flash_attn_v100.py @@ -533,6 +533,7 @@ def _sm70_profile_trace(message: str, *args: object) -> None: _logged_prefill_prefix_splitkv = False _logged_prefill_paged_cache = False _logged_prefill_smallq_decode = False +_logged_prefill_prefix_decode_rows = False _logged_prefill_smallq_decode_xqa = False _logged_prefill_smallq_grouped_verify = False _logged_prefill_smallq_grouped_verify_gate = False @@ -7822,6 +7823,219 @@ def _should_use_prefill_gather_dense( ) return eligible + def _prefill_prefix_decode_rows_allowed( + self, + *, + causal: bool, + anchor_lens: torch.Tensor | None, + num_seqs: int, + query: torch.Tensor, + window_size: tuple[int, int], + ) -> bool: + return ( + envs.VLLM_FLASH_V100_PREFILL_PREFIX_DECODE_ROWS + and causal + and anchor_lens is None + and num_seqs > 1 + and self.use_flash_v100_decode + and self.use_flash_v100_prefill_paged + and not self.use_decode_paged_prefill + and not self.use_decode_dense_cache + and not self.use_decode_dense_reference + and window_size == (-1, -1) + and not _is_cuda_graph_capturing(query) + ) + + def _run_prefill_prefix_decode_rows( + self, + layer: torch.nn.Module, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + attn_metadata: TritonAttentionMetadata, + out_view: torch.Tensor, + query_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + window_size: tuple[int, int], + ) -> set[int]: + """Run the small-query rows of a mixed batch as one paged-decode batch. + + Inside a chunked-prefill batch every row takes the prefill route, and + ``prefill_paged_fwd`` gives a small-q row one CTA per query head for + the whole context (kernel/fused_mha_api.cpp launches ``grid(ceil(q/BM), + 1, B*H)``). A resident decoder at 240K pays ~58 ms per layer that way + versus ~1.8 ms on the partitioned decode kernel (1CatAI/1Cat-vLLM#490), + and an MTP/DFlash verify row (q = K+1) has the same grid. Every query + token of a selected row becomes one decode row whose visible KV length + grows by one, the expansion _flash_v100_small_query_prefill_as_decode + uses for the verifier, so the causal mask is preserved. Returns the row + indices consumed here; the caller's per-sequence loop skips them. + """ + global _logged_prefill_prefix_decode_rows + num_seqs = len(query_start_loc) - 1 + qsl = query_start_loc[: num_seqs + 1].tolist() + seq_lens_host = seq_lens[:num_seqs].tolist() + max_q = max(1, int(self.smallq_decode_max_query_len)) + rows = [ + i + for i in range(num_seqs) + if 1 <= qsl[i + 1] - qsl[i] <= max_q + and int(seq_lens_host[i]) > qsl[i + 1] - qsl[i] + ] + if not rows or len(rows) == num_seqs: + return set() + + token_idx: list[int] = [] + token_rows: list[int] = [] + token_seq_lens: list[int] = [] + for i in rows: + q_len = qsl[i + 1] - qsl[i] + seq_len = int(seq_lens_host[i]) + for j in range(q_len): + token_idx.append(qsl[i] + j) + token_rows.append(i) + token_seq_lens.append(seq_len - q_len + 1 + j) + device = query.device + start_idx = torch.tensor(token_idx, device=device, dtype=torch.long) + q_rows = query.index_select(0, start_idx) + out_rows = torch.empty_like(q_rows) + block_table = attn_metadata.block_table.index_select( + 0, torch.tensor(token_rows, device=device, dtype=torch.long) + ) + seq_lens_rows = torch.tensor( + token_seq_lens, device=device, dtype=attn_metadata.seq_lens.dtype + ) + max_seq_len_hint = max(token_seq_lens) + max_query_len_rows = max(qsl[i + 1] - qsl[i] for i in rows) + + num_kv_heads = int(key_cache.shape[2]) + q_per_kv = ( + int(q_rows.shape[1]) // num_kv_heads + if num_kv_heads > 0 and int(q_rows.shape[1]) % num_kv_heads == 0 + else 0 + ) + fp16_kv = ( + self.kv_cache_dtype in ("auto", "float16", "bfloat16") + and key_cache.dtype == torch.float16 + and value_cache.dtype == torch.float16 + ) + fp8_e5m2_kv = ( + self.kv_cache_dtype == "fp8_e5m2" + and key_cache.dtype == torch.uint8 + and value_cache.dtype == torch.uint8 + ) + fp8_e4m3_kv = ( + self.kv_cache_dtype in ("fp8", "fp8_e4m3") + and key_cache.dtype == torch.uint8 + and value_cache.dtype == torch.uint8 + # The E4M3 XQA wave route retains half partials; a DFlash2 target + # keeps the FP32 state policy of the scalar route. + and not getattr(attn_metadata, "is_dflash_selector_target", False) + ) + # Same selection as the uniform-decode path (_flash_v100_decode), with + # the sequence hint taken from this batch's rows because build() only + # attaches decode shape hints when max_query_len == 1. + use_xqa = ( + self.use_decode_xqa + and self.flash_attn_decode_paged_xqa is not None + and (fp16_kv or fp8_e5m2_kv or fp8_e4m3_kv) + and int(q_rows.shape[2]) == 256 + and ( + q_per_kv in (6, 8) + or (q_per_kv == 4 and max_seq_len_hint >= _decode_xqa_q4_min_seq_len()) + ) + and ( + not fp8_e4m3_kv + or ( + q_per_kv == 6 + and (q_rows.shape[0] == 1 or _e4m3_batch_xqa_allowed(q_rows)) + ) + ) + and ( + not fp8_e5m2_kv + or (q_per_kv != 4 and max_seq_len_hint >= _decode_fp8_xqa_min_seq_len()) + ) + ) + partition_size_hint = ( + _g6_aligned_page_partition_size_hint( + q_rows, key_cache, value_cache, self.kv_cache_dtype + ) + if use_xqa + else None + ) + k_scale = float(layer._k_scale_float) + v_scale = float(layer._v_scale_float) + if use_xqa: + route = "prefill_prefix_decode_rows_xqa" + + def run() -> torch.Tensor: + self.flash_attn_decode_paged_xqa( + q_rows, + key_cache, + value_cache, + block_table, + seq_lens_rows, + softmax_scale=self.scale, + out=out_rows, + kv_cache_dtype=self.kv_cache_dtype, + k_scale=k_scale, + v_scale=v_scale, + window_size=window_size, + max_seq_len_hint=max_seq_len_hint, + partition_size_hint=partition_size_hint, + ) + return out_rows + + else: + route = "prefill_prefix_decode_rows_scalar" + + def run() -> torch.Tensor: + self._call_flash_attn_decode_paged( + q_rows, + key_cache, + value_cache, + block_table, + seq_lens_rows, + softmax_scale=self.scale, + out=out_rows, + kv_cache_dtype=self.kv_cache_dtype, + k_scale=k_scale, + v_scale=v_scale, + window_size=window_size, + max_seq_len_hint=max_seq_len_hint, + ) + return out_rows + + if not _logged_prefill_prefix_decode_rows: + logger.info( + "FLASH_ATTN_V100 mixed-batch small-query rows take the paged " + "decode route (%s, rows=%d of %d, max_q=%d, max_seq_len=%d).", + route, + len(rows), + num_seqs, + max_query_len_rows, + max_seq_len_hint, + ) + _logged_prefill_prefix_decode_rows = True + self._run_prefill_paged_call( + route=route, + q_len=max_query_len_rows, + seq_len=max_seq_len_hint, + heads_q=int(q_rows.shape[1]), + heads_kv=num_kv_heads, + head_dim=int(q_rows.shape[2]), + block_size=int(key_cache.shape[1]), + fn=run, + ) + _log_fp8_kv_cache_route( + "decode", + self.kv_cache_dtype, + "xqa_paged" if use_xqa else "scalar_paged", + ) + _record_route(route) + out_view.index_copy_(0, start_idx, out_rows) + return set(rows) + def _run_prefill_paged_call( self, *, @@ -7981,7 +8195,29 @@ def _flash_v100_prefill_with_prefix( seq_lens, ) + decode_rows: set[int] = set() + if self._prefill_prefix_decode_rows_allowed( + causal=causal, + anchor_lens=anchor_lens, + num_seqs=num_seqs, + query=query, + window_size=window_size, + ): + decode_rows = self._run_prefill_prefix_decode_rows( + layer, + query, + key_cache, + value_cache, + attn_metadata, + out_view, + query_start_loc, + seq_lens, + window_size, + ) + for i in range(num_seqs): + if i in decode_rows: + continue start = int(query_start_loc[i].item()) end = int(query_start_loc[i + 1].item()) if end <= start: diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 17d9ea05ad..e81d3b47c2 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -8,6 +8,7 @@ from dataclasses import replace from typing import Any +from vllm import envs from vllm.compilation.cuda_graph import CUDAGraphStat from vllm.config import VllmConfig from vllm.distributed.ec_transfer.ec_connector.base import ( @@ -203,6 +204,8 @@ def __init__( # requests skipped in waiting flow due async deps or constraints. self.skipped_waiting = create_request_queue(self.policy) self.running: list[Request] = [] + # Experimental (1CatAI/1Cat-vLLM#490): see VLLM_1CAT_PREFILL_PACE_STEPS. + self._prefill_pace_steps = max(0, int(envs.VLLM_1CAT_PREFILL_PACE_STEPS)) # The request IDs that are finished in between the previous and the # current steps. This is used to notify the workers about the finished @@ -507,6 +510,12 @@ def schedule(self) -> SchedulerOutput: self.kv_cache_manager.new_step_starts() + # Prefill pacing applies only while some running request is decoding; + # a lone prefill keeps its full budget. + pace_prefills = self._prefill_pace_steps > 0 and any( + r.num_computed_tokens >= r.num_prompt_tokens for r in self.running + ) + # First, schedule the RUNNING requests. req_index = 0 while req_index < len(self.running) and token_budget > 0: @@ -717,6 +726,15 @@ def schedule(self) -> SchedulerOutput: num_scheduled_tokens[request_id] = num_new_tokens token_budget -= num_new_tokens req_index += 1 + if ( + pace_prefills + and request.num_computed_tokens + num_new_tokens + < request.num_prompt_tokens + ): + # Still prefilling after this chunk: skip the next N-1 steps. + request.next_decode_eligible_step = ( + self.current_step + self._prefill_pace_steps + ) # Speculative decode related. if request.spec_token_ids: