diff --git a/benchmarks/benchmark_sm70_decode.py b/benchmarks/benchmark_sm70_decode.py index 93f4596100..6a482ac926 100644 --- a/benchmarks/benchmark_sm70_decode.py +++ b/benchmarks/benchmark_sm70_decode.py @@ -1746,6 +1746,7 @@ def _write_results( "engine_kwargs": llm_kwargs, "sampling_params": first_case["sampling_params"], "cuda_profile_repeat": args.cuda_profile_repeat, + "cuda_profile_case": args.cuda_profile_case, "reset_prefix_cache_before_case": args.reset_prefix_cache_before_case, "reset_prefix_cache_between_runs": args.reset_prefix_cache_between_runs, "prompt": first_case["prompt"], @@ -1873,6 +1874,14 @@ def _parse_args() -> argparse.Namespace: action="store_true", help="Wrap only measured repeats in cudaProfilerStart/Stop.", ) + parser.add_argument( + "--cuda-profile-case", + help=( + "When profiling measured repeats, capture only the named case. " + "Other cases remain in the same loaded engine but outside the " + "cudaProfilerStart/Stop range." + ), + ) parser.add_argument( "--require-sm70-fa2-d256-prefill", action="store_true", @@ -1895,6 +1904,8 @@ def main() -> int: raise ValueError( "Choose only one prefix-cache reset mode: before-case or between-runs." ) + if args.cuda_profile_case and not args.cuda_profile_repeat: + raise ValueError("--cuda-profile-case requires --cuda-profile-repeat") import torch from transformers import AutoTokenizer @@ -1942,6 +1953,12 @@ def main() -> int: ), } ) + if args.cuda_profile_case and not any( + case["name"] == args.cuda_profile_case for case in cases + ): + raise ValueError( + f"--cuda-profile-case did not match a case: {args.cuda_profile_case}" + ) engine_kwargs = _parse_extra_engine_args(args.engine_arg) llm_kwargs: dict[str, Any] = { @@ -1998,6 +2015,18 @@ def main() -> int: } ) + def start_cuda_profile() -> None: + try: + llm.start_profile() + except Exception as exc: + raise RuntimeError( + "--cuda-profile-repeat requires a worker profiler for " + "multiprocess TP runs. Pass " + '--engine-arg \'profiler_config={"profiler":"cuda"}\' ' + "so Nsight Compute/Systems capture the TP worker kernels." + ) from exc + + profile_active = False if args.cuda_profile_repeat: # Warm every case before starting the profiler so the capture contains # measured repeats only. Normal sweeps warm each case immediately @@ -2014,15 +2043,9 @@ def main() -> int: ) for _ in range(case["warmup"]) ] - try: - llm.start_profile() - except Exception as exc: - raise RuntimeError( - "--cuda-profile-repeat requires a worker profiler for " - "multiprocess TP runs. Pass " - '--engine-arg \'profiler_config={"profiler":"cuda"}\' ' - "so Nsight Compute/Systems capture the TP worker kernels." - ) from exc + if args.cuda_profile_case is None: + start_cuda_profile() + profile_active = True try: for case, result in zip(cases, case_results): if not args.cuda_profile_repeat: @@ -2037,15 +2060,26 @@ def main() -> int: ) for _ in range(case["warmup"]) ] - repeats = [ - _run_once( - llm, - case["prompt"], - case["sampling_params"], - reset_prefix_cache=args.reset_prefix_cache_between_runs, - ) - for _ in range(case["repeat"]) - ] + profile_this_case = ( + args.cuda_profile_repeat and args.cuda_profile_case == case["name"] + ) + if profile_this_case: + start_cuda_profile() + profile_active = True + try: + repeats = [ + _run_once( + llm, + case["prompt"], + case["sampling_params"], + reset_prefix_cache=args.reset_prefix_cache_between_runs, + ) + for _ in range(case["repeat"]) + ] + finally: + if profile_this_case and profile_active: + llm.stop_profile() + profile_active = False result["repeats"] = repeats result["summary"] = _summarize(repeats) if args.checkpoint_after_case: @@ -2058,7 +2092,7 @@ def main() -> int: case_results, ) finally: - if args.cuda_profile_repeat: + if profile_active: llm.stop_profile() payload = _write_results( diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 0144df53bf..e3ff61461f 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -130,6 +130,108 @@ Goal: [long-context result](https://github.com/yangzhuxinyzx/1Cat-vLLM-private/pull/23#issuecomment-5437096030) comments. +### Grouped Page4 QSA follow-up, 2026-08-28 + +- Public PR #378 landed the Page4 precursor. The exact grouped follow-up is + Draft PR #387 on + `codex/v100-qwen38-grouped-page4-prefill-20260828-144759`, rebuilt from + public `main@03c04da68e`. The frozen endpoint route remains + TP4/V2/ModelOpt NVFP4/FP16 activations and KV/no MTP, with 8192-token chunked + prefill. The retained BN16 endpoint baseline is + `4532.07/4446.64/4108.16 tok/s` at 32K/64K/131K; do not compare short route + probes or TTFT-inclusive numbers against it. +- The retained 8192-token Nsight Systems capture makes selected QSA the first + optimization target: `55.151 ms` per layer/rank and `36.42%` of kernel time. + GEMM is about `35.8%`, TP collectives about `12.8%`, and the main GDN kernel + `3.11%`. Nsight Compute counters remain unavailable to this user with + `ERR_NVGPUCTRPERM`; do not repeat the counter attempt or change host policy + for this campaign. +- Upstream audits found no drop-in SM70 kernel. FlashInfer PRs + [#4474](https://github.com/flashinfer-ai/flashinfer/pull/4474) and + [#4689](https://github.com/flashinfer-ai/flashinfer/pull/4689) provide the + useful BSR/KV256 and weight-stationary scheduling ideas, but target SM75+ + and assume one semantic query block shares one sparse row. SGLang PR + [#36497](https://github.com/sgl-project/sglang/pull/36497) and vLLM PR + [#53896](https://github.com/vllm-project/vllm/pull/53896) retain a direct + per-row Triton QSA prefill kernel. Qwen TP4 has six query heads, D256, and + different top-k sets per row, so exact cross-query masks are required. +- The retained production-layout Page4 partition sweep admits P1024 without + changing the 192-thread/two-CTA-per-SM resource contract. P256/P512/P1024 + measure `30.742/29.361/28.769 ms`; P1024 is `6.4%` faster than P256 with + maximum absolute error `0.000244140625` and relative L2 about `3.95e-4`. + This is the fallback for large rows that cannot use grouped planning. +- The grouped candidate assigns one CTA to eight adjacent query rows and six + heads, reusing each unioned FP16 K/V Page4 block. A 32-bit mask carries four + token bits for each of eight queries, preserving different top-k sets and + causal tails exactly. Real retained selection diagnostics are `81.92%` + adjacent-row overlap and `54.88%` at gap four; the admission microbenchmark + uses a 1163-page union matching both values, rather than assuming one common + set across all eight rows. +- The GPU planner uses a per-group 8192-entry shared-memory hash table. It + resolves physical pages for every request independently, unions all rows, + and groups output by the seven possible active WMMA row-tile masks. Category + boundaries are padded to eight Page4 blocks so a 32-token kernel tile never + mixes categories. A deterministic block-wide prefix orders entries by hash + slot; four replays produce bitwise-identical tables and outputs. The earlier + shared-atomic ordering is rejected despite a slightly lower median because + it varied by one FP16 ULP across replays. +- After the latest-main repair and QSA cache-table canonicalization, grouped + planning also consumes per-row query positions and per-request live sequence + lengths. It truncates early prompt rows, validates partial tails, request and + logical-page bounds, and rejects stale physical pages before hashing. A + 4096-row hybrid Page16/Page4 test spans 1/2/3/4-token early rows, the + 2048/2049/2050/2051 boundaries, mature rows, and a nonmonotonic physical page + table; it matches Triton at relative L2 `2.434e-4` and cosine `1.0`. An all + invalid/padded-group test produces bitwise-equal zero output instead of + leaving the caller's output buffer untouched. +- The final isolated route, including Python dispatch, GPU planning, padding, + and grouped attention, measures `16.132 ms` median (`15.588 ms` minimum) + versus Triton `55.940 ms`, a `3.468x` speedup. Maximum absolute error is + `0.0001220703125`, relative L2 `3.535e-4`, and cosine + `0.9999998807907104`. The standalone planner is about `0.32 ms`; its exact + group-0 block/mask dictionary matches the reference. Focused Python tests + are `10 passed`, and the SM70 extension builds successfully with zero local + memory for the grouped kernel. Focused latest-source coverage is `19 passed, + 16 warnings`. +- The TP4 endpoint gate route-hits grouped Page4 on full 8192-row chunks and on + the 8120-row 131K tail. Pure-prefill throughput is `5998.65`, `5777.43`, and + `5450.92 tok/s` at 32K/64K/131K, versus the matching retained + `4532.07/4446.64/4108.16` baseline: `1.3236x/1.2993x/1.3269x`. An exact 8K + case measures `6394.74 tok/s`. Arithmetic, Chinese, and all long-context + token hashes are bitwise identical to the baseline, including both repeats + at 32K and 64K. The recovery endpoint run overlapped host/disk activity from + an unrelated GPU4-7 model load beginning at `15:23:39`, so treat these as + conservative endpoint values rather than a clean-host ceiling. +- The clean double-locked 8192-token Nsight capture reports `5094.991 ms` of + aggregate kernel service over four ranks. Grouped attention is + `9.632 ms` per QSA layer/rank and its planner is `0.362 ms`, or `9.994 ms` + together: `5.518x` faster and `81.88%` lower than the old `55.151 ms` QSA + path. QSA plus planning is now `9.42%` of aggregate service instead of + `36.42%`. Per-GPU kernel-union duty is `98.27-98.61%` over a `1293.2 ms` + capture span; only `18.0-22.4 ms` is between kernels. The representative + NVML sample is `99-100%` GPU busy at `1530 MHz` and `268-314 W`. These are + scheduling-duty and power observations, not achieved Tensor-Core or HBM + counters; NCU remains unavailable under `ERR_NVGPUCTRPERM`. +- The post-QSA first hotspot is the routed NVFP4 MoE chain. Its two grouped + GEMMs consume `25.12%`; materializing the top-10 input permutation consumes + `5.61%`; deterministic unpermute/finalize consumes `2.46%`. These three + stages alone are `33.19%`, before routing, sorting, and activation. The + NVFP4 kernel uses 128 threads, 122 registers/thread, and 16.4 KiB shared + memory, giving a four-CTA/SM, 16-warp/SM (`25%`) static occupancy ceiling. + Useful W13/W2 rates are only about `28.8-31.9 TFLOP/s`. HC combine-norm and + gate-mix are `6.44%/3.85%`, NCCL all-reduce is `8.45%`, and the main GDN + kernel is `4.43%`. GDN is separately grid/resource limited (48 CTAs for 80 + SMs, 253 registers/thread, 91 KiB dynamic shared memory), but its absolute + opportunity is smaller than routed MoE. +- Do not spend another full-model startup rechecking the same grouped-QSA + route. The next bounded experiment is an indexed-A NVFP4 W13 operator gate: + retain the existing deterministic expert sort and inverse map, but avoid + materializing 8192 rows into 81920 top-10 rows and let TurboMind's existing + indexed SM70 A iterator read the original token-major activation. Admit it + only if the complete sort/map/W13 chain wins and remains bitwise or within + the existing FP16 accumulation tolerance; a noncontiguous indexed read that + merely moves the copy cost into GEMM is rejected before any endpoint run. + ## Active MRV2 DFlash2 campaign, 2026-08-20 - Integration base: `onecat/main@7aede2cf010d92815c9d7bff25867b4fa009b6cb`. diff --git a/flash-attention-v100/include/fused_mha.h b/flash-attention-v100/include/fused_mha.h index 0987368cba..6e80b06e2f 100644 --- a/flash-attention-v100/include/fused_mha.h +++ b/flash-attention-v100/include/fused_mha.h @@ -57,6 +57,19 @@ at::Tensor flash_attention_grouped_verify_paged( const std::string& kv_cache_dtype, const float k_scale, const float v_scale, const bool one_pass); +at::Tensor flash_attention_grouped_sparse_page4( + const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& out_, const at::Tensor& block_table, + const at::Tensor& token_masks, const at::Tensor& seq_lens, at::Tensor& lse, + const float softmax_scale); + +at::Tensor flash_attention_grouped_sparse_page4_plan( + const at::Tensor& logical_indices, const at::Tensor& block_table, + const at::Tensor& token_to_req, const at::Tensor& query_positions, + const at::Tensor& sequence_lengths, at::Tensor& output_blocks, + at::Tensor& output_masks, at::Tensor& output_seq_lens, const int page_size, + const int physical_page_stride, const int num_cache_blocks); + at::Tensor flash_attention_decode_paged_wmma( const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, std::optional& out_, const at::Tensor& block_table, diff --git a/flash-attention-v100/kernel/flash_decode_paged.cu b/flash-attention-v100/kernel/flash_decode_paged.cu index 224653d046..8c65757ea6 100644 --- a/flash-attention-v100/kernel/flash_decode_paged.cu +++ b/flash-attention-v100/kernel/flash_decode_paged.cu @@ -1855,6 +1855,7 @@ struct alignas(256) GroupedVerifySmem { alignas(16) float row_sum[kGroupedVerifyRows]; alignas(16) float row_scale[kGroupedVerifyRows]; alignas(16) int page_ids[kGroupedVerifyPageIdsCapacity]; + alignas(16) uint32_t sparse_token_masks[kGroupedVerifyBlockN / 4]; }; static_assert(sizeof(GroupedVerifySmem) <= 64 * 1024, @@ -1880,13 +1881,17 @@ __device__ __forceinline__ int grouped_verify_active_splits( __device__ __forceinline__ void grouped_verify_qk( const __half* __restrict__ shared_q, const __half* __restrict__ shared_k, - float* __restrict__ shared_scores, const float qk_scale) { + float* __restrict__ shared_scores, const float qk_scale, + const int active_m_tiles) { const int warp_id = threadIdx.x / kWarpSize; if (warp_id >= kGroupedVerifyQKWarps) { return; } const int m_tile = warp_id / (kGroupedVerifyBlockN / 16); + if ((active_m_tiles & (1 << m_tile)) == 0) { + return; + } const int n_tile = warp_id % (kGroupedVerifyBlockN / 16); volta::fragment q_fragment; @@ -1933,9 +1938,31 @@ __device__ __forceinline__ void grouped_verify_scale_output_fragment( fragment.x[7] *= second_scale; } +template +__device__ __forceinline__ bool grouped_verify_key_visible( + const uint32_t* __restrict__ sparse_tile_masks, const int token_idx, + const int query_len, const int head_idx, const int kv_idx, + const int valid_k_rows, const int lane_or_col, const int prefix_kv_len) { + const bool row_valid = token_idx < query_len && + head_idx < kGroupedVerifyHeads && + lane_or_col < valid_k_rows; + if (!row_valid) { + return false; + } + if constexpr (SPARSE_PAGE4) { + const uint32_t token_mask = + sparse_tile_masks[(kv_idx & (kGroupedVerifyBlockN - 1)) >> 2]; + const int mask_bit = token_idx * 4 + (kv_idx & 3); + return (token_mask & (1u << mask_bit)) != 0; + } + return kv_idx <= prefix_kv_len + token_idx; +} + template + bool STAGE_PARTITION_PAGE_IDS = false, + int KV_DTYPE = flash_v100::KV_CACHE_DTYPE_FP8_E5M2, + bool SPARSE_PAGE4 = false> __global__ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_e5m2_partial_kernel( const __half* __restrict__ q, const void* __restrict__ k_cache, @@ -1946,22 +1973,44 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ const int64_t k_block_stride, const int64_t k_token_stride, const int64_t k_head_stride, const int64_t v_block_stride, const int64_t v_token_stride, const int64_t v_head_stride, - const float qk_scale, const float v_scale) { + const float qk_scale, const float v_scale, + const uint32_t* __restrict__ sparse_token_masks = nullptr, + const int num_groups = 1) { using Traits = GroupedVerifyTraits; const int head_group = blockIdx.x; const int split_id = blockIdx.y; - (void)max_num_blocks; + const int group_idx = SPARSE_PAGE4 ? blockIdx.z : 0; if (head_group >= Traits::kHeadGroups || split_id >= Traits::kSplits || - query_len <= 0 || query_len > MAX_QUERY_TOKENS) { + group_idx >= num_groups || query_len <= 0 || + query_len > MAX_QUERY_TOKENS) { return; } - const int total_kv = seq_lens[0]; + const int total_kv = seq_lens[SPARSE_PAGE4 ? group_idx : 0]; if (total_kv <= 0) { + if constexpr (SPARSE_PAGE4) { + constexpr int kGroupOutputElements = + kGroupedVerifyQ8MaxQ * kGroupedVerifyHeads * kGroupedVerifyHeadDim; + const int64_t group_output_offset = + static_cast(group_idx) * kGroupOutputElements; + for (int idx = threadIdx.x; idx < kGroupOutputElements; + idx += kGroupedVerifyThreads) { + partial_out[group_output_offset + idx] = __float2half_rn(0.0f); + } + constexpr int kGroupLseElements = + kGroupedVerifyQ8MaxQ * kGroupedVerifyHeads; + if (threadIdx.x < kGroupLseElements) { + partial_lse[static_cast(group_idx) * kGroupLseElements + + threadIdx.x] = kXQANegInf; + } + } return; } const int active_splits = - grouped_verify_active_splits(total_kv); + SPARSE_PAGE4 + ? 1 + : grouped_verify_active_splits( + total_kv); if (split_id >= active_splits) { return; } @@ -1988,7 +2037,8 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ __half* shared_kv = smem.storage.compute.kv; float* shared_scores = smem.storage.compute.scores; __half* shared_probs = smem.storage.compute.probs; - const int* page_ids = block_table; + const int* page_ids = + block_table + static_cast(group_idx) * max_num_blocks; int split_page_offset = 0; bool use_staged_page_ids = false; if constexpr (STAGE_PARTITION_PAGE_IDS) { @@ -2023,8 +2073,10 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ const int local_head = row % Traits::kHeadsPerCta; const int head_idx = head_start + local_head; if (token_idx < query_len && head_idx < kGroupedVerifyHeads) { + const int64_t query_row = + static_cast(group_idx) * MAX_QUERY_TOKENS + token_idx; shared_q_vec[row * kSharedQVecsPerRow + vec_col] = __ldg( - q_vec + (token_idx * kGroupedVerifyHeads + head_idx) * kVecsPerRow + + q_vec + (query_row * kGroupedVerifyHeads + head_idx) * kVecsPerRow + vec_col); } else { shared_q_vec[row * kSharedQVecsPerRow + vec_col] = make_uint4(0, 0, 0, 0); @@ -2056,12 +2108,23 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ tile_start += kGroupedVerifyBlockN) { const int valid_k_rows = min(kGroupedVerifyBlockN, split_end - tile_start); + if constexpr (SPARSE_PAGE4) { + const int valid_sparse_pages = (valid_k_rows + 3) / 4; + if (tid < kGroupedVerifyBlockN / 4) { + smem.sparse_token_masks[tid] = + tid < valid_sparse_pages + ? __ldg(sparse_token_masks + + static_cast(group_idx) * max_num_blocks + + (tile_start >> 2) + tid) + : 0; + } + } const int tile_page_offset = use_staged_page_ids ? split_page_offset + tile_start - split_start : tile_start; load_xqa_tc_kv_panel( + kGroupedVerifyThreads, KV_DTYPE, + KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E5M2>( shared_kv, k_cache, page_ids, valid_k_rows, kPanelStrideVec, kSharedStrideVec, tile_page_offset, 0, page_block_size, 0, k_block_stride, k_token_stride, k_head_stride, 0); @@ -2072,7 +2135,26 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ } __syncthreads(); - grouped_verify_qk(shared_q, shared_kv, shared_scores, qk_scale); + int active_m_tiles = 0x7; + if constexpr (SPARSE_PAGE4) { + uint32_t active_query_nibbles = 0; +#pragma unroll + for (int page = 0; page < kGroupedVerifyBlockN / 4; ++page) { + active_query_nibbles |= smem.sparse_token_masks[page]; + } + active_m_tiles = 0; +#pragma unroll + for (int token = 0; token < kGroupedVerifyQ8MaxQ; ++token) { + if ((active_query_nibbles & (0xFu << (token * 4))) != 0) { + const int first_row = token * kGroupedVerifyHeads; + const int last_row = first_row + kGroupedVerifyHeads - 1; + active_m_tiles |= 1 << (first_row / 16); + active_m_tiles |= 1 << (last_row / 16); + } + } + } + grouped_verify_qk(shared_q, shared_kv, shared_scores, qk_scale, + active_m_tiles); __syncthreads(); #pragma unroll @@ -2082,9 +2164,9 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ const int local_head = row % Traits::kHeadsPerCta; const int head_idx = head_start + local_head; const int kv_idx = tile_start + lane_id; - const bool visible = - token_idx < query_len && head_idx < kGroupedVerifyHeads && - lane_id < valid_k_rows && kv_idx <= prefix_kv_len + token_idx; + const bool visible = grouped_verify_key_visible( + smem.sparse_token_masks, token_idx, query_len, head_idx, kv_idx, + valid_k_rows, lane_id, prefix_kv_len); const float score = visible ? shared_scores[row * kGroupedVerifyScoreStride + lane_id] : kXQANegInf; @@ -2114,12 +2196,23 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ for (int tile_start = split_start; tile_start < split_end; tile_start += kGroupedVerifyBlockN) { const int valid_k_rows = min(kGroupedVerifyBlockN, split_end - tile_start); + if constexpr (SPARSE_PAGE4) { + const int valid_sparse_pages = (valid_k_rows + 3) / 4; + if (tid < kGroupedVerifyBlockN / 4) { + smem.sparse_token_masks[tid] = + tid < valid_sparse_pages + ? __ldg(sparse_token_masks + + static_cast(group_idx) * max_num_blocks + + (tile_start >> 2) + tid) + : 0; + } + } const int tile_page_offset = use_staged_page_ids ? split_page_offset + tile_start - split_start : tile_start; load_xqa_tc_kv_panel( + kGroupedVerifyThreads, KV_DTYPE, + KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E5M2>( shared_kv, k_cache, page_ids, valid_k_rows, kPanelStrideVec, kSharedStrideVec, tile_page_offset, 0, page_block_size, 0, k_block_stride, k_token_stride, k_head_stride, 0); @@ -2130,7 +2223,26 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ } __syncthreads(); - grouped_verify_qk(shared_q, shared_kv, shared_scores, qk_scale); + int active_m_tiles = 0x7; + if constexpr (SPARSE_PAGE4) { + uint32_t active_query_nibbles = 0; +#pragma unroll + for (int page = 0; page < kGroupedVerifyBlockN / 4; ++page) { + active_query_nibbles |= smem.sparse_token_masks[page]; + } + active_m_tiles = 0; +#pragma unroll + for (int token = 0; token < kGroupedVerifyQ8MaxQ; ++token) { + if ((active_query_nibbles & (0xFu << (token * 4))) != 0) { + const int first_row = token * kGroupedVerifyHeads; + const int last_row = first_row + kGroupedVerifyHeads - 1; + active_m_tiles |= 1 << (first_row / 16); + active_m_tiles |= 1 << (last_row / 16); + } + } + } + grouped_verify_qk(shared_q, shared_kv, shared_scores, qk_scale, + active_m_tiles); __syncthreads(); if constexpr (TWO_PASS) { @@ -2143,8 +2255,9 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ const int head_idx = head_start + local_head; const int kv_idx = tile_start + col; const bool visible = - token_idx < query_len && head_idx < kGroupedVerifyHeads && - col < valid_k_rows && kv_idx <= prefix_kv_len + token_idx && + grouped_verify_key_visible( + smem.sparse_token_masks, token_idx, query_len, head_idx, kv_idx, + valid_k_rows, col, prefix_kv_len) && smem.row_sum[row] > 0.0f; const float probability = visible ? __expf(fmaxf( @@ -2164,9 +2277,9 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ const int local_head = row % Traits::kHeadsPerCta; const int head_idx = head_start + local_head; const int kv_idx = tile_start + lane_id; - const bool visible = - token_idx < query_len && head_idx < kGroupedVerifyHeads && - lane_id < valid_k_rows && kv_idx <= prefix_kv_len + token_idx; + const bool visible = grouped_verify_key_visible( + smem.sparse_token_masks, token_idx, query_len, head_idx, kv_idx, + valid_k_rows, lane_id, prefix_kv_len); const float score = visible ? shared_scores[row * kGroupedVerifyScoreStride + lane_id] : kXQANegInf; @@ -2202,8 +2315,8 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ } load_xqa_tc_kv_panel( + kGroupedVerifyThreads, KV_DTYPE, + KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E5M2>( shared_kv, v_cache, page_ids, valid_k_rows, kPanelStrideVec, kSharedStrideVec, tile_page_offset, 0, page_block_size, 0, v_block_stride, v_token_stride, v_head_stride, 0); @@ -2220,6 +2333,11 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ const int output_tile = warp_id + fragment_idx * kGroupedVerifyWarps; const int m_tile = output_tile / (kGroupedVerifyHeadDim / 16); const int d_tile = output_tile % (kGroupedVerifyHeadDim / 16); + if constexpr (SPARSE_PAGE4) { + if ((active_m_tiles & (1 << m_tile)) == 0) { + continue; + } + } volta::fragment probability_fragment; volta::fragment @@ -2268,12 +2386,21 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ if (token_idx < query_len && head_idx < kGroupedVerifyHeads) { const float sum = smem.row_sum[row]; const float scale = sum > 0.0f ? v_scale / sum : 0.0f; - const int64_t output_idx = - (((static_cast(split_id) * MAX_QUERY_TOKENS + token_idx) * - kGroupedVerifyHeads + - head_idx) * - kGroupedVerifyHeadDim + - d); + int64_t output_idx; + if constexpr (SPARSE_PAGE4) { + const int64_t global_token_idx = + static_cast(group_idx) * MAX_QUERY_TOKENS + token_idx; + output_idx = (global_token_idx * kGroupedVerifyHeads + head_idx) * + kGroupedVerifyHeadDim + + d; + } else { + output_idx = + (((static_cast(split_id) * MAX_QUERY_TOKENS + token_idx) * + kGroupedVerifyHeads + + head_idx) * + kGroupedVerifyHeadDim + + d); + } partial_out[output_idx] = __float2half_rn(shared_output[idx] * scale); } } @@ -2283,10 +2410,17 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ const int head_idx = head_start + local_head; if (token_idx < query_len && head_idx < kGroupedVerifyHeads) { const float sum = smem.row_sum[tid]; - const int64_t lse_idx = - (static_cast(split_id) * MAX_QUERY_TOKENS + token_idx) * - kGroupedVerifyHeads + - head_idx; + int64_t lse_idx; + if constexpr (SPARSE_PAGE4) { + const int64_t global_token_idx = + static_cast(group_idx) * MAX_QUERY_TOKENS + token_idx; + lse_idx = global_token_idx * kGroupedVerifyHeads + head_idx; + } else { + lse_idx = + (static_cast(split_id) * MAX_QUERY_TOKENS + token_idx) * + kGroupedVerifyHeads + + head_idx; + } partial_lse[lse_idx] = sum > 0.0f ? smem.row_max[tid] + logf(sum) : kXQANegInf; } @@ -3520,8 +3654,381 @@ void launch_flash_attention_decode_qk_scores( k_cache.stride(1), k_cache.stride(2), softmax_scale, k_scale); } +constexpr int kGroupedSparseQueries = 8; +constexpr int kGroupedSparsePlannerThreads = 512; +constexpr int kGroupedSparseHashCapacity = 8192; +constexpr unsigned long long kGroupedSparseEmptyEntry = 0x00000000ffffffffULL; + +__device__ __forceinline__ void grouped_sparse_hash_insert( + unsigned long long* __restrict__ hash_table, const int physical_microblock, + const uint32_t token_mask) { + if (physical_microblock < 0 || token_mask == 0) { + return; + } + int slot = (static_cast(physical_microblock) * 2654435761u) & + (kGroupedSparseHashCapacity - 1); + const unsigned long long desired = + (static_cast(token_mask) << 32) | + static_cast(physical_microblock); +#pragma unroll 1 + for (int probe = 0; probe < kGroupedSparseHashCapacity; ++probe) { + const unsigned long long old = + atomicCAS(hash_table + slot, kGroupedSparseEmptyEntry, desired); + if (old == kGroupedSparseEmptyEntry) { + return; + } + if (static_cast(old) == + static_cast(physical_microblock)) { + atomicOr(hash_table + slot, static_cast(token_mask) + << 32); + return; + } + slot = (slot + 1) & (kGroupedSparseHashCapacity - 1); + } +} + +__device__ __forceinline__ int grouped_sparse_physical_microblock( + const int token, const int request_idx, + const int* __restrict__ request_block_table, + const int64_t request_block_table_stride, const int block_table_width, + const int page_size, const int physical_page_stride, + const int num_cache_blocks) { + if (token < 0) { + return -1; + } + const int logical_page = token / page_size; + if (logical_page < 0 || logical_page >= block_table_width) { + return -1; + } + const int page_offset = token - logical_page * page_size; + const int physical_page = + __ldg(request_block_table + + static_cast(request_idx) * request_block_table_stride + + logical_page); + if (physical_page < 0 || physical_page >= num_cache_blocks) { + return -1; + } + return physical_page * physical_page_stride + page_offset / 4; +} + +__device__ __forceinline__ int grouped_sparse_active_m_tiles( + const uint32_t token_mask) { + int active_m_tiles = 0; +#pragma unroll + for (int query = 0; query < kGroupedSparseQueries; ++query) { + if ((token_mask & (0xFu << (query * 4))) != 0) { + const int first_row = query * kGroupedVerifyHeads; + const int last_row = first_row + kGroupedVerifyHeads - 1; + active_m_tiles |= 1 << (first_row / 16); + active_m_tiles |= 1 << (last_row / 16); + } + } + return active_m_tiles; +} + +__global__ +__launch_bounds__(kGroupedSparsePlannerThreads, 1) void grouped_sparse_page4_plan_kernel( + const int* __restrict__ logical_indices, + const int* __restrict__ request_block_table, + const int* __restrict__ token_to_req, + const int64_t* __restrict__ query_positions, + const int* __restrict__ sequence_lengths, int* __restrict__ output_blocks, + uint32_t* __restrict__ output_masks, int* __restrict__ output_seq_lens, + const int selection_width, const int64_t logical_indices_stride, + const int64_t request_block_table_stride, const int num_requests, + const int block_table_width, const int output_width, const int page_size, + const int physical_page_stride, const int num_cache_blocks) { + const int group_idx = blockIdx.x; + const int tid = threadIdx.x; + __shared__ int category_counts[8]; + __shared__ int category_offsets[8]; + __shared__ int category_cursors[8]; + __shared__ int + warp_category_prefix[(kGroupedSparsePlannerThreads / kWarpSize) * 8]; + extern __shared__ unsigned long long hash_table[]; + for (int slot = tid; slot < kGroupedSparseHashCapacity; + slot += kGroupedSparsePlannerThreads) { + hash_table[slot] = kGroupedSparseEmptyEntry; + } + __syncthreads(); + + const int full_page4_count = selection_width / 4; + for (int selected_page = tid; selected_page < full_page4_count; + selected_page += kGroupedSparsePlannerThreads) { +#pragma unroll + for (int query = 0; query < kGroupedSparseQueries; ++query) { + const int row = group_idx * kGroupedSparseQueries + query; + const int request_idx = __ldg(token_to_req + row); + if (request_idx < 0 || request_idx >= num_requests) { + continue; + } + const int sequence_length = __ldg(sequence_lengths + request_idx); + const int64_t query_visible_tokens = __ldg(query_positions + row) + 1; + const int visible_tokens = + query_visible_tokens <= 0 + ? 0 + : (query_visible_tokens < sequence_length + ? static_cast(query_visible_tokens) + : max(sequence_length, 0)); + const int row_complete_page4_count = + min(min(visible_tokens / 4, sequence_length / 4), full_page4_count); + if (selected_page >= row_complete_page4_count) { + continue; + } + const int* selected = logical_indices + + static_cast(row) * logical_indices_stride + + selected_page * 4; + const int first_token = __ldg(selected); + if (first_token < 0) { + continue; + } + const bool full_page4 = __ldg(selected + 1) == first_token + 1 && + __ldg(selected + 2) == first_token + 2 && + __ldg(selected + 3) == first_token + 3 && + (first_token & 3) == 0 && + first_token + 3 < sequence_length; + if (full_page4) { + const int physical_microblock = grouped_sparse_physical_microblock( + first_token, request_idx, request_block_table, + request_block_table_stride, block_table_width, page_size, + physical_page_stride, num_cache_blocks); + grouped_sparse_hash_insert(hash_table, physical_microblock, + 0xFu << (query * 4)); + } else { +#pragma unroll + for (int token_offset = 0; token_offset < 4; ++token_offset) { + const int token = __ldg(selected + token_offset); + if (token >= 0 && token < sequence_length) { + const int physical_microblock = grouped_sparse_physical_microblock( + token, request_idx, request_block_table, + request_block_table_stride, block_table_width, page_size, + physical_page_stride, num_cache_blocks); + grouped_sparse_hash_insert(hash_table, physical_microblock, + 1u << (query * 4 + (token & 3))); + } + } + } + } + } + if (tid < kGroupedSparseQueries) { + const int query = tid; + const int row = group_idx * kGroupedSparseQueries + query; + const int request_idx = __ldg(token_to_req + row); + if (request_idx >= 0 && request_idx < num_requests) { + const int sequence_length = __ldg(sequence_lengths + request_idx); + const int64_t query_visible_tokens = __ldg(query_positions + row) + 1; + const int visible_tokens = + query_visible_tokens <= 0 + ? 0 + : (query_visible_tokens < sequence_length + ? static_cast(query_visible_tokens) + : max(sequence_length, 0)); + const int complete_page4_count = + min(min(visible_tokens / 4, sequence_length / 4), full_page4_count); + const int tail_count = visible_tokens & 3; + const int tail_index = complete_page4_count * 4; + const int selected_tail_token = + tail_index < selection_width + ? __ldg(logical_indices + + static_cast(row) * logical_indices_stride + + tail_index) + : -1; + const int expected_tail_token = (visible_tokens / 4) * 4; + if (tail_count > 0 && selected_tail_token == expected_tail_token && + selected_tail_token < sequence_length) { + const int physical_microblock = grouped_sparse_physical_microblock( + selected_tail_token, request_idx, request_block_table, + request_block_table_stride, block_table_width, page_size, + physical_page_stride, num_cache_blocks); + const uint32_t tail_mask = ((1u << tail_count) - 1) << (query * 4); + grouped_sparse_hash_insert(hash_table, physical_microblock, tail_mask); + } + } + } + __syncthreads(); + + if (tid < 8) { + category_counts[tid] = 0; + category_offsets[tid] = 0; + category_cursors[tid] = 0; + } + __syncthreads(); + for (int slot = tid; slot < kGroupedSparseHashCapacity; + slot += kGroupedSparsePlannerThreads) { + const unsigned long long entry = hash_table[slot]; + if (static_cast(entry) != 0xffffffffu) { + const int category = + grouped_sparse_active_m_tiles(static_cast(entry >> 32)); + atomicAdd(category_counts + category, 1); + } + } + __syncthreads(); + if (tid == 0) { + int padded_offset = 0; +#pragma unroll + for (int category = 1; category < 8; ++category) { + category_offsets[category] = padded_offset; + padded_offset += (category_counts[category] + 7) & ~7; + } + category_offsets[0] = padded_offset; + } + __syncthreads(); + constexpr int kPlannerWarps = kGroupedSparsePlannerThreads / kWarpSize; + const int lane = tid & (kWarpSize - 1); + const int warp = tid / kWarpSize; + for (int chunk_start = 0; chunk_start < kGroupedSparseHashCapacity; + chunk_start += kGroupedSparsePlannerThreads) { + const unsigned long long entry = hash_table[chunk_start + tid]; + const uint32_t physical_microblock = static_cast(entry); + int category = 0; + if (physical_microblock != 0xffffffffu) { + const uint32_t token_mask = static_cast(entry >> 32); + category = grouped_sparse_active_m_tiles(token_mask); + } + unsigned category_lanes = 0; +#pragma unroll + for (int scan_category = 1; scan_category < 8; ++scan_category) { + const unsigned lanes = + __ballot_sync(0xffffffffu, category == scan_category); + if (category == scan_category) { + category_lanes = lanes; + } + if (lane == 0) { + warp_category_prefix[warp * 8 + scan_category] = __popc(lanes); + } + } + __syncthreads(); + if (tid > 0 && tid < 8) { + int prefix = category_cursors[tid]; +#pragma unroll + for (int scan_warp = 0; scan_warp < kPlannerWarps; ++scan_warp) { + const int count = warp_category_prefix[scan_warp * 8 + tid]; + warp_category_prefix[scan_warp * 8 + tid] = prefix; + prefix += count; + } + category_cursors[tid] = prefix; + } + __syncthreads(); + if (category != 0) { + const int category_rank = warp_category_prefix[warp * 8 + category] + + __popc(category_lanes & ((1u << lane) - 1)); + const int output_idx = category_offsets[category] + category_rank; + if (output_idx < output_width) { + output_blocks[static_cast(group_idx) * output_width + + output_idx] = static_cast(physical_microblock); + output_masks[static_cast(group_idx) * output_width + + output_idx] = static_cast(entry >> 32); + } + } + __syncthreads(); + } + __syncthreads(); + if (tid > 0 && tid < 8) { + const int category = tid; + const int padded_count = (category_counts[category] + 7) & ~7; + for (int local_idx = category_counts[category]; local_idx < padded_count; + ++local_idx) { + const int output_idx = category_offsets[category] + local_idx; + if (output_idx < output_width) { + output_blocks[static_cast(group_idx) * output_width + + output_idx] = 0; + output_masks[static_cast(group_idx) * output_width + + output_idx] = 0; + } + } + } + if (tid == 0) { + output_seq_lens[group_idx] = min(category_offsets[0], output_width) * 4; + } +} + } // namespace +at::Tensor flash_attention_grouped_sparse_page4_plan( + const at::Tensor& logical_indices, const at::Tensor& block_table, + const at::Tensor& token_to_req, const at::Tensor& query_positions, + const at::Tensor& sequence_lengths, at::Tensor& output_blocks, + at::Tensor& output_masks, at::Tensor& output_seq_lens, const int page_size, + const int physical_page_stride, const int num_cache_blocks) { + TORCH_CHECK(logical_indices.is_cuda() && block_table.is_cuda() && + token_to_req.is_cuda() && query_positions.is_cuda() && + sequence_lengths.is_cuda() && output_blocks.is_cuda() && + output_masks.is_cuda() && output_seq_lens.is_cuda(), + "grouped sparse page4 planner tensors must be CUDA tensors"); + TORCH_CHECK(logical_indices.dtype() == torch::kInt32 && + block_table.dtype() == torch::kInt32 && + token_to_req.dtype() == torch::kInt32 && + query_positions.dtype() == torch::kInt64 && + sequence_lengths.dtype() == torch::kInt32 && + output_blocks.dtype() == torch::kInt32 && + output_masks.scalar_type() == at::ScalarType::UInt32 && + output_seq_lens.dtype() == torch::kInt32, + "grouped sparse page4 planner requires int32/uint32 metadata"); + TORCH_CHECK(logical_indices.dim() == 2 && logical_indices.size(0) > 0 && + logical_indices.size(0) % kGroupedSparseQueries == 0 && + logical_indices.size(1) == 2051, + "grouped sparse page4 planner requires [8*N, 2051] indices"); + const int64_t num_groups = logical_indices.size(0) / kGroupedSparseQueries; + TORCH_CHECK( + block_table.dim() == 2 && + token_to_req.sizes() == at::IntArrayRef({logical_indices.size(0)}), + "grouped sparse page4 planner request metadata is invalid"); + TORCH_CHECK( + query_positions.sizes() == at::IntArrayRef({logical_indices.size(0)}) && + sequence_lengths.sizes() == at::IntArrayRef({block_table.size(0)}), + "grouped sparse page4 planner visibility metadata is invalid"); + TORCH_CHECK(output_blocks.dim() == 2 && output_blocks.size(0) == num_groups && + output_blocks.size(1) >= 4160 && + output_masks.sizes() == output_blocks.sizes() && + output_seq_lens.sizes() == at::IntArrayRef({num_groups}), + "grouped sparse page4 planner outputs must be [groups, >=4160]"); + TORCH_CHECK( + logical_indices.is_contiguous() && block_table.is_contiguous() && + token_to_req.is_contiguous() && query_positions.is_contiguous() && + sequence_lengths.is_contiguous() && output_blocks.is_contiguous() && + output_masks.is_contiguous() && output_seq_lens.is_contiguous(), + "grouped sparse page4 planner metadata must be contiguous"); + TORCH_CHECK(page_size > 0 && page_size % 4 == 0 && physical_page_stride > 0 && + num_cache_blocks > 0, + "grouped sparse page4 planner requires page_size divisible by 4"); + TORCH_CHECK(logical_indices.device() == block_table.device() && + logical_indices.device() == token_to_req.device() && + logical_indices.device() == query_positions.device() && + logical_indices.device() == sequence_lengths.device() && + logical_indices.device() == output_blocks.device() && + logical_indices.device() == output_masks.device() && + logical_indices.device() == output_seq_lens.device(), + "grouped sparse page4 planner tensors must share one device"); + + c10::cuda::CUDAGuard device_guard(logical_indices.device()); + const auto* properties = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(properties->major == 7 && properties->minor == 0, + "grouped sparse page4 planner supports SM70 only"); + cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + constexpr size_t kPlannerSharedMemory = + kGroupedSparseHashCapacity * sizeof(unsigned long long); + const cudaError_t smem_status = cudaFuncSetAttribute( + grouped_sparse_page4_plan_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, kPlannerSharedMemory); + TORCH_CHECK(smem_status == cudaSuccess, + "Failed to set grouped sparse page4 planner shared memory: ", + cudaGetErrorString(smem_status)); + grouped_sparse_page4_plan_kernel<<(num_groups), + kGroupedSparsePlannerThreads, + kPlannerSharedMemory, stream>>>( + logical_indices.data_ptr(), block_table.data_ptr(), + token_to_req.data_ptr(), query_positions.data_ptr(), + sequence_lengths.data_ptr(), output_blocks.data_ptr(), + output_masks.data_ptr(), output_seq_lens.data_ptr(), + static_cast(logical_indices.size(1)), logical_indices.stride(0), + block_table.stride(0), static_cast(block_table.size(0)), + static_cast(block_table.size(1)), + static_cast(output_blocks.size(1)), page_size, physical_page_stride, + num_cache_blocks); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return output_blocks; +} + at::Tensor flash_attention_grouped_verify_paged( const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, std::optional& out_, const at::Tensor& block_table, @@ -3640,7 +4147,7 @@ at::Tensor flash_attention_grouped_verify_paged( static_cast(k_cache.size(1)), k_cache.stride(0), \ k_cache.stride(1), k_cache.stride(2), v_cache.stride(0), \ v_cache.stride(1), v_cache.stride(2), softmax_scale * k_scale, \ - v_scale); \ + v_scale, nullptr, 1); \ } while (0) #define DISPATCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, \ @@ -3721,6 +4228,95 @@ at::Tensor flash_attention_grouped_verify_paged( return out; } +at::Tensor flash_attention_grouped_sparse_page4( + const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& out_, const at::Tensor& block_table, + const at::Tensor& token_masks, const at::Tensor& seq_lens, at::Tensor& lse, + const float softmax_scale) { + constexpr int kQueriesPerGroup = kGroupedVerifyQ8MaxQ; + TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), + "grouped sparse page4 q/K/V must be CUDA tensors"); + TORCH_CHECK(block_table.is_cuda() && token_masks.is_cuda() && + seq_lens.is_cuda() && lse.is_cuda(), + "grouped sparse page4 metadata must be CUDA tensors"); + TORCH_CHECK(q.dtype() == torch::kFloat16 && + k_cache.dtype() == torch::kFloat16 && + v_cache.dtype() == torch::kFloat16, + "grouped sparse page4 requires fp16 q/K/V"); + TORCH_CHECK(block_table.dtype() == torch::kInt32 && + seq_lens.dtype() == torch::kInt32 && + token_masks.scalar_type() == at::ScalarType::UInt32, + "grouped sparse page4 tables must be int32/uint32"); + TORCH_CHECK(q.dim() == 3 && q.size(0) > 0 && + q.size(0) % kQueriesPerGroup == 0 && + q.size(1) == kGroupedVerifyHeads && + q.size(2) == kGroupedVerifyHeadDim, + "grouped sparse page4 q must have shape [8*N, 6, 256]"); + TORCH_CHECK(k_cache.dim() == 4 && v_cache.dim() == 4 && + k_cache.sizes() == v_cache.sizes() && k_cache.size(1) == 4 && + k_cache.size(2) == 1 && + k_cache.size(3) == kGroupedVerifyHeadDim, + "grouped sparse page4 KV must have shape [blocks, 4, 1, 256]"); + const int64_t num_groups = q.size(0) / kQueriesPerGroup; + TORCH_CHECK(block_table.dim() == 2 && block_table.size(0) == num_groups && + token_masks.sizes() == block_table.sizes(), + "grouped sparse page4 block IDs/masks must be [groups, pages]"); + TORCH_CHECK(seq_lens.sizes() == at::IntArrayRef({num_groups}), + "grouped sparse page4 seq_lens must have shape [groups]"); + TORCH_CHECK(q.is_contiguous() && block_table.is_contiguous() && + token_masks.is_contiguous() && seq_lens.is_contiguous(), + "grouped sparse page4 q/metadata must be contiguous"); + TORCH_CHECK(k_cache.stride(-1) == 1 && v_cache.stride(-1) == 1, + "grouped sparse page4 KV head dimension must be contiguous"); + TORCH_CHECK( + lse.sizes() == at::IntArrayRef({q.size(0), kGroupedVerifyHeads}) && + lse.dtype() == torch::kFloat32 && lse.is_contiguous(), + "grouped sparse page4 lse must be contiguous [rows, 6] fp32"); + + at::Tensor out = out_.has_value() ? out_.value() : torch::empty_like(q); + TORCH_CHECK(out.is_cuda() && out.device() == q.device() && + out.dtype() == torch::kFloat16 && out.sizes() == q.sizes() && + out.is_contiguous(), + "grouped sparse page4 out must be contiguous fp16 and q-shaped"); + TORCH_CHECK(q.device() == k_cache.device() && + q.device() == v_cache.device() && + q.device() == block_table.device() && + q.device() == token_masks.device() && + q.device() == seq_lens.device() && + q.device() == lse.device() && q.device() == out.device(), + "all grouped sparse page4 tensors must be on the same device"); + + c10::cuda::CUDAGuard device_guard(q.device()); + const auto* properties = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(properties->major == 7 && properties->minor == 0, + "grouped sparse page4 supports SM70 only"); + cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + const dim3 grid(1, 1, static_cast(num_groups)); + const size_t shared_mem = sizeof(GroupedVerifySmem); + auto kernel = (void*)flash_attention_grouped_verify_e5m2_partial_kernel< + kQueriesPerGroup, false, 4, false, false, false, + flash_v100::KV_CACHE_DTYPE_FP16, true>; + const cudaError_t smem_status = cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_mem); + TORCH_CHECK(smem_status == cudaSuccess, + "Failed to set grouped sparse page4 shared memory: ", + cudaGetErrorString(smem_status)); + flash_attention_grouped_verify_e5m2_partial_kernel< + kQueriesPerGroup, false, 4, false, false, false, + flash_v100::KV_CACHE_DTYPE_FP16, true> + <<>>( + reinterpret_cast(q.data_ptr()), k_cache.data_ptr(), + v_cache.data_ptr(), block_table.data_ptr(), + seq_lens.data_ptr(), reinterpret_cast<__half*>(out.data_ptr()), + lse.data_ptr(), kQueriesPerGroup, + static_cast(block_table.size(1)), 4, k_cache.stride(0), + k_cache.stride(1), k_cache.stride(2), v_cache.stride(0), + v_cache.stride(1), v_cache.stride(2), softmax_scale, 1.0f, + token_masks.data_ptr(), static_cast(num_groups)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + at::Tensor flash_attention_decode_paged( const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, std::optional& out_, const at::Tensor& block_table, @@ -4465,10 +5061,12 @@ at::Tensor flash_attention_decode_paged_xqa( (batch_context_route == XQABatchContextRoute::kDualCta || batch_context_route == XQABatchContextRoute::kDualCtaSplit) && xqa_e5m2_batch_wide_load_enabled(); - const bool use_qsa_page4 = - q.size(0) >= 4096 && q_per_kv == 6 && partition_size == 256 && - k_cache.size(1) == 4 && k_cache.size(2) == 1 && - k_cache.scalar_type() == at::kHalf && block_table.size(1) == 513; + const bool use_qsa_page4 = q.size(0) >= 4096 && q_per_kv == 6 && + (partition_size == 256 || partition_size == 512 || + partition_size == 1024) && + k_cache.size(1) == 4 && k_cache.size(2) == 1 && + k_cache.scalar_type() == at::kHalf && + block_table.size(1) == 513; const bool use_g6_dual_cta = use_qsa_page4 || use_g6_p1024_auto || use_g6_p1024_sawtooth || use_mtp5_dual_cta || use_e5m2_g6_dual_cta || @@ -4480,8 +5078,7 @@ at::Tensor flash_attention_decode_paged_xqa( batch_context_route == XQABatchContextRoute::kDualCtaSplit || (use_g6_dual_cta && xqa_split_reduce_enabled()); const bool supports_block16_index = use_g6_dual_cta && k_cache.size(1) == 16; - const bool use_block4_index = - use_g6_dual_cta && partition_size == 256 && k_cache.size(1) == 4; + const bool use_block4_index = use_g6_dual_cta && k_cache.size(1) == 4; const bool supports_block16_contiguous_layout = supports_block16_index && k_cache.size(2) == 1 && k_cache.stride(0) == 4096 && k_cache.stride(1) == 256 && @@ -4824,12 +5421,28 @@ at::Tensor flash_attention_decode_paged_xqa( split_reduce_dim_tile, stream); } } else if (use_block4_index) { - launch_flash_attention_decode_paged_xqa_tc_256_wide< - 256, 6, true, kXQATCG6DualCtaThreads, 2, 4, false>( - q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, - exp_sums, active_num_partitions, softmax_scale, k_scale, v_scale, - launch_num_partitions, use_split_reduce, split_reduce_dim_tile, - stream); + if (partition_size == 256) { + launch_flash_attention_decode_paged_xqa_tc_256_wide< + 256, 6, true, kXQATCG6DualCtaThreads, 2, 4, false>( + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, + max_logits, exp_sums, active_num_partitions, softmax_scale, k_scale, + v_scale, launch_num_partitions, use_split_reduce, + split_reduce_dim_tile, stream); + } else if (partition_size == 512) { + launch_flash_attention_decode_paged_xqa_tc_256_wide< + 512, 6, true, kXQATCG6DualCtaThreads, 2, 4, false>( + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, + max_logits, exp_sums, active_num_partitions, softmax_scale, k_scale, + v_scale, launch_num_partitions, use_split_reduce, + split_reduce_dim_tile, stream); + } else { + launch_flash_attention_decode_paged_xqa_tc_256_wide< + 1024, 6, true, kXQATCG6DualCtaThreads, 2, 4, false>( + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, + max_logits, exp_sums, active_num_partitions, softmax_scale, k_scale, + v_scale, launch_num_partitions, use_split_reduce, + split_reduce_dim_tile, stream); + } } else if (block16_layout_mode == 2) { launch_flash_attention_decode_paged_xqa_tc_256_wide< 256, 6, true, kXQATCG6DualCtaThreads, 2, 16, true>( diff --git a/flash-attention-v100/kernel/fused_mha_api.cpp b/flash-attention-v100/kernel/fused_mha_api.cpp index ec1f942353..89c436316d 100644 --- a/flash-attention-v100/kernel/fused_mha_api.cpp +++ b/flash-attention-v100/kernel/fused_mha_api.cpp @@ -24,6 +24,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Staged FlashAttention XQA decode over paged KV cache (Volta)"); m.def("grouped_verify_paged_fwd", &flash_attention_grouped_verify_paged, "Exact grouped DFlash2 verification over paged KV cache (Volta)"); + m.def("grouped_sparse_page4_fwd", &flash_attention_grouped_sparse_page4, + "Grouped exact QSA page4 attention over paged KV cache (Volta)"); + m.def("grouped_sparse_page4_plan_fwd", + &flash_attention_grouped_sparse_page4_plan, + "Build grouped exact QSA page4 tables over paged KV cache (Volta)"); m.def("decode_paged_wmma_fwd", &flash_attention_decode_paged_wmma, "FlashAttention single-query decode through paged-prefill WMMA order " "(Volta)"); diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index ce64cc09da..0a626edf8e 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -31,13 +31,28 @@ _SM70_QSA_XQA_PAGE4_MIN_ROWS = int( os.getenv("VLLM_SM70_QSA_XQA_PAGE4_MIN_ROWS", "4096") ) -_SM70_QSA_XQA_PAGE4_PARTITION = 256 +_SM70_QSA_XQA_PAGE4_PARTITION = 1024 _SM70_QSA_XQA_PAGE4_PAGES = 513 _SM70_QSA_XQA_PAGE4_MARKER = 1 << 30 +_SM70_QSA_GROUPED_PAGE4 = os.getenv("VLLM_SM70_QSA_GROUPED_PAGE4", "1") == "1" +_SM70_QSA_GROUPED_PAGE4_QUERIES = 8 +_SM70_QSA_GROUPED_PAGE4_OUTPUT_PAGES = ( + _SM70_QSA_XQA_PAGE4_PAGES * _SM70_QSA_GROUPED_PAGE4_QUERIES + 56 +) _SM70_QSA_XQA_PAGE4_WORKSPACES: dict[ tuple[int, int, int, int, int], tuple[int, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], ] = {} +_SM70_QSA_GROUPED_PAGE4_WORKSPACES: dict[ + tuple[int, int], + tuple[ + int, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ], +] = {} @triton.jit @@ -1427,6 +1442,131 @@ def _qsa_xqa_page4_workspace( ) +def _qsa_grouped_page4_workspace( + q: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + groups = q.shape[0] // _SM70_QSA_GROUPED_PAGE4_QUERIES + device_index = q.device.index if q.device.index is not None else -1 + stream_id = int(torch.cuda.current_stream(q.device).cuda_stream) + key = (device_index, stream_id) + workspace = _SM70_QSA_GROUPED_PAGE4_WORKSPACES.get(key) + if workspace is None or workspace[0] < groups: + capacity = 1 << (groups - 1).bit_length() + grouped_pages = torch.empty( + (capacity, _SM70_QSA_GROUPED_PAGE4_OUTPUT_PAGES), + dtype=torch.int32, + device=q.device, + ) + token_masks = torch.empty( + (capacity, _SM70_QSA_GROUPED_PAGE4_OUTPUT_PAGES), + dtype=torch.uint32, + device=q.device, + ) + grouped_sequence_lengths = torch.empty( + (capacity,), dtype=torch.int32, device=q.device + ) + lse = torch.empty( + (capacity * _SM70_QSA_GROUPED_PAGE4_QUERIES, q.shape[1]), + dtype=torch.float32, + device=q.device, + ) + workspace = ( + capacity, + grouped_pages, + token_masks, + grouped_sequence_lengths, + lse, + ) + _SM70_QSA_GROUPED_PAGE4_WORKSPACES[key] = workspace + _, grouped_pages, token_masks, grouped_sequence_lengths, lse = workspace + return ( + grouped_pages[:groups], + token_masks[:groups], + grouped_sequence_lengths[:groups], + lse[: q.shape[0]], + ) + + +def _qsa_xqa_page4_physical_kv( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + microblock_stride = 4 * q.shape[2] + if k_cache.stride(0) == k_cache.shape[1] * q.shape[2]: + microblocks_per_cache_block = k_cache.shape[1] // 4 + physical_k_cache = k_cache.view( + k_cache.shape[0] * microblocks_per_cache_block, + 4, + 1, + q.shape[2], + ) + physical_v_cache = v_cache.view_as(physical_k_cache) + else: + # The local FlashAttention ABI interleaves K and V inside every + # physical cache block. The virtual page IDs carry that doubled block + # stride, while this narrow view exposes a four-token page stride. + physical_shape = (k_cache.shape[0], 4, 1, q.shape[2]) + physical_strides = ( + microblock_stride, + q.shape[2], + q.shape[2], + 1, + ) + physical_k_cache = k_cache.as_strided(physical_shape, physical_strides) + physical_v_cache = v_cache.as_strided(physical_shape, physical_strides) + return physical_k_cache, physical_v_cache + + +def _qsa_sparse_paged_attention_sm70_grouped_page4( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + logical_indices: torch.Tensor, + block_table: torch.Tensor, + token_to_req: torch.Tensor, + query_positions: torch.Tensor, + sequence_lengths: torch.Tensor, + out: torch.Tensor, + flash_attn_v100_cuda, +) -> torch.Tensor: + grouped_pages, token_masks, grouped_sequence_lengths, lse = ( + _qsa_grouped_page4_workspace(q) + ) + physical_page_stride = k_cache.stride(0) // (4 * q.shape[2]) + flash_attn_v100_cuda.grouped_sparse_page4_plan_fwd( + logical_indices, + block_table, + token_to_req, + query_positions, + sequence_lengths, + grouped_pages, + token_masks, + grouped_sequence_lengths, + k_cache.shape[1], + physical_page_stride, + k_cache.shape[0], + ) + physical_k_cache, physical_v_cache = _qsa_xqa_page4_physical_kv(q, k_cache, v_cache) + flash_attn_v100_cuda.grouped_sparse_page4_fwd( + q, + physical_k_cache, + physical_v_cache, + out, + grouped_pages, + token_masks, + grouped_sequence_lengths, + lse, + q.shape[2] ** -0.5, + ) + logger.info_once( + "Using SM70 grouped QSA Flash-V100 page4 prefill route (rows=%d, groups=%d).", + q.shape[0], + q.shape[0] // _SM70_QSA_GROUPED_PAGE4_QUERIES, + ) + return out + + def _qsa_sparse_paged_attention_sm70_xqa_page4( q: torch.Tensor, k_cache: torch.Tensor, @@ -1453,6 +1593,27 @@ def _qsa_sparse_paged_attention_sm70_xqa_page4( ) return None + grouped_bindings_available = hasattr( + flash_attn_v100_cuda, "grouped_sparse_page4_plan_fwd" + ) and hasattr(flash_attn_v100_cuda, "grouped_sparse_page4_fwd") + if ( + _SM70_QSA_GROUPED_PAGE4 + and q.shape[0] % _SM70_QSA_GROUPED_PAGE4_QUERIES == 0 + and grouped_bindings_available + ): + return _qsa_sparse_paged_attention_sm70_grouped_page4( + q, + k_cache, + v_cache, + logical_indices, + block_table, + token_to_req, + query_positions, + sequence_lengths, + out, + flash_attn_v100_cuda, + ) + virtual_block_table, xqa_sequence_lengths = _qsa_xqa_page4_block_table( logical_indices, block_table, @@ -1467,30 +1628,7 @@ def _qsa_sparse_paged_attention_sm70_xqa_page4( temporary_output, max_logits, exp_sums, active_num_partitions = ( _qsa_xqa_page4_workspace(q, num_partitions) ) - microblock_stride = 4 * q.shape[2] - if k_cache.stride(0) == k_cache.shape[1] * q.shape[2]: - microblocks_per_cache_block = k_cache.shape[1] // 4 - physical_k_cache = k_cache.view( - k_cache.shape[0] * microblocks_per_cache_block, - 4, - 1, - q.shape[2], - ) - physical_v_cache = v_cache.view_as(physical_k_cache) - else: - # The local FlashAttention ABI interleaves K and V inside every - # physical cache block. The virtual page IDs carry that doubled block - # stride, while this narrow view exposes a four-token page stride to - # XQA. No cache data is copied or rearranged. - physical_shape = (k_cache.shape[0], 4, 1, q.shape[2]) - physical_strides = ( - microblock_stride, - q.shape[2], - q.shape[2], - 1, - ) - physical_k_cache = k_cache.as_strided(physical_shape, physical_strides) - physical_v_cache = v_cache.as_strided(physical_shape, physical_strides) + physical_k_cache, physical_v_cache = _qsa_xqa_page4_physical_kv(q, k_cache, v_cache) flash_attn_v100_cuda.decode_paged_xqa_fwd( q, physical_k_cache,