From f2b2b567aacc08a81121ecaed76cf23c739478e8 Mon Sep 17 00:00:00 2001 From: SII-yangdian <84713386+yangdian96@users.noreply.github.com> Date: Fri, 14 Aug 2026 10:00:22 +0800 Subject: [PATCH] perf(jit_kernel/deepseek_v4): optimize paged_mqa_metadata (#25855) Co-authored-by: SII-yangdian --- benchmark/kernels/bench_paged_mqa_metadata.py | 57 +++ .../csrc/deepseek_v4/paged_mqa_metadata.cuh | 354 +++++++++++++++--- .../sglang/kernels/ops/attention/dsv4/attn.py | 11 +- .../ops/attention/test_paged_mqa_metadata.py | 277 ++++++++++++++ 4 files changed, 641 insertions(+), 58 deletions(-) create mode 100644 benchmark/kernels/bench_paged_mqa_metadata.py create mode 100644 test/registered/kernels/ops/attention/test_paged_mqa_metadata.py diff --git a/benchmark/kernels/bench_paged_mqa_metadata.py b/benchmark/kernels/bench_paged_mqa_metadata.py new file mode 100644 index 000000000..49102c369 --- /dev/null +++ b/benchmark/kernels/bench_paged_mqa_metadata.py @@ -0,0 +1,57 @@ +"""Benchmark paged_mqa_metadata JIT kernel. + +Reports per-shape median latency in µs via ``marker.do_bench`` (CUDA-graph +timing). + +Shape axes: + - ``bs``: dense sweep from single-request decode (1) to large multi-block + batch (32768). Covers the three internal dispatch paths + (tiny ``bs<=64`` / small ``bs<=2048`` / multi-block ``bs>2048``). + - ``max_ctx``: two extremes (2048, 32768). The kernel is value-invariant + (cost is O(bs) regardless of seq_lens values); sweeping both bookends + makes that empirically visible. + +Constants: ``num_sm`` queried from the active GPU; ``page_size = 64``. + +Local run: + python benchmark/kernels/bench_paged_mqa_metadata.py +""" + +import torch + +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.ops.attention.dsv4 import get_paged_mqa_logits_metadata + +NUM_SM = ( + torch.cuda.get_device_properties(0).multi_processor_count + if torch.cuda.is_available() + else 132 +) +PAGE_SIZE = 64 +DEVICE = "cuda" + + +def _make_seq_lens(bs: int, max_ctx: int, seed: int = 0) -> torch.Tensor: + g = torch.Generator(device=DEVICE).manual_seed(seed) + return torch.randint( + 1, max_ctx + 1, (bs,), dtype=torch.int32, device=DEVICE, generator=g + ) + + +@marker.parametrize( + "bs", + [1, 8, 16, 32, 64, 128, 256, 384, 512, 1024, 2048, 4096, 8192, 16384, 32768], + [128, 2048], +) +@marker.parametrize("max_ctx", [2048, 32768], [8192]) +@marker.benchmark("impl", ["jit"]) +def benchmark(bs: int, max_ctx: int, impl: str): + seq_lens = _make_seq_lens(bs, max_ctx) + return marker.do_bench( + get_paged_mqa_logits_metadata, + input_args=(seq_lens, PAGE_SIZE, NUM_SM), + ) + + +if __name__ == "__main__": + benchmark.run() diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/paged_mqa_metadata.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/paged_mqa_metadata.cuh index b52ab9ff1..0b25c1137 100644 --- a/python/sglang/kernels/jit/csrc/deepseek_v4/paged_mqa_metadata.cuh +++ b/python/sglang/kernels/jit/csrc/deepseek_v4/paged_mqa_metadata.cuh @@ -1,93 +1,314 @@ +// paged_mqa_metadata: batch-size-adaptive dispatch. +// +// Replaces upstream's single-block kernel (grid=1, Phase-3 lane-serial +// advance, O(bs) dependent loads on the critical path) with three internal +// kernels dispatched by batch_size, all sharing the same Phase-1/2 prefix +// sum and a `num_sm + 1`-thread parallel upper_bound for Phase 3. + #include #include #include -#include +#include #include #include +#include + namespace sglang { -constexpr uint32_t kBlockSize = 1024; constexpr uint32_t kSplitKV = 256; // const for both SM90 and SM100 +constexpr uint32_t kTinyBlock = 256; +constexpr uint32_t kTinyMax = 64; + +constexpr uint32_t kSmallBlock = 256; +constexpr uint32_t kSmallMax = 2048; +constexpr uint32_t kSmallItemsPerThread = 8; +static_assert(kSmallBlock * kSmallItemsPerThread == kSmallMax); + +constexpr uint32_t kMBTileSize = 4096; +constexpr uint32_t kMBBlockSize = 1024; +constexpr uint32_t kMBItemsPerThread = 4; +constexpr uint32_t kKernelBThreads = 256; +static_assert(kMBBlockSize * kMBItemsPerThread == kMBTileSize); + struct MetadataParams { - /// NOTE: batch_size > 0 uint32_t batch_size; uint32_t num_sm; const uint32_t* __restrict__ context_lens; uint32_t* __restrict__ schedule_metadata; - bool use_smem = true; }; -__global__ __launch_bounds__(kBlockSize, 1) // - void smxx_paged_mqa_logits_metadata(const MetadataParams params) { - using namespace device; - extern __shared__ uint32_t s_length[]; - static constexpr auto kNumWarps = kBlockSize / kWarpThreads; - static_assert(kNumWarps == kWarpThreads); +// bs <= 64. Warp-0 inclusive scan, 256 B static smem. +__global__ __launch_bounds__(kTinyBlock, 1) // + void paged_mqa_metadata_tiny_kernel(const MetadataParams params) { + __shared__ uint32_t s_prefix[kTinyMax]; + __shared__ uint32_t s_global_sum; - const auto tx = threadIdx.x; - const auto lane_id = tx % kWarpThreads; - const auto warp_id = tx / kWarpThreads; + const uint32_t tx = threadIdx.x; + const uint32_t bs = params.batch_size; + const uint32_t num_sm = params.num_sm; - __shared__ uint32_t s_warp_sum[kNumWarps]; - - uint32_t local_sum = 0; - for (uint32_t i = tx; i < params.batch_size; i += kBlockSize) { - const auto length = params.context_lens[i]; - local_sum += (length + kSplitKV - 1) / kSplitKV; - if (params.use_smem) s_length[i] = length; + if (tx < 32) { + uint32_t running = 0; +#pragma unroll + for (uint32_t base = 0; base < kTinyMax; base += 32) { + const uint32_t idx = base + tx; + uint32_t v = 0; + if (idx < bs) { + const uint32_t length = params.context_lens[idx]; + v = (length + kSplitKV - 1) >> 8; + } +#pragma unroll + for (int o = 1; o < 32; o <<= 1) { + uint32_t y = __shfl_up_sync(0xffffffff, v, o); + if (tx >= static_cast(o)) v += y; + } + v += running; + if (idx < bs) s_prefix[idx] = v; + running = __shfl_sync(0xffffffff, v, 31); + } + if (tx == 0) s_global_sum = running; } - - s_warp_sum[warp_id] = warp::reduce_sum(local_sum); __syncthreads(); - const auto global_sum = warp::reduce_sum(s_warp_sum[lane_id]); - if (lane_id != 0) return; + const uint32_t global_sum = s_global_sum; + const uint32_t avg = global_sum / num_sm; + const uint32_t ret = global_sum % num_sm; + const uint32_t pivot = num_sm - ret; - const auto length_ptr = params.use_smem ? s_length : params.context_lens; + // Stride loop so num_sm > blockDim.x - 1 is fully written. + for (uint32_t i = tx; i <= num_sm; i += blockDim.x) { + // Match DeepGEMM's reversed remainder allocation: leading SMs get + // `avg` work and the final `ret` SMs get `avg + 1`. When global_sum is + // smaller than num_sm, empty SMs stay at the valid (q=0, offset=0) + // boundary instead of starting at q=batch_size. + const uint32_t target = i * avg + (i > pivot ? i - pivot : 0); - const auto avg = global_sum / params.num_sm; - const auto ret = global_sum % params.num_sm; - uint32_t q = 0; - uint32_t num_work = (length_ptr[0] + kSplitKV - 1) / kSplitKV; - uint32_t sum_work = num_work; - for (auto i = warp_id; i <= params.num_sm; i += kNumWarps) { - const auto target = i * avg + min(i, ret); - while (sum_work <= target) { - if (++q >= params.batch_size) break; - num_work = (length_ptr[q] + kSplitKV - 1) / kSplitKV; - sum_work += num_work; + uint32_t lo = 0; + uint32_t hi = bs; + while (lo < hi) { + const uint32_t mid = (lo + hi) >> 1; + if (s_prefix[mid] <= target) + lo = mid + 1; + else + hi = mid; } - if (q >= params.batch_size) { - params.schedule_metadata[2 * i + 0] = params.batch_size; + const uint32_t q = lo; + + if (q >= bs) { + params.schedule_metadata[2 * i + 0] = bs; params.schedule_metadata[2 * i + 1] = 0; } else { - // sum > target && (sum - length) <= target + const uint32_t prefix_prev = (q == 0) ? 0u : s_prefix[q - 1]; params.schedule_metadata[2 * i + 0] = q; - params.schedule_metadata[2 * i + 1] = target - (sum_work - num_work); + params.schedule_metadata[2 * i + 1] = target - prefix_prev; } } } -template -void setup_kernel_smem_once(host::DebugInfo where = {}) { - [[maybe_unused]] - static const auto result = [] { - const auto fptr = std::bit_cast(f); - return ::cudaFuncSetAttribute(fptr, ::cudaFuncAttributeMaxDynamicSharedMemorySize, kMaxDynamicSMEM); - }(); - host::RuntimeDeviceCheck(result, where); +// 64 < bs <= 2048. CUB BlockScan, 8 KB static smem. +__global__ __launch_bounds__(kSmallBlock, 1) // + void paged_mqa_metadata_small_kernel(const MetadataParams params) { + using BlockScan = cub::BlockScan; + __shared__ uint32_t s_prefix[kSmallMax]; + __shared__ typename BlockScan::TempStorage temp_storage; + __shared__ uint32_t s_global_sum; + + const uint32_t tx = threadIdx.x; + const uint32_t bs = params.batch_size; + const uint32_t num_sm = params.num_sm; + + uint32_t thread_items[kSmallItemsPerThread]; +#pragma unroll + for (uint32_t k = 0; k < kSmallItemsPerThread; ++k) { + const uint32_t i = tx * kSmallItemsPerThread + k; + if (i < bs) { + const uint32_t length = params.context_lens[i]; + thread_items[k] = (length + kSplitKV - 1) >> 8; + } else { + thread_items[k] = 0; + } + } + + uint32_t block_aggregate; + BlockScan(temp_storage).InclusiveSum(thread_items, thread_items, block_aggregate); + + if (tx == 0) s_global_sum = block_aggregate; + +#pragma unroll + for (uint32_t k = 0; k < kSmallItemsPerThread; ++k) { + const uint32_t i = tx * kSmallItemsPerThread + k; + if (i < bs) s_prefix[i] = thread_items[k]; + } + __syncthreads(); + + const uint32_t global_sum = s_global_sum; + const uint32_t avg = global_sum / num_sm; + const uint32_t ret = global_sum % num_sm; + const uint32_t pivot = num_sm - ret; + + // Stride loop so num_sm > blockDim.x - 1 is fully written. + for (uint32_t i = tx; i <= num_sm; i += blockDim.x) { + const uint32_t target = i * avg + (i > pivot ? i - pivot : 0); + + uint32_t lo = 0; + uint32_t hi = bs; + while (lo < hi) { + const uint32_t mid = (lo + hi) >> 1; + if (s_prefix[mid] <= target) + lo = mid + 1; + else + hi = mid; + } + const uint32_t q = lo; + + if (q >= bs) { + params.schedule_metadata[2 * i + 0] = bs; + params.schedule_metadata[2 * i + 1] = 0; + } else { + const uint32_t prefix_prev = (q == 0) ? 0u : s_prefix[q - 1]; + params.schedule_metadata[2 * i + 0] = q; + params.schedule_metadata[2 * i + 1] = target - prefix_prev; + } + } +} + +// bs > 2048, Phase 1: ceil(bs / kMBTileSize) blocks each emit an in-tile +// inclusive prefix into scratch_prefix and a per-tile sum into tile_sums. +__global__ __launch_bounds__(kMBBlockSize, 1) // + void phase1_tile_scan_kernel( + const MetadataParams params, uint32_t* __restrict__ scratch_prefix, uint32_t* __restrict__ tile_sums) { + using TileBlockScan = cub::BlockScan; + __shared__ typename TileBlockScan::TempStorage temp_storage; + + const uint32_t bs = params.batch_size; + const uint32_t tile_idx = blockIdx.x; + const uint32_t tile_base = tile_idx * kMBTileSize; + const uint32_t tx = threadIdx.x; + + uint32_t thread_items[kMBItemsPerThread]; +#pragma unroll + for (uint32_t k = 0; k < kMBItemsPerThread; ++k) { + const uint32_t i = tile_base + tx * kMBItemsPerThread + k; + if (i < bs) { + const uint32_t length = params.context_lens[i]; + thread_items[k] = (length + kSplitKV - 1) >> 8; + } else { + thread_items[k] = 0; + } + } + + uint32_t block_aggregate; + TileBlockScan(temp_storage).InclusiveSum(thread_items, thread_items, block_aggregate); + +#pragma unroll + for (uint32_t k = 0; k < kMBItemsPerThread; ++k) { + const uint32_t i = tile_base + tx * kMBItemsPerThread + k; + if (i < bs) scratch_prefix[i] = thread_items[k]; + } + + if (tx == 0) tile_sums[tile_idx] = block_aggregate; +} + +// bs > 2048, Phase 2/3: one block, kKernelBThreads threads. Warp-0 scans +// tile_sums into s_tile_prefix; then num_sm+1 threads do tile-level +// upper_bound + within-tile upper_bound to recover (batch_idx, offset). +__global__ __launch_bounds__(kKernelBThreads, 1) // + void schedule_from_tiles_kernel( + const MetadataParams params, + const uint32_t* __restrict__ scratch_prefix, + const uint32_t* __restrict__ tile_sums, + uint32_t num_tiles) { + extern __shared__ uint32_t s_tile_prefix[]; + __shared__ uint32_t s_global_sum; + + const uint32_t tx = threadIdx.x; + const uint32_t bs = params.batch_size; + const uint32_t num_sm = params.num_sm; + + if (tx < 32) { + uint32_t running = 0; + for (uint32_t base = 0; base < num_tiles; base += 32) { + const uint32_t idx = base + tx; + uint32_t v = (idx < num_tiles) ? tile_sums[idx] : 0; +#pragma unroll + for (int o = 1; o < 32; o <<= 1) { + uint32_t y = __shfl_up_sync(0xffffffff, v, o); + if (tx >= static_cast(o)) v += y; + } + v += running; + if (idx < num_tiles) s_tile_prefix[idx] = v; + running = __shfl_sync(0xffffffff, v, 31); + } + if (tx == 0) s_global_sum = running; + } + __syncthreads(); + + const uint32_t global_sum = s_global_sum; + const uint32_t avg = global_sum / num_sm; + const uint32_t ret = global_sum % num_sm; + const uint32_t pivot = num_sm - ret; + + // Stride loop so num_sm > blockDim.x - 1 is fully written. `continue` + // replaces the original early-out `return` so later strided targets + // still get processed. + for (uint32_t i = tx; i <= num_sm; i += blockDim.x) { + const uint32_t target = i * avg + (i > pivot ? i - pivot : 0); + + uint32_t t_lo = 0; + uint32_t t_hi = num_tiles; + while (t_lo < t_hi) { + const uint32_t mid = (t_lo + t_hi) >> 1; + if (s_tile_prefix[mid] <= target) + t_lo = mid + 1; + else + t_hi = mid; + } + const uint32_t tile = t_lo; + + if (tile >= num_tiles) { + params.schedule_metadata[2 * i + 0] = bs; + params.schedule_metadata[2 * i + 1] = 0; + continue; + } + + const uint32_t tile_offset = (tile == 0) ? 0u : s_tile_prefix[tile - 1]; + const uint32_t tile_start = tile * kMBTileSize; + uint32_t tile_end = tile_start + kMBTileSize; + if (tile_end > bs) tile_end = bs; + const uint32_t local_target = target - tile_offset; + + uint32_t lo = tile_start; + uint32_t hi = tile_end; + while (lo < hi) { + const uint32_t mid = (lo + hi) >> 1; + if (scratch_prefix[mid] <= local_target) + lo = mid + 1; + else + hi = mid; + } + const uint32_t q = lo; + + if (q >= bs) { + params.schedule_metadata[2 * i + 0] = bs; + params.schedule_metadata[2 * i + 1] = 0; + } else { + const uint32_t prefix_prev = (q == tile_start) ? tile_offset : (scratch_prefix[q - 1] + tile_offset); + params.schedule_metadata[2 * i + 0] = q; + params.schedule_metadata[2 * i + 1] = target - prefix_prev; + } + } } struct IndexerMetadataKernel { - static constexpr auto kMaxBatchSizeInSmem = 16384 * 2; // 128 KB smeme - static void run(tvm::ffi::TensorView seq_lens, tvm::ffi::TensorView metadata) { + static void run(tvm::ffi::TensorView seq_lens, tvm::ffi::TensorView metadata, tvm::ffi::TensorView workspace) { using namespace host; auto N = SymbolicSize{"batch_size"}; auto M = SymbolicSize{"num_sm"}; + auto W = SymbolicSize{"workspace"}; auto device = SymbolicDevice{}; device.set_options(); TensorMatcher({N}) // @@ -98,21 +319,40 @@ struct IndexerMetadataKernel { .with_dtype() .with_device(device) .verify(metadata); + TensorMatcher({W}) // + .with_dtype() + .with_device(device) + .verify(workspace); + const auto batch_size = static_cast(N.unwrap()); const auto num_sm = static_cast(M.unwrap()) - 1; - RuntimeCheck(num_sm <= 1024); - const auto use_smem = batch_size <= kMaxBatchSizeInSmem; + RuntimeCheck(num_sm >= 1 && num_sm <= 1024); + const auto params = MetadataParams{ .batch_size = batch_size, .num_sm = num_sm, .context_lens = static_cast(seq_lens.data_ptr()), .schedule_metadata = static_cast(metadata.data_ptr()), - .use_smem = use_smem, }; - constexpr auto kernel = smxx_paged_mqa_logits_metadata; - setup_kernel_smem_once(); - const auto smem = use_smem ? (batch_size + 1) * sizeof(uint32_t) : 0; - LaunchKernel(1, kBlockSize, device.unwrap(), smem)(kernel, params); + + const auto dl_device = device.unwrap(); + + if (batch_size <= kTinyMax) { + LaunchKernel(1, kTinyBlock, dl_device)(paged_mqa_metadata_tiny_kernel, params); + } else if (batch_size <= kSmallMax) { + LaunchKernel(1, kSmallBlock, dl_device)(paged_mqa_metadata_small_kernel, params); + } else { + const auto num_tiles = (batch_size + kMBTileSize - 1) / kMBTileSize; + const auto required = static_cast(batch_size) + num_tiles; + RuntimeCheck(static_cast(W.unwrap()) >= required, "workspace too small for multi-block path"); + auto* scratch_prefix = static_cast(workspace.data_ptr()); + auto* tile_sums = scratch_prefix + batch_size; + + LaunchKernel(num_tiles, kMBBlockSize, dl_device)(phase1_tile_scan_kernel, params, scratch_prefix, tile_sums); + const auto kb_smem_bytes = static_cast(num_tiles) * sizeof(uint32_t); + LaunchKernel(1, kKernelBThreads, dl_device, kb_smem_bytes)( + schedule_from_tiles_kernel, params, scratch_prefix, tile_sums, num_tiles); + } } }; diff --git a/python/sglang/kernels/ops/attention/dsv4/attn.py b/python/sglang/kernels/ops/attention/dsv4/attn.py index 6f5b95841..f973bf916 100644 --- a/python/sglang/kernels/ops/attention/dsv4/attn.py +++ b/python/sglang/kernels/ops/attention/dsv4/attn.py @@ -45,9 +45,18 @@ def _jit_fused_store_module( def get_paged_mqa_logits_metadata(seq_lens: torch.Tensor, page_size: int, num_sm: int): assert page_size == 64 seq_lens = seq_lens.view(-1).to(torch.int32) + bs = int(seq_lens.shape[0]) metadata = seq_lens.new_empty(num_sm + 1, 2) + # Workspace for the multi-block path; kMBTileSize must match the .cuh. + if bs > 2048: + kMBTileSize = 4096 + workspace = seq_lens.new_empty( + bs + (bs + kMBTileSize - 1) // kMBTileSize, dtype=torch.int32 + ) + else: + workspace = seq_lens.new_empty(0, dtype=torch.int32) module = _jit_metadata_module() - module.run(seq_lens, metadata) + module.run(seq_lens, metadata, workspace) return metadata diff --git a/test/registered/kernels/ops/attention/test_paged_mqa_metadata.py b/test/registered/kernels/ops/attention/test_paged_mqa_metadata.py new file mode 100644 index 000000000..d16f09b08 --- /dev/null +++ b/test/registered/kernels/ops/attention/test_paged_mqa_metadata.py @@ -0,0 +1,277 @@ +"""Unit tests for paged_mqa_metadata JIT kernel. + +Verifies byte-equal correctness against a pure-PyTorch reference oracle +across the shape envelope. Output is int32 ``[num_sm + 1, 2]`` — a +deterministic partition table — so equality is strict (``torch.equal``, +no atol/rtol). + +Test groups: + +1. ``test_matches_pytorch_ref`` — random-input envelope sweep over + ``bs x max_ctx`` (powers-of-2 + off-by-one for ``bs`` to stress the + ``ret`` branch where ``bs % num_sm != 0``; kSplitKV=256 boundary + values for ``max_ctx``). + +2. ``test_matches_pytorch_ref_at_ksplitkv_boundary`` — hand-crafted + ``seq_lens`` straddling the internal ``kSplitKV=256`` boundary + (catches off-by-one in ``ceil(len/256)``). + +3. ``test_byte_equal_at_correctness_floor`` — ``bs`` above the smem-path + ceiling (``bs > 32768``); exercises the multi-block gmem path. Catches + regressions where a future kernel adds a ``batch_size`` upper bound for + smem convenience. + +4. ``test_matches_pytorch_ref_at_large_num_sm`` — ``num_sm in [1, 1024]`` + contract; guards against per-block thread-guard truncation across the + three dispatch paths. + +5. ``test_matches_deep_gemm`` — byte-equality against the production + ``deep_gemm`` oracle; auto-skips at ``bs >= 16384`` where deep_gemm + exceeds sm_90's smem cap. +""" + +import itertools + +import pytest +import torch + +from sglang.kernels.jit.utils import get_ci_test_range +from sglang.kernels.ops.attention.dsv4 import get_paged_mqa_logits_metadata +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="1-gpu-large") + + +KSPLITKV = 256 # internal kernel constant (note: public API page_size=64 is unrelated) +NUM_SM = ( + 132 # H200 reference; CI box may have fewer SMs but the kernel is SM-count agnostic +) +PAGE_SIZE = 64 +DEVICE = "cuda" + + +# Algorithmic spec for the int32 [num_sm+1, 2] schedule. Ground-truth +# correctness is gated by ``test_matches_deep_gemm`` below; this ref is +# used by the wider envelope tests where deep_gemm exceeds the sm_90 +# smem cap. +def paged_mqa_metadata_ref( + seq_lens: torch.Tensor, num_sm: int, page_size: int +) -> torch.Tensor: + assert page_size == 64, f"page_size must be 64, got {page_size}" + assert ( + seq_lens.dtype == torch.int32 + ), f"seq_lens dtype must be int32, got {seq_lens.dtype}" + assert ( + seq_lens.dim() == 1 + ), f"seq_lens must be 1-D, got shape {tuple(seq_lens.shape)}" + + device = seq_lens.device + batch_size = int(seq_lens.shape[0]) + + work_per_batch = (seq_lens.to(torch.int64) + KSPLITKV - 1) // KSPLITKV + global_sum = int(work_per_batch.sum().item()) + avg = global_sum // num_sm + ret = global_sum % num_sm + pivot = num_sm - ret + + schedule_metadata = torch.empty((num_sm + 1, 2), dtype=torch.int32, device=device) + work = work_per_batch.tolist() + q = 0 + sum_work = work[0] if batch_size > 0 else 0 + for i in range(num_sm + 1): + # Match DeepGEMM's reversed allocation: the final ``ret`` SMs get + # one extra unit of work. This keeps leading empty SMs at q=0 when + # there is less total work than available SMs. + target = i * avg + max(i - pivot, 0) + while sum_work <= target: + q += 1 + if q >= batch_size: + break + sum_work += work[q] + if q >= batch_size: + schedule_metadata[i, 0] = batch_size + schedule_metadata[i, 1] = 0 + else: + schedule_metadata[i, 0] = q + schedule_metadata[i, 1] = target - (sum_work - work[q]) + return schedule_metadata + + +# ----------------------------------------------------------------------------- +# Shape envelope +# +# bs values: +# 1 single-request decode (smallest realistic input) +# 17, 129, 257 non-power-of-2; Phase 3 q-advance hits `ret` branch +# (bs % num_sm != 0 → uneven work split) +# 1025 ditto, large-bs ret-branch stressor +# 32, 128, 512, 1024, 2048 DSv4 decode/prefill realistic batches +# 4096..32768 multi-block path (bs > kSmallMax = 2048) +# +# max_ctx values: +# 1 degenerate (all seq_lens == 1 → work_per_batch == 1) +# 255 just under kSplitKV=256 boundary (ceil(255/256) = 1) +# 256 exactly at kSplitKV boundary (ceil(256/256) = 1) +# 257 just over boundary (ceil(257/256) = 2) +# 2048, 8192 realistic decode/short-prefill contexts +# 32768 long-context upper bound +# ----------------------------------------------------------------------------- +BS_FULL = [1, 17, 32, 128, 129, 257, 512, 1024, 1025, 2048, 4096, 8192, 16384, 32768] +MAX_CTX_FULL = [1, 255, 256, 257, 2048, 8192, 32768] + +# bs values above the in-smem path's ceiling (kSmallMax = 2048; multi-block +# path takes over). Tested separately to keep the CI parametrize matrix +# small while still guarding the gmem path. +BS_CORRECTNESS_FLOOR = [65536, 131072] + +BS_LIST = get_ci_test_range( + full_range=BS_FULL, + ci_range=[1, 128, 1025], # tiny + typical decode + large ret-branch +) +MAX_CTX_LIST = get_ci_test_range( + full_range=MAX_CTX_FULL, + ci_range=[256, 8192], # kSplitKV boundary + typical decode +) + + +def _make_seq_lens(bs: int, max_ctx: int, seed: int = 0) -> torch.Tensor: + g = torch.Generator(device=DEVICE).manual_seed(seed) + return torch.randint( + 1, max_ctx + 1, (bs,), dtype=torch.int32, device=DEVICE, generator=g + ) + + +# ----------------------------------------------------------------------------- +# Tests +# ----------------------------------------------------------------------------- +@pytest.mark.parametrize("bs,max_ctx", list(itertools.product(BS_LIST, MAX_CTX_LIST))) +def test_matches_pytorch_ref(bs: int, max_ctx: int): + """Kernel output bit-exact vs PyTorch reference across the shape envelope.""" + seq_lens = _make_seq_lens(bs, max_ctx) + got = get_paged_mqa_logits_metadata(seq_lens, PAGE_SIZE, NUM_SM) + ref = paged_mqa_metadata_ref(seq_lens, NUM_SM, PAGE_SIZE) + assert torch.equal(got, ref), ( + f"kernel != ref for bs={bs} max_ctx={max_ctx}\n" + f" kernel first row: {got[0].tolist()}\n" + f" ref first row: {ref[0].tolist()}" + ) + + +_KSPLITKV_BOUNDARY_LENS = [ + [1], # minimum + [256], # exact kSplitKV multiple + [255, 256, 257], # straddle boundary + [256] * 132, # all-equal at boundary, bs == num_sm + [1] * 131 + [32768], # one giant + rest minimum (skewed) + [1, 256, 512, 768, 1024], # exact multiples + [255, 511, 767, 1023, 1279], # one below each multiple + [257, 513, 769, 1025, 1281], # one above each multiple +] + + +@pytest.mark.parametrize("seq_lens_data", _KSPLITKV_BOUNDARY_LENS) +def test_matches_pytorch_ref_at_ksplitkv_boundary(seq_lens_data): + """Bit-exact vs ref on hand-crafted kSplitKV=256 boundary inputs. + + Catches off-by-one in ceil(len/256) (Phase 1) and uneven work-split + in Phase 3's advance loop. + """ + seq_lens = torch.tensor(seq_lens_data, dtype=torch.int32, device=DEVICE) + got = get_paged_mqa_logits_metadata(seq_lens, PAGE_SIZE, NUM_SM) + ref = paged_mqa_metadata_ref(seq_lens, NUM_SM, PAGE_SIZE) + assert torch.equal(got, ref), ( + f"kernel != ref for seq_lens={seq_lens_data}\n" + f" kernel: {got.tolist()}\n ref: {ref.tolist()}" + ) + + +@pytest.mark.parametrize("bs", BS_CORRECTNESS_FLOOR) +@pytest.mark.parametrize("max_ctx", [8192]) +def test_byte_equal_at_correctness_floor(bs: int, max_ctx: int): + """bs above smem-path ceiling: multi-block path must remain byte-equal. + + Guards against a future kernel adding a ``batch_size`` upper bound for + smem convenience and breaking the gmem fallback. + """ + seq_lens = _make_seq_lens(bs, max_ctx) + got = get_paged_mqa_logits_metadata(seq_lens, PAGE_SIZE, NUM_SM) + ref = paged_mqa_metadata_ref(seq_lens, NUM_SM, PAGE_SIZE) + assert torch.equal(got, ref), f"kernel != ref at correctness-floor bs={bs}" + + +# ----------------------------------------------------------------------------- +# Defensive: kernel must handle the full num_sm range [1, 1024] guaranteed +# by the public API contract. The internal dispatch uses smaller per-block +# thread counts on some paths, so the per-target write loop must stride +# rather than use a `tx <= num_sm` thread guard. +# ----------------------------------------------------------------------------- +@pytest.mark.parametrize( + "bs,num_sm", + [ + # bs spans the tiny / small / multi-block paths. + # num_sm in {256, 257, 500, 1024} crosses the boundary where a + # 256-thread block would silently truncate schedule_metadata. + (64, 256), + (64, 257), + (64, 1024), + (128, 257), + (128, 1024), + (8192, 257), + (8192, 1024), + ], +) +def test_matches_pytorch_ref_at_large_num_sm(bs: int, num_sm: int): + """schedule_metadata[0 .. num_sm] must be fully populated for any + num_sm in [1, 1024], regardless of which dispatch path bs selects.""" + seq_lens = _make_seq_lens(bs, max_ctx=8192) + got = get_paged_mqa_logits_metadata(seq_lens, PAGE_SIZE, num_sm) + ref = paged_mqa_metadata_ref(seq_lens, num_sm, PAGE_SIZE) + assert torch.equal(got, ref), ( + f"kernel != ref at bs={bs} num_sm={num_sm}; " + f"last-5 rows: got={got[-5:].tolist()} ref={ref[-5:].tolist()}" + ) + + +def _to_2d_context_lens(seq_lens: torch.Tensor) -> torch.Tensor: + return seq_lens.contiguous().view(-1, 1) + + +def _load_deep_gemm(): + try: + import deep_gemm # noqa: PLC0415 + + return deep_gemm + except Exception as e: # noqa: BLE001 + pytest.skip(f"deep_gemm unavailable: {type(e).__name__}: {e}") + + +@pytest.mark.parametrize("bs", [64, 128, 1025, 2048, 4096, 8192, 32768, 65536]) +@pytest.mark.parametrize("max_ctx", [256, 32768]) +def test_matches_deep_gemm(bs: int, max_ctx: int): + """Byte-equal against the production deep_gemm reference. Auto-skips + at large bs where deep_gemm exceeds sm_90 smem cap.""" + deep_gemm = _load_deep_gemm() + seq_lens = _make_seq_lens(bs, max_ctx) + got = get_paged_mqa_logits_metadata(seq_lens, PAGE_SIZE, NUM_SM) + try: + dg = deep_gemm.get_paged_mqa_logits_metadata( + _to_2d_context_lens(seq_lens), PAGE_SIZE, NUM_SM + ) + except RuntimeError as e: + msg = str(e) + if "smem" in msg.lower() or "capacity" in msg.lower(): + pytest.skip( + f"deep_gemm smem cap exceeded at bs={bs}: {msg.splitlines()[0]}" + ) + raise + assert torch.equal(got, dg), ( + f"kernel != deep_gemm for bs={bs} max_ctx={max_ctx}\n" + f" kernel first row: {got[0].tolist()}\n" + f" dg first row: {dg[0].tolist()}" + ) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__]))