From 8d75652949948da1abc3e41fccbf0828638623a6 Mon Sep 17 00:00:00 2001 From: claude Date: Fri, 14 Aug 2026 07:48:52 +0000 Subject: [PATCH] feat: import CUDA kernels from xllm/CCCL/FLA upstream repos MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Sources cloned and tree'd (no --depth): - jd-opensource/xllm: ILU kernels, CUDA kernels, MoE kernels - NVIDIA/cccl: CUB tuning/dispatch headers (block-level primitives) - fla-org/flash-linear-attention: Triton GDN kernels - NVIDIA/cutlass: grouped GEMM reference (read, not copied) - Dao-AILab/flash-attention: attention kernel reference (SM80+, read only) New CUDA kernels (from xllm, SM-agnostic, portable to BI-V100): ex_engine/xllm_kernels/cuda/activation.cu (188 lines) — silu_and_mul, gelu ex_engine/xllm_kernels/cuda/norm.cu (600 lines) — rms_norm, fused_add_rms_norm ex_engine/xllm_kernels/cuda/rope.cu (258 lines) — rotary_embedding ex_engine/xllm_kernels/cuda/block_copy.cu (209 lines) — copy_blocks, swap_blocks ex_engine/xllm_kernels/cuda/reshape_paged_cache.cu (101 lines) — KV cache ops ex_engine/xllm_kernels/cuda/headers/ (5 headers for compilation) ILU bridge kernel sources (from xllm, verified SAME as upstream): ex_engine/xllm_kernels/ilu/ (10 files, 925 lines total) — activation.cpp, attention.cpp, fused_moe.cpp, group_gemm.cpp, matmul.cpp, norm.cpp, rope.cpp, ilu_ops_api.h, ixformer.h, utils.h FLA Triton GDN kernels (for GatedDeltaNet without SM90+ FlashQLA): ex_engine/fla_kernels/gated_delta_rule/ (7 files, 2370 lines) — chunk_fwd.py (428), chunk.py (487), wy_fast.py (409), fused_recurrent.py (392), naive.py (161), gate.py (380) CCCL sync (12 tuning + 14 dispatch headers updated from NVIDIA/cccl): cccl_upstream/cub/cub/device/dispatch/tuning/ — 12 changed files synced cccl_upstream/cub/cub/device/dispatch/ — 14 changed dispatch files synced Compilation targets for real machine (ivcore10): 1. CUDA kernels: --cuda-gpu-arch=ivcore10 via corex clang/16 2. ILU bridges: torch.utils.cpp_extension linking ixformer .so 3. FLA kernels: Triton JIT (if Triton works on BI-V100) --- .../dispatch/dispatch_adjacent_difference.cuh | 34 +- .../device/dispatch/dispatch_batch_memcpy.cuh | 4 +- .../device/dispatch/dispatch_batched_topk.cuh | 1244 ++++++-- .../device/dispatch/dispatch_copy_mdspan.cuh | 85 +- .../device/dispatch/dispatch_histogram.cuh | 12 +- .../dispatch/dispatch_reduce_by_key.cuh | 32 +- .../cub/cub/device/dispatch/dispatch_scan.cuh | 138 +- .../device/dispatch/dispatch_scan_by_key.cuh | 54 +- .../dispatch_segmented_radix_sort.cuh | 2 +- .../dispatch/dispatch_segmented_reduce.cuh | 32 +- .../dispatch/dispatch_segmented_scan.cuh | 22 +- .../device/dispatch/dispatch_select_if.cuh | 17 +- .../dispatch/dispatch_three_way_partition.cuh | 24 +- .../dispatch_transform_tile_config.cuh | 2 +- .../cub/cub/device/dispatch/tuning/common.cuh | 5 +- .../dispatch/tuning/tuning_batch_memcpy.cuh | 4 +- .../dispatch/tuning/tuning_batched_topk.cuh | 323 ++- .../device/dispatch/tuning/tuning_merge.cuh | 23 +- .../dispatch/tuning/tuning_reduce_by_key.cuh | 4 +- .../dispatch/tuning/tuning_rle_encode.cuh | 4 +- .../tuning/tuning_rle_non_trivial_runs.cuh | 4 +- .../device/dispatch/tuning/tuning_scan.cuh | 6 +- .../dispatch/tuning/tuning_scan_by_key.cuh | 4 +- .../dispatch/tuning/tuning_segmented_sort.cuh | 28 +- .../dispatch/tuning/tuning_select_if.cuh | 22 +- .../tuning/tuning_three_way_partition.cuh | 4 +- .../fla_kernels/gated_delta_rule/__init__.py | 17 + .../fla_kernels/gated_delta_rule/chunk.py | 591 ++++ .../fla_kernels/gated_delta_rule/chunk_fwd.py | 428 +++ .../gated_delta_rule/fused_recurrent.py | 478 +++ .../fla_kernels/gated_delta_rule/gate.py | 344 +++ .../fla_kernels/gated_delta_rule/naive.py | 161 ++ .../fla_kernels/gated_delta_rule/wy_fast.py | 351 +++ ex_engine/fla_kernels/utils/__init__.py | 65 + ex_engine/fla_kernels/utils/cache.py | 449 +++ ex_engine/fla_kernels/utils/op.py | 101 + ex_engine/xllm_kernels/cuda/activation.cu | 188 ++ ex_engine/xllm_kernels/cuda/block_copy.cu | 209 ++ .../xllm_kernels/cuda/headers/cuda_ops_api.h | 306 ++ .../cuda/headers/device_utils.cuh | 116 + .../cuda/headers/fp8_quant_utils.cuh | 239 ++ .../cuda/headers/type_convert.cuh | 231 ++ ex_engine/xllm_kernels/cuda/headers/utils.h | 163 ++ ex_engine/xllm_kernels/cuda/norm.cu | 600 ++++ .../xllm_kernels/cuda/reshape_paged_cache.cu | 101 + ex_engine/xllm_kernels/cuda/rope.cu | 258 ++ ex_engine/xllm_kernels/ilu/activation.cpp | 32 + ex_engine/xllm_kernels/ilu/attention.cpp | 163 ++ ex_engine/xllm_kernels/ilu/fused_moe.cpp | 99 + ex_engine/xllm_kernels/ilu/group_gemm.cpp | 39 + ex_engine/xllm_kernels/ilu/ilu_ops_api.h | 153 + ex_engine/xllm_kernels/ilu/ixformer.h | 147 + ex_engine/xllm_kernels/ilu/matmul.cpp | 73 + ex_engine/xllm_kernels/ilu/norm.cpp | 51 + ex_engine/xllm_kernels/ilu/rope.cpp | 31 + ex_engine/xllm_kernels/ilu/utils.h | 63 + .../zh/getting_started/quick_start_GLM5.md | 1166 ++++---- .../layers/quantization/utils/__init__.py | 6 +- vllm_overrides/core/block/block_table.py | 652 ++--- .../core/block/cpu_gpu_block_allocator.py | 762 ++--- .../core/block/prefix_caching_block.py | 1658 +++++------ vllm_overrides/core/block_manager_v2.py | 852 +++--- vllm_overrides/core/evictor_v2.py | 220 +- .../model_executor/layers/sampler.py | 2558 ++++++++--------- .../model_executor/sampling_metadata.py | 1096 +++---- vllm_overrides/sampling_params.py | 942 +++--- 66 files changed, 12834 insertions(+), 5458 deletions(-) create mode 100644 ex_engine/fla_kernels/gated_delta_rule/__init__.py create mode 100644 ex_engine/fla_kernels/gated_delta_rule/chunk.py create mode 100644 ex_engine/fla_kernels/gated_delta_rule/chunk_fwd.py create mode 100644 ex_engine/fla_kernels/gated_delta_rule/fused_recurrent.py create mode 100644 ex_engine/fla_kernels/gated_delta_rule/gate.py create mode 100644 ex_engine/fla_kernels/gated_delta_rule/naive.py create mode 100644 ex_engine/fla_kernels/gated_delta_rule/wy_fast.py create mode 100644 ex_engine/fla_kernels/utils/__init__.py create mode 100644 ex_engine/fla_kernels/utils/cache.py create mode 100644 ex_engine/fla_kernels/utils/op.py create mode 100644 ex_engine/xllm_kernels/cuda/activation.cu create mode 100644 ex_engine/xllm_kernels/cuda/block_copy.cu create mode 100644 ex_engine/xllm_kernels/cuda/headers/cuda_ops_api.h create mode 100644 ex_engine/xllm_kernels/cuda/headers/device_utils.cuh create mode 100644 ex_engine/xllm_kernels/cuda/headers/fp8_quant_utils.cuh create mode 100644 ex_engine/xllm_kernels/cuda/headers/type_convert.cuh create mode 100644 ex_engine/xllm_kernels/cuda/headers/utils.h create mode 100644 ex_engine/xllm_kernels/cuda/norm.cu create mode 100644 ex_engine/xllm_kernels/cuda/reshape_paged_cache.cu create mode 100644 ex_engine/xllm_kernels/cuda/rope.cu create mode 100644 ex_engine/xllm_kernels/ilu/activation.cpp create mode 100644 ex_engine/xllm_kernels/ilu/attention.cpp create mode 100644 ex_engine/xllm_kernels/ilu/fused_moe.cpp create mode 100644 ex_engine/xllm_kernels/ilu/group_gemm.cpp create mode 100644 ex_engine/xllm_kernels/ilu/ilu_ops_api.h create mode 100644 ex_engine/xllm_kernels/ilu/ixformer.h create mode 100644 ex_engine/xllm_kernels/ilu/matmul.cpp create mode 100644 ex_engine/xllm_kernels/ilu/norm.cpp create mode 100644 ex_engine/xllm_kernels/ilu/rope.cpp create mode 100644 ex_engine/xllm_kernels/ilu/utils.h diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_adjacent_difference.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_adjacent_difference.cuh index 69480e69..150cab2c 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_adjacent_difference.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_adjacent_difference.cuh @@ -246,14 +246,15 @@ struct CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceAdjacentDifference" error = CubDebug( THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( num_tiles, AdjacentDifferencePolicyT::BLOCK_THREADS, 0, stream) - .doit(detail::adjacent_difference::DeviceAdjacentDifferenceDifferenceKernel < KernelPolicySelector, - InputIteratorT, - OutputIteratorT, - DifferenceOpT, - OffsetT, - InputT, - AliasOpt == MayAlias::Yes, - ReadOpt == ReadOption::Left >, + .doit(detail::adjacent_difference::DeviceAdjacentDifferenceDifferenceKernel< + KernelPolicySelector, + InputIteratorT, + OutputIteratorT, + DifferenceOpT, + OffsetT, + InputT, + AliasOpt == MayAlias::Yes, + ReadOpt == ReadOption::Left>, d_input, first_tile_previous, d_output, @@ -438,14 +439,15 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto dispatch( if (const auto error = CubDebug( THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(num_tiles, active_policy.threads_per_block, 0, stream) - .doit(DeviceAdjacentDifferenceDifferenceKernel < policy_selector_t, - InputIteratorT, - OutputIteratorT, - DifferenceOpT, - offset_t, - input_t, - AliasOpt == MayAlias::Yes, - ReadOpt == ReadOption::Left >, + .doit(DeviceAdjacentDifferenceDifferenceKernel< + policy_selector_t, + InputIteratorT, + OutputIteratorT, + DifferenceOpT, + offset_t, + input_t, + AliasOpt == MayAlias::Yes, + ReadOpt == ReadOption::Left>, d_input, first_tile_previous, d_output, diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_batch_memcpy.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_batch_memcpy.cuh index b86dd7a8..642ba8a3 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_batch_memcpy.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_batch_memcpy.cuh @@ -152,8 +152,8 @@ __launch_bounds__(int(current_policy().lookback.large_buffer.thr { if (thread_offset < buffer_sizes[buffer_id]) { - const auto value = read_item < MemcpyOpt == CopyAlg::Memcpy, AliasT, - InputBufferT > (input_buffer_it[buffer_id], thread_offset); + const auto value = + read_item(input_buffer_it[buffer_id], thread_offset); write_item( output_buffer_it[buffer_id], thread_offset, value); } diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_batched_topk.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_batched_topk.cuh index f95538f0..ce4b5ab3 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_batched_topk.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_batched_topk.cuh @@ -2,8 +2,8 @@ // Apache-2.0 WITH LLVM-exception //! @file -//! cub::DeviceTopK provides device-wide, parallel operations for finding the K largest (or smallest) items from -//! sequences of unordered data items residing within device-accessible memory. +//! Internal device-wide dispatch for cub::DeviceBatchedTopK: selects between the baseline (worker-per-segment) and +//! cluster (SM 9.0+) backends and launches them through a single kernel symbol. #pragma once @@ -17,14 +17,17 @@ # pragma system_header #endif // no system header +#include +#include #include -#include +#include #include #include #include #include #include #include +#include #include #include #include @@ -32,16 +35,32 @@ #include #include +#include +#include +#include #include #include +#include #include +#include +#include +#include +#include #include +#include #include +#include #include #include +#include +#include +#include #include +#include #include +#include + CUB_NAMESPACE_BEGIN namespace detail::batched_topk @@ -53,7 +72,7 @@ namespace detail::batched_topk // The selection direction is compile-time only: callers pass `::cuda::args::constant`, which maps to a // value-less static_discrete_param. Because the direction is fixed at compile time and carries no runtime value, it // can never disagree with its only supported option, so dispatch can never silently degrade to a no-op. -template +template [[nodiscard]] _CCCL_HOST_DEVICE auto wrap_select_direction(::cuda::args::constant) { return params::static_discrete_param{}; @@ -84,7 +103,7 @@ template // per-segment tile counts that we exclusive-scan to obtain per-segment tile // offsets. // ----------------------------------------------------------------------------- -template +template struct segment_size_to_tile_count_op { SegmentSizeParameterT segment_sizes; @@ -93,49 +112,824 @@ struct segment_size_to_tile_count_op template _CCCL_HOST_DEVICE _CCCL_FORCEINLINE constexpr TotalNumItemsValueType operator()(SegmentIndexT segment_id) const { - return static_cast( - ::cuda::ceil_div(params::get_param(segment_sizes, segment_id), large_segment_agent_tile_size)); + return static_cast(::cuda::ceil_div( + params::__get_and_clamp_param_to_nonnegative(segment_sizes, segment_id), large_segment_agent_tile_size)); } }; // ----------------------------------------------------------------------------- -// Segmented Top-K Dispatch +// Automatic backend selector // ----------------------------------------------------------------------------- +// Stateless selector built purely from the compile-time request facts. It owns the entire backend decision, including +// computing `baseline_can_cover` from the concrete agent types -- the reason it lives here (where +// `baseline_can_cover_v` and the baseline agent are visible) rather than in the tuning header. +template +struct policy_selector_from_types +{ + // TODO(bgruber): to let the baseline policy vary per CC, move this coverage check into operator() and evaluate it for + // the passed CC. Only the check is hard: it instantiates the agent for sizeof(TempStorage), so it needs the CC as a + // compile-time constant, whereas operator()'s `cc` is a runtime parameter (building the policy itself is just the + // value make_baseline_policy(cc)). Recover the compile-time CC by folding over + // ::cuda::__target_compute_capabilities() (as detail::dispatch_to_cc_list does) and evaluate baseline_can_cover_v for + // the matching CC. That also removes the invariant below, since coverage and the returned baseline would then derive + // from the same cc. -//! @param d_temp_storage Device-accessible allocation of temporary storage. When `nullptr`, the required allocation -//! size is written to `temp_storage_bytes` and no work is done. -//! @param temp_storage_bytes Reference to size in bytes of `d_temp_storage` allocation -//! @param d_key_segments_it d_key_segments_it[segment_index] -> iterator to the input sequence of key data for segment -//! `segment_index` -//! @param d_key_segments_out_it d_key_segments_out_it[segment_index] -> iterator to the output sequence of key data for -//! segment `segment_index` -//! @param d_value_segments_it d_value_segments_it[segment_index] -> iterator to the input sequence of associated value -//! items for segment `segment_index`. When cub::NullType**, only keys are provided. -//! @param d_value_segments_out_it d_value_segments_out_it[segment_index] -> iterator to the output sequence of -//! associated value items for segment `segment_index` -//! @param segment_sizes Parameter providing segment sizes for each segment -//! @param k Parameter providing K for each segment -//! @param select_directions Parameter providing the selection direction for each segment -//! @param num_segments Number of segments -//! @param total_num_items_guarantee Allows the user to provide a guarantee on the upper bound of the total number of -//! items -template topk_policy + { + return topk_policy{topk_algorithm::baseline, baseline_policy, {}}; + } + }; + + // Whether a one-worker-per-segment (default baseline) policy fits the static max segment size in shared memory; feeds + // the backend decision below. + static constexpr bool baseline_can_cover = baseline_can_cover_v< + policy_getter_17, + SegmentSizeParameterT, + KeyInputItItT, + KeyOutputItItT, + ValueInputItItT, + ValueOutputItItT, + SegmentSizeParameterT, + KParameterT, + SelectDirectionParameterT, + NumSegmentsParameterT, + LargeSegmentTileOffsetT>; + + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const -> topk_policy + { + constexpr bool deterministic = (Determinism != ::cuda::execution::determinism::__determinism_t::__not_guaranteed) + || (TieBreak != ::cuda::execution::tie_break::__tie_break_t::__unspecified); + + topk_algorithm backend = topk_algorithm::unsupported; + if (deterministic || !baseline_can_cover) + { + // A deterministic result set / concrete tie-break preference, or a segment too large for the single-block + // baseline, is served only by the cluster backend (SM 9.0+); otherwise the request cannot run here. + backend = cluster_capable(cc) ? topk_algorithm::cluster : topk_algorithm::unsupported; + } + else + { + // Baseline can cover: use the cluster backend only where it is measured to win. The size crossover is a fixed + // selector constant (not read from the tunable cluster policy), so tuning the cluster policy never shifts the + // backend choice. The threshold is applied on every cluster-capable architecture, not gated to a minimum CC. + const bool beneficial = StaticMaxSegSize >= cluster_beneficial_min_segment_size; + backend = (cluster_capable(cc) && beneficial) ? topk_algorithm::cluster : topk_algorithm::baseline; + } + return topk_policy{backend, baseline_policy, make_cluster_policy()}; + } +}; + +// ----------------------------------------------------------------------------- +// Dispatch (both backends behind one kernel symbol) +// ----------------------------------------------------------------------------- +// The dispatch is host-only: it launches the single kernel symbol (`device_batched_topk_kernel`, in +// kernel_batched_topk.cuh) via the CUDA runtime. The algorithm does not support device-side (CDP) launch. + +// Corrected form of `launcher_factory.max_dynamic_smem_size_for` (host path: `cub::MaxPotentialDynamicSmemBytes`), +// returning the usable dynamic budget as `opt-in - static footprint`. That facility currently subtracts the per-block +// reserved shared memory a second time even though `cudaDevAttrMaxSharedMemoryPerBlockOptin` already excludes it, +// under-reporting the budget by ~`reserved` (~1 KiB) -- enough to drop the cluster kernel's top table tier (see the +// TODO in MaxPotentialDynamicSmemBytes). TODO: once that facility is fixed, delete this and call +// `launcher_factory.max_dynamic_smem_size_for(...)` directly. +template +_CCCL_HOST_API cudaError_t max_dynamic_smem_size_for_fixed(int& max_dynamic_smem_bytes, KernelPtr kernel_ptr) +{ + max_dynamic_smem_bytes = -1; + int device_id = 0; + if (const auto error = CubDebug(cudaGetDevice(&device_id))) + { + return error; + } + int max_smem_optin_bytes = 0; + if (const auto error = + CubDebug(cudaDeviceGetAttribute(&max_smem_optin_bytes, cudaDevAttrMaxSharedMemoryPerBlockOptin, device_id))) + { + return error; + } + cudaFuncAttributes kernel_attrs{}; + if (const auto error = CubDebug(cudaFuncGetAttributes(&kernel_attrs, kernel_ptr))) + { + return error; + } + const int static_smem_bytes = static_cast(kernel_attrs.sharedSizeBytes); + max_dynamic_smem_bytes = (max_smem_optin_bytes > static_smem_bytes) ? max_smem_optin_bytes - static_smem_bytes : 0; + return cudaSuccess; +} + +// Largest number of CTA blocks per cluster the kernel/architecture admits at `dynamic_smem_bytes` of dynamic SMEM. The +// config's cluster dimension is ignored by the query (placeholder here); the non-portable opt-in must already be set +// for it to report sizes beyond the portable ceiling. +template +_CCCL_HOST_API ::cuda::std::expected +probe_max_cluster_blocks(KernelPtr kernel_ptr, cudaStream_t stream, int threads_per_block, int dynamic_smem_bytes) +{ + ::cudaLaunchAttribute cluster_attr{}; + cluster_attr.id = ::cudaLaunchAttributeClusterDimension; + cluster_attr.val.clusterDim = {1, 1, 1}; + + ::cudaLaunchConfig_t cfg{}; + cfg.gridDim = dim3(1); + cfg.blockDim = dim3(static_cast(threads_per_block)); + cfg.dynamicSmemBytes = static_cast<::cuda::std::size_t>(dynamic_smem_bytes); + cfg.stream = stream; + cfg.attrs = &cluster_attr; + cfg.numAttrs = 1; + + int cluster_blocks = 0; + if (const auto error = CubDebug( + ::cudaOccupancyMaxPotentialClusterSize(&cluster_blocks, reinterpret_cast(kernel_ptr), &cfg))) + { + return ::cuda::std::unexpected(error); + } + return cluster_blocks; +} + +// Device-wide count of `cluster_blocks`-CTA clusters that can be co-resident at `dynamic_smem_bytes` of dynamic SMEM +// (clusters per wave). `cudaOccupancyMaxActiveClusters` rejects a grid that is not a multiple of the cluster, so the +// grid is set to exactly one cluster; the returned capacity is independent of the actual grid size. +template +_CCCL_HOST_API ::cuda::std::expected probe_clusters_per_wave( + KernelPtr kernel_ptr, cudaStream_t stream, int threads_per_block, int cluster_blocks, int dynamic_smem_bytes) +{ + ::cudaLaunchAttribute cluster_attr{}; + cluster_attr.id = ::cudaLaunchAttributeClusterDimension; + cluster_attr.val.clusterDim = {static_cast(cluster_blocks), 1, 1}; + + ::cudaLaunchConfig_t cfg{}; + cfg.gridDim = dim3(static_cast(cluster_blocks)); + cfg.blockDim = dim3(static_cast(threads_per_block)); + cfg.dynamicSmemBytes = static_cast<::cuda::std::size_t>(dynamic_smem_bytes); + cfg.stream = stream; + cfg.attrs = &cluster_attr; + cfg.numAttrs = 1; + + int clusters_per_wave = 0; + if (const auto error = + CubDebug(::cudaOccupancyMaxActiveClusters(&clusters_per_wave, reinterpret_cast(kernel_ptr), &cfg))) + { + return ::cuda::std::unexpected(error); + } + return clusters_per_wave; +} + +// The cluster backend's launch shape: CTA blocks per cluster and the dynamic-SMEM bytes to launch with. +struct cluster_launch_shape +{ + int cluster_blocks = 0; + int dynamic_smem_bytes = 0; +}; + +// Chooses the cluster launch shape for the statically-bounded max segment size. Probes occupancy through the CUDA +// runtime; the caller has already set the kernel's dynamic-SMEM opt-in to the maximum, so every probed config and the +// final launch run under one consistent opt-in. +template +_CCCL_HOST_API ::cuda::std::expected select_cluster_launch_shape( + ::cuda::std::uint64_t max_segment_size, + ::cuda::std::uint64_t num_segments, + int max_dynamic_smem_bytes, + cluster_topk_policy policy, + cudaStream_t stream, + KernelPtr kernel_ptr) +{ + using layout_t = LayoutT; + + const int threads_per_block = policy.threads_per_block; + + // Computed before any occupancy query so the single-CTA fast path below can skip one -- that driver query + // otherwise dominates the runtime of tiny launches. + const int max_block_resident_items = static_cast(layout_t::max_block_resident_items(max_dynamic_smem_bytes)); + if (max_block_resident_items <= 0) + { + // Not even one load-aligned chunk fits in the opt-in budget; the kernel cannot run. + return ::cuda::std::unexpected(cudaErrorInvalidValue); + } + + // Smallest cluster block count for full residency: at the largest SMEM each CTA holds `max_block_resident_items` + // items. 64-bit to match the launch-shape arithmetic below; the value is small (`max_segment_size <= 2^21`). + const auto min_blocks_per_segment = + ::cuda::ceil_div(max_segment_size, static_cast<::cuda::std::uint64_t>(max_block_resident_items)); + + int cluster_blocks = 0; + int dynamic_smem_bytes = 0; + + if (batched_topk_cluster::is_single_cta_eligible( + static_cast<::cuda::std::uint32_t>(max_segment_size), + static_cast<::cuda::std::uint32_t>(max_block_resident_items), + policy.single_block_max_seg_size)) + { + // Single-CTA fast path: the segment fits resident in one CTA and is small enough that the agent's + // cluster-barrier-free path beats spreading it across more CTAs. One CTA at in-budget SMEM is always launchable, + // so the occupancy probe is skipped. Larger fully-resident segments fall through to the wave-aware search below. + cluster_blocks = 1; + dynamic_smem_bytes = layout_t::min_smem_bytes_from_num_items(max_segment_size); + } + else + { + // Hardware cluster ceiling (max blocks per cluster), queried at runtime (not hardcoded) so a future device with + // larger non-portable clusters is not capped. Probed at zero dynamic SMEM for the arch/kernel ceiling alone; each + // candidate is re-validated against its own SMEM below. + const auto hw_cluster_ceiling = + probe_max_cluster_blocks(kernel_ptr, stream, threads_per_block, /*dynamic_smem_bytes=*/0); + if (!hw_cluster_ceiling) + { + return ::cuda::std::unexpected(hw_cluster_ceiling.error()); + } + if (*hw_cluster_ceiling <= 0) + { + return ::cuda::std::unexpected(cudaErrorInvalidValue); + } + // `max_blocks_per_cluster == 0` -> the full hardware ceiling; a non-zero knob narrows it, clamped to that ceiling. + // A cap narrower than a segment needs pushes it into the oversize/streaming fallback below. + const int eff_max_blocks_per_cluster = + (policy.max_blocks_per_cluster == 0) + ? *hw_cluster_ceiling + : (::cuda::std::min) (policy.max_blocks_per_cluster, *hw_cluster_ceiling); + + // Wave-aware selection: the free variable is the cluster block count (one cluster per segment), paired with the + // smallest SMEM that keeps the segment fully resident (fewer blocks = more SMEM/fewer clusters-per-wave, more = + // the reverse). Pick the count minimizing waves, ties toward the largest (smallest SMEM, most L1 -- the profiled + // fast configs). Enumerated analytically, so a register-limited occupancy cannot collapse the candidate set. + if (min_blocks_per_segment <= static_cast<::cuda::std::uint64_t>(eff_max_blocks_per_cluster)) + { + // Full residency achievable: `max_segment_size <= min_blocks_per_segment * max_block_resident_items` and + // `min_blocks_per_segment <= eff_max_blocks_per_cluster`, so every per-CTA capacity below fits `int`. + + // Cluster blocks the max segment actually needs (shared with the device so the launch is never wider than + // necessary). At `min_chunks_per_block == 1` this equals the segment's chunk count; a larger knob shrinks it. + const auto desired_cluster_blocks = ::cuda::narrow(batched_topk_cluster::compute_num_logical_cluster_blocks( + static_cast<::cuda::std::uint32_t>(layout_t::num_chunks_from_num_items(max_segment_size)), + policy.min_chunks_per_block, + ::cuda::narrow<::cuda::std::uint32_t>(eff_max_blocks_per_cluster))); + + // Scan `[min_candidate_blocks, max_candidate_blocks]` for the min-waves block count, tie-breaking largest. + // `max_candidate_blocks == max(desired_cluster_blocks, min(min_candidate_blocks, eff_max_blocks_per_cluster))`: + // the segment-needed count `desired_cluster_blocks` (<= `eff_max_blocks_per_cluster`, capped in + // `compute_num_logical_cluster_blocks`), floored at `min_candidate_blocks`. The `clamp` operands are ordered so + // `lo <= hi` holds even when `eff_max_blocks_per_cluster == 1` forces `min_candidate_blocks (== 2) > eff_max`; + // there `max_candidate_blocks == 1` empties the scan and the single-CTA fallback below runs (that edge is a + // one-CTA-resident segment with the single-CTA path disabled, so `min_blocks_per_segment == 1`). + const auto min_candidate_blocks = (::cuda::std::max) (2, static_cast(min_blocks_per_segment)); + const auto max_candidate_blocks = + ::cuda::std::clamp(min_candidate_blocks, desired_cluster_blocks, eff_max_blocks_per_cluster); + auto best_waves = (::cuda::std::numeric_limits<::cuda::std::uint64_t>::max)(); + for (int candidate_blocks = min_candidate_blocks; candidate_blocks <= max_candidate_blocks; ++candidate_blocks) + { + const auto num_block_items = ::cuda::ceil_div(max_segment_size, candidate_blocks); + const int resident_smem_bytes = layout_t::min_smem_bytes_from_num_items(num_block_items); + if (resident_smem_bytes > max_dynamic_smem_bytes) + { + // Unreachable for candidate_blocks >= min_blocks_per_segment, but guards the SMEM budget regardless. + continue; + } + + const auto clusters_per_wave = + probe_clusters_per_wave(kernel_ptr, stream, threads_per_block, candidate_blocks, resident_smem_bytes); + if (!clusters_per_wave) + { + return ::cuda::std::unexpected(clusters_per_wave.error()); + } + if (*clusters_per_wave <= 0) + { + continue; // cluster blocks not launchable at this SMEM. + } + + const auto waves = ::cuda::ceil_div(num_segments, *clusters_per_wave); + // Min waves, tie-break largest count: the loop ascends, so `<=` keeps the largest at equal waves (`best_waves` + // starts at `UINT64_MAX`, so the first launchable count always wins). + if (waves <= best_waves) + { + best_waves = waves; + cluster_blocks = candidate_blocks; + dynamic_smem_bytes = resident_smem_bytes; + } + } + + if (cluster_blocks == 0 && min_blocks_per_segment == 1) + { + // No multi-CTA config was launchable; fall back to single-CTA full residency. Slower for large segments, but + // `min_blocks_per_segment == 1` guarantees the resident SMEM fits the budget and one CTA is always launchable. + cluster_blocks = 1; + dynamic_smem_bytes = layout_t::min_smem_bytes_from_num_items(max_segment_size); + } + } + + if (cluster_blocks == 0) + { + // Oversize (`min_blocks_per_segment > eff_max_blocks_per_cluster`) or nothing launchable: full residency + // is impossible, so maximize residency with the largest launchable cluster at the largest SMEM and stream the + // overflow. + const auto hw_max_cluster_blocks = + probe_max_cluster_blocks(kernel_ptr, stream, threads_per_block, max_dynamic_smem_bytes); + if (!hw_max_cluster_blocks) + { + return ::cuda::std::unexpected(hw_max_cluster_blocks.error()); + } + cluster_blocks = (::cuda::std::min) (*hw_max_cluster_blocks, eff_max_blocks_per_cluster); + if (cluster_blocks <= 0) + { + return ::cuda::std::unexpected(cudaErrorInvalidValue); + } + dynamic_smem_bytes = max_dynamic_smem_bytes; + } + } + + return cluster_launch_shape{cluster_blocks, dynamic_smem_bytes}; +} + +// Cluster arm of the dispatch (host-only): after the shared query-pass / CC-guard setup, launches the single kernel +// symbol via `cudaLaunchKernelEx` using the resolved-CC cluster policy and geometry from `policy_getter`. +// `select_directions` arrives already wrapped; the cluster tuning comes from `policy_getter` (the resolved-CC policy) +// and the requested `Determinism`/`TieBreak` from the dispatch. The kernel launch goes through `launcher_factory`; the +// cluster occupancy / shared-memory setup queries still use the CUDA runtime directly. +template >, - it_value_t>, - ::cuda::std::int64_t, - ::cuda::args::__traits::highest>> -#if _CCCL_HAS_CONCEPTS() - requires batched_topk_policy_selector -#endif // _CCCL_HAS_CONCEPTS() -CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( + typename KernelLauncherFactory> +_CCCL_HOST_API cudaError_t launch_cluster_arm( + PolicyGetter policy_getter, + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputItItT d_key_segments_it, + KeyOutputItItT d_key_segments_out_it, + ValueInputItItT d_value_segments_it, + ValueOutputItItT d_value_segments_out_it, + SegmentSizeParameterT segment_sizes, + KParameterT k_param, + SelectDirectionParameterT select_directions, + NumSegmentsParameterT num_segments, + cudaStream_t stream, + KernelLauncherFactory launcher_factory) +{ + // A tie-break preference is only meaningful once the result set itself is deterministic. + static_assert(Determinism != ::cuda::execution::determinism::__determinism_t::__not_guaranteed + || TieBreak == ::cuda::execution::tie_break::__tie_break_t::__unspecified, + "A tie-break preference requires a deterministic execution requirement"); + + // The cluster arm needs no temporary storage; report a positive size so the two-phase protocol proceeds. + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + + // A `tune`d override (`UserProvidedTuning`) can force the cluster backend on a device that cannot run it: return + // cudaErrorNotSupported rather than launch a cluster kernel the device lacks (the deferred-mode runtime behavior + // tests and benchmarks rely on). The automatic selector never routes here below SM 9.0, so its instantiation drops + // this check. `PtxComputeCap` is the running code's capability (never above the hardware SM), so it also rejects an + // SM 9.0+ build on older hardware. + if constexpr (UserProvidedTuning) + { + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) + { + return error; + } + if (cc < ::cuda::compute_capability{9, 0}) + { + return cudaErrorNotSupported; + } + } + + // Single kernel symbol; its cluster vs baseline arm is selected device-side via `current_policy()`. + // Taking its address here ODR-uses the `__global__` template, which is what drives its emission and registration. + // Not `constexpr`: MSVC (C2326) rejects a `constexpr` local captured and ODR-used inside the lambdas below. + auto kernel_ptr = &device_batched_topk_kernel< + PolicySelector, + KeyInputItItT, + KeyOutputItItT, + ValueInputItItT, + ValueOutputItItT, + SegmentSizeParameterT, + KParameterT, + SelectDirectionParameterT, + NumSegmentsParameterT, + LargeSegmentTileOffsetT, + Determinism, + TieBreak>; + + // Cluster sub-policy for the *resolved* architecture -- exactly what the device kernel instantiates via + // `current_policy()`, so the host launch config (block size, shared-memory math) stays in lock-step + // with the device policy per CC. `policy_getter()` is a constant expression, so `policy` is a non-type template arg. + constexpr cluster_topk_policy policy = policy_getter().cluster; + constexpr int threads_per_block = policy.threads_per_block; + constexpr int chunk_bytes = policy.chunk_bytes; + constexpr int load_align_bytes = policy.load_align_bytes; + constexpr int max_chunk_slots_per_block = policy.max_chunk_slots_per_block; + static_assert(policy.max_blocks_per_cluster >= 0, + "max_blocks_per_cluster must be 0 (unrestricted) or a positive cluster block count"); + static_assert(max_chunk_slots_per_block >= 0, + "max_chunk_slots_per_block must be 0 (unrestricted) or a positive count"); + + using key_it_t = it_value_t; + using key_t = it_value_t; + using layout_t = batched_topk_cluster::smem_block_tile_layout; + static_assert(is_valid_cluster_policy(policy)); + static_assert(load_align_bytes % int{sizeof(key_t)} == 0); + + // Tightest upper bound the segment-size argument carries -- for a static-bounded per-segment sequence a loose type + // max, not the actual runtime maximum across segments. + const auto max_seg_size = ::cuda::args::__highest_(segment_sizes); + using num_segments_val_t = typename ::cuda::args::__traits::element_type; + // `num_segments > 0` and `max_seg_size > 0` here: the generic `dispatch` returns for the empty-batch cases (no + // segments, or a non-positive max segment size) before invoking this launch arm. + const auto num_seg_val = detail::params::get_param(num_segments, num_segments_val_t{0}); + + // Opt in to non-portable cluster blocks (>8 on Hopper). + if (const auto error = CubDebug(::cudaFuncSetAttribute( + reinterpret_cast(kernel_ptr), cudaFuncAttributeNonPortableClusterSizeAllowed, 1))) + { + return error; + } + + // Usable dynamic shared-memory budget (opt-in minus the kernel's static footprint); the policy slot cap may narrow + // it further into `max_dynamic_smem_bytes` below. + int hw_dynamic_smem_bytes = 0; + if (const auto error = max_dynamic_smem_size_for_fixed(hw_dynamic_smem_bytes, kernel_ptr)) + { + return error; + } + // Optional policy cap on resident chunk slots per block (`max_chunk_slots_per_block == 0` -> unrestricted, i.e. + // the full hardware budget). Expressed as the SMEM those slots need, then clamped to the hardware budget: a cap the + // hardware cannot satisfy is a no-op (hardware wins). Fewer slots lowers every CTA's resident dynamic shared-memory + // request, so a smaller segment overflows into streaming -- useful to leave shared memory free for a concurrent + // kernel (or to reach the streaming / schedule paths at a small footprint in tests). A cap below one slot trips the + // `max_block_resident_items <= 0` guard below. + const int max_dynamic_smem_bytes = + (max_chunk_slots_per_block == 0) + ? hw_dynamic_smem_bytes + : (::cuda::std::min) (hw_dynamic_smem_bytes, layout_t::min_smem_bytes_from_num_chunks(max_chunk_slots_per_block)); + + // Set the kernel's dynamic-SMEM opt-in once, to the per-symbol maximum, before any occupancy probe or launch. + // `max_dynamic_smem_bytes` is fixed by the compile-time policy and the device, so every thread sharing this kernel + // symbol writes the identical value: the process-global attribute cannot be raced to a lower value that would fail a + // concurrent launch. It also covers every launch shape the search below can pick (all `<= max_dynamic_smem_bytes`) + // and keeps the occupancy probes and the final launch on one consistent opt-in. + if (const auto error = launcher_factory.set_max_dynamic_smem_size_for(kernel_ptr, max_dynamic_smem_bytes)) + { + return error; + } + + // Resolve the launch shape (cluster blocks + dynamic SMEM) for the max segment size. + const auto shape = select_cluster_launch_shape( + static_cast<::cuda::std::uint64_t>(max_seg_size), + static_cast<::cuda::std::uint64_t>(num_seg_val), + max_dynamic_smem_bytes, + policy, + stream, + kernel_ptr); + if (!shape) + { + return shape.error(); + } + + const int cluster_blocks = shape->cluster_blocks; + const int dynamic_smem_bytes = shape->dynamic_smem_bytes; + const auto max_block_resident_items = layout_t::max_block_resident_items(dynamic_smem_bytes); + + // One cluster per segment, its CTAs stacked in the grid's y-dimension so the x-extent stays `num_segments`: a + // flattened x == num_segments * cluster_blocks would overrun the 2^31-1 grid-x limit for a multi-CTA cluster well + // before `num_segments` reached its INT_MAX maximum (already <= INT_MAX by the entry check in `dispatch`; + // `cluster_blocks` is far below the y-dimension limit). The device reads the segment id from clusterid.x (segments + // stay the x-extent) and the CTA rank from cluster_ctarank (linearized across the cluster's dims), so this needs no + // agent change. + const dim3 grid_dim{static_cast(num_seg_val), static_cast(cluster_blocks), 1u}; + const dim3 cluster_dim{1u, static_cast(cluster_blocks), 1u}; + + // The cluster dimension routes the host launch through `cudaLaunchKernelEx`. + if (const auto error = CubDebug( + launcher_factory(grid_dim, + dim3{static_cast(threads_per_block)}, + static_cast<::cuda::std::size_t>(dynamic_smem_bytes), + stream, + /*dependent_launch=*/false, + cluster_dim) + .doit(kernel_ptr, + d_key_segments_it, + d_key_segments_out_it, + d_value_segments_it, + d_value_segments_out_it, + segment_sizes, + k_param, + select_directions, + num_segments, + baseline_kernel_args{}, + cluster_kernel_args{static_cast<::cuda::std::uint32_t>(max_block_resident_items)}))) + { + return error; + } + + return CubDebug(detail::DebugSyncStream(stream)); +} + +// Baseline host-launch arm of the dispatch. Launches the single kernel symbol +// (`device_batched_topk_kernel`, packing the large-segment bookkeeping into `baseline_kernel_args` and passing an empty +// `cluster_kernel_args`). `select_directions` arrives already wrapped and the baseline tuning is taken from the +// `PolicySelector`. All kernel launches, memsets and nested scans go through +// `launcher_factory`. +template +_CCCL_HOST_API cudaError_t launch_baseline_arm( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputItItT d_key_segments_it, + KeyOutputItItT d_key_segments_out_it, + ValueInputItItT d_value_segments_it, + ValueOutputItItT d_value_segments_out_it, + SegmentSizeParameterT segment_sizes, + KParameterT k, + SelectDirectionParameterT select_directions, + NumSegmentsParameterT num_segments, + cudaStream_t stream, + KernelLauncherFactory launcher_factory) +{ + // Whether some one-worker-per-segment policy covers the static max segment size within the shared-memory limit. + // Computed from this call's concrete agent types (a tuning override exposes no `baseline_can_cover` member). The + // automatic selector never routes here when this is false; only a trusted `tune`d override forces the baseline + // backend on an oversize segment. Strict mode rejects that at compile time (the static_assert below, mirroring the + // arch-unsupported one in `dispatch`); deferred mode keeps the two-phase runtime cudaErrorNotSupported path. + constexpr bool baseline_can_cover = baseline_can_cover_v< + PolicyGetter, + SegmentSizeParameterT, + KeyInputItItT, + KeyOutputItItT, + ValueInputItItT, + ValueOutputItItT, + SegmentSizeParameterT, + KParameterT, + SelectDirectionParameterT, + NumSegmentsParameterT, + LargeSegmentTileOffsetT>; +#if !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC) \ + && !defined(CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT) + static_assert( + baseline_can_cover, + "cub::DeviceBatchedTopK: the forced baseline backend cannot cover the static maximum segment size within the " + "shared-memory limit. Force the cluster backend, lower the segment-size bound, or define " + "CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT to defer the diagnosis to runtime (cudaErrorNotSupported)."); +#endif // !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC) + // && !defined(CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT) + if constexpr (!baseline_can_cover) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + return cudaErrorNotSupported; + } + else + { + using large_segment_tile_offset_t = LargeSegmentTileOffsetT; + + // Determine which one-worker-per-segment policy covers the segment-size range and k. Resolve from the handed + // resolved-CC `PolicyGetter` (not `PolicySelector`) so the host picks the same baseline policy the device kernel + // instantiates for the resolved CC -- they would otherwise diverge once baseline tuning becomes CC-dependent. + constexpr auto policy = find_smallest_covering_policy_for_getter< + PolicyGetter, + SegmentSizeParameterT, + KeyInputItItT, + KeyOutputItItT, + ValueInputItItT, + ValueOutputItItT, + SegmentSizeParameterT, + KParameterT, + SelectDirectionParameterT, + NumSegmentsParameterT, + large_segment_tile_offset_t>::policy; + constexpr worker_policy worker_per_segment_policy = policy.worker_per_segment_policy; + constexpr multi_worker_policy multi_worker_per_segment_policy = policy.multi_worker_per_segment_policy; + + static constexpr int worker_per_segment_tile_size = + worker_per_segment_policy.threads_per_block * worker_per_segment_policy.items_per_thread; + static constexpr bool any_small_segments = + ::cuda::args::__traits::lowest <= worker_per_segment_tile_size; + static constexpr bool only_small_segments = + ::cuda::args::__traits::highest <= worker_per_segment_tile_size; + + // Allocation layout: + // only_small_segments: [0] dummy. + // any_small_segments && !only_small_segments (mixed): [0] tile offsets, [1] counters struct, + // [2] large-segment ids. + // !any_small_segments (large-only): [0] tile offsets, [1] segment-size transform-scan temp storage. + static constexpr int allocations_array_size = only_small_segments ? 1 : (any_small_segments ? 3 : 2); + size_t allocation_sizes[allocations_array_size] = {1}; + + using num_segments_val_t = typename ::cuda::args::__traits::element_type; + using counters_t = batched_topk_counters; + using segment_size_scan_offset_t = detail::choose_offset_t; + using segment_size_scan_input_op_t = + segment_size_to_tile_count_op; + static constexpr auto multi_worker_per_segment_tile_size = + multi_worker_per_segment_policy.threads_per_block * multi_worker_per_segment_policy.items_per_thread; + const segment_size_scan_input_op_t segment_size_scan_input_op{segment_sizes, multi_worker_per_segment_tile_size}; + // Transform iterator over [0, num_segments) producing the tile-count for each segment. + [[maybe_unused]] const auto segment_size_scan_input_it = ::cuda::transform_iterator( + ::cuda::counting_iterator{num_segments_val_t{0}}, segment_size_scan_input_op); + + if constexpr (!only_small_segments) + { + const auto num_segments_val = params::get_param(num_segments, 0); + // TODO(topk): the baseline large-segment (multi-CTA) path is WIP. Completing it requires: (1) guarding the + // `num_segments_val * sizeof(...)` byte counts below against size_t overflow (safe today only because the entry + // bounds num_segments_val to <= INT_MAX); (2) making the baseline tunable by populating its `epilogue` and + // `multi_worker_per_segment_policy` sub-policies and adding matching knobs to the segmented_topk benchmarks, + // which leave them zero-initialized today so baseline sweeps are not yet meaningful. + allocation_sizes[0] = num_segments_val * sizeof(large_segment_tile_offset_t); + if constexpr (any_small_segments) + { + allocation_sizes[1] = sizeof(counters_t); + allocation_sizes[2] = num_segments_val * sizeof(num_segments_val_t); + } + else + { + // Query the temporary storage requirement of the segment-size transform-scan. + if (const auto error = CubDebug(detail::scan::dispatch( + nullptr, + allocation_sizes[1], + segment_size_scan_input_it, + static_cast(nullptr), + ::cuda::std::plus<>{}, + detail::InputValue(large_segment_tile_offset_t{0}), + static_cast(num_segments_val), + stream, + {}, + {}, + launcher_factory))) + { + return error; + } + } + } + + void* allocations[allocations_array_size] = {}; + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + if (d_temp_storage == nullptr) + { + return cudaSuccess; + } + + // `num_segments > 0` and the max segment size > 0 here: the generic `dispatch` returns for the empty-batch cases + // (no segments, or a non-positive max segment size) before invoking this launch arm. + + if constexpr (any_small_segments) + { + if constexpr (!only_small_segments) + { + // Zero-initialize the counters struct read by the agent's atomics. + if (const auto error = CubDebug(launcher_factory.MemsetAsync(allocations[1], 0, sizeof(counters_t), stream))) + { + return error; + } + } + const int grid_dim = static_cast(params::get_param(num_segments, 0)); + constexpr int block_dim = worker_per_segment_policy.threads_per_block; + if (const auto error = CubDebug( + launcher_factory(grid_dim, block_dim, 0, stream, /*dependent_launch=*/false) + .doit( + device_batched_topk_kernel< + PolicySelector, + KeyInputItItT, + KeyOutputItItT, + ValueInputItItT, + ValueOutputItItT, + SegmentSizeParameterT, + KParameterT, + SelectDirectionParameterT, + NumSegmentsParameterT, + large_segment_tile_offset_t, + Determinism, + TieBreak>, + d_key_segments_it, + d_key_segments_out_it, + d_value_segments_it, + d_value_segments_out_it, + segment_sizes, + k, + select_directions, + num_segments, + baseline_kernel_args{ + only_small_segments ? nullptr : static_cast(allocations[1]), + only_small_segments ? nullptr : static_cast(allocations[2]), + only_small_segments ? nullptr : static_cast(allocations[0])}, + cluster_kernel_args{}))) + { + return error; + } + } + else + { + // No small segments: compute the per-segment tile offsets directly via a transform-scan over all segment sizes. + if (const auto error = CubDebug(detail::scan::dispatch( + allocations[1], + allocation_sizes[1], + segment_size_scan_input_it, + static_cast(allocations[0]), + ::cuda::std::plus<>{}, + detail::InputValue(large_segment_tile_offset_t{0}), + static_cast(params::get_param(num_segments, 0)), + stream, + {}, + {}, + launcher_factory))) + { + return error; + } + } + + return CubDebug(detail::DebugSyncStream(stream)); + } +} + +#if !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC) +// Returns true if at least one architecture this translation unit targets (the compile target list exposed as +// `::cuda::__target_compute_capabilities()`) resolves to the `unsupported` backend for `PolicySelector` -- e.g. a +// deterministic request while a pre-SM90 target is present in the list. Used to turn a would-be runtime +// `cudaErrorNotSupported` into a compile-time diagnostic (see the static_assert in `dispatch`). +template +[[nodiscard]] _CCCL_HOST_DEVICE_API _CCCL_CONSTEVAL bool any_target_cc_unsupported() +{ + bool any = false; + for (const auto cc : ::cuda::__target_compute_capabilities()) + { + any = any || (PolicySelector{}(cc).backend == topk_algorithm::unsupported); + } + return any; +} +#endif // !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC) + +// Internal entry point: the single dispatch that replaces the standalone baseline / cluster dispatches. It resolves the +// runtime compute capability, then uses `dispatch_compute_cap` to pick, per architecture, the backend chosen by the +// resolved policy selector (deterministic -> cluster; otherwise the arch+size crossover). Both host arms launch the +// same kernel symbol. `Determinism`/`TieBreak` are compile-time selection inputs. +// +// `tuning_env` carries an optional `tune`d policy selector (keyed on `topk_policy`): when present it fully replaces the +// automatic selector -- its `.backend` chooses the arm and its `.baseline`/`.cluster` carry the tunings. Matching +// DeviceScan/DeviceTransform, the tuned backend choice is trusted; only the determinism/tie-break guard below still +// applies. `launcher_factory` routes the kernel launches, memsets, nested scans and the routing CC query (the cluster +// arm's occupancy / shared-memory queries still call the CUDA runtime directly). +template < + ::cuda::execution::determinism::__determinism_t Determinism = + ::cuda::execution::determinism::__determinism_t::__not_guaranteed, + ::cuda::execution::tie_break::__tie_break_t TieBreak = ::cuda::execution::tie_break::__tie_break_t::__unspecified, + typename KeyInputItItT, + typename KeyOutputItItT, + typename ValueInputItItT, + typename ValueOutputItItT, + typename SegmentSizeParameterT, + typename KParameterT, + typename SelectDirectionT, + typename NumSegmentsParameterT, + typename TotalNumItemsGuaranteeT, + typename TuningEnvT = ::cuda::std::execution::env<>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +_CCCL_HOST_API cudaError_t dispatch( void* d_temp_storage, size_t& temp_storage_bytes, KeyInputItItT d_key_segments_it, @@ -147,209 +941,222 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( SelectDirectionT select_direction, NumSegmentsParameterT num_segments, [[maybe_unused]] TotalNumItemsGuaranteeT total_num_items_guarantee, - cudaStream_t stream = nullptr, - [[maybe_unused]] PolicySelector policy_selector = {}) + cudaStream_t stream, + const TuningEnvT& = {}, + KernelLauncherFactory launcher_factory = {}) { - using large_segment_tile_offset_t = typename ::cuda::args::__traits::element_type; + // Both arms resolve `num_segments` on the host (allocation sizing, grid extent, empty-batch guard), so it must be a + // host-known single value; device-resident counts are future work. Defensive: the public entry checks this too, but + // `dispatch` is also called directly (tests / benchmarks). + static_assert(::cuda::args::__traits::is_single_value + && !::cuda::args::__traits::is_deferred, + "cub::DeviceBatchedTopK requires a host-known uniform number of segments (constant, immediate, or a " + "plain integral value)."); - // Wrap the raw enum into the internal discrete param type - auto select_directions = wrap_select_direction(select_direction); - using SelectDirectionParameterT = decltype(select_directions); + // The selection direction is a compile-time constant carried as `::cuda::args::constant`. Wrap it into the + // internal discrete param the kernel/agent expect (both host arms take the wrapped form). + // Type derived from the parameter type rather than `decltype(select_directions)`: GCC 7 rejects the latter ("use of + // 'select_directions' before deduction of 'auto'") when it feeds the `constexpr baseline_can_cover` initializer + // below. Declaring `select_directions` with the alias keeps its (const-qualified) type single-sourced. + using SelectDirectionParameterT = const decltype(wrap_select_direction(::cuda::std::declval())); + SelectDirectionParameterT select_directions = wrap_select_direction(select_direction); - // Helper that determines (a) whether there's any one-worker-per-segment policy supporting the range of segment - // sizes and k, and (b) if so, which set of one-worker-per-segment policies to use - constexpr auto policy = find_smallest_covering_policy< - PolicySelector, + using key_t = it_value_t>; + using value_t = it_value_t>; + using LargeSegmentTileOffsetT = typename ::cuda::args::__traits::element_type; + + constexpr ::cuda::std::int64_t max_k = ::cuda::args::__traits::highest; + constexpr ::cuda::std::int64_t static_max_seg = ::cuda::args::__traits::highest; + + // Default automatic selector from the compile-time inputs; it computes its own baseline coverage. A `tune`d selector + // in the environment (keyed on `topk_policy`) replaces it wholesale. + using default_policy_selector_t = policy_selector_from_types< + key_t, + value_t, + max_k, + static_max_seg, + Determinism, + TieBreak, SegmentSizeParameterT, KeyInputItItT, KeyOutputItItT, ValueInputItItT, ValueOutputItItT, - SegmentSizeParameterT, KParameterT, SelectDirectionParameterT, NumSegmentsParameterT, - large_segment_tile_offset_t>::policy; - constexpr worker_policy worker_per_segment_policy = policy.worker_per_segment_policy; - constexpr multi_worker_policy multi_worker_per_segment_policy = policy.multi_worker_per_segment_policy; + LargeSegmentTileOffsetT>; - static constexpr int worker_per_segment_tile_size = - worker_per_segment_policy.threads_per_block * worker_per_segment_policy.items_per_thread; - static constexpr bool any_small_segments = - ::cuda::args::__traits::lowest <= worker_per_segment_tile_size; - static constexpr bool only_small_segments = - ::cuda::args::__traits::highest <= worker_per_segment_tile_size; + // Type derived from the query-result trait rather than `decltype(policy_selector)`: GCC 7 rejects the latter ("use of + // 'policy_selector' before deduction of 'auto'") when `policy_selector_t` is later named inside the dispatch lambda. + using policy_selector_t = + ::cuda::std::execution::__query_result_or_t; +#if _CCCL_HAS_CONCEPTS() + static_assert(topk_policy_selector, + "Invalid policy selector for cub::DeviceBatchedTopK::dispatch"); +#endif // _CCCL_HAS_CONCEPTS() - // Allocation layout: - // only_small_segments: [0] dummy. - // any_small_segments && !only_small_segments (mixed): [0] tile offsets, [1] counters struct, - // [2] large-segment ids. - // !any_small_segments (large-only): [0] tile offsets, [1] segment-size transform-scan temp storage. - static constexpr int allocations_array_size = only_small_segments ? 1 : (any_small_segments ? 3 : 2); - size_t allocation_sizes[allocations_array_size] = {1}; +#if !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC) \ + && !defined(CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT) + // Strict mode (default): fail at compile time if the request cannot be served on *any* architecture this translation + // unit targets. Two causes reach here: a deterministic / large-segment request while a pre-SM90 target is present + // (the cluster backend requires SM90+), or _CCCL_DISABLE_DYNAMIC_CLUSTER_LAUNCH disabling the cluster backend on all + // architectures. This is the least-surprising UX for callers whose build targets multiple architectures. Define + // `CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT` to defer the diagnosis to runtime instead (the dispatch then returns + // `cudaErrorNotSupported` on unsupported devices); CUB's own tests and benchmarks do this so they can compile the + // full configuration space across all target architectures and skip at runtime where unsupported. + static_assert( + !any_target_cc_unsupported(), + "cub::DeviceBatchedTopK: the requested top-k configuration cannot be served on at least one architecture this " + "translation unit targets. The deterministic / large-segment path requires the cluster backend (SM90+), which is " + "unavailable either because a pre-SM90 architecture is targeted or because _CCCL_DISABLE_DYNAMIC_CLUSTER_LAUNCH is " + "defined (which disables the cluster backend on all architectures). To fix: target only SM90+ and leave " + "_CCCL_DISABLE_DYNAMIC_CLUSTER_LAUNCH undefined, relax the request (non-deterministic and small enough for the " + "baseline backend), or define CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT to defer the diagnosis to runtime " + "(cudaErrorNotSupported)."); +#endif // !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC) + // && !defined(CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT) - using num_segments_val_t = typename ::cuda::args::__traits::element_type; - using counters_t = batched_topk_counters; - using segment_size_scan_offset_t = detail::choose_offset_t; - using segment_size_scan_input_op_t = - segment_size_to_tile_count_op; - static constexpr auto multi_worker_per_segment_tile_size = - multi_worker_per_segment_policy.threads_per_block * multi_worker_per_segment_policy.items_per_thread; - const segment_size_scan_input_op_t segment_size_scan_input_op{segment_sizes, multi_worker_per_segment_tile_size}; - // Transform iterator over [0, num_segments) producing the tile-count for each segment. - [[maybe_unused]] const auto segment_size_scan_input_it = ::cuda::transform_iterator( - ::cuda::counting_iterator{num_segments_val_t{0}}, segment_size_scan_input_op); + // The supported maximum segment size (2^21) is enforced at compile time at the public entry; a statically negative + // lower bound is allowed and negative runtime sizes are clamped to 0 (see + // detail::params::__get_and_clamp_param_to_nonnegative). A per-segment value outside its declared bound is a caller + // error (UB): the statically declared bounds are validated at compile time, while the argument values are + // bounds-checked only by assertions active in assertion-enabled (e.g. debug) builds -- host-side for a host-known + // immediate value and device-side for values read from a deferred / deferred_sequence handle. - if constexpr (!only_small_segments) - { - const auto num_segments_val = params::get_param(num_segments, 0); - // Scan output - allocation_sizes[0] = num_segments_val * sizeof(large_segment_tile_offset_t); - if constexpr (any_small_segments) - { - allocation_sizes[1] = sizeof(counters_t); - // Large segment ids for indirectly accessing the large segment parameters - allocation_sizes[2] = num_segments_val * sizeof(num_segments_val_t); - } - else - { - // Query the temporary storage requirement of the segment-size transform-scan. - if (const auto error = CubDebug(detail::scan::dispatch( - nullptr, - allocation_sizes[1], - segment_size_scan_input_it, - static_cast(nullptr), - ::cuda::std::plus<>{}, - detail::InputValue(large_segment_tile_offset_t{0}), - static_cast(num_segments_val), - stream))) - { - return error; - } - } - } - - // Compute allocation pointers into the single storage blob (or compute the necessary size of the blob) - void* allocations[allocations_array_size] = {}; - if (const auto error = - CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) { return error; } - if (d_temp_storage == nullptr) + // `num_segments` maps to the grid's x-extent in both host launch arms (the baseline arm launches one block per + // segment; the cluster arm launches one cluster per segment, stacking the cluster's CTAs in the grid's y-dimension), + // so it must fit a positive 32-bit grid dimension. A count above INT_MAX cannot, so reject it as an out-of-contract + // value at this single host boundary: otherwise the baseline arm would silently narrow it to `int` and the cluster + // arm would build an out-of-range grid.x. + using num_segments_val_t = typename ::cuda::args::__traits::element_type; + const num_segments_val_t num_segments_val = detail::params::get_param(num_segments, num_segments_val_t{0}); + // Unary `+` integer-promotes the count to a standard integer type so the sign-safe `cmp_*` comparators accept it: + // they are constrained to `__cccl_is_integer_v`, which excludes the character count types the public API permits. + if (::cuda::std::cmp_greater(+num_segments_val, ::cuda::std::numeric_limits::max())) { + return cudaErrorInvalidValue; + } + // A negative count is no work (like a zero count), matching DeviceSegmentedReduce. Short-circuit here, before + // `dispatch_compute_cap`, so the query pass cannot fall into the baseline arm where `num_segments_val * sizeof(...)` + // would cast the negative count to a huge `size_t`. (Zero is handled by `empty_batch_no_launch` below, which keeps + // the arch-gated no-op semantics.) TODO(topk): file an issue to unify the negative-`num_segments` contract across + // CUB device algorithms. + if (::cuda::std::cmp_less(+num_segments_val, 0)) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + } return cudaSuccess; } - // TODO (elstehle): support number of segments provided by device-accessible iterator - // Only uniform number of segments are supported (i.e., we need to resolve the number of segments on the host) - static_assert(::cuda::args::__traits::is_single_value, - "Only uniform segment sizes are currently supported."); + // Empty batch = no work to launch: no segments, or a non-positive tightest max segment size (every segment empty, + // e.g. a uniform negative size clamped to 0). `== 0` suffices for `num_segments`: a negative count is already + // short-circuited as no work above. Consulted only on the launch (`d_temp_storage != nullptr`) of a *supported* arm + // below: the query pass falls through to size `temp_storage_bytes`, and the unsupported arm ignores it so an + // unavailable request still fails with cudaErrorNotSupported rather than being masked into success. + const auto empty_batch_no_launch = [&] { + return d_temp_storage != nullptr + && (detail::params::get_param(num_segments, 0) == 0 || ::cuda::args::__highest_(segment_sizes) <= 0); + }; - if constexpr (any_small_segments) - { - if constexpr (!only_small_segments) + return detail::dispatch_compute_cap(policy_selector_t{}, cc, [&](auto policy_getter) -> cudaError_t { + constexpr topk_policy active_policy = policy_getter(); +#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + NV_IF_TARGET(NV_IS_HOST, ({ + ::std::stringstream ss; + ss << active_policy; + _CubLog("Dispatching DeviceBatchedTopK to compute capability %d.%d with tuning: %s\n", + cc.major_cap(), + cc.minor_cap(), + ss.str().c_str()); + })) +#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + if constexpr (active_policy.backend == topk_algorithm::baseline) { - // Zero-initialize the counters struct that holds the large-segment queue length and the block retirement - // counter; both are read by the agent's atomic operations and must start at 0. - if (const auto error = CubDebug(cudaMemsetAsync(allocations[1], 0, sizeof(counters_t), stream))) + // Computed from the template parameters, not a captured function-scope constant: MSVC rejects the latter as + // non-constant inside this lambda's `if constexpr`. + constexpr bool deterministic = (Determinism != ::cuda::execution::determinism::__determinism_t::__not_guaranteed) + || (TieBreak != ::cuda::execution::tie_break::__tie_break_t::__unspecified); + if constexpr (deterministic) { - return error; + // A `tune`d selector forced the baseline backend for a deterministic / tie-break request it cannot serve (only + // the SM 9.0+ cluster backend is deterministic). Mirror the arch-unsupported / oversize-baseline failure model: + // a hard compile error by default, deferred to a runtime cudaErrorNotSupported only under the escape hatches. +#if !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC) \ + && !defined(CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT) + static_assert( + !deterministic, + "cub::DeviceBatchedTopK: a tuned policy selector forced the baseline backend for a deterministic " + "/ tie-break request it cannot serve (only the SM 9.0+ cluster backend is deterministic). Drop " + "the override, relax the determinism / tie-break requirement, or define " + "CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT to defer the diagnosis to runtime " + "(cudaErrorNotSupported)."); +#endif // !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC) + // && !defined(CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT) + // Report a positive temp-storage size so the two-phase protocol proceeds, then fail the launch explicitly. + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + return cudaErrorNotSupported; + } + else + { + if (empty_batch_no_launch()) + { + return cudaSuccess; + } + return launch_baseline_arm( + d_temp_storage, + temp_storage_bytes, + d_key_segments_it, + d_key_segments_out_it, + d_value_segments_it, + d_value_segments_out_it, + segment_sizes, + k, + select_directions, + num_segments, + stream, + launcher_factory); } } - const int grid_dim = static_cast(params::get_param(num_segments, 0)); - constexpr int block_dim = worker_per_segment_policy.threads_per_block; - if (const auto error = CubDebug( - THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(grid_dim, block_dim, 0, stream) - .doit( - device_segmented_topk_kernel< - PolicySelector, - KeyInputItItT, - KeyOutputItItT, - ValueInputItItT, - ValueOutputItItT, - SegmentSizeParameterT, - KParameterT, - SelectDirectionParameterT, - NumSegmentsParameterT, - large_segment_tile_offset_t>, - d_key_segments_it, - d_key_segments_out_it, - d_value_segments_it, - d_value_segments_out_it, - segment_sizes, - k, - select_directions, - num_segments, - only_small_segments ? nullptr : static_cast(allocations[1]), - only_small_segments ? nullptr : static_cast(allocations[2]), - only_small_segments ? nullptr : static_cast(allocations[0])))) + else if constexpr (active_policy.backend == topk_algorithm::cluster) { - return error; - } - } - else - { - // No small segments: the small-kernel epilogue (which would otherwise produce the per-segment tile offsets) does - // not run. Compute the per-segment tile offsets directly via a transform-scan over all segment sizes. - // The large segment agent will either consume these offsets directly (segment_id -> tile offset) or, when going - // through the large-segment queue, via a transform iterator over `d_large_segments_ids` (level of indirection). - if (const auto error = CubDebug(detail::scan::dispatch( - allocations[1], - allocation_sizes[1], - segment_size_scan_input_it, - static_cast(allocations[0]), - ::cuda::std::plus<>{}, - detail::InputValue(large_segment_tile_offset_t{0}), - static_cast(params::get_param(num_segments, 0)), - stream))) - { - return error; - } - } - - if constexpr (!only_small_segments) - { - // TODO (elstehle): support larger number of segments through multiple kernel launches - // Depending on any_small_segments, we need to either: - // - Indirectly get the large segment parameters via the queued large segment IDs - // - Directly take the segment parameters since all segments are large - } - return CubDebug(detail::DebugSyncStream(stream)); -} -// Env-based dispatch function handling memory allocation as well. This is usually done by the device-layer, but there -// is no public API for segmented topk yet. -template > -[[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch_with_env( - KeyInputItItT d_key_segments_it, - KeyOutputItItT d_key_segments_out_it, - ValueInputItItT d_value_segments_it, - ValueOutputItItT d_value_segments_out_it, - SegmentSizeParameterT segment_sizes, - KParameterT k, - SelectDirectionParameterT select_directions, - NumSegmentsParameterT num_segments, - TotalNumItemsGuaranteeT total_num_items_guarantee, - const EnvT& env = {}) -{ - using default_policy_selector = - policy_selector_from_types>, - it_value_t>, - ::cuda::std::int64_t, - ::cuda::args::__traits::highest>; - return detail::dispatch_with_env_and_tuning( - env, [&](auto policy_selector, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { - return dispatch( +#if !_CCCL_HAS_DYNAMIC_CLUSTER_LAUNCH() + // The automatic selector never picks the cluster backend when dynamic cluster launches are disabled (see + // cluster_capable), so reaching here means a `tune`d selector forced it. The kernel would launch without its + // cluster extent (triple_chevron drops it), so reject the contradiction at compile time rather than run wrong. + static_assert(active_policy.backend != topk_algorithm::cluster, + "cub::DeviceBatchedTopK: a tuned policy selector forced the cluster backend, but " + "_CCCL_DISABLE_DYNAMIC_CLUSTER_LAUNCH is defined. Drop the override or the macro."); +#endif // !_CCCL_HAS_DYNAMIC_CLUSTER_LAUNCH() + if (empty_batch_no_launch()) + { + return cudaSuccess; + } + // `UserProvidedTuning`: false for the automatic selector, which returns `cluster` solely for a + // `cluster_capable(cc)` and so needs no runtime re-check; a `tune`d override is a different type and keeps it. + // Inlined as a type trait rather than a function-scope constexpr, which MSVC rejects inside this lambda. + return launch_cluster_arm>( + policy_getter, d_temp_storage, temp_storage_bytes, d_key_segments_it, @@ -360,10 +1167,21 @@ template #include +#include #include +#include +#include +#include #include +#include +#include #include CUB_NAMESPACE_BEGIN @@ -47,6 +53,46 @@ struct copy_mdspan_t } }; +template +[[nodiscard]] _CCCL_HOST_API ::cudaError_t +__copy_mdspan_bytes(::cuda::stream_ref __stream, _MDSpanIn&& __mdspan_in, _MDSpanOut&& __mdspan_out) +{ + _CCCL_TRY + { + ::cuda::copy_bytes(__stream, __mdspan_in, __mdspan_out); + } + _CCCL_CATCH (const ::cuda::cuda_error& __e) + { + return __e.status(); + } +#if _CCCL_HOSTED() + _CCCL_CATCH (const ::std::invalid_argument& __e) + { + static_cast(__e); + return ::cudaErrorInvalidValue; + } +#endif // _CCCL_HOSTED + _CCCL_CATCH_ALL + { + return ::cudaErrorUnknown; + } + + return ::cudaSuccess; +} + +template +[[nodiscard]] CUB_RUNTIME_FUNCTION ::cudaError_t +__transform_copy(_MDSpanIn&& __mdspan_in, _MDSpanOut&& __mdspan_out, const _Env& __env) +{ + return CUB_NS_QUALIFIER::DeviceTransform::__transform_internal( + ::cuda::std::make_tuple(__mdspan_in.data_handle()), + __mdspan_out.data_handle(), + __mdspan_in.size(), + ::cuda::always_true{}, + ::cuda::std::identity{}, + __env); +} + template mdspan_in, if (mdspan_in.is_exhaustive() && mdspan_out.is_exhaustive() && detail::have_same_strides(mdspan_in.mapping(), mdspan_out.mapping())) { - return cub::DeviceTransform::__transform_internal( - ::cuda::std::make_tuple(mdspan_in.data_handle()), - mdspan_out.data_handle(), - mdspan_in.size(), - ::cuda::always_true{}, - ::cuda::std::identity{}, - env); + // NOLINTBEGIN(bugprone-branch-clone) + if constexpr (::cuda::std::same_as + && ::cuda::__detail::__can_mdspan_copy_bytes + && ::cuda::std::__is_callable_v<::cuda::get_stream_t, const EnvT&>) + { + NV_IF_TARGET( + NV_IS_HOST, + ({ + auto __stream = ::cuda::get_stream(env); + + // cuda::copy_bytes() builds an __ensure_current_context(stream_ref), which calls + // cuStreamGetCtx(). That driver call rejects the NULL stream with + // CUDA_ERROR_INVALID_VALUE. Use the transform path, which goes through the runtime + // API and accepts the NULL stream. + // + // Likewise, we cannot retrieve the context for a stream that is capturings so we + // need to call the kernel. + if (__stream.get() == nullptr + || (::cuda::__driver::__streamIsCapturing(__stream.get()) == ::CU_STREAM_CAPTURE_STATUS_ACTIVE)) + { + return CUB_NS_QUALIFIER::detail::copy_mdspan::__transform_copy(mdspan_in, mdspan_out, env); + } + + return CUB_NS_QUALIFIER::detail::copy_mdspan::__copy_mdspan_bytes(__stream, mdspan_in, mdspan_out); + }), + (return CUB_NS_QUALIFIER::detail::copy_mdspan::__transform_copy(mdspan_in, mdspan_out, env);)) + } + else + { + return CUB_NS_QUALIFIER::detail::copy_mdspan::__transform_copy(mdspan_in, mdspan_out, env); + } + // NOLINTEND(bugprone-branch-clone) } // TODO (fbusato): add ForEachInLayout when mdspan_in and mdspan_out have compatible layouts // Compatible layouts could use more efficient iteration patterns diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_histogram.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_histogram.cuh index 704c57ad..673be75f 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_histogram.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_histogram.cuh @@ -216,12 +216,12 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( if constexpr (IsDeviceInit) { return kernel_source.template HistogramSweepKernelDeviceInit< - PolicySelector, - PRIVATIZED_SMEM_BINS, - FirstLevelArrayT, - SecondLevelArrayT, - IsEven, - IsByteSample>(); + PolicySelector, + PRIVATIZED_SMEM_BINS, + FirstLevelArrayT, + SecondLevelArrayT, + IsEven, + IsByteSample>(); } else { diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_reduce_by_key.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_reduce_by_key.cuh index 91560416..926f40ef 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_reduce_by_key.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_reduce_by_key.cuh @@ -208,14 +208,14 @@ __launch_bounds__(int(current_policy().lookback.threads_per_bloc { static constexpr ReduceByKeyPolicy policy = current_policy(); using AgentReduceByKeyPolicyT = agent_reduce_by_key_policy< - policy.lookback.threads_per_block, - policy.lookback.items_per_thread, - policy.lookback.load_algorithm, - policy.lookback.load_modifier, - policy.lookback.scan_algorithm, - delay_constructor_t>; + policy.lookback.threads_per_block, + policy.lookback.items_per_thread, + policy.lookback.load_algorithm, + policy.lookback.load_modifier, + policy.lookback.scan_algorithm, + delay_constructor_t>; using vsmem_helper_t = vsmem_helper_default_fallback_policy_t< AgentReduceByKeyPolicyT, @@ -656,14 +656,14 @@ _CCCL_HOST_DEVICE_API auto determine_threads_items_vsmem(PolicyGetter policy_get // TODO(bgruber): refactor this in the future constexpr ReduceByKeyPolicy policy = policy_getter(); using Policy = agent_reduce_by_key_policy< - policy.lookback.threads_per_block, - policy.lookback.items_per_thread, - policy.lookback.load_algorithm, - policy.lookback.load_modifier, - policy.lookback.scan_algorithm, - delay_constructor_t>; + policy.lookback.threads_per_block, + policy.lookback.items_per_thread, + policy.lookback.load_algorithm, + policy.lookback.load_modifier, + policy.lookback.scan_algorithm, + delay_constructor_t>; using vsmem_helper_t = vsmem_helper_default_fallback_policy_t; return ::cuda::std::tuple{vsmem_helper_t::agent_policy_t::BLOCK_THREADS, vsmem_helper_t::agent_policy_t::ITEMS_PER_THREAD, diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_scan.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_scan.cuh index bb3762ef..a62f5e81 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_scan.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_scan.cuh @@ -220,11 +220,11 @@ template < typename ScanOpT, typename InitValueT, typename OffsetT, - typename AccumT = ::cuda::std::__accumulator_t, - ::cuda::std::_If<::cuda::std::is_same_v, - cub::detail::it_value_t, - typename InitValueT::value_type>>, + typename AccumT = ::cuda::std::__accumulator_t, + ::cuda::std::_If<::cuda::std::is_same_v, + cub::detail::it_value_t, + typename InitValueT::value_type>>, ForceInclusive EnforceInclusive = ForceInclusive::No, typename PolicyHub = detail::scan:: policy_hub, detail::it_value_t, AccumT, OffsetT, ScanOpT>, @@ -552,38 +552,44 @@ struct CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceScan") DispatchScan int smem_size = smem_size_1_stage; // When launched from the host, maximize the number of stages that we can fit inside the shared memory. - NV_IF_TARGET(NV_IS_HOST, ({ - // number of stages to have an even workload across all SMs (improves small problem sizes), assuming - // 1 CTA per SM +1 since it tends to improve performance - // TODO(bgruber): make the +1 a tuning parameter - const int max_stages_for_even_workload = static_cast( - ::cuda::ceil_div(num_items, static_cast(sm_count * lookahead_policy.tile_size())) + 1); + NV_IF_TARGET( + NV_IS_HOST, ({ + // number of stages to have an even workload across all SMs (improves small problem sizes), assuming + // 1 CTA per SM +1 since it tends to improve performance + // TODO(bgruber): make the +1 a tuning parameter + const int max_stages_for_even_workload = static_cast( + ::cuda::ceil_div(num_items, static_cast(sm_count * lookahead_policy.tile_size())) + 1); - while (num_stages <= max_stages_for_even_workload) - { - const int next_smem_size = detail::scan::smem_for_stages( - lookahead_policy, - num_stages + 1, - static_cast(kernel_source.InputSize()), - static_cast(kernel_source.InputAlign()), - static_cast(kernel_source.OutputAlign()), - static_cast(kernel_source.AccumSize()), - static_cast(kernel_source.AccumAlign())); - if (next_smem_size > max_dynamic_smem_size) - { - // This number of stages failed, so stay at the current settings - break; - } + while (num_stages <= max_stages_for_even_workload) + { + const int next_smem_size = detail::scan::smem_for_stages( + lookahead_policy, + num_stages + 1, + static_cast(kernel_source.InputSize()), + static_cast(kernel_source.InputAlign()), + static_cast(kernel_source.OutputAlign()), + static_cast(kernel_source.AccumSize()), + static_cast(kernel_source.AccumAlign())); + if (next_smem_size > max_dynamic_smem_size) + { + // This number of stages failed, so stay at the current settings + break; + } - smem_size = next_smem_size; - ++num_stages; - } + smem_size = next_smem_size; + ++num_stages; + } - if (const auto error = launcher_factory.set_max_dynamic_smem_size_for(scan_kernel, smem_size)) - { - return error; - } - })) + // Set scan kernel's max shared memory limit to the max smem value. We might not use all of it, but it prevents + // multiple kernels from overwriting the max shared memory limit by different values. + // + // TODO: Since CTK 13.2 we can use CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE to allow non-portable shared memory + // sizes, however we need something that works even with older CTKs. + if (const auto error = launcher_factory.set_max_dynamic_smem_size_for(scan_kernel, max_dynamic_smem_size)) + { + return error; + } + })) // Invoke init kernel { @@ -1163,38 +1169,44 @@ CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t invoke_lookahead( int smem_size = smem_size_1_stage; // When launched from the host, maximize the number of stages that we can fit inside the shared memory. - NV_IF_TARGET(NV_IS_HOST, ({ - // number of stages to have an even workload across all SMs (improves small problem sizes), assuming - // 1 CTA per SM +1 since it tends to improve performance - // TODO(bgruber): make the +1 a tuning parameter - const int max_stages_for_even_workload = static_cast( - ::cuda::ceil_div(num_items, static_cast(sm_count * lookahead_policy.tile_size())) + 1); + NV_IF_TARGET( + NV_IS_HOST, ({ + // number of stages to have an even workload across all SMs (improves small problem sizes), assuming + // 1 CTA per SM +1 since it tends to improve performance + // TODO(bgruber): make the +1 a tuning parameter + const int max_stages_for_even_workload = static_cast( + ::cuda::ceil_div(num_items, static_cast(sm_count * lookahead_policy.tile_size())) + 1); - while (num_stages <= max_stages_for_even_workload) - { - const int next_smem_size = detail::scan::smem_for_stages( - lookahead_policy, - num_stages + 1, - static_cast(kernel_source.InputSize()), - static_cast(kernel_source.InputAlign()), - static_cast(kernel_source.OutputAlign()), - static_cast(kernel_source.AccumSize()), - static_cast(kernel_source.AccumAlign())); - if (next_smem_size > max_dynamic_smem_size) - { - // This number of stages failed, so stay at the current settings - break; - } + while (num_stages <= max_stages_for_even_workload) + { + const int next_smem_size = detail::scan::smem_for_stages( + lookahead_policy, + num_stages + 1, + static_cast(kernel_source.InputSize()), + static_cast(kernel_source.InputAlign()), + static_cast(kernel_source.OutputAlign()), + static_cast(kernel_source.AccumSize()), + static_cast(kernel_source.AccumAlign())); + if (next_smem_size > max_dynamic_smem_size) + { + // This number of stages failed, so stay at the current settings + break; + } - smem_size = next_smem_size; - ++num_stages; - } + smem_size = next_smem_size; + ++num_stages; + } - if (const auto error = launcher_factory.set_max_dynamic_smem_size_for(scan_kernel, smem_size)) - { - return error; - } - })) + // Set scan kernel's max shared memory limit to the max smem value. We might not use all of it, but it prevents + // multiple kernels from overwriting the max shared memory limit by different values. + // + // TODO: Since CTK 13.2 we can use CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE to allow non-portable shared memory + // sizes, however we need something that works even with older CTKs. + if (const auto error = launcher_factory.set_max_dynamic_smem_size_for(scan_kernel, max_dynamic_smem_size)) + { + return error; + } + })) // Invoke init kernel { diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_scan_by_key.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_scan_by_key.cuh index bb316143..3ee74db8 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_scan_by_key.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_scan_by_key.cuh @@ -238,15 +238,15 @@ template < typename PolicyHub = policy_hub, ScanOpT>, typename PolicySelector = policy_selector_from_hub, typename KernelSource = DeviceScanByKeyKernelSource< - PolicySelector, - KeysInputIteratorT, - ValuesInputIteratorT, - ValuesOutputIteratorT, - EqualityOp, - ScanOpT, - InitValueT, - OffsetT, - AccumT>, + PolicySelector, + KeysInputIteratorT, + ValuesInputIteratorT, + ValuesOutputIteratorT, + EqualityOp, + ScanOpT, + InitValueT, + OffsetT, + AccumT>, typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> struct dispatch_scan_by_key { @@ -672,15 +672,15 @@ template < cub::detail::it_value_t, ScanOpT>, typename KernelSource = DeviceScanByKeyKernelSource< - PolicySelector, - KeysInputIteratorT, - ValuesInputIteratorT, - ValuesOutputIteratorT, - EqualityOp, - ScanOpT, - InitValueT, - OffsetT, - AccumT>, + PolicySelector, + KeysInputIteratorT, + ValuesInputIteratorT, + ValuesOutputIteratorT, + EqualityOp, + ScanOpT, + InitValueT, + OffsetT, + AccumT>, typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> #if _CCCL_HAS_CONCEPTS() requires scan_by_key_policy_selector @@ -875,15 +875,15 @@ template < detail::scan_by_key::policy_hub, ScanOpT>, typename PolicySelector = detail::scan_by_key::policy_selector_from_hub, typename KernelSource = detail::scan_by_key::DeviceScanByKeyKernelSource< - PolicySelector, - KeysInputIteratorT, - ValuesInputIteratorT, - ValuesOutputIteratorT, - EqualityOp, - ScanOpT, - InitValueT, - OffsetT, - AccumT>, + PolicySelector, + KeysInputIteratorT, + ValuesInputIteratorT, + ValuesOutputIteratorT, + EqualityOp, + ScanOpT, + InitValueT, + OffsetT, + AccumT>, typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> using DispatchScanByKey CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceScan") = detail::scan_by_key::dispatch_scan_by_key< diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_radix_sort.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_radix_sort.cuh index 2c9d6fbe..ce33983c 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_radix_sort.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_radix_sort.cuh @@ -911,7 +911,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( { using default_policy_selector_t = policy_selector_from_types; using policy_selector_t = ::cuda::std::decay_t< - ::cuda::std::execution::__query_result_or_t>; + ::cuda::std::execution::__query_result_or_t>; #if _CCCL_HAS_CONCEPTS() static_assert(segmented_radix_sort_policy_selector); #endif // _CCCL_HAS_CONCEPTS() diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_reduce.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_reduce.cuh index 35c432b4..a01cfe45 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_reduce.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_reduce.cuh @@ -503,15 +503,15 @@ template < decltype(select_segmented_accum_t(static_cast(nullptr))), typename PolicySelector = policy_selector_from_types, typename KernelSource = DeviceSegmentedReduceKernelSource< - PolicySelector, - InputIteratorT, - OutputIteratorT, - BeginOffsetIteratorT, - EndOffsetIteratorT, - OffsetT, - ReductionOpT, - InitValueT, - AccumT>, + PolicySelector, + InputIteratorT, + OutputIteratorT, + BeginOffsetIteratorT, + EndOffsetIteratorT, + OffsetT, + ReductionOpT, + InitValueT, + AccumT>, typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> #if _CCCL_HAS_CONCEPTS() requires segmented_reduce_policy_selector @@ -711,13 +711,13 @@ template (nullptr))), typename PolicySelector = policy_selector_from_types, typename KernelSource = DeviceFixedSizeSegmentedReduceKernelSource< - PolicySelector, - InputIteratorT, - OutputIteratorT, - OffsetT, - ReductionOpT, - InitValueT, - AccumT>, + PolicySelector, + InputIteratorT, + OutputIteratorT, + OffsetT, + ReductionOpT, + InitValueT, + AccumT>, typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY, ::cuda::std::enable_if_t<::cuda::std::is_arithmetic_v, int> = 0> #if _CCCL_HAS_CONCEPTS() diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_scan.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_scan.cuh index c8dd92d2..673c9b7f 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_scan.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_scan.cuh @@ -96,17 +96,17 @@ template < common_iterator_value_t, typename PolicySelector = policy_selector_from_types, typename KernelSource = device_segmented_scan_kernel_source< - PolicySelector, - InputIteratorT, - OutputIteratorT, - BeginOffsetIteratorInputT, - EndOffsetIteratorInputT, - BeginOffsetIteratorOutputT, - OffsetT, - ScanOpT, - InitValueT, - AccumT, - EnforceInclusive>, + PolicySelector, + InputIteratorT, + OutputIteratorT, + BeginOffsetIteratorInputT, + EndOffsetIteratorInputT, + BeginOffsetIteratorOutputT, + OffsetT, + ScanOpT, + InitValueT, + AccumT, + EnforceInclusive>, typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> #if _CCCL_HAS_CONCEPTS() requires segmented_scan_policy_selector diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_select_if.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_select_if.cuh index fa8bc5c8..3e845587 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_select_if.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_select_if.cuh @@ -216,14 +216,15 @@ struct make_vsmem_helper { static constexpr SelectPolicy active_policy = DefaultPolicyGetter{}(); using agent_policy_t = detail::agent_select_if_policy< - active_policy.lookback.threads_per_block, - active_policy.lookback.items_per_thread, - active_policy.lookback.load_algorithm, - active_policy.lookback.load_modifier, - active_policy.lookback.scan_algorithm, - delay_constructor_t>; + active_policy.lookback.threads_per_block, + active_policy.lookback.items_per_thread, + active_policy.lookback.load_algorithm, + active_policy.lookback.load_modifier, + active_policy.lookback.scan_algorithm, + delay_constructor_t, + active_policy.lookback._load_prefetch>; using type = vsmem_helper_default_fallback_policy_t< agent_policy_t, bind_selection_opt::template agent_t, diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_three_way_partition.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_three_way_partition.cuh index 13ab3b78..61c0d352 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_three_way_partition.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_three_way_partition.cuh @@ -444,18 +444,18 @@ template , per_partition_offset_t>, typename KernelSource = DeviceThreeWayPartitionKernelSource< - PolicySelector, - InputIteratorT, - FirstOutputIteratorT, - SecondOutputIteratorT, - UnselectedOutputIteratorT, - NumSelectedIteratorT, - ScanTileStateT, - SelectFirstPartOp, - SelectSecondPartOp, - per_partition_offset_t, - streaming_context_t, - OffsetT>, + PolicySelector, + InputIteratorT, + FirstOutputIteratorT, + SecondOutputIteratorT, + UnselectedOutputIteratorT, + NumSelectedIteratorT, + ScanTileStateT, + SelectFirstPartOp, + SelectSecondPartOp, + per_partition_offset_t, + streaming_context_t, + OffsetT>, typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> #if _CCCL_HAS_CONCEPTS() requires three_way_partition_policy_selector diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_transform_tile_config.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_transform_tile_config.cuh index 3130fb6e..7a3327eb 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_transform_tile_config.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_transform_tile_config.cuh @@ -26,7 +26,7 @@ # pragma system_header #endif // no system header -#define _CCCL_CUB_HAS_TILE_TRANSFORM() _CCCL_TILE_COMPILATION() +#define _CCCL_CUB_HAS_TILE_TRANSFORM() _CCCL_TILE_COMPILATION() && _CCCL_STD_VER >= 2020 #if _CCCL_CUB_HAS_TILE_TRANSFORM() && defined(_CCCL_ENABLE_EXPERIMENTAL_TILE_TRANSFORM_DISPATCH) # define _CCCL_CUB_TILE_TRANSFORM_DISPATCH_ENABLED() 1 diff --git a/cccl_upstream/cub/cub/device/dispatch/tuning/common.cuh b/cccl_upstream/cub/cub/device/dispatch/tuning/common.cuh index 0e27f17c..a0518e6f 100644 --- a/cccl_upstream/cub/cub/device/dispatch/tuning/common.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/tuning/common.cuh @@ -13,13 +13,14 @@ # pragma system_header #endif // no system header +#include #include #include -#include #include #include +#include #include #include #include @@ -130,7 +131,7 @@ template return iterator_info{ static_cast(size_of), static_cast(align_of), - THRUST_NS_QUALIFIER::is_trivially_relocatable_v, + ::cuda::is_trivially_copyable_v, THRUST_NS_QUALIFIER::is_contiguous_iterator_v}; } diff --git a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_batch_memcpy.cuh b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_batch_memcpy.cuh index 37072f3c..6cb318aa 100644 --- a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_batch_memcpy.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_batch_memcpy.cuh @@ -158,13 +158,13 @@ struct BatchedCopyPolicy BatchedCopyLookbackPolicy lookback; //!< The policy for the batched-copy algorithm based on decoupled-lookback. Only //!< used when @p algorithm is @p lookback. - [[nodiscard]] _CCCL_API friend constexpr bool + [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator==(const BatchedCopyPolicy& lhs, const BatchedCopyPolicy& rhs) noexcept { return lhs.algorithm == rhs.algorithm && lhs.lookback == rhs.lookback; } - [[nodiscard]] _CCCL_API friend constexpr bool + [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator!=(const BatchedCopyPolicy& lhs, const BatchedCopyPolicy& rhs) noexcept { return !(lhs == rhs); diff --git a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_batched_topk.cuh b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_batched_topk.cuh index 48f97651..679dbbfc 100644 --- a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_batched_topk.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_batched_topk.cuh @@ -16,20 +16,27 @@ #include #include #include +#include +#include #include +#include +#include #include #include +#include CUB_NAMESPACE_BEGIN namespace detail::batched_topk { +//! Sub-policy for the compaction epilogue shared by the baseline @ref DeviceBatchedTopK workers: it scans the radix +//! histogram and writes out the selected keys. struct epilogue_policy { - int items_per_thread; - BlockLoadAlgorithm load_algorithm; - BlockStoreAlgorithm store_algorithm; - BlockScanAlgorithm scan_algorithm; + int items_per_thread; //!< Keys each thread loads/stores per tile in the epilogue. + BlockLoadAlgorithm load_algorithm; //!< Block load algorithm used to read keys back in the epilogue. + BlockStoreAlgorithm store_algorithm; //!< Block store algorithm used to write the selected keys. + BlockScanAlgorithm scan_algorithm; //!< Block scan algorithm used for the histogram prefix sum. _CCCL_HOST_DEVICE_API friend constexpr bool operator==(const epilogue_policy& lhs, const epilogue_policy& rhs) { @@ -42,24 +49,26 @@ struct epilogue_policy return !(lhs == rhs); } -#if !_CCCL_COMPILER(NVRTC) +#if _CCCL_HOSTED() friend ::std::ostream& operator<<(::std::ostream& os, const epilogue_policy& p) { return os << "epilogue_policy { .items_per_thread = " << p.items_per_thread << ", .load_algorithm = " << p.load_algorithm << ", .store_algorithm = " << p.store_algorithm << ", .scan_algorithm = " << p.scan_algorithm << " }"; } -#endif // !_CCCL_COMPILER(NVRTC) +#endif // _CCCL_HOSTED() }; +//! Per-segment worker sub-policy for the baseline backend: one thread block cooperatively computes the top-k of a +//! single segment. @ref baseline_topk_policy holds several of these, ordered by decreasing tile size. struct worker_policy { - int threads_per_block; - int items_per_thread; - BlockLoadAlgorithm load_algorithm; - BlockStoreAlgorithm store_algorithm; + int threads_per_block; //!< Number of threads in a CUDA block. + int items_per_thread; //!< Keys each thread loads/processes per tile (with `threads_per_block` sets the tile size). + BlockLoadAlgorithm load_algorithm; //!< Block load algorithm used to read the segment's keys. + BlockStoreAlgorithm store_algorithm; //!< Block store algorithm used to write the selected keys. - epilogue_policy epilogue; + epilogue_policy epilogue; //!< Sub-policy for the compaction epilogue. _CCCL_HOST_DEVICE_API friend constexpr bool operator==(const worker_policy& lhs, const worker_policy& rhs) { @@ -80,13 +89,15 @@ struct worker_policy << ", .items_per_thread = " << p.items_per_thread << ", .load_algorithm = " << p.load_algorithm << ", .store_algorithm = " << p.store_algorithm << ", .epilogue = " << p.epilogue << " }"; } -#endif // !_CCCL_COMPILER(NVRTC) +#endif // _CCCL_HOSTED() }; +//! Sub-policy for the baseline backend's multiple-blocks-per-segment worker path, used for segments too large for a +//! single worker block. struct multi_worker_policy { - int threads_per_block; - int items_per_thread; + int threads_per_block; //!< Number of threads in a CUDA block. + int items_per_thread; //!< Keys each thread loads/processes per tile. _CCCL_HOST_DEVICE_API friend constexpr bool operator==(const multi_worker_policy& lhs, const multi_worker_policy& rhs) { @@ -98,7 +109,7 @@ struct multi_worker_policy return !(lhs == rhs); } -#if !_CCCL_COMPILER(NVRTC) +#if _CCCL_HOSTED() friend ::std::ostream& operator<<(::std::ostream& os, const multi_worker_policy& p) { return os << "multi_worker_policy { .threads_per_block = " << p.threads_per_block @@ -107,28 +118,31 @@ struct multi_worker_policy #endif // _CCCL_HOSTED() }; -struct batched_topk_policy +//! Sub-policy for the baseline (worker-per-segment) backend of @ref DeviceBatchedTopK. +struct baseline_topk_policy { - // The list of per-segment agent policies is ordered by decreasing tile size. At compile time, the smallest policy - // whose tile size still covers the upper bound of the segment size is selected. + //! Per-segment worker policies ordered by decreasing tile size. At compile time the smallest policy whose tile size + //! still covers the upper bound of the segment size is selected. ::cuda::std::array worker_per_segment_policies; - multi_worker_policy multi_worker_per_segment_policy; + multi_worker_policy multi_worker_per_segment_policy; //!< Worker policy for segments too large for a single block. - _CCCL_HOST_DEVICE_API friend constexpr bool operator==(const batched_topk_policy& lhs, const batched_topk_policy& rhs) + _CCCL_HOST_DEVICE_API friend constexpr bool + operator==(const baseline_topk_policy& lhs, const baseline_topk_policy& rhs) { return lhs.worker_per_segment_policies == rhs.worker_per_segment_policies && lhs.multi_worker_per_segment_policy == rhs.multi_worker_per_segment_policy; } - _CCCL_HOST_DEVICE_API friend constexpr bool operator!=(const batched_topk_policy& lhs, const batched_topk_policy& rhs) + _CCCL_HOST_DEVICE_API friend constexpr bool + operator!=(const baseline_topk_policy& lhs, const baseline_topk_policy& rhs) { return !(lhs == rhs); } #if _CCCL_HOSTED() - friend ::std::ostream& operator<<(::std::ostream& os, const batched_topk_policy& p) + friend ::std::ostream& operator<<(::std::ostream& os, const baseline_topk_policy& p) { - os << "batched_topk_policy { .worker_per_segment_policies = { "; + os << "baseline_topk_policy { .worker_per_segment_policies = { "; for (::cuda::std::size_t i = 0; i < p.worker_per_segment_policies.size(); ++i) { if (i != 0) @@ -142,45 +156,242 @@ struct batched_topk_policy #endif // _CCCL_HOSTED() }; +// Default baseline sub-policy. Tuning is currently CC-independent. +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto make_baseline_policy() -> baseline_topk_policy +{ + constexpr auto load_alg = BLOCK_LOAD_WARP_TRANSPOSE; + constexpr auto store_alg = BLOCK_STORE_WARP_TRANSPOSE; + constexpr auto scan_alg = BLOCK_SCAN_WARP_SCANS; + constexpr auto epilogue = epilogue_policy{16, load_alg, store_alg, scan_alg}; + return baseline_topk_policy{ + {{ + worker_policy{256, 64, load_alg, store_alg, epilogue}, + worker_policy{256, 32, load_alg, store_alg, epilogue}, + worker_policy{256, 16, load_alg, store_alg, epilogue}, + worker_policy{256, 8, load_alg, store_alg, epilogue}, + worker_policy{256, 4, load_alg, store_alg, epilogue}, + worker_policy{128, 2, load_alg, store_alg, epilogue}, + }}, + multi_worker_policy{256, 64}}; +} + +// Largest maximum segment size (in keys) the baseline (worker-per-segment) backend can cover: the largest worker tile +// (threads_per_block * items_per_thread) in `policy`. A larger statically-known maximum segment size makes the baseline +// backend ineligible (the selector then picks the cluster backend where supported, otherwise `unsupported`). This is +// only the tile-based necessary condition; the exact predicate `baseline_can_cover_v` also checks the agent's +// shared-memory fit (which needs the concrete agent types). +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr ::cuda::std::int64_t +baseline_max_covered_segment_size(const baseline_topk_policy& policy) +{ + ::cuda::std::int64_t max_tile_size = 0; + for (const auto& worker : policy.worker_per_segment_policies) + { + const ::cuda::std::int64_t tile_size = ::cuda::std::int64_t{worker.threads_per_block} * worker.items_per_thread; + if (tile_size > max_tile_size) + { + max_tile_size = tile_size; + } + } + return max_tile_size; +} + +//! Execution shape for the thread-block-cluster backend of @ref DeviceBatchedTopK. The dispatch picks the number of +//! cluster blocks and the dynamic shared-memory block_tile capacity at runtime (occupancy / wave-aware), so this policy +//! mostly carries per-block tuning knobs; the two trailing `max_*` fields are optional launch-geometry caps that bound +//! that runtime choice. +struct cluster_topk_policy +{ + int threads_per_block; //!< Number of threads in a CUDA block. + int min_blocks_per_sm; //!< Minimum resident blocks per SM, forwarded as the kernel launch-bounds occupancy hint. + + int min_chunks_per_block; //!< Minimum number of chunks a block must own to join a segment's effective cluster (the + //!< divisor mapping a segment's chunk count to its cluster width). Must be >= 1. + + int chunk_bytes; //!< Size in bytes of one block_tile chunk -- the granularity of the async-copy load pipeline. + int load_align_bytes; //!< Load / bulk-copy alignment in bytes. Must be a power of two and >= 16 + //!< (`detail::bulk_copy_min_align`), and `chunk_bytes` must be a multiple of it (see + //!< `is_valid_cluster_policy`). + int pipeline_stages; //!< Depth of the async-copy (mbarrier) pipeline that stages chunks into shared memory. + + int single_block_max_seg_size; //!< Largest segment size, in keys, still eligible for the single-block fast path + //!< (kept out of the byte-unit loading group above as it is measured in items). + int bits_per_pass; //!< Radix digit width per pass; each pass' histogram spans `1 << bits_per_pass` buckets. Together + //!< with `threads_per_block` this implicitly fixes the histogram block-scan's items per thread, + //!< `ceil_div(1 << bits_per_pass, threads_per_block)` buckets scanned per thread. + int histogram_items_per_thread; //!< Keys each thread accumulates per tile during the radix histogram passes. + int tie_break_items_per_thread; //!< Keys each thread processes per tile during the final tie-break / filter phase. + int copy_items_per_thread; //!< Keys each thread copies per tile on the select-all (k >= segment size) fast path. + + // Launch-geometry caps that bound the otherwise heuristic / hardware-derived cluster width and resident shared-memory + // footprint. Both default to 0 (= unrestricted) and are deliberately not auto-tuned. Beyond deterministically + // steering tests onto the streaming / cluster paths at a small footprint, they let a caller trade top-k throughput + // for resources it wants to leave free -- e.g. capping resident slots to fit a shared-memory carveout reserved for a + // concurrently running kernel, or narrowing the cluster width to co-schedule other work. The algorithm stays correct + // at any cap: a segment that no longer fits resident simply streams the remainder from global memory. + int max_blocks_per_cluster; //!< Upper bound on the launched cluster width (CTAs per segment); 0 = unrestricted (the + //!< hardware cluster-width ceiling, queried from the runtime). Non-zero is additionally + //!< clamped to that same ceiling. A cap narrower than a segment needs pushes it into the + //!< streaming fallback (cap 1 -> single-CTA streaming). + int max_chunk_slots_per_block; //!< Upper bound on resident chunk slots per block; 0 = unrestricted (the full + //!< shared-memory budget: the hardware opt-in budget). A smaller cap shrinks each + //!< CTA's resident capacity (and thus its dynamic shared-memory request), so a smaller + //!< segment overflows into the streaming path. + + // Equality/streaming make this a regular type (required by the `policy_selector` concept / `dispatch_compute_cap`). + _CCCL_HOST_DEVICE_API friend constexpr bool operator==(const cluster_topk_policy& lhs, const cluster_topk_policy& rhs) + { + return lhs.threads_per_block == rhs.threads_per_block && lhs.min_blocks_per_sm == rhs.min_blocks_per_sm + && lhs.min_chunks_per_block == rhs.min_chunks_per_block && lhs.chunk_bytes == rhs.chunk_bytes + && lhs.load_align_bytes == rhs.load_align_bytes && lhs.pipeline_stages == rhs.pipeline_stages + && lhs.single_block_max_seg_size == rhs.single_block_max_seg_size && lhs.bits_per_pass == rhs.bits_per_pass + && lhs.histogram_items_per_thread == rhs.histogram_items_per_thread + && lhs.tie_break_items_per_thread == rhs.tie_break_items_per_thread + && lhs.copy_items_per_thread == rhs.copy_items_per_thread + && lhs.max_blocks_per_cluster == rhs.max_blocks_per_cluster + && lhs.max_chunk_slots_per_block == rhs.max_chunk_slots_per_block; + } + + _CCCL_HOST_DEVICE_API friend constexpr bool operator!=(const cluster_topk_policy& lhs, const cluster_topk_policy& rhs) + { + return !(lhs == rhs); + } + +#if _CCCL_HOSTED() + friend ::std::ostream& operator<<(::std::ostream& os, const cluster_topk_policy& p) + { + return os + << "cluster_topk_policy { .threads_per_block = " << p.threads_per_block + << ", .min_blocks_per_sm = " << p.min_blocks_per_sm << ", .min_chunks_per_block = " << p.min_chunks_per_block + << ", .chunk_bytes = " << p.chunk_bytes << ", .load_align_bytes = " << p.load_align_bytes + << ", .pipeline_stages = " << p.pipeline_stages + << ", .single_block_max_seg_size = " << p.single_block_max_seg_size << ", .bits_per_pass = " << p.bits_per_pass + << ", .histogram_items_per_thread = " << p.histogram_items_per_thread << ", .tie_break_items_per_thread = " + << p.tie_break_items_per_thread << ", .copy_items_per_thread = " << p.copy_items_per_thread + << ", .max_blocks_per_cluster = " << p.max_blocks_per_cluster + << ", .max_chunk_slots_per_block = " << p.max_chunk_slots_per_block << " }"; + } +#endif // _CCCL_HOSTED() +}; + +// Default cluster sub-policy. Tuning is currently CC-independent. +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto make_cluster_policy() -> cluster_topk_policy +{ + return cluster_topk_policy{ + /*threads_per_block=*/512, + /*min_blocks_per_sm=*/1, + /*min_chunks_per_block=*/1, + /*chunk_bytes=*/16 * 1024, + /*load_align_bytes=*/128, + /*pipeline_stages=*/8, + /*single_block_max_seg_size=*/8 * 1024, + /*bits_per_pass=*/11, + /*histogram_items_per_thread=*/8, + /*tie_break_items_per_thread=*/8, + /*copy_items_per_thread=*/8, + /*max_blocks_per_cluster=*/0, + /*max_chunk_slots_per_block=*/0}; +} + +// Hard constraints a cluster sub-policy must satisfy, mirroring the agent's compile-time invariants so a bad policy +// (e.g. from a `tune` override) trips `launch_cluster_arm`'s `static_assert(is_valid_cluster_policy(policy))` with a +// clear message instead of a cryptic failure deep in the agent (or a host-side divide-by-zero in the launch-shape +// math). The block_tile byte geometry constraints stem from the aligned bulk-copy (TMA) load path, which addresses +// gmem/smem in `load_align_bytes`-sized, aligned units. +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto is_valid_cluster_policy(cluster_topk_policy policy) -> bool +{ + return policy.chunk_bytes > 0 && policy.load_align_bytes >= bulk_copy_min_align + && ::cuda::is_power_of_two(policy.load_align_bytes) && policy.chunk_bytes % policy.load_align_bytes == 0 + && policy.threads_per_block > 0 && policy.threads_per_block % warp_threads == 0 && policy.min_blocks_per_sm >= 0 + && policy.pipeline_stages >= 1 && policy.pipeline_stages <= 32 && policy.min_chunks_per_block >= 1 + && policy.bits_per_pass >= 1 && policy.bits_per_pass <= 16 && policy.histogram_items_per_thread > 0 + && policy.tie_break_items_per_thread > 0 && policy.copy_items_per_thread > 0 + && policy.single_block_max_seg_size >= 0 && policy.max_blocks_per_cluster >= 0 + && policy.max_chunk_slots_per_block >= 0; +} + +static_assert(is_valid_cluster_policy(make_cluster_policy())); + +// ----------------------------------------------------------------------------- +// Backend selection +// ----------------------------------------------------------------------------- +//! Backend algorithms for @ref DeviceBatchedTopK. Both backends are launched through a single kernel symbol; which one +//! runs is decided per architecture by `policy_selector` below, whose result also drives the device-side agent +//! selection (via `current_policy`). +enum class topk_algorithm +{ + baseline, //!< worker-per-segment backend (single thread block per segment) + cluster, //!< thread-block-cluster backend (SM 9.0+) + unsupported //!< no backend can serve the request on the target architecture; dispatch returns cudaErrorNotSupported +}; + +#if _CCCL_HOSTED() +[[nodiscard]] inline ::std::ostream& operator<<(::std::ostream& os, topk_algorithm backend) +{ + switch (backend) + { + case topk_algorithm::baseline: + return os << "baseline"; + case topk_algorithm::cluster: + return os << "cluster"; + default: + return os << "unsupported"; + } +} +#endif // _CCCL_HOSTED() + +//! The tuning policy for all backends of @ref DeviceBatchedTopK. It carries the selected backend plus both backends' +//! sub-policies; the kernel instantiates only the arm named by @p backend (chosen device-side via `current_policy`). +//! +//! This is a regular type: `detail::dispatch_compute_cap` (and the `policy_selector` concept) require the selector's +//! result to be `::cuda::std::regular`, hence the equality/streaming operators below. +struct topk_policy +{ + topk_algorithm backend; //!< Backend the dispatch selected, i.e. the kernel arm that runs. + baseline_topk_policy baseline; //!< Sub-policy used when @p backend is @p topk_algorithm::baseline. + cluster_topk_policy cluster; //!< Sub-policy used when @p backend is @p topk_algorithm::cluster. + + _CCCL_HOST_DEVICE_API friend constexpr bool operator==(const topk_policy& lhs, const topk_policy& rhs) + { + return lhs.backend == rhs.backend && lhs.baseline == rhs.baseline && lhs.cluster == rhs.cluster; + } + + _CCCL_HOST_DEVICE_API friend constexpr bool operator!=(const topk_policy& lhs, const topk_policy& rhs) + { + return !(lhs == rhs); + } + +#if _CCCL_HOSTED() + friend ::std::ostream& operator<<(::std::ostream& os, const topk_policy& p) + { + return os << "topk_policy { .backend = " << p.backend << ", .baseline = " << p.baseline + << ", .cluster = " << p.cluster << " }"; + } +#endif // _CCCL_HOSTED() +}; + #if _CCCL_HAS_CONCEPTS() template -concept batched_topk_policy_selector = policy_selector; +concept topk_policy_selector = policy_selector; #endif // _CCCL_HAS_CONCEPTS() -struct policy_selector -{ - [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const -> batched_topk_policy - { - constexpr auto load_alg = BLOCK_LOAD_WARP_TRANSPOSE; - constexpr auto store_alg = BLOCK_STORE_WARP_TRANSPOSE; - constexpr auto scan_alg = BLOCK_SCAN_WARP_SCANS; - constexpr auto epilogue = epilogue_policy{16, load_alg, store_alg, scan_alg}; - return batched_topk_policy{ - {{ - worker_policy{256, 64, load_alg, store_alg, epilogue}, - worker_policy{256, 32, load_alg, store_alg, epilogue}, - worker_policy{256, 16, load_alg, store_alg, epilogue}, - worker_policy{256, 8, load_alg, store_alg, epilogue}, - worker_policy{256, 4, load_alg, store_alg, epilogue}, - worker_policy{128, 2, load_alg, store_alg, epilogue}, - }}, - multi_worker_policy{256, 64}}; - } -}; +// Crossover knobs (TODO: tune via SM100 benchmarks). +//! Clusters require SM 9.0+. +inline constexpr int cluster_min_cc_major = 9; +//! Smallest statically-known maximum segment size at which the cluster backend starts to win (measured on B200). This +//! is the backend crossover threshold and is intentionally part of the selector -- not a tunable policy field -- so +//! that tuning the cluster policy (e.g. its single-CTA threshold) does not silently shift which backend is chosen. +inline constexpr ::cuda::std::int64_t cluster_beneficial_min_segment_size = 8 * 1024; -template -struct policy_selector_from_types +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr bool cluster_capable([[maybe_unused]] ::cuda::compute_capability cc) { - [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const - -> batched_topk_policy - { - return policy_selector{}(cc); - } -}; - -#if _CCCL_HAS_CONCEPTS() -static_assert(batched_topk_policy_selector); -#endif // _CCCL_HAS_CONCEPTS() +#if _CCCL_HAS_DYNAMIC_CLUSTER_LAUNCH() + return cc >= ::cuda::compute_capability{cluster_min_cc_major, 0}; +#else // ^^^ dynamic cluster launches enabled ^^^ / vvv dynamic cluster launches disabled vvv + // The cluster backend launches with a runtime cluster width, which _CCCL_DISABLE_DYNAMIC_CLUSTER_LAUNCH compiles out; + // reporting no architecture as cluster-capable makes the selector fall back to baseline (or report unsupported). + return false; +#endif // _CCCL_HAS_DYNAMIC_CLUSTER_LAUNCH() +} } // namespace detail::batched_topk CUB_NAMESPACE_END diff --git a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_merge.cuh b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_merge.cuh index e44a9bdb..c16725b6 100644 --- a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_merge.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_merge.cuh @@ -20,9 +20,9 @@ #include #include -#include #include +#include #include #include #include @@ -43,10 +43,15 @@ struct MergePolicy [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator==(const MergePolicy& lhs, const MergePolicy& rhs) noexcept { + // gcc 8 folds comparisons of adjacent bool members within one expression into a BIT_FIELD_REF, which its + // constexpr evaluator cannot handle (ICE in cxx_eval_bit_field_ref, fixed in gcc 9). Keep each bool + // comparison in a separate statement to avoid the fold. + const bool same_bulk_copy_for_keys = lhs.use_bulk_copy_for_keys == rhs.use_bulk_copy_for_keys; + const bool same_bulk_copy_for_values = lhs.use_bulk_copy_for_values == rhs.use_bulk_copy_for_values; + const bool same_unroll = lhs.unroll == rhs.unroll; return lhs.threads_per_block == rhs.threads_per_block && lhs.items_per_thread == rhs.items_per_thread && lhs.load_modifier == rhs.load_modifier && lhs.store_algorithm == rhs.store_algorithm - && lhs.use_bulk_copy_for_keys == rhs.use_bulk_copy_for_keys - && lhs.use_bulk_copy_for_values == rhs.use_bulk_copy_for_values && lhs.unroll == rhs.unroll; + && same_bulk_copy_for_keys && same_bulk_copy_for_values && same_unroll; } [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool @@ -90,10 +95,10 @@ struct policy_selector [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const -> MergePolicy { - const int tune_type_size = key_size + value_size; - const int ipt_800_plus = nominal_4B_items_to_items(15, tune_type_size); - const bool can_bulk_keys = (key_size == key_align) && key_is_trivially_relocatable && key_iterators_are_contiguous - && key_iterator_value_types_are_the_same; + const int tune_type_size = key_size + value_size; + const int ipt_800_plus = nominal_4B_items_to_items(15, tune_type_size); + const bool can_bulk_keys = (key_size == key_align) && key_is_trivially_relocatable && key_iterators_are_contiguous + && key_iterator_value_types_are_the_same; const bool can_bulk_values = (value_size == value_align) && value_is_trivially_relocatable && value_iterators_are_contiguous && value_iterator_value_types_are_the_same; @@ -163,12 +168,12 @@ struct policy_selector_from_types return policy_selector{ int{sizeof(key_t)}, int{alignof(key_t)}, - THRUST_NS_QUALIFIER::is_trivially_relocatable_v, + ::cuda::is_trivially_copyable_v, THRUST_NS_QUALIFIER::is_contiguous_iterator_v && THRUST_NS_QUALIFIER::is_contiguous_iterator_v, ::cuda::std::is_same_v>, ::cuda::std::is_same_v ? 0 : int{sizeof(item_t)}, int{alignof(item_t)}, - THRUST_NS_QUALIFIER::is_trivially_relocatable_v, + ::cuda::is_trivially_copyable_v, THRUST_NS_QUALIFIER::is_contiguous_iterator_v && THRUST_NS_QUALIFIER::is_contiguous_iterator_v, ::cuda::std::is_same_v>, diff --git a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_reduce_by_key.cuh b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_reduce_by_key.cuh index 9f2ccb15..402e8cab 100644 --- a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_reduce_by_key.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_reduce_by_key.cuh @@ -99,13 +99,13 @@ struct ReduceByKeyPolicy ReduceByKeyLookbackPolicy lookback; //!< The policy for the reduce-by-key algorithm based on decoupled-lookback. Only //!< used when @p algorithm is @lookback. - [[nodiscard]] _CCCL_API friend constexpr bool + [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator==(const ReduceByKeyPolicy& lhs, const ReduceByKeyPolicy& rhs) noexcept { return lhs.algorithm == rhs.algorithm && lhs.lookback == rhs.lookback; } - [[nodiscard]] _CCCL_API friend constexpr bool + [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator!=(const ReduceByKeyPolicy& lhs, const ReduceByKeyPolicy& rhs) noexcept { return !(lhs == rhs); diff --git a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_rle_encode.cuh b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_rle_encode.cuh index 7f1c9d17..1ea80618 100644 --- a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_rle_encode.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_rle_encode.cuh @@ -100,13 +100,13 @@ struct RleEncodePolicy RleAlgorithm algorithm = RleAlgorithm::lookback; //!< The RLE-encode algorithm to use RleLookbackPolicy lookback; //!< The lookback policy - [[nodiscard]] _CCCL_API friend constexpr bool + [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator==(const RleEncodePolicy& lhs, const RleEncodePolicy& rhs) noexcept { return lhs.algorithm == rhs.algorithm && lhs.lookback == rhs.lookback; } - [[nodiscard]] _CCCL_API friend constexpr bool + [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator!=(const RleEncodePolicy& lhs, const RleEncodePolicy& rhs) noexcept { return !(lhs == rhs); diff --git a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_rle_non_trivial_runs.cuh b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_rle_non_trivial_runs.cuh index 495d708f..10d6353e 100644 --- a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_rle_non_trivial_runs.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_rle_non_trivial_runs.cuh @@ -104,13 +104,13 @@ struct RleNonTrivialRunsPolicy RleNonTrivialRunsLookbackPolicy lookback; //!< The policy for the non-trivial-runs algorithm based on //!< decoupled-lookback. Only used when @p algorithm is @lookback. - [[nodiscard]] _CCCL_API friend constexpr bool + [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator==(const RleNonTrivialRunsPolicy& lhs, const RleNonTrivialRunsPolicy& rhs) noexcept { return lhs.algorithm == rhs.algorithm && lhs.lookback == rhs.lookback; } - [[nodiscard]] _CCCL_API friend constexpr bool + [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator!=(const RleNonTrivialRunsPolicy& lhs, const RleNonTrivialRunsPolicy& rhs) noexcept { return !(lhs == rhs); diff --git a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_scan.cuh b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_scan.cuh index 93948713..756ac570 100644 --- a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_scan.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_scan.cuh @@ -187,12 +187,14 @@ struct ScanPolicy ScanLookbackPolicy lookback; //!< The look-back scan policy (used when algorithm is @p lookback, otherwise ignored) ScanLookaheadPolicy lookahead; //!< The lookahead scan policy (used when algorithm is @p lookahead, otherwise ignored) - [[nodiscard]] _CCCL_API friend constexpr bool operator==(const ScanPolicy& lhs, const ScanPolicy& rhs) noexcept + [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool + operator==(const ScanPolicy& lhs, const ScanPolicy& rhs) noexcept { return lhs.lookback == rhs.lookback && lhs.lookahead == rhs.lookahead && lhs.algorithm == rhs.algorithm; } - [[nodiscard]] _CCCL_API friend constexpr bool operator!=(const ScanPolicy& lhs, const ScanPolicy& rhs) noexcept + [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool + operator!=(const ScanPolicy& lhs, const ScanPolicy& rhs) noexcept { return !(lhs == rhs); } diff --git a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_scan_by_key.cuh b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_scan_by_key.cuh index 04905fc2..015aa095 100644 --- a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_scan_by_key.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_scan_by_key.cuh @@ -103,13 +103,13 @@ struct ScanByKeyPolicy ScanByKeyLookbackPolicy lookback; //!< The policy for the scan-by-key algorithm based on decoupled-lookback. Only used //!< when @p algorithm is @lookback. - [[nodiscard]] _CCCL_API friend constexpr bool + [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator==(const ScanByKeyPolicy& lhs, const ScanByKeyPolicy& rhs) noexcept { return lhs.algorithm == rhs.algorithm && lhs.lookback == rhs.lookback; } - [[nodiscard]] _CCCL_API friend constexpr bool + [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator!=(const ScanByKeyPolicy& lhs, const ScanByKeyPolicy& rhs) noexcept { return !(lhs == rhs); diff --git a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh index 6f730b7b..5b9773d4 100644 --- a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh @@ -575,13 +575,13 @@ struct policy_hub static constexpr int BLOCK_THREADS = 256; static constexpr int PARTITIONING_THRESHOLD = 500; using LargeSegmentPolicy = detail::agent_radix_sort_downsweep_policy< - BLOCK_THREADS, - 23, - DominantT, - BLOCK_LOAD_TRANSPOSE, - LOAD_DEFAULT, - RADIX_RANK_MEMOIZE, - BLOCK_SCAN_WARP_SCANS, + BLOCK_THREADS, + 23, + DominantT, + BLOCK_LOAD_TRANSPOSE, + LOAD_DEFAULT, + RADIX_RANK_MEMOIZE, + BLOCK_SCAN_WARP_SCANS, (sizeof(KeyT) > 1) ? 6 : 4>; static constexpr int ITEMS_PER_SMALL_THREAD = Nominal4BItemsToItems(9); @@ -606,13 +606,13 @@ struct policy_hub static constexpr int BLOCK_THREADS = 256; static constexpr int PARTITIONING_THRESHOLD = 500; using LargeSegmentPolicy = detail::agent_radix_sort_downsweep_policy< - BLOCK_THREADS, - 23, - DominantT, - BLOCK_LOAD_TRANSPOSE, - LOAD_DEFAULT, - RADIX_RANK_MEMOIZE, - BLOCK_SCAN_WARP_SCANS, + BLOCK_THREADS, + 23, + DominantT, + BLOCK_LOAD_TRANSPOSE, + LOAD_DEFAULT, + RADIX_RANK_MEMOIZE, + BLOCK_SCAN_WARP_SCANS, (sizeof(KeyT) > 1) ? 6 : 4>; static constexpr bool LARGE_ITEMS = sizeof(DominantT) > 4; diff --git a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_select_if.cuh b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_select_if.cuh index 05e64ff3..5aaa1aac 100644 --- a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_select_if.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_select_if.cuh @@ -18,6 +18,7 @@ #include #include #include +#include #include #include #include @@ -43,13 +44,15 @@ struct SelectLookbackPolicy CacheLoadModifier load_modifier; //!< The @ref CacheLoadModifier used for loading items from global memory BlockScanAlgorithm scan_algorithm; //!< The @ref BlockScanAlgorithm used for scanning LookbackDelayPolicy lookback_delay; //!< The policy configuring the delay used in decoupled lookback + detail::LoadPrefetch _load_prefetch = detail::LoadPrefetch::none; //!< Implementation detail; do not use directly [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator==(const SelectLookbackPolicy& lhs, const SelectLookbackPolicy& rhs) noexcept { return lhs.threads_per_block == rhs.threads_per_block && lhs.items_per_thread == rhs.items_per_thread && lhs.load_algorithm == rhs.load_algorithm && lhs.load_modifier == rhs.load_modifier - && lhs.scan_algorithm == rhs.scan_algorithm && lhs.lookback_delay == rhs.lookback_delay; + && lhs.scan_algorithm == rhs.scan_algorithm && lhs.lookback_delay == rhs.lookback_delay + && lhs._load_prefetch == rhs._load_prefetch; } [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool @@ -62,9 +65,10 @@ struct SelectLookbackPolicy friend ::std::ostream& operator<<(::std::ostream& os, const SelectLookbackPolicy& p) { return os - << "SelectLookbackPolicy { .threads_per_block = " << p.threads_per_block << ", .items_per_thread = " - << p.items_per_thread << ", .load_algorithm = " << p.load_algorithm << ", .load_modifier = " << p.load_modifier - << ", .scan_algorithm = " << p.scan_algorithm << ", .lookback_delay = " << p.lookback_delay << " }"; + << "SelectLookbackPolicy { .threads_per_block = " << p.threads_per_block + << ", .items_per_thread = " << p.items_per_thread << ", .load_algorithm = " << p.load_algorithm + << ", .load_modifier = " << p.load_modifier << ", .scan_algorithm = " << p.scan_algorithm + << ", .lookback_delay = " << p.lookback_delay << ", ._load_prefetch = " << p._load_prefetch << " }"; } #endif // _CCCL_HOSTED() }; @@ -102,12 +106,14 @@ struct SelectPolicy SelectLookbackPolicy lookback; //!< The policy for the selection algorithm based on decoupled-lookback. Only used when //!< @p algorithm is @lookback. - [[nodiscard]] _CCCL_API friend constexpr bool operator==(const SelectPolicy& lhs, const SelectPolicy& rhs) noexcept + [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool + operator==(const SelectPolicy& lhs, const SelectPolicy& rhs) noexcept { return lhs.algorithm == rhs.algorithm && lhs.lookback == rhs.lookback; } - [[nodiscard]] _CCCL_API friend constexpr bool operator!=(const SelectPolicy& lhs, const SelectPolicy& rhs) noexcept + [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool + operator!=(const SelectPolicy& lhs, const SelectPolicy& rhs) noexcept { return !(lhs == rhs); } @@ -190,13 +196,13 @@ struct PartitionPolicy PartitionLookbackPolicy lookback; //!< The policy for the partition algorithm based on decoupled-lookback. Only used //!< when algorithm is @lookback. - [[nodiscard]] _CCCL_API friend constexpr bool + [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator==(const PartitionPolicy& lhs, const PartitionPolicy& rhs) noexcept { return lhs.algorithm == rhs.algorithm && lhs.lookback == rhs.lookback; } - [[nodiscard]] _CCCL_API friend constexpr bool + [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator!=(const PartitionPolicy& lhs, const PartitionPolicy& rhs) noexcept { return !(lhs == rhs); diff --git a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_three_way_partition.cuh b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_three_way_partition.cuh index ae3d4d1c..2099f08d 100644 --- a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_three_way_partition.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_three_way_partition.cuh @@ -98,13 +98,13 @@ struct ThreeWayPartitionPolicy ThreeWayPartitionLookbackPolicy lookback; //!< The policy for the three-way partition algorithm based on //!< decoupled-lookback. Only used when @p algorithm is @lookback. - [[nodiscard]] _CCCL_API friend constexpr bool + [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator==(const ThreeWayPartitionPolicy& lhs, const ThreeWayPartitionPolicy& rhs) noexcept { return lhs.algorithm == rhs.algorithm && lhs.lookback == rhs.lookback; } - [[nodiscard]] _CCCL_API friend constexpr bool + [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator!=(const ThreeWayPartitionPolicy& lhs, const ThreeWayPartitionPolicy& rhs) noexcept { return !(lhs == rhs); diff --git a/ex_engine/fla_kernels/gated_delta_rule/__init__.py b/ex_engine/fla_kernels/gated_delta_rule/__init__.py new file mode 100644 index 00000000..7e65713b --- /dev/null +++ b/ex_engine/fla_kernels/gated_delta_rule/__init__.py @@ -0,0 +1,17 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +from .chunk import chunk_gated_delta_rule, chunk_gdn +from .fused_recurrent import fused_recurrent_gated_delta_rule, fused_recurrent_gdn +from .naive import naive_chunk_gated_delta_rule, naive_recurrent_gated_delta_rule + +__all__ = [ + "chunk_gated_delta_rule", "chunk_gdn", + "fused_recurrent_gated_delta_rule", "fused_recurrent_gdn", + "naive_chunk_gated_delta_rule", + "naive_recurrent_gated_delta_rule", +] diff --git a/ex_engine/fla_kernels/gated_delta_rule/chunk.py b/ex_engine/fla_kernels/gated_delta_rule/chunk.py new file mode 100644 index 00000000..576278e9 --- /dev/null +++ b/ex_engine/fla_kernels/gated_delta_rule/chunk.py @@ -0,0 +1,591 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +import warnings + +import torch + +from fla.modules.l2norm import l2norm_bwd, l2norm_fwd +from fla.ops.backends import dispatch +from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h +from fla.ops.common.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o +from fla.ops.common.gate import fused_beta_sigmoid, fused_beta_sigmoid_bwd +from fla.ops.cp import FLACPContext +from fla.ops.cp.chunk_delta_h import ( + chunk_gated_delta_rule_bwd_dhu_pre_process, + chunk_gated_delta_rule_fwd_h_pre_process, + compress_h0, + expand_h0, +) +from fla.ops.gated_delta_rule.chunk_fwd import chunk_gated_delta_rule_fwd_intra +from fla.ops.gated_delta_rule.gate import gdn_gate_bwd, gdn_gate_chunk_cumsum +from fla.ops.gated_delta_rule.wy_fast import prepare_wy_repr_bwd, recompute_w_u_fwd +from fla.ops.utils import chunk_local_cumsum +from fla.ops.utils.constant import RCP_LN2 +from fla.ops.utils.index import prepare_chunk_indices +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + + +def chunk_gated_delta_rule_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + chunk_indices: torch.LongTensor | None = None, + use_gate_in_kernel: bool = False, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + chunk_size: int = 64, +): + g_input = g if use_gate_in_kernel else None + if use_gate_in_kernel: + g = gdn_gate_chunk_cumsum( + g=g, + A_log=A_log, + chunk_size=chunk_size, + scale=RCP_LN2, + dt_bias=dt_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + else: + g = chunk_local_cumsum( + g, + chunk_size=chunk_size, + scale=RCP_LN2, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + # obtain WY representation. u is actually the new v. + # fused kkt + solve_tril + recompute_w_u + w, u, A = chunk_gated_delta_rule_fwd_intra( + k=k, + v=v, + g=g, + beta=beta, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + + if cp_context is not None: + initial_state = chunk_gated_delta_rule_fwd_h_pre_process( + k=k, + w=w, + u=u, + g=g, + cu_seqlens=cu_seqlens, + initial_state=initial_state, + context=cp_context, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + g=g, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + + if cp_context is not None: + initial_state = compress_h0(initial_state, context=cp_context) + + o = chunk_fwd_o( + q=q, + k=k, + v=v_new, + h=h, + g=g, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + return g, o, A, final_state, initial_state, g_input + + +def chunk_gated_delta_rule_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + chunk_indices: torch.LongTensor | None = None, + use_gate_in_kernel: bool = False, + g_input: torch.Tensor | None = None, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + chunk_size: int = 64, +): + w, u = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + g=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + + if cp_context is not None: + initial_state = expand_h0(initial_state, context=cp_context) + + h, v_new, _ = chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + g=g, + initial_state=initial_state, + output_final_state=False, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + dv = chunk_bwd_dv_local( + q=q, + k=k, + g=g, + do=do, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + + if cp_context is not None: + # initial_state is None in the CP mode + # We only need to compute dht of current rank and pass it to the backward kernel + dht, initial_state = chunk_gated_delta_rule_bwd_dhu_pre_process( + q=q, + k=k, + w=w, + do=do, + dv=dv, + g=g, + scale=scale, + cu_seqlens=cu_seqlens, + dht=dht, + initial_state=initial_state, + context=cp_context, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + + dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( + q=q, + k=k, + w=w, + g=g, + h0=initial_state, + dht=dht, + do=do, + dv=dv, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + dq, dk, dw, dg = chunk_bwd_dqkwg( + q=q, + k=k, + v=v_new, + w=w, + g=g, + h=h, + dv=dv, + do=do, + dh=dh, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + dk2, dv, db, dg2 = prepare_wy_repr_bwd( + k=k, + v=v, + beta=beta, + g=g, + A=A, + dw=dw, + du=dv, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + dk.add_(dk2) + dg.add_(dg2) + dg = chunk_local_cumsum(dg, chunk_size=chunk_size, reverse=True, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices) + dA_log, ddt_bias = None, None + if use_gate_in_kernel: + dg, dA_log, ddt_bias = gdn_gate_bwd(g=g_input, A_log=A_log, dt_bias=dt_bias, dyg=dg) + return dq, dk, dv, db, dg, dh0, dA_log, ddt_bias + + +class ChunkGatedDeltaRuleFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + use_qk_l2norm_in_kernel: bool = False, + use_gate_in_kernel: bool = False, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + cp_context: FLACPContext | None = None, + chunk_size: int = 64, + ): + q_rstd, k_rstd = None, None + if use_qk_l2norm_in_kernel: + q, q_rstd = l2norm_fwd(q) + k, k_rstd = l2norm_fwd(k) + + beta_raw = beta + if use_beta_sigmoid_in_kernel: + beta = fused_beta_sigmoid(beta_raw, scale=2.0 if allow_neg_eigval else 1.0) + + chunk_indices = None + if cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) + g, o, A, final_state, initial_state, g_input = chunk_gated_delta_rule_fwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + cp_context=cp_context, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + use_gate_in_kernel=use_gate_in_kernel, + A_log=A_log, + dt_bias=dt_bias, + chunk_size=chunk_size, + ) + ctx.save_for_backward( + q, + q_rstd, + k, + k_rstd, + v, + g, + beta_raw, + beta, + A, + initial_state, + cu_seqlens, + chunk_indices, + g_input, + A_log, + dt_bias, + ) + ctx.scale = scale + ctx.chunk_size = chunk_size + ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel + ctx.use_beta_sigmoid_in_kernel = use_beta_sigmoid_in_kernel + ctx.allow_neg_eigval = allow_neg_eigval + ctx.cp_context = cp_context + ctx.state_v_first = state_v_first + ctx.use_gate_in_kernel = use_gate_in_kernel + return o.to(q.dtype), final_state + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward( + ctx, + do: torch.Tensor, + dht: torch.Tensor, + ): + ( + q, + q_rstd, + k, + k_rstd, + v, + g, + beta_raw, + beta, + A, + initial_state, + cu_seqlens, + chunk_indices, + g_input, + A_log, + dt_bias, + ) = ctx.saved_tensors + dq, dk, dv, db, dg, dh0, dA_log, ddt_bias = chunk_gated_delta_rule_bwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A=A, + scale=ctx.scale, + initial_state=initial_state, + do=do, + dht=dht, + cu_seqlens=cu_seqlens, + cp_context=ctx.cp_context, + chunk_indices=chunk_indices, + state_v_first=ctx.state_v_first, + use_gate_in_kernel=ctx.use_gate_in_kernel, + g_input=g_input, + A_log=A_log, + dt_bias=dt_bias, + chunk_size=ctx.chunk_size, + ) + if ctx.use_qk_l2norm_in_kernel: + dq = l2norm_bwd(q, q_rstd, dq) + dk = l2norm_bwd(k, k_rstd, dk) + if ctx.use_beta_sigmoid_in_kernel: + db = fused_beta_sigmoid_bwd(beta_raw, db, scale=2.0 if ctx.allow_neg_eigval else 1.0) + return ( + dq.to(q), dk.to(k), dv.to(v), dg.to(g), db.to(beta_raw), + None, dh0, None, None, None, None, None, None, dA_log, ddt_bias, + None, None, None, None, + ) + + +@dispatch('gated_delta_rule') +@torch.compiler.disable +def chunk_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + **kwargs, +): + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, HV, V]`. + GVA (Grouped Value Attention) is applied if `HV > H`, where `HV` must be divisible by `H`. + g (torch.Tensor): + (forget) gating tensor of shape `[B, T, HV]`. + When `use_gate_in_kernel=False` (default), `g` should be in log space (pre-computed decay). + When `use_gate_in_kernel=True`, `g` is the raw input before gate activation; + the kernel fuses `-exp(A_log) * softplus(g + dt_bias)` + chunk cumsum internally. + beta (torch.Tensor): + betas of shape `[B, T, HV]`. + scale (Optional[float]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, HV, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`. + use_qk_l2norm_in_kernel (bool): + Whether to apply L2norm to the q/k tensor internally. Default: `False`. + use_gate_in_kernel (bool): + Whether to compute the log-space GDN decay internally. + When `True`, the passed `g` is the raw input, and `A_log` must be provided. + The kernel fuses gate activation + chunk cumsum in a single pass. + Default: `False`. + A_log (Optional[torch.Tensor]): + Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`. + dt_bias (Optional[torch.Tensor]): + Bias added to `g` before activation, of shape `[HV]`. + Only used when `use_gate_in_kernel=True`. + use_beta_sigmoid_in_kernel (bool): + Whether to apply `torch.sigmoid(beta)` before launching the chunk kernel. + - If `True`, the passed `beta` acts as the raw beta logits. + - If `False`, `beta` is expected to already be in post-sigmoid space. + Default: `False`. + allow_neg_eigval (bool): + Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`. + Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case + the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`. + state_v_first (Optional[bool]): + Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + cp_context (Optional[FLACPContext]): + Context parallel context for distributed training across multiple devices. + When provided, `initial_state` and `output_final_state` are not supported, + and `cu_seqlens` will be overridden by the context. Default: `None`. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, HV, V]`. + final_state (torch.Tensor): + Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.gated_delta_rule import chunk_gated_delta_rule + # inputs with equal lengths + >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512 + >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') + >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1) + >>> v = torch.randn(B, T, HV, V, dtype=torch.bfloat16, device='cuda') + >>> beta = torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda').sigmoid() + >>> g = F.logsigmoid(torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda')) + >>> h0 = torch.randn(B, HV, K, V, dtype=torch.bfloat16, device='cuda') + >>> o, ht = chunk_gated_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o, ht = chunk_gated_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + if 'transpose_state_layout' in kwargs: + if state_v_first: + raise ValueError("Cannot pass both `state_v_first` and the deprecated `transpose_state_layout`.") + warnings.warn( + "`transpose_state_layout` is deprecated and renamed to `state_v_first`.", + DeprecationWarning, + stacklevel=2, + ) + state_v_first = kwargs.pop('transpose_state_layout') + + # Validate head dimensions + if q.shape[2] != k.shape[2]: + raise ValueError( + f"q and k must have the same number of heads, " + f"but got q.shape[2]={q.shape[2]} and k.shape[2]={k.shape[2]}" + ) + H, HV = q.shape[2], v.shape[2] + if HV % H != 0: + raise ValueError( + f"For GVA, num_v_heads (HV={HV}) must be evenly divisible by " + f"num_heads (H={H}), but got HV % H = {HV % H}" + ) + + if 'head_first' in kwargs: + raise DeprecationWarning( + "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", + ) + + chunk_size = kwargs.pop('chunk_size', 64) + if chunk_size not in (16, 32, 64): + raise ValueError(f"`chunk_size` must be 16, 32, or 64 for Gated Delta Rule, got {chunk_size}.") + + if cp_context is not None: + assert initial_state is None, "Initial state is not supported for CP" + assert output_final_state is False, "Output final state is not supported for CP" + assert cp_context.cu_seqlens is not None, "cu_seqlens is required for CP" + cu_seqlens = cp_context.cu_seqlens + if cp_context.cu_seqlens_cpu is not None: + cu_seqlens_cpu = cp_context.cu_seqlens_cpu + + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + use_gate_in_kernel = kwargs.get('use_gate_in_kernel', False) + A_log = kwargs.get('A_log') + dt_bias = kwargs.get('dt_bias') + if use_gate_in_kernel: + assert A_log is not None, "A_log must be provided when use_gate_in_kernel=True." + if allow_neg_eigval and not use_beta_sigmoid_in_kernel: + raise ValueError("`allow_neg_eigval=True` requires `use_beta_sigmoid_in_kernel=True`.") + + if scale is None: + scale = k.shape[-1] ** -0.5 + o, final_state = ChunkGatedDeltaRuleFunction.apply( + q, + k, + v, + g, + beta, + scale, + initial_state, + output_final_state, + state_v_first, + cu_seqlens, + cu_seqlens_cpu, + use_qk_l2norm_in_kernel, + use_gate_in_kernel, + A_log, + dt_bias, + use_beta_sigmoid_in_kernel, + allow_neg_eigval, + cp_context, + chunk_size, + ) + return o, final_state + + +chunk_gdn = chunk_gated_delta_rule diff --git a/ex_engine/fla_kernels/gated_delta_rule/chunk_fwd.py b/ex_engine/fla_kernels/gated_delta_rule/chunk_fwd.py new file mode 100644 index 00000000..76824219 --- /dev/null +++ b/ex_engine/fla_kernels/gated_delta_rule/chunk_fwd.py @@ -0,0 +1,428 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +import torch +import triton +import triton.language as tl + +from fla.ops.backends import dispatch +from fla.ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd +from fla.ops.gated_delta_rule.wy_fast import recompute_w_u_fwd +from fla.ops.utils import prepare_chunk_indices, solve_tril +from fla.ops.utils.cache import fla_cache_autotune +from fla.ops.utils.op import exp2 +from fla.utils import IS_INTEL, IS_TF32_SUPPORTED, autotune_cache_kwargs + +if IS_TF32_SUPPORTED: + SOLVE_TRIL_DOT_PRECISION = tl.constexpr('tf32') +else: + SOLVE_TRIL_DOT_PRECISION = tl.constexpr('ieee') + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@fla_cache_autotune( + configs=[ + triton.Config({'BK': BK}, num_warps=num_warps) + for BK in [32, 64] + for num_warps in [1, 2, 4] + ], + key=['H', 'HV', 'K', 'BC'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_gated_delta_rule_fwd_kkt_solve_kernel( + k, + g, + beta, + A, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + USE_G: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + """ + Fused kernel: compute beta * K @ K^T (lower triangular) + solve_tril (I+A)^{-1} in one pass. + + This kernel fuses chunk_scaled_dot_kkt_fwd and solve_tril into a single kernel, + avoiding the HBM round-trip for the intermediate A matrix. + + Steps: + 1. Compute all 10 lower-triangular [BC, BC] blocks of beta * K @ K^T in registers + 2. Apply gate and beta scaling + 3. Forward substitution on diagonal blocks + 4. Block merge to get full (I+A)^{-1} + 5. Write result to A (output) + """ + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + i_b, i_h = i_bh // HV, i_bh % HV + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT >= T: + return + + i_tc0 = i_t * BT + i_tc1 = i_t * BT + BC + i_tc2 = i_t * BT + 2 * BC + i_tc3 = i_t * BT + 3 * BC + + k += (bos * H + i_h // (HV // H)) * K + A += (bos * HV + i_h) * BT + + o_i = tl.arange(0, BC) + m_tc0 = (i_tc0 + o_i) < T + m_tc1 = (i_tc1 + o_i) < T + m_tc2 = (i_tc2 + o_i) < T + m_tc3 = (i_tc3 + o_i) < T + + # load beta for each sub-chunk + p_b0 = beta + bos * HV + i_h + (i_tc0 + o_i) * HV + p_b1 = beta + bos * HV + i_h + (i_tc1 + o_i) * HV + p_b2 = beta + bos * HV + i_h + (i_tc2 + o_i) * HV + p_b3 = beta + bos * HV + i_h + (i_tc3 + o_i) * HV + b_b0 = tl.load(p_b0, mask=m_tc0, other=0.0).to(tl.float32) + b_b1 = tl.load(p_b1, mask=m_tc1, other=0.0).to(tl.float32) + b_b2 = tl.load(p_b2, mask=m_tc2, other=0.0).to(tl.float32) + b_b3 = tl.load(p_b3, mask=m_tc3, other=0.0).to(tl.float32) + + # load gate if used + if USE_G: + p_g0 = g + bos * HV + i_h + (i_tc0 + o_i) * HV + p_g1 = g + bos * HV + i_h + (i_tc1 + o_i) * HV + p_g2 = g + bos * HV + i_h + (i_tc2 + o_i) * HV + p_g3 = g + bos * HV + i_h + (i_tc3 + o_i) * HV + + b_g0 = tl.load(p_g0, mask=m_tc0, other=0.0).to(tl.float32) + b_g1 = tl.load(p_g1, mask=m_tc1, other=0.0).to(tl.float32) + b_g2 = tl.load(p_g2, mask=m_tc2, other=0.0).to(tl.float32) + b_g3 = tl.load(p_g3, mask=m_tc3, other=0.0).to(tl.float32) + + ############################################################################ + # Step 1: compute all 10 lower-triangular [BC, BC] blocks of K @ K^T + ############################################################################ + + # 4 diagonal blocks + b_A00 = tl.zeros([BC, BC], dtype=tl.float32) + b_A11 = tl.zeros([BC, BC], dtype=tl.float32) + b_A22 = tl.zeros([BC, BC], dtype=tl.float32) + b_A33 = tl.zeros([BC, BC], dtype=tl.float32) + + # 6 off-diagonal blocks + b_A10 = tl.zeros([BC, BC], dtype=tl.float32) + b_A20 = tl.zeros([BC, BC], dtype=tl.float32) + b_A21 = tl.zeros([BC, BC], dtype=tl.float32) + b_A30 = tl.zeros([BC, BC], dtype=tl.float32) + b_A31 = tl.zeros([BC, BC], dtype=tl.float32) + b_A32 = tl.zeros([BC, BC], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + p_k0 = k + (i_tc0 + o_i)[:, None] * (H*K) + o_k[None, :] + b_k0 = tl.load(p_k0, mask=m_tc0[:, None] & (o_k[None, :] < K), other=0.0) + # diagonal block 0 + b_A00 += tl.dot(b_k0, tl.trans(b_k0)) + + if i_tc1 < T: + p_k1 = k + (i_tc1 + o_i)[:, None] * (H*K) + o_k[None, :] + b_k1 = tl.load(p_k1, mask=m_tc1[:, None] & (o_k[None, :] < K), other=0.0) + # diagonal block 1 + b_A11 += tl.dot(b_k1, tl.trans(b_k1)) + # off-diagonal (1,0) + b_A10 += tl.dot(b_k1, tl.trans(b_k0)) + + if i_tc2 < T: + p_k2 = k + (i_tc2 + o_i)[:, None] * (H*K) + o_k[None, :] + b_k2 = tl.load(p_k2, mask=m_tc2[:, None] & (o_k[None, :] < K), other=0.0) + # diagonal block 2 + b_A22 += tl.dot(b_k2, tl.trans(b_k2)) + # off-diagonal (2,0), (2,1) + b_A20 += tl.dot(b_k2, tl.trans(b_k0)) + b_A21 += tl.dot(b_k2, tl.trans(b_k1)) + + if i_tc3 < T: + p_k3 = k + (i_tc3 + o_i)[:, None] * (H*K) + o_k[None, :] + b_k3 = tl.load(p_k3, mask=m_tc3[:, None] & (o_k[None, :] < K), other=0.0) + # diagonal block 3 + b_A33 += tl.dot(b_k3, tl.trans(b_k3)) + # off-diagonal (3,0), (3,1), (3,2) + b_A30 += tl.dot(b_k3, tl.trans(b_k0)) + b_A31 += tl.dot(b_k3, tl.trans(b_k1)) + b_A32 += tl.dot(b_k3, tl.trans(b_k2)) + + ############################################################################ + # Step 2: apply gate and beta scaling + ############################################################################ + + # apply gate, beta scaling, and masking + # m_d: strictly lower triangular mask for diagonal blocks + # m_tc: boundary mask to prevent NaN from 0 * inf (IEEE 754) when + # out-of-bounds g loads as 0 via boundary_check and exp2(0 - g_inbounds) overflows + m_d = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + if USE_G: + b_A00 *= tl.where(m_d & m_tc0[:, None] & m_tc0[None, :], exp2(b_g0[:, None] - b_g0[None, :]), 0.) + b_A11 *= tl.where(m_d & m_tc1[:, None] & m_tc1[None, :], exp2(b_g1[:, None] - b_g1[None, :]), 0.) + b_A22 *= tl.where(m_d & m_tc2[:, None] & m_tc2[None, :], exp2(b_g2[:, None] - b_g2[None, :]), 0.) + b_A33 *= tl.where(m_d & m_tc3[:, None] & m_tc3[None, :], exp2(b_g3[:, None] - b_g3[None, :]), 0.) + + b_A10 *= tl.where(m_tc1[:, None] & m_tc0[None, :], exp2(b_g1[:, None] - b_g0[None, :]), 0.) + b_A20 *= tl.where(m_tc2[:, None] & m_tc0[None, :], exp2(b_g2[:, None] - b_g0[None, :]), 0.) + b_A21 *= tl.where(m_tc2[:, None] & m_tc1[None, :], exp2(b_g2[:, None] - b_g1[None, :]), 0.) + b_A30 *= tl.where(m_tc3[:, None] & m_tc0[None, :], exp2(b_g3[:, None] - b_g0[None, :]), 0.) + b_A31 *= tl.where(m_tc3[:, None] & m_tc1[None, :], exp2(b_g3[:, None] - b_g1[None, :]), 0.) + b_A32 *= tl.where(m_tc3[:, None] & m_tc2[None, :], exp2(b_g3[:, None] - b_g2[None, :]), 0.) + else: + b_A00 = tl.where(m_d, b_A00, 0.) + b_A11 = tl.where(m_d, b_A11, 0.) + b_A22 = tl.where(m_d, b_A22, 0.) + b_A33 = tl.where(m_d, b_A33, 0.) + + # diagonal blocks: scaled by beta + b_A00 = b_A00 * b_b0[:, None] + b_A11 = b_A11 * b_b1[:, None] + b_A22 = b_A22 * b_b2[:, None] + b_A33 = b_A33 * b_b3[:, None] + + # off-diagonal blocks: full block, scaled by beta + b_A10 = b_A10 * b_b1[:, None] + b_A20 = b_A20 * b_b2[:, None] + b_A21 = b_A21 * b_b2[:, None] + b_A30 = b_A30 * b_b3[:, None] + b_A31 = b_A31 * b_b3[:, None] + b_A32 = b_A32 * b_b3[:, None] + + ############################################################################ + # Step 3: forward substitution on diagonal blocks -> (I + A_diag)^{-1} + # + # Same algorithm as solve_tril, but rows are extracted from in-register + # [BC, BC] tensor via tl.sum(tl.where(mask, tensor, 0), 0) instead of + # tl.load from HBM. + ############################################################################ + + b_Ai00 = -b_A00 + b_Ai11 = -b_A11 + b_Ai22 = -b_A22 + b_Ai33 = -b_A33 + + for i in range(2, min(BC, T - i_tc0)): + b_a00 = tl.sum(tl.where((o_i == i)[:, None], -b_A00, 0.), 0) + b_a00 = tl.where(o_i < i, b_a00, 0.) + b_a00 = b_a00 + tl.sum(b_a00[:, None] * b_Ai00, 0) + b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00) + for i in range(2, min(BC, T - i_tc1)): + b_a11 = tl.sum(tl.where((o_i == i)[:, None], -b_A11, 0.), 0) + b_a11 = tl.where(o_i < i, b_a11, 0.) + b_a11 = b_a11 + tl.sum(b_a11[:, None] * b_Ai11, 0) + b_Ai11 = tl.where((o_i == i)[:, None], b_a11, b_Ai11) + for i in range(2, min(BC, T - i_tc2)): + b_a22 = tl.sum(tl.where((o_i == i)[:, None], -b_A22, 0.), 0) + b_a22 = tl.where(o_i < i, b_a22, 0.) + b_a22 = b_a22 + tl.sum(b_a22[:, None] * b_Ai22, 0) + b_Ai22 = tl.where((o_i == i)[:, None], b_a22, b_Ai22) + for i in range(2, min(BC, T - i_tc3)): + b_a33 = tl.sum(tl.where((o_i == i)[:, None], -b_A33, 0.), 0) + b_a33 = tl.where(o_i < i, b_a33, 0.) + b_a33 = b_a33 + tl.sum(b_a33[:, None] * b_Ai33, 0) + b_Ai33 = tl.where((o_i == i)[:, None], b_a33, b_Ai33) + + b_Ai00 += m_I + b_Ai11 += m_I + b_Ai22 += m_I + b_Ai33 += m_I + + ############################################################################ + # Step 4: block merge -> full (I + A)^{-1} + ############################################################################ + + b_Ai10 = -tl.dot( + tl.dot(b_Ai11, b_A10, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai00, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai21 = -tl.dot( + tl.dot(b_Ai22, b_A21, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai11, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai32 = -tl.dot( + tl.dot(b_Ai33, b_A32, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai22, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + + b_Ai20 = -tl.dot( + b_Ai22, + tl.dot(b_A20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_A21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + b_Ai31 = -tl.dot( + b_Ai33, + tl.dot(b_A31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_A32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + b_Ai30 = -tl.dot( + b_Ai33, + tl.dot(b_A30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_A31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_A32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + + ############################################################################ + # Step 5: store full (I + A)^{-1} to output A + ############################################################################ + + p_A00 = A + (i_tc0 + o_i)[:, None] * (HV*BT) + o_i[None, :] + p_A10 = A + (i_tc1 + o_i)[:, None] * (HV*BT) + o_i[None, :] + p_A11 = A + (i_tc1 + o_i)[:, None] * (HV*BT) + (BC + o_i)[None, :] + p_A20 = A + (i_tc2 + o_i)[:, None] * (HV*BT) + o_i[None, :] + p_A21 = A + (i_tc2 + o_i)[:, None] * (HV*BT) + (BC + o_i)[None, :] + p_A22 = A + (i_tc2 + o_i)[:, None] * (HV*BT) + (2*BC + o_i)[None, :] + p_A30 = A + (i_tc3 + o_i)[:, None] * (HV*BT) + o_i[None, :] + p_A31 = A + (i_tc3 + o_i)[:, None] * (HV*BT) + (BC + o_i)[None, :] + p_A32 = A + (i_tc3 + o_i)[:, None] * (HV*BT) + (2*BC + o_i)[None, :] + p_A33 = A + (i_tc3 + o_i)[:, None] * (HV*BT) + (3*BC + o_i)[None, :] + + m_A0 = m_tc0[:, None] & (o_i[None, :] < BT) + m_A1 = m_tc1[:, None] & (o_i[None, :] < BT) + m_A2 = m_tc2[:, None] & (o_i[None, :] < BT) + m_A3 = m_tc3[:, None] & (o_i[None, :] < BT) + m_A11 = m_tc1[:, None] & ((BC + o_i)[None, :] < BT) + m_A21 = m_tc2[:, None] & ((BC + o_i)[None, :] < BT) + m_A22 = m_tc2[:, None] & ((2*BC + o_i)[None, :] < BT) + m_A31 = m_tc3[:, None] & ((BC + o_i)[None, :] < BT) + m_A32 = m_tc3[:, None] & ((2*BC + o_i)[None, :] < BT) + m_A33 = m_tc3[:, None] & ((3*BC + o_i)[None, :] < BT) + + tl.store(p_A00, b_Ai00.to(A.dtype.element_ty), mask=m_A0) + tl.store(p_A10, b_Ai10.to(A.dtype.element_ty), mask=m_A1) + tl.store(p_A11, b_Ai11.to(A.dtype.element_ty), mask=m_A11) + tl.store(p_A20, b_Ai20.to(A.dtype.element_ty), mask=m_A2) + tl.store(p_A21, b_Ai21.to(A.dtype.element_ty), mask=m_A21) + tl.store(p_A22, b_Ai22.to(A.dtype.element_ty), mask=m_A22) + tl.store(p_A30, b_Ai30.to(A.dtype.element_ty), mask=m_A3) + tl.store(p_A31, b_Ai31.to(A.dtype.element_ty), mask=m_A31) + tl.store(p_A32, b_Ai32.to(A.dtype.element_ty), mask=m_A32) + tl.store(p_A33, b_Ai33.to(A.dtype.element_ty), mask=m_A33) + + +@dispatch('gated_delta_rule') +def chunk_gated_delta_rule_fwd_intra( + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + r""" + GDN intra-chunk forward: fused or unfused kkt + solve_tril + recompute_w_u. + + For ``chunk_size == 64``, this uses the fused kkt + solve_tril path. For + other supported chunk sizes, it computes the mathematically equivalent + representation with ``chunk_scaled_dot_kkt_fwd`` followed by ``solve_tril``. + + Args: + k (torch.Tensor): + The key tensor of shape `[B, T, H, K]`. + v (torch.Tensor): + The value tensor of shape `[B, T, HV, V]`. + g (torch.Tensor): + The cumulative sum of the gate tensor of shape `[B, T, HV]`. Default: `None`. + beta (torch.Tensor): + The beta tensor of shape `[B, T, HV]`. + cu_seqlens (torch.LongTensor): + The cumulative sequence lengths. Default: `None`. + chunk_size (int): + The chunk size. Default: 64. + chunk_indices (torch.LongTensor): + Precomputed chunk indices. Default: `None`. + + Returns: + w (torch.Tensor): shape `[B, T, HV, K]` + u (torch.Tensor): shape `[B, T, HV, V]` + A (torch.Tensor): shape `[B, T, HV, BT]`, the solved (I+A)^{-1} matrix + """ + if chunk_size not in (16, 32, 64): + raise ValueError(f"`chunk_size` must be 16, 32, or 64, got {chunk_size}.") + + B, T, H, K, HV = *k.shape, beta.shape[2] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + + # The fused kernel keeps ten [BC, BC] fp32 accumulators live across the K loop. + # That fits NVIDIA's register file but spills on Intel GPUs, where the unfused + # two-kernel path measures 2.3-3.0x faster despite the extra HBM round-trip. + if BT == 64 and not IS_INTEL: + # Step 1: fused kkt + solve_tril + BC = 16 + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + A = torch.zeros(B, T, HV, BT, device=k.device, dtype=k.dtype) + chunk_gated_delta_rule_fwd_kkt_solve_kernel[(NT, B * HV)]( + k=k, + g=g, + beta=beta, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + HV=HV, + K=K, + BT=BT, + BC=BC, + ) + else: + # Step 1: mathematically equivalent unfused kkt + solve_tril + A = chunk_scaled_dot_kkt_fwd( + k=k, + g=g, + beta=beta, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=BT, + output_dtype=torch.float32, + ) + A = solve_tril( + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + output_dtype=k.dtype, + ) + + # Step 2: recompute_w_u + w, u = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + g=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + return w, u, A diff --git a/ex_engine/fla_kernels/gated_delta_rule/fused_recurrent.py b/ex_engine/fla_kernels/gated_delta_rule/fused_recurrent.py new file mode 100644 index 00000000..0207dc9e --- /dev/null +++ b/ex_engine/fla_kernels/gated_delta_rule/fused_recurrent.py @@ -0,0 +1,478 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +import warnings + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.op import exp +from fla.ops.utils.softplus import softplus +from fla.utils import input_guard + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'USE_GK': lambda args: args['gk'] is not None, + 'USE_GV': lambda args: args['gv'] is not None, + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, + 'USE_GATE_IN_KERNEL': lambda args: args['A_log'] is not None, + 'HAS_DT_BIAS': lambda args: args['dt_bias'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def fused_recurrent_gated_delta_rule_fwd_kernel( + q, + k, + v, + g, + gk, + gv, + beta, + A_log, + dt_bias, + o, + h0, + ht, + cu_seqlens, + scale, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_GV: tl.constexpr, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + IS_BETA_HEADWISE: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + STATE_V_FIRST: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_GATE_IN_KERNEL: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + APPLY_BETA_SIGMOID: tl.constexpr, + ALLOW_NEG_EIGVAL: tl.constexpr, +): + pid = tl.program_id(0) + NV = tl.cdiv(V, BV) + i_v, i_nh = pid % NV, (pid // NV).to(tl.int64) + i_n, i_hv = i_nh // HV, i_nh % HV + i_h = i_hv // (HV // H) + + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + + p_q = q + (bos * H + i_h) * K + o_k + p_k = k + (bos * H + i_h) * K + o_k + p_v = v + (bos * HV + i_hv) * V + o_v + if USE_G: + p_g = g + bos * HV + i_hv + if USE_GK: + p_gk = gk + (bos * HV + i_hv) * K + o_k + if USE_GV: + p_gv = gv + (bos * HV + i_hv) * V + o_v + if IS_BETA_HEADWISE: + p_beta = beta + bos * HV + i_hv + else: + p_beta = beta + (bos * HV + i_hv) * V + o_v + + p_o = o + (bos * HV + i_hv) * V + o_v + + mask_k = o_k < K + mask_v = o_v < V + if STATE_V_FIRST: + mask_h = mask_v[:, None] & mask_k[None, :] + else: + mask_h = mask_k[:, None] & mask_v[None, :] + + if STATE_V_FIRST: + b_h = tl.zeros([BV, BK], dtype=tl.float32) + else: + b_h = tl.zeros([BK, BV], dtype=tl.float32) + if USE_INITIAL_STATE: + if STATE_V_FIRST: + p_h0 = h0 + i_nh * K*V + o_v[:, None] * K + o_k[None, :] + else: + p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + for _ in tl.range(0, T): + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + b_q = b_q * scale + if IS_BETA_HEADWISE: + b_beta = tl.load(p_beta).to(tl.float32) + else: + b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) + if APPLY_BETA_SIGMOID: + b_beta = tl.sigmoid(b_beta) + if ALLOW_NEG_EIGVAL: + b_beta = b_beta * 2 + + if USE_G: + b_g = tl.load(p_g).to(tl.float32) + if USE_GATE_IN_KERNEL: + b_A = tl.load(A_log + i_hv).to(tl.float32) + if HAS_DT_BIAS: + b_g = b_g + tl.load(dt_bias + i_hv).to(tl.float32) + b_g = -exp(b_A) * softplus(b_g) + b_h *= exp(b_g) + + if USE_GK: + b_gk = tl.load(p_gk).to(tl.float32) + if STATE_V_FIRST: + b_h *= exp(b_gk[None, :]) + else: + b_h *= exp(b_gk[:, None]) + + if USE_GV: + b_gv = tl.load(p_gv).to(tl.float32) + if STATE_V_FIRST: + b_h *= exp(b_gv[:, None]) + else: + b_h *= exp(b_gv[None, :]) + + if STATE_V_FIRST: + b_v = b_beta * (b_v - tl.sum(b_h * b_k[None, :], 1)) + b_h += b_v[:, None] * b_k[None, :] + b_o = tl.sum(b_h * b_q[None, :], 1) + else: + b_v = b_beta * (b_v - tl.sum(b_h * b_k[:, None], 0)) + b_h += b_k[:, None] * b_v + b_o = tl.sum(b_h * b_q[:, None], 0) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + + p_q += H*K + p_k += H*K + p_v += HV*V + if USE_G: + p_g += HV + if USE_GK: + p_gk += HV*K + if USE_GV: + p_gv += HV*V + p_beta += HV * (1 if IS_BETA_HEADWISE else V) + p_o += HV*V + + if STORE_FINAL_STATE: + if STATE_V_FIRST: + p_ht = ht + i_nh * K*V + o_v[:, None] * K + o_k[None, :] + else: + p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) + + +def fused_recurrent_gated_delta_rule_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + gv: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + HV = v.shape[2] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + BK = triton.next_power_of_2(K) + BV = min(8, triton.next_power_of_2(V)) if gv is None else triton.next_power_of_2(V) + NV = triton.cdiv(V, BV) + + o = torch.empty_like(v) + if output_final_state: + if state_v_first: + final_state = q.new_empty(N, HV, V, K, dtype=torch.float32) + else: + final_state = q.new_empty(N, HV, K, V, dtype=torch.float32) + else: + final_state = None + + grid = (NV * N * HV,) + fused_recurrent_gated_delta_rule_fwd_kernel[grid]( + q=q, + k=k, + v=v, + g=g, + gk=gk, + gv=gv, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + o=o, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + scale=scale, + T=T, + H=H, + HV=HV, + K=K, + V=V, + BK=BK, + BV=BV, + IS_BETA_HEADWISE=beta.ndim != v.ndim, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + APPLY_BETA_SIGMOID=use_beta_sigmoid_in_kernel, + ALLOW_NEG_EIGVAL=allow_neg_eigval, + STATE_V_FIRST=state_v_first, + num_warps=1, + num_stages=3, + ) + return o, final_state + + +class FusedRecurrentFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + gv: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + ): + o, final_state = fused_recurrent_gated_delta_rule_fwd( + q=q, + k=k, + v=v, + g=g, + gk=gk, + gv=gv, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel, + allow_neg_eigval=allow_neg_eigval, + state_v_first=state_v_first, + cu_seqlens=cu_seqlens, + ) + + return o, final_state + + @staticmethod + @input_guard + def backward(ctx, do, dht): + raise NotImplementedError( + "Backward pass is not implemented yet and we do not have plans to implement it " + "because we haven't figured out how to compute dg without materializing the full " + "hidden states for all time steps.", + ) + + +def fused_recurrent_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + gv: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + use_gate_in_kernel: bool = False, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + **kwargs, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, HV, V]`. + GVA (Grouped Value Attention) is applied if `HV > H`, where `HV` must be divisible by `H`. + g (torch.Tensor): + g (decays) of shape `[B, T, HV]`. Default: `None`. + When `use_gate_in_kernel=False` (default), `g` must be in log space (pre-computed decay). + When `use_gate_in_kernel=True`, `g` is the raw pre-activation input; the kernel fuses + `-exp(A_log) * softplus(g + dt_bias)` internally per step. + gk (torch.Tensor): + gk (decays) of shape `[B, T, HV, K]`. Default: `None`. + gv (torch.Tensor): + gv (decays) of shape `[B, T, HV, V]`. Default: `None`. + beta (torch.Tensor): + betas of shape `[B, T, HV]`. + scale (Optional[float]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, HV, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`. + use_qk_l2norm_in_kernel (Optional[bool]): + Whether to use L2 normalization in the kernel. Default: `False`. + use_gate_in_kernel (bool): + Whether to compute the log-space GDN decay internally. + When `True`, `g` is the raw input and `A_log` must be provided; the kernel fuses + gate activation into the recurrence. Default: `False`. + A_log (Optional[torch.Tensor]): + Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`. + dt_bias (Optional[torch.Tensor]): + Bias added to `g` before activation, of shape `[HV]`. + Only used when `use_gate_in_kernel=True`. + use_beta_sigmoid_in_kernel (Optional[bool]): + Whether to apply `torch.sigmoid(beta)` inside the kernel. + - If `True`, the passed `beta` acts as the raw beta logits. + - If `False`, `beta` is expected to already be in post-sigmoid space. + Default: `False`. + allow_neg_eigval (Optional[bool]): + Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`. + Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case + the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`. + state_v_first (Optional[bool]): + Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, HV, V]`. + final_state (torch.Tensor): + Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule + # inputs with equal lengths + >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512 + >>> q = torch.randn(B, T, H, K, device='cuda') + >>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1) + >>> v = torch.randn(B, T, HV, V, device='cuda') + >>> g = F.logsigmoid(torch.rand(B, T, HV, device='cuda')) + >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid() + >>> h0 = torch.randn(B, HV, K, V, device='cuda') + >>> o, ht = fused_gated_recurrent_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g, beta)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o, ht = fused_gated_recurrent_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + if 'transpose_state_layout' in kwargs: + if state_v_first: + raise ValueError("Cannot pass both `state_v_first` and the deprecated `transpose_state_layout`.") + warnings.warn( + "`transpose_state_layout` is deprecated and renamed to `state_v_first`.", + DeprecationWarning, + stacklevel=2, + ) + state_v_first = kwargs.pop('transpose_state_layout') + + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + if beta is None: + beta = torch.ones_like(q[..., 0]) + if use_gate_in_kernel: + if A_log is None: + raise ValueError("`A_log` must be provided when `use_gate_in_kernel=True`.") + if g is None: + raise ValueError("`g` (raw pre-activation) must be provided when `use_gate_in_kernel=True`.") + else: + A_log = None + dt_bias = None + if allow_neg_eigval and not use_beta_sigmoid_in_kernel: + raise ValueError("`allow_neg_eigval=True` requires `use_beta_sigmoid_in_kernel=True`.") + + o, final_state = FusedRecurrentFunction.apply( + q, + k, + v, + g, + gk, + gv, + beta, + A_log, + dt_bias, + scale, + initial_state, + output_final_state, + use_qk_l2norm_in_kernel, + use_beta_sigmoid_in_kernel, + allow_neg_eigval, + state_v_first, + cu_seqlens, + ) + return o, final_state + + +fused_recurrent_gdn = fused_recurrent_gated_delta_rule diff --git a/ex_engine/fla_kernels/gated_delta_rule/gate.py b/ex_engine/fla_kernels/gated_delta_rule/gate.py new file mode 100644 index 00000000..564177e1 --- /dev/null +++ b/ex_engine/fla_kernels/gated_delta_rule/gate.py @@ -0,0 +1,344 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from fla.ops.backends import dispatch +from fla.ops.utils.cache import fla_cache_autotune +from fla.ops.utils.index import prepare_chunk_indices +from fla.ops.utils.op import exp +from fla.ops.utils.softplus import softplus +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard + + +def naive_gdn_gate( + g: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None = None, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """ + Torch reference implementation for GDN gate computation. + + Computes: ``g = -A_log.exp() * softplus(g + dt_bias)`` + + Args: + g (torch.Tensor): + Input tensor of shape `[..., HV]`. + A_log (torch.Tensor): + Decay parameter tensor with `HV` elements. + dt_bias (torch.Tensor | None): + Optional bias tensor added to `g` before activation, shape `[HV]`. + + Returns: + Output tensor of shape `[..., HV]`. + """ + g = g.float() + if dt_bias is not None: + g = g + dt_bias.float() + return (-A_log.float().exp() * F.softplus(g)).to(output_dtype) + + +@triton.heuristics({ + 'HAS_BIAS': lambda args: args['dt_bias'] is not None, + 'HAS_SCALE': lambda args: args['scale'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@fla_cache_autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8] + ], + key=['H', 'BT', 'IS_VARLEN', 'REVERSE'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def gdn_gate_chunk_cumsum_scalar_kernel( + g, + A_log, + dt_bias, + o, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + REVERSE: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_SCALE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_g = g + bos * H + i_h + o_t * H + p_o = o + bos * H + i_h + o_t * H + + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) + if HAS_BIAS: + b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32) + b_A = tl.load(A_log + i_h).to(tl.float32) + b_gate = -exp(b_A) * softplus(b_g) + + b_o = tl.cumsum(b_gate, axis=0) + if REVERSE: + b_z = tl.sum(b_gate, axis=0) + b_o = -b_o + b_z[None] + b_gate + if HAS_SCALE: + b_o *= scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_t) + + +@triton.heuristics({ + 'HAS_BIAS': lambda args: args['dt_bias'] is not None, +}) +@fla_cache_autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8] + ], + key=['H', 'BT'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def gdn_gate_bwd_kernel( + g, + A_log, + dt_bias, + dyg, + dg, + dA, + T, + H: tl.constexpr, + BT: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1) + + b_A = tl.load(A_log + i_h).to(tl.float32) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_g = g + i_h + o_t * H + p_dg = dg + i_h + o_t * H + p_dyg = dyg + i_h + o_t * H + + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) + b_dyg = tl.load(p_dyg, mask=m_t, other=0.0).to(tl.float32) + + if HAS_BIAS: + b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32) + + # gate = -exp(A_log) * softplus(g + bias) + # d(gate)/d(g) = -exp(A_log) * sigmoid(g + bias) (softplus' = sigmoid) + # d(gate)/d(A_log) = -exp(A_log) * softplus(g + bias) = gate + b_neg_expA = -exp(b_A) + b_yg = b_neg_expA * softplus(b_g) + b_dg = b_neg_expA * (b_dyg * tl.sigmoid(b_g)) + b_dA = tl.sum(b_dyg * b_yg, 0) + + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_t) + tl.store(dA + i_t * H + i_h, b_dA) + + +@input_guard +@dispatch('gated_delta_rule') +def gdn_gate_chunk_cumsum( + g: torch.Tensor, + A_log: torch.Tensor, + chunk_size: int, + scale: float = None, + dt_bias: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + output_dtype: torch.dtype | None = torch.float, +) -> torch.Tensor: + B, T, H = g.shape + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + o = torch.empty_like(g, dtype=output_dtype or g.dtype) + gdn_gate_chunk_cumsum_scalar_kernel[(NT, B * H)]( + g=g, + A_log=A_log, + dt_bias=dt_bias, + o=o, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + BT=BT, + REVERSE=False, + ) + return o + + +@dispatch('gated_delta_rule') +def gdn_gate_bwd( + g: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None, + dyg: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + H = g.shape[-1] + T = g.numel() // H + BT = 32 + NT = triton.cdiv(T, BT) + + dg = torch.empty_like(g, dtype=torch.float32) + dA = A_log.new_empty(NT, H, dtype=torch.float32) + + gdn_gate_bwd_kernel[(NT, H)]( + g=g, + A_log=A_log, + dt_bias=dt_bias, + dyg=dyg, + dg=dg, + dA=dA, + T=T, + H=H, + BT=BT, + ) + + dg = dg.view_as(g).type_as(g) + dA = dA.sum(0).view_as(A_log).type_as(A_log) + dbias = dg.view(-1, H).sum(0).to(dt_bias) if dt_bias is not None else None + + return dg, dA, dbias + + +@triton.heuristics({ + 'HAS_BIAS': lambda args: args['dt_bias'] is not None, +}) +@fla_cache_autotune( + configs=[ + triton.Config({'BT': BT}, num_warps=num_warps, num_stages=num_stages) + for BT in [32, 64, 128] + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3] + ], + key=['H'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def gdn_gate_fwd_kernel( + g, + A_log, + dt_bias, + yg, + T, + H: tl.constexpr, + BT: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1) + + b_A = tl.load(A_log + i_h).to(tl.float32) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_g = g + i_h + o_t * H + p_yg = yg + i_h + o_t * H + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) + if HAS_BIAS: + b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32) + b_yg = -exp(b_A) * softplus(b_g) + tl.store(p_yg, b_yg.to(p_yg.dtype.element_ty), mask=m_t) + + +@dispatch('gated_delta_rule') +def gdn_gate_fwd( + g: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None = None, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + H = g.shape[-1] + T = g.numel() // H + + yg = torch.empty_like(g, dtype=output_dtype) + + def grid(meta): + return (triton.cdiv(T, meta['BT']), H) + + gdn_gate_fwd_kernel[grid]( + g=g, + A_log=A_log, + dt_bias=dt_bias, + yg=yg, + T=T, + H=H, + ) + return yg + + +class GDNGateFunction(torch.autograd.Function): + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + g: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None = None, + output_dtype: torch.dtype = torch.float32, + ) -> torch.Tensor: + yg = gdn_gate_fwd(g=g, A_log=A_log, dt_bias=dt_bias, output_dtype=output_dtype) + ctx.save_for_backward(g, A_log, dt_bias) + return yg + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, dyg: torch.Tensor): + g, A_log, dt_bias = ctx.saved_tensors + dg, dA, dbias = gdn_gate_bwd(g=g, A_log=A_log, dt_bias=dt_bias, dyg=dyg) + return dg, dA, dbias, None + + +@torch.compiler.disable +def fused_gdn_gate( + g: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None = None, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + r""" + Fused GDN gate computation with autograd support. + + Computes: ``g = -A_log.exp() * softplus(g + dt_bias)`` + + Args: + g (torch.Tensor): + Input tensor of shape `[..., HV]`. + A_log (torch.Tensor): + Decay parameter tensor with `HV` elements. + dt_bias (torch.Tensor | None): + Optional bias tensor added to `g` before activation, shape `[HV]`. + output_dtype (torch.dtype): + The dtype of the output tensor. Default: `torch.float32`. + + Returns: + Output tensor of shape `[..., HV]`. + """ + return GDNGateFunction.apply(g, A_log, dt_bias, output_dtype) diff --git a/ex_engine/fla_kernels/gated_delta_rule/naive.py b/ex_engine/fla_kernels/gated_delta_rule/naive.py new file mode 100644 index 00000000..cd0cf0d1 --- /dev/null +++ b/ex_engine/fla_kernels/gated_delta_rule/naive.py @@ -0,0 +1,161 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +import torch +import torch.nn.functional as F +from einops import rearrange + + +def naive_recurrent_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + g: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, +): + """ + Reference PyTorch implementation of recurrent gated delta rule. + + Args: + q: [B, T, H, K] + k: [B, T, H, K] + v: [B, T, H, V] + beta: [B, T, H] + g: [B, T, H] + scale: float, optional + initial_state: [B, H, K, V], optional + output_final_state: bool + + Returns: + o: [B, T, H, V] + final_state: [B, H, K, V] if output_final_state else None + """ + q, k, v, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g]) + B, H, T, K, V = *k.shape, v.shape[-1] + o = torch.zeros(B, H, T, V).to(v) + h = torch.zeros(B, H, K, V).to(v) + if initial_state is not None: + h = initial_state.to(torch.float32) + if scale is None: + scale = 1 / (q.shape[-1] ** 0.5) + q = q * scale + + for i in range(T): + b_q = q[:, :, i] + b_k = k[:, :, i] + b_v = v[:, :, i].clone() + h = h.clone() * g[:, :, i].exp()[..., None, None] + b_beta = beta[:, :, i] + b_v = b_v - (h.clone() * b_k[..., None]).sum(-2) + b_v = b_v * b_beta[..., None] + h = h.clone() + b_k.unsqueeze(-1) * b_v.unsqueeze(-2) + o[:, :, i] = torch.einsum('bhd,bhdm->bhm', b_q, h) + + if not output_final_state: + h = None + o = o.transpose(1, 2).contiguous() + return o, h + + +def naive_chunk_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + chunk_size: int = 64, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, +): + """ + Reference PyTorch implementation of chunk gated delta rule. + + Args: + q: [B, T, H, K] + k: [B, T, H, K] + v: [B, T, H, V] + g: [B, T, H] + beta: [B, T, H] + chunk_size: int + scale: float, optional + initial_state: [B, H, K, V], optional + output_final_state: bool + + Returns: + o: [B, T, H, V] + final_state: [B, H, K, V] if output_final_state else None + """ + BT = chunk_size + if scale is None: + scale = 1 / (q.shape[-1] ** 0.5) + + q, k, v, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g]) + + T = q.shape[-2] + pad_len = (BT - (T % BT)) % BT + if pad_len > 0: + q = F.pad(q, (0, 0, 0, pad_len)) + k = F.pad(k, (0, 0, 0, pad_len)) + v = F.pad(v, (0, 0, 0, pad_len)) + beta = F.pad(beta, (0, pad_len)) + g = F.pad(g, (0, pad_len)) + + q, k, v, beta, g = map(lambda x: x.to(torch.float32), [q, k, v, beta, g]) + decay = g + chunk_size = BT + b, h, l, d_k = q.shape + d_v = v.shape[-1] + q = q * scale + v = v * beta[..., None] + k_beta = k * beta[..., None] + assert l % chunk_size == 0 + + # note that diagonal is masked. + mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0) + q, k, v, k_beta, decay = map( + lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size), + [q, k, v, k_beta, decay.unsqueeze(-1)], + ) + decay = decay.squeeze(-1).cumsum(-1) + decay_exp = decay.exp()[..., None] + L_mask = ((decay.unsqueeze(-1) - decay.unsqueeze(-2)).tril().exp().float()).tril() + attn = -((k_beta @ k.transpose(-1, -2)) * L_mask).masked_fill(mask, 0) + for i in range(1, chunk_size): + attn[..., i, :i] = attn[..., i, :i].clone() + (attn[..., i, :i, None].clone() * attn[..., :i, :i].clone()).sum(-2) + attn = attn + torch.eye(chunk_size, dtype=torch.float, device=q.device) + attn = attn + k_cumsum = attn @ v + k_cumdecay = attn @ (k_beta * decay_exp) + v = k_cumsum + + S = k.new_zeros(b, h, d_k, d_v) + if initial_state is not None: + S = initial_state.to(torch.float32) + + o = torch.zeros_like(v) + mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1) + for i in range(0, l // chunk_size): + q_i, k_i, v_i = q[:, :, i], k[:, :, i], v[:, :, i] + attn = (q_i @ k_i.transpose(-1, -2) * L_mask[:, :, i]).masked_fill_(mask, 0) + v_prime = (k_cumdecay[:, :, i]) @ S + v_new = v_i - v_prime + o_inter = (q_i * decay[:, :, i, :, None].exp()) @ S + o[:, :, i] = o_inter + attn @ v_new + S = S * decay[:, :, i, -1, None, None].exp() + (k_i * (decay[:, :, i, -1, None] - decay[:, :, i]).exp() + [..., None]).transpose(-1, -2) @ v_new + if not output_final_state: + S = None + + # unpad + o = rearrange(o, 'b h n c d -> b h (n c) d') + o = o[:, :, :T] + o = o.transpose(1, 2) + return o, S diff --git a/ex_engine/fla_kernels/gated_delta_rule/wy_fast.py b/ex_engine/fla_kernels/gated_delta_rule/wy_fast.py new file mode 100644 index 00000000..4cbfe1b5 --- /dev/null +++ b/ex_engine/fla_kernels/gated_delta_rule/wy_fast.py @@ -0,0 +1,351 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +import torch +import triton +import triton.language as tl + +from fla.ops.backends import dispatch +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.cache import fla_cache_autotune +from fla.ops.utils.op import exp2 +from fla.utils import IS_INTEL, IS_NVIDIA_BLACKWELL, autotune_cache_kwargs, check_shared_mem + +# Blackwell can select unstable Triton configs for prepare_wy_repr_bwd_kernel +# during autotuning (see #913). Restrict it to the config that has been +# validated on B200 until the wider config space is re-validated. +PREPARE_WY_REPR_BWD_NUM_WARPS = [2] if IS_NVIDIA_BLACKWELL else [2, 4] +PREPARE_WY_REPR_BWD_NUM_STAGES = [4] if IS_NVIDIA_BLACKWELL else [2, 3, 4] + +# Intel keeps scaling past the warp counts NVIDIA prefers: 16 warps is ~1.3x faster +# than 8 for recompute_w_u. +RECOMPUTE_W_U_NUM_WARPS = [2, 4, 8, 16] if IS_INTEL else [2, 4, 8] + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@fla_cache_autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in RECOMPUTE_W_U_NUM_WARPS + for num_stages in [2, 3, 4] + ], + key=['H', 'HV', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def recompute_w_u_fwd_kernel( + k, + v, + beta, + w, + u, + A, + g, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + i_b, i_h = i_bh // HV, i_bh % HV + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + o_t = i_t * BT + tl.arange(0, BT) + o_A = tl.arange(0, BT) + m_t = o_t < T + m_A = m_t[:, None] & (o_A[None, :] < BT) + p_b = beta + bos*HV + i_h + o_t * HV + b_b = tl.load(p_b, mask=m_t, other=0.0) + + p_A = A + (bos*HV + i_h) * BT + o_t[:, None] * (HV*BT) + o_A[None, :] + b_A = tl.load(p_A, mask=m_A, other=0.0) + + for i_v in range(tl.cdiv(V, BV)): + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :] + p_u = u + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_u = tl.dot(b_A, b_vb, allow_tf32=False) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_v) + + if USE_G: + p_g = g + (bos*HV + i_h) + o_t * HV + b_g = exp2(tl.load(p_g, mask=m_t, other=0.0)) + + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = m_t[:, None] & (o_k[None, :] < K) + p_k = k + (bos*H + i_h // (HV // H)) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_w = w + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_k, other=0.0) + b_kb = b_k * b_b[:, None] + if USE_G: + b_kb *= b_g[:, None] + b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_k) + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@fla_cache_autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in PREPARE_WY_REPR_BWD_NUM_WARPS + for num_stages in PREPARE_WY_REPR_BWD_NUM_STAGES + ], + key=['H', 'HV', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def prepare_wy_repr_bwd_kernel( + k, + v, + beta, + g, + A, + dw, + du, + dk, + dv, + db, + dg, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + i_b, i_h = i_bh // HV, i_bh % HV + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_t = i_t * BT + tl.arange(0, BT) + o_A = tl.arange(0, BT) + m_t = o_t < T + m_AT = (o_A[:, None] < BT) & m_t[None, :] + p_b = beta + (bos*HV + i_h) + o_t * HV + p_db = db + (bos*HV + i_h) + o_t * HV + p_A = A + (bos*HV + i_h) * BT + o_A[:, None] + o_t[None, :] * (HV*BT) + + b_b = tl.load(p_b, mask=m_t, other=0.0) + b_db = tl.zeros([BT], dtype=tl.float32) + b_A = tl.load(p_A, mask=m_AT, other=0.0) + b_dA = tl.zeros([BT, BT], dtype=tl.float32) + + if USE_G: + p_g = g + (bos*HV + i_h) + o_t * HV + b_g = tl.load(p_g, mask=m_t, other=0.0) + b_g_exp = exp2(b_g) + b_dg = tl.zeros([BT], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = m_t[:, None] & (o_k[None, :] < K) + p_k = k + (bos*H + i_h // (HV // H)) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_dk = dk + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :] + p_dw = dw + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :] + # [BT, BK] + b_k = tl.load(p_k, mask=m_k, other=0.0) + if USE_G: + b_kbg = b_k * (b_b * b_g_exp)[:, None] + else: + b_kbg = b_k * b_b[:, None] + b_dw = tl.load(p_dw, mask=m_k, other=0.0) + + b_dA += tl.dot(b_dw, tl.trans(b_kbg).to(b_dw.dtype)) + b_dkbg = tl.dot(b_A, b_dw) + if USE_G: + b_dk = b_dkbg * (b_g_exp * b_b)[:, None] + b_db += tl.sum(b_dkbg * b_k * b_g_exp[:, None], 1) + b_dg += tl.sum(b_dkbg * b_kbg, 1) + else: + b_dk = b_dkbg * b_b[:, None] + b_db += tl.sum(b_dkbg * b_k, 1) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_k) + + for i_v in range(tl.cdiv(V, BV)): + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :] + p_dv = dv + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :] + p_du = du + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_du = tl.load(p_du, mask=m_v, other=0.0) + b_dA += tl.dot(b_du, tl.trans(b_vb)) + b_dvb = tl.dot(b_A, b_du) + b_dv = b_dvb * b_b[:, None] + b_db += tl.sum(b_dvb * b_v, 1) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_v) + + m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) + b_dA = tl.where(m_A, b_dA, 0) + b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) + b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) + + if USE_G: + b_dA *= exp2(b_g[:, None] - b_g[None, :]) + + b_A = tl.zeros([BT, BT], dtype=tl.float32) + b_dA = tl.where(m_A, -b_dA, 0).to(k.dtype.element_ty) + + tl.debug_barrier() + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = m_t[:, None] & (o_k[None, :] < K) + p_k = k + (bos*H + i_h // (HV // H)) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_dk = dk + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_k, other=0.0) + b_kt = tl.trans(b_k) + b_kb = b_k * b_b[:, None] + + b_A += tl.dot(b_k, b_kt) + b_dkb = tl.dot(b_dA, b_k) + b_db += tl.sum(b_dkb * b_k, 1) + b_dk = b_dkb * b_b[:, None] + tl.trans(tl.dot(tl.trans(b_kb).to(b_dA.dtype), b_dA)) + b_dk += tl.load(p_dk, mask=m_k, other=0.0) + + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_k) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), mask=m_t) + + b_A *= b_b[:, None] + if USE_G: + b_AdA = b_dA * b_A + p_dg = dg + (bos*HV + i_h) + o_t * HV + b_dg += tl.sum(b_AdA, axis=1) - tl.sum(b_AdA, axis=0) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_t) + + +@dispatch('gated_delta_rule') +def recompute_w_u_fwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + g: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V, HV = *k.shape, v.shape[-1], v.shape[2] + BT = A.shape[-1] + BK = 64 + BV = 64 + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + w = k.new_empty(B, T, HV, K) + u = torch.empty_like(v) + recompute_w_u_fwd_kernel[(NT, B*HV)]( + k=k, + v=v, + beta=beta, + w=w, + u=u, + A=A, + g=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + HV=HV, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return w, u + + +@dispatch('gated_delta_rule') +def prepare_wy_repr_bwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + dw: torch.Tensor, + du: torch.Tensor, + g: torch.Tensor = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K, V, HV = *k.shape, v.shape[-1], v.shape[2] + BT = A.shape[-1] + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + CONST_TILING = 64 if check_shared_mem() else 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + + dk = k.new_empty(B, T, HV, K) + dv = torch.empty_like(v) + dg = torch.empty_like(g) if g is not None else None + db = torch.empty_like(beta) + prepare_wy_repr_bwd_kernel[(NT, B * HV)]( + k=k, + v=v, + beta=beta, + g=g, + A=A, + dw=dw, + du=du, + dk=dk, + dv=dv, + db=db, + dg=dg, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + HV=HV, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + if H != HV: + dk = dk.view(B, T, H, HV // H, K).sum(3) + return dk, dv, db, dg + + +fwd_recompute_w_u = recompute_w_u_fwd +bwd_prepare_wy_repr = prepare_wy_repr_bwd diff --git a/ex_engine/fla_kernels/utils/__init__.py b/ex_engine/fla_kernels/utils/__init__.py new file mode 100644 index 00000000..88acd8b9 --- /dev/null +++ b/ex_engine/fla_kernels/utils/__init__.py @@ -0,0 +1,65 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +from .csr import prepare_block_csr +from .cumsum import ( + chunk_global_cumsum, + chunk_global_cumsum_scalar, + chunk_global_cumsum_vector, + chunk_local_cumsum, + chunk_local_cumsum_scalar, + chunk_local_cumsum_vector, +) +from .index import ( + get_max_num_splits, + prepare_chunk_indices, + prepare_chunk_offsets, + prepare_cu_seqlens_from_lens, + prepare_cu_seqlens_from_mask, + prepare_lens, + prepare_lens_from_mask, + prepare_position_ids, + prepare_sequence_ids, + prepare_token_indices, +) +from .logsumexp import logsumexp_fwd +from .matmul import addmm, matmul +from .pack import pack_sequence, unpack_sequence +from .pooling import mean_pooling +from .softmax import softmax_bwd, softmax_fwd +from .softplus import softplus +from .solve_tril import solve_tril + +__all__ = [ + "addmm", + "chunk_global_cumsum", + "chunk_global_cumsum_scalar", + "chunk_global_cumsum_vector", + "chunk_local_cumsum", + "chunk_local_cumsum_scalar", + "chunk_local_cumsum_vector", + "get_max_num_splits", + "logsumexp_fwd", + "matmul", + "mean_pooling", + "pack_sequence", + "prepare_block_csr", + "prepare_chunk_indices", + "prepare_chunk_offsets", + "prepare_cu_seqlens_from_lens", + "prepare_cu_seqlens_from_mask", + "prepare_lens", + "prepare_lens_from_mask", + "prepare_position_ids", + "prepare_sequence_ids", + "prepare_token_indices", + "softmax_bwd", + "softmax_fwd", + "softplus", + "solve_tril", + "unpack_sequence", +] diff --git a/ex_engine/fla_kernels/utils/cache.py b/ex_engine/fla_kernels/utils/cache.py new file mode 100644 index 00000000..40ab2dbc --- /dev/null +++ b/ex_engine/fla_kernels/utils/cache.py @@ -0,0 +1,449 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +import dataclasses +import enum +import json +import logging +import os +import re +from functools import cache, lru_cache +from pathlib import Path +from typing import Any + +import torch +import triton +from packaging import version +from triton.runtime.autotuner import Autotuner + +TRITON_ABOVE_3_5_1 = version.parse(triton.__version__) >= version.parse("3.5.1") +TRITON_ABOVE_3_4_0 = version.parse(triton.__version__) >= version.parse("3.4.0") + + +class FlaCacheMode(enum.Enum): + """Controls how FLA loads kernel configs from its config cache (FLA_CACHE_MODE env var). + + DISABLED — skip all cache lookups, always fall back to Triton autotune (default when FLA_CACHE_MODE is unset) + STRICT — exact key match only; falls back to Triton autotune if no match + FUZZY — exact key match → fuzzy key match; falls back to Triton autotune if no match + FULL — exact key match → fuzzy key match → default_config fallback + DEFAULT — use only the top-level default_config field, skip key-based lookup + ALWAYS — like DEFAULT, but re-reads config files on every kernel call; + useful for debugging: edit default_config in a JSON file and the next + kernel call picks it up without restarting the process + """ + DISABLED = "disabled" + STRICT = "strict" + FUZZY = "fuzzy" + FULL = "full" + DEFAULT = "default" + ALWAYS = "always" + + def uses_default_config(self) -> bool: + """Return True for modes that may fall back to default_config (FULL, DEFAULT, ALWAYS).""" + return self in (FlaCacheMode.FULL, FlaCacheMode.DEFAULT, FlaCacheMode.ALWAYS) + + @classmethod + def from_env(cls) -> "FlaCacheMode": + mode_str = os.environ.get("FLA_CACHE_MODE", cls.DISABLED.value) + try: + return cls(mode_str) + except ValueError: + valid = [m.value for m in cls] + raise ValueError( + f"Invalid FLA_CACHE_MODE={mode_str!r}. Valid values: {valid}" + ) from None + + +FLA_CACHE_MODE: FlaCacheMode = FlaCacheMode.from_env() +logger = logging.getLogger(__name__) + + +def sanitize_gpu_name(gpu_name: str) -> str: + sanitized = re.sub(r"[^0-9A-Za-z]+", "_", gpu_name) + sanitized = sanitized.strip("_") + return sanitized or "unknown_gpu" + + +@lru_cache(maxsize=1) +def get_gpu_info(): + """Get GPU model information. + + This function detects the GPU model and returns a sanitized string identifier. + It prioritizes FLA_GPU_NAME environment variable if set, then detects from + available hardware (CUDA, ROCm, Intel GPU, or CPU). + """ + # Check if GPU name is overridden via environment variable + gpu_name = None + # Check if GPU name is overridden via environment variable + if "FLA_GPU_NAME" in os.environ: + gpu_name = os.environ["FLA_GPU_NAME"] + # Try to get device name based on availability + elif torch.cuda.is_available(): + # Works for both NVIDIA and AMD GPUs (ROCm) + gpu_name = torch.cuda.get_device_name(0) + elif hasattr(torch, 'xpu') and torch.xpu.is_available(): + gpu_name = torch.xpu.get_device_name(0) + + if gpu_name: + return sanitize_gpu_name(gpu_name) + + # Default to CPU if no GPU available + return "cpu" + + +def get_fla_config_dir() -> Path: + """Get FLA's configs directory. + + The directory can be overridden by setting the FLA_CONFIG_DIR environment variable. + If set, configs will be loaded directly from $FLA_CONFIG_DIR/. Otherwise FLA + falls back to the default fla/configs/{GPU}/ directory in the project. + """ + # Check if custom config dir is set via environment variable + if "FLA_CONFIG_DIR" in os.environ: + return Path(os.environ["FLA_CONFIG_DIR"]) + + # Default: project_dir/fla/configs/{GPU}/ + project_dir = Path(__file__).parent.parent.parent + return project_dir / "configs" / get_gpu_info() + + +@dataclasses.dataclass(frozen=True) +class AutotuneKey: + """Autotune key with exact/fuzzy matching, serialization, and construction helpers.""" + autotune_key: tuple[Any, ...] + + @staticmethod + def normalize_autotune_key(value: Any) -> Any: + if isinstance(value, (list, tuple)): + return [AutotuneKey.normalize_autotune_key(v) for v in value] + if isinstance(value, dict): + return {k: AutotuneKey.normalize_autotune_key(v) for k, v in value.items()} + return value + + @staticmethod + def serialize(key: Any) -> str: + return json.dumps(AutotuneKey.normalize_autotune_key(key), separators=(",", ":"), sort_keys=True) + + @staticmethod + def key_hash(key: Any) -> str: + import hashlib + return hashlib.md5(AutotuneKey.serialize(key).encode()).hexdigest() + + @staticmethod + def is_numeric(value: Any) -> bool: + return isinstance(value, (int, float)) and not isinstance(value, bool) + + @staticmethod + def keys_fuzzy_match(cached_key: Any, requested_key: Any) -> bool: + # Fuzzy match: numeric leaves are compatible regardless of their actual numeric values + # (e.g. a config tuned for seq_len=1024 can apply to seq_len=2048). + # Structure (type, length, dict keys) must still match exactly. + if AutotuneKey.is_numeric(cached_key) and AutotuneKey.is_numeric(requested_key): + return True + if isinstance(cached_key, (list, tuple)) and isinstance(requested_key, (list, tuple)): + return len(cached_key) == len(requested_key) and all( + AutotuneKey.keys_fuzzy_match(c, r) for c, r in zip(cached_key, requested_key) + ) + if isinstance(cached_key, dict) and isinstance(requested_key, dict): + return cached_key.keys() == requested_key.keys() and all( + AutotuneKey.keys_fuzzy_match(cached_key[k], requested_key[k]) for k in cached_key + ) + return cached_key == requested_key + + @classmethod + def build( + cls, + arg_names: list[str], + key_names: list[str], + positional_args: tuple[Any, ...], + runtime_kwargs: dict[str, Any], + ) -> "AutotuneKey": + named_args = dict(zip(arg_names, positional_args)) + all_args = {**named_args, **runtime_kwargs} + tracked_args = {k: v for (k, v) in all_args.items() if k in arg_names} + tuning_key = [tracked_args[name] for name in key_names if name in tracked_args] + for arg in tracked_args.values(): + if hasattr(arg, "dtype"): + tuning_key.append(str(arg.dtype)) + return cls(autotune_key=tuple(tuning_key)) + + def exact_matches(self, entry_key: Any) -> bool: + return self.serialize(self.autotune_key) == self.serialize(entry_key) + + def fuzzy_matches(self, entry_key: Any) -> bool: + self_normalized = self.normalize_autotune_key(self.autotune_key) + entry_normalized = self.normalize_autotune_key(entry_key) + return ( + isinstance(self_normalized, list) + and isinstance(entry_normalized, list) + and len(self_normalized) == len(entry_normalized) + and AutotuneKey.keys_fuzzy_match(self_normalized, entry_normalized) + ) + + +@dataclasses.dataclass(frozen=True) +class KernelConfigFile: + """Validated in-memory representation of a {kernel_name}.json config file.""" + kernel_name: str | None + triton_version: str | None + autotune_entries: dict[str, dict[str, Any]] | None + default_config: dict[str, Any] | None + + @classmethod + def from_dict(cls, config_file: Path, data: Any) -> "KernelConfigFile | None": + """Parse and validate a raw JSON dict. Returns None (with a warning) if malformed.""" + def fail(msg, *args): + logger.warning(msg, *args) + raise ValueError + + try: + if not isinstance(data, dict): + fail("Malformed config %s: root is %s, expected dict", config_file, type(data).__name__) + raw_entries = data.get("autotune_entries") + entries: dict[str, dict[str, Any]] | None = None + if raw_entries is not None: + if not isinstance(raw_entries, dict): + fail("Malformed config %s: 'autotune_entries' is %s, expected dict", + config_file, type(raw_entries).__name__) + for h, entry in raw_entries.items(): + if not isinstance(entry, dict): + fail("Malformed config %s: autotune_entries[%r] is %s, expected dict", + config_file, h, type(entry).__name__) + if not isinstance(entry.get("config"), dict): + fail("Malformed config %s: autotune_entries[%r] missing valid 'config' field", config_file, h) + entries = raw_entries + default_config = data.get("default_config") + if default_config is not None and not isinstance(default_config, dict): + fail("Malformed config %s: 'default_config' is %s, expected dict", config_file, type(default_config).__name__) + return cls( + kernel_name=data.get("kernel_name"), + triton_version=data.get("triton_version"), + autotune_entries=entries, + default_config=default_config, + ) + except ValueError: + return None + + @classmethod + def from_file(cls, config_file: Path) -> "KernelConfigFile | None": + """Read and validate a config file. Returns None if the file is missing or malformed.""" + config_data = read_config_file(config_file) + if config_data is None: + return None + return cls.from_dict(config_file, config_data) + + def lookup_exact(self, key: AutotuneKey) -> dict[str, Any] | None: + if self.autotune_entries is None: + return None + return self.autotune_entries.get(AutotuneKey.key_hash(key.autotune_key)) + + def lookup_fuzzy(self, key: AutotuneKey) -> dict[str, Any] | None: + if self.autotune_entries is None: + return None + for entry in self.autotune_entries.values(): + if key.fuzzy_matches(entry.get("autotune_key")): + return entry + return None + + +@cache +def load_config_file(config_file: Path) -> dict[str, Any] | None: + try: + with open(config_file) as f: + return json.load(f) + except Exception as e: + logger.warning("Error reading config file %s: %s", config_file, e) + return None + + +def read_config_file(config_file: Path) -> dict[str, Any] | None: + """Read a config file, bypassing the in-process cache in ALWAYS mode.""" + if FLA_CACHE_MODE is FlaCacheMode.ALWAYS: + return load_config_file.__wrapped__(config_file) + return load_config_file(config_file) + + +def load_cached_config(kernel_name: str, autotune_key: AutotuneKey | None = None) -> dict[str, Any] | None: + """ + Load cached best config for a kernel from FLA configs directory. + + This function loads the cached best configuration for a given kernel name + from get_fla_config_dir()/{kernel_name}.json. + + Cache files may contain multiple autotune entries keyed by Triton's + runtime tuning key plus a top-level default config. + + If the config file is not found or cannot be loaded, a warning is printed + and None is returned, allowing fallback to Triton's autotune. + + The lookup mode is controlled by the FLA_CACHE_MODE environment variable (see FlaCacheMode). + + Args: + kernel_name: Name of the kernel (e.g., "causal_conv1d_fwd_kernel") + autotune_key: Triton autotune key for the current invocation + + Returns: + Best config dictionary or None if not found or disabled + """ + if FLA_CACHE_MODE is FlaCacheMode.DISABLED: + return None + + config_dir = get_fla_config_dir() + config_file = config_dir / f"{kernel_name}.json" + + if not config_file.exists(): + return None + + config_data = read_config_file(config_file) + if config_data is None: + return None + config = KernelConfigFile.from_dict(config_file, config_data) + if config is None: + return None + + if FLA_CACHE_MODE is FlaCacheMode.DEFAULT or FLA_CACHE_MODE is FlaCacheMode.ALWAYS: + return config.default_config + + # STRICT mode: exact match only, no fuzzy fallback + if FLA_CACHE_MODE is FlaCacheMode.STRICT: + if autotune_key is not None: + entry = config.lookup_exact(autotune_key) + if entry is not None: + return entry["config"] + return None + + # FULL and FUZZY modes: try exact key match first, then fuzzy match + if autotune_key is not None: + entry = config.lookup_exact(autotune_key) or config.lookup_fuzzy(autotune_key) + if entry is not None: + return entry["config"] + + if FLA_CACHE_MODE is FlaCacheMode.FUZZY: + return None + + # FULL mode: fall back to default_config, then legacy raw config (no autotune_entries) + if config.default_config is not None: + return config.default_config + if config.autotune_entries is not None: + return None + return config_data + + +class CachedAutotuner(Autotuner): + """ + A modified autotuner that loads best config from FLA's config directory. + + This class extends Triton's Autotuner but overrides the run method to + try loading cached configuration first before falling back to autotune. + """ + + def __init__(self, fn, arg_names, configs, key, reset_to_zero, restore_value, **kwargs): + super().__init__(fn, arg_names, configs, key, reset_to_zero, restore_value, **kwargs) + self.kernel_name = fn.fn.__name__ if hasattr(fn, 'fn') else fn.__name__ + + # None-safe pre/post hooks: Triton's defaults crash when a restore_value / reset_to_zero arg + # is None (idiomatic for optional pointers gated by a tl.constexpr flag). + # Fixed upstream in triton-lang/triton#10295 — remove this override once FLA's minimum Triton version has it. + if not self.user_defined_pre_hook and (self.reset_to_zero or self.restore_value): + def _pre_hook(kw, reset_only=False): + for n in self.reset_to_zero: + if kw[n] is not None: + kw[n].zero_() + if not reset_only: + self.restore_copies = {n: kw[n].clone() for n in self.restore_value if kw[n] is not None} + self.pre_hook = _pre_hook + if not self.user_defined_post_hook and self.restore_value: + def _post_hook(kw, exception): + for n, copy in self.restore_copies.items(): + kw[n].copy_(copy) + self.restore_copies = {} + self.post_hook = _post_hook + + def should_check_fla_cache(self, key: AutotuneKey) -> bool: + if FLA_CACHE_MODE is FlaCacheMode.DISABLED: + return False + if FLA_CACHE_MODE is FlaCacheMode.ALWAYS: + return True + return key.autotune_key not in self.cache + + def run(self, *args, **kwargs): + key = AutotuneKey.build(self.arg_names, self.keys, args, kwargs) + if self.should_check_fla_cache(key): + self.maybe_load_cached_config(key) + return super().run(*args, **kwargs) + + def maybe_load_cached_config(self, key: AutotuneKey): + best_config = load_cached_config(self.kernel_name, key) + + if best_config is not None: + kw = best_config["kwargs"] + num_warps = best_config["num_warps"] + num_stages = best_config["num_stages"] + + extra = { + "num_ctas": best_config["num_ctas"], + "maxnreg": best_config.get("maxnreg"), + "pre_hook": None, + "ir_override": best_config.get("ir_override"), + } if TRITON_ABOVE_3_5_1 else {} + cfg = triton.Config(kw, num_warps=num_warps, num_stages=num_stages, **extra) + + self.cache[key.autotune_key] = cfg + else: + logger.debug( + "No cached config found for kernel %s and key %s; falling back to Triton autotune", + self.kernel_name, + list(key.autotune_key), + ) + + +def fla_cache_autotune(configs, key=None, prune_configs_by=None, reset_to_zero=None, restore_value=None, + pre_hook=None, post_hook=None, warmup=None, rep=None, use_cuda_graph=False, + do_bench=None, cache_results=False): + """ + Decorator for auto-tuning a :code:`triton.jit`'d function with FLA config support. + + Extends Triton's autotune to load best configurations from FLA's config directory + (default: fla/configs/{GPU}/, or FLA_CONFIG_DIR/ when overridden), keyed by kernel + name from {kernel_name}.json. Lookup behaviour is controlled by FLA_CACHE_MODE. + Falls back to normal Triton autotuning when no cached config is found. + """ + # key can be None when we want to use cache only (no fallback autotune) + if key is None: + key = [] + + def decorator(fn): + kwargs = {} + if TRITON_ABOVE_3_4_0: + kwargs = {"cache_results": cache_results} + + return CachedAutotuner(fn, fn.arg_names, configs, key, reset_to_zero, restore_value, + pre_hook=pre_hook, post_hook=post_hook, + prune_configs_by=prune_configs_by, warmup=warmup, rep=rep, + use_cuda_graph=use_cuda_graph, do_bench=do_bench, + **kwargs, + ) + + return decorator + + +def configure_fla_cache_autotune(): + triton.autotune = fla_cache_autotune + logger.info( + "configure_fla_cache_autotune() is enabling FLA fla_cache_autotune; " + "triton.autotune will be replaced with fla_cache_autotune." + ) + + +def restore_autotune_backend(): + from triton.runtime.autotuner import autotune as original_autotune + triton.autotune = original_autotune + logger.info( + "restore_autotune_backend() is restoring Triton's original autotune; " + "triton.autotune will be replaced with triton.runtime.autotuner.autotune." + ) diff --git a/ex_engine/fla_kernels/utils/op.py b/ex_engine/fla_kernels/utils/op.py new file mode 100644 index 00000000..10e5b300 --- /dev/null +++ b/ex_engine/fla_kernels/utils/op.py @@ -0,0 +1,101 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +import os + +import triton +import triton.language as tl +import triton.language.extra.libdevice as tldevice + +from fla.utils import IS_GATHER_SUPPORTED, IS_NVIDIA_BLACKWELL + +if os.environ.get('FLA_USE_FAST_OPS', '0') == '1': + @triton.jit + def exp(x): return tldevice.fast_expf(x.to(tl.float32)) + @triton.jit + def exp2(x): return tldevice.exp2(x.to(tl.float32)) + @triton.jit + def log(x): return tldevice.fast_logf(x.to(tl.float32)) + @triton.jit + def log2(x): return tldevice.fast_log2f(x.to(tl.float32)) + @triton.jit + def tanh(x): return tldevice.fast_tanhf(x.to(tl.float32)) +else: + @triton.jit + def exp(x): return tl.exp(x.to(tl.float32)) + @triton.jit + def exp2(x): return tl.math.exp2(x.to(tl.float32)) + @triton.jit + def log(x): return tl.log(x.to(tl.float32)) + @triton.jit + def log2(x): return tl.log2(x.to(tl.float32)) + @triton.jit + def tanh(x): return tldevice.tanh(x.to(tl.float32)) + + +if IS_NVIDIA_BLACKWELL: + """ + Compute tl.dot with Blackwell workaround. + + On SM100 datacenter and SM120 consumer Blackwell GPUs, wraps the result in + inline assembly to prevent the TritonGPUHoistTMEMAlloc pass from incorrectly + fusing add and dot operations. + See: https://github.com/fla-org/flash-linear-attention/issues/638 + + TODO: Remove this workaround once the Triton compiler bug is fixed. + Track upstream issue at: https://github.com/triton-lang/triton/issues/8695 + """ + @triton.jit + def safe_dot(a, b, allow_tf32: tl.constexpr = None): + return tl.inline_asm_elementwise( + asm="mov.f32 $0, $1;", + constraints="=r,r", + args=[tl.dot(a, b, allow_tf32=allow_tf32)], + dtype=tl.float32, + is_pure=True, + pack=1, + ) +else: + @triton.jit + def safe_dot(a, b, allow_tf32: tl.constexpr = None): + return tl.dot(a, b, allow_tf32=allow_tf32) + + +if not IS_GATHER_SUPPORTED: + @triton.jit + def gather(src, index, axis, _builder=None): + """ + Gather operation that works when tl.gather is not supported. + This is a fallback implementation that returns None. + Just to make triton compiler happy. + """ + return None +else: + gather = tl.gather + + +if hasattr(triton.language, '_experimental_make_tensor_descriptor'): + # For Triton 3.3.x + make_tensor_descriptor = triton.language._experimental_make_tensor_descriptor +elif hasattr(triton.language, 'make_tensor_descriptor'): + # For Triton 3.4.x and later + make_tensor_descriptor = triton.language.make_tensor_descriptor +else: + """ + Fallback implementation when TMA is not supported. + Returns None to indicate TMA descriptors are unavailable. + Just make triton compiler happy. + """ + @triton.jit + def make_tensor_descriptor( + base, + shape, + strides, + block_shape, + _builder=None, + ): + return None diff --git a/ex_engine/xllm_kernels/cuda/activation.cu b/ex_engine/xllm_kernels/cuda/activation.cu new file mode 100644 index 00000000..1f4a71e5 --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/activation.cu @@ -0,0 +1,188 @@ +/* Copyright 2025 The vLLM Authors and The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ +#include +#include + +#include + +#include "cuda_ops_api.h" +#include "device_utils.cuh" + +// ref to: +// https://github.com/vllm-project/vllm/blob/main/csrc/activation_kernels.cu + +namespace { + +using ::xllm::kernel::cuda::xllm_ldg; + +template +__device__ __forceinline__ scalar_t compute(const scalar_t& x, + const scalar_t& y) { + return act_first ? ACT_FN(x) * y : x * ACT_FN(y); +} + +// Check if pointer is 16-byte aligned for int4 vectorized access +__device__ __forceinline__ bool is_16byte_aligned(const void* ptr) { + return (reinterpret_cast(ptr) & 15) == 0; +} + +// Activation and gating kernel template with 128-bit vectorized access +// optimization. +template +__global__ void XLLM_KERNEL_ATTR(1024) + act_and_mul_kernel(scalar_t* __restrict__ out, // [..., d] + const scalar_t* __restrict__ input, // [..., 2, d] + const int d) { + constexpr int kVecSize = 16 / sizeof(scalar_t); + const int64_t token_idx = blockIdx.x; + const scalar_t* x_ptr = input + token_idx * 2 * d; + const scalar_t* y_ptr = x_ptr + d; + scalar_t* out_ptr = out + token_idx * d; + + // Check alignment for 128-bit vectorized access. + // All three pointers must be 16-byte aligned for safe int4 operations. + const bool aligned = is_16byte_aligned(x_ptr) && is_16byte_aligned(y_ptr) && + is_16byte_aligned(out_ptr); + + if (aligned && d >= kVecSize) { + // Fast path: 128-bit vectorized loop + const int4* x_vec = reinterpret_cast(x_ptr); + const int4* y_vec = reinterpret_cast(y_ptr); + int4* out_vec = reinterpret_cast(out_ptr); + const int num_vecs = d / kVecSize; + const int vec_end = num_vecs * kVecSize; + + for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) { + int4 x = xllm_ldg(&x_vec[i]), y = xllm_ldg(&y_vec[i]), r; + auto* xp = reinterpret_cast(&x); + auto* yp = reinterpret_cast(&y); + auto* rp = reinterpret_cast(&r); +#pragma unroll + for (int j = 0; j < kVecSize; j++) { + rp[j] = compute(xp[j], yp[j]); + } + out_vec[i] = r; + } + // Scalar cleanup for remaining elements + for (int i = vec_end + threadIdx.x; i < d; i += blockDim.x) { + out_ptr[i] = compute(xllm_ldg(&x_ptr[i]), + xllm_ldg(&y_ptr[i])); + } + } else { + // Scalar fallback for unaligned data or small d + for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) { + const scalar_t x = xllm_ldg(&x_ptr[idx]); + const scalar_t y = xllm_ldg(&y_ptr[idx]); + out_ptr[idx] = compute(x, y); + } + } +} + +template +__device__ __forceinline__ T silu_kernel(const T& x) { + // x * sigmoid(x) + const float f = static_cast(x); + return static_cast(f / (1.0f + expf(-f))); +} + +template +__device__ __forceinline__ T gelu_kernel(const T& x) { + // Equivalent to PyTorch GELU with 'none' approximation. + // Refer to: + // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38 + const float f = static_cast(x); + constexpr float kAlpha = M_SQRT1_2; + return static_cast(f * 0.5f * (1.0f + ::erf(f * kAlpha))); +} + +template +__device__ __forceinline__ T gelu_tanh_kernel(const T& x) { + // Equivalent to PyTorch GELU with 'tanh' approximation. + // Refer to: + // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30 + const float f = static_cast(x); + constexpr float kBeta = M_SQRT2 * M_2_SQRTPI * 0.5f; + constexpr float kKappa = 0.044715; + float x_cube = f * f * f; + float inner = kBeta * (f + kKappa * x_cube); + return static_cast(0.5f * f * (1.0f + ::tanhf(inner))); +} + +#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, ACT_FIRST) \ + int d = input.size(-1) / 2; \ + int64_t num_tokens = input.numel() / input.size(-1); \ + dim3 grid(num_tokens); \ + dim3 block(std::min(d, 1024)); \ + if (num_tokens == 0) { \ + return; \ + } \ + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \ + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \ + DISPATCH_FLOATING_TYPES(input.scalar_type(), "act_and_mul_kernel", [&] { \ + act_and_mul_kernel, ACT_FIRST> \ + <<>>( \ + out.data_ptr(), input.data_ptr(), d); \ + }); + +void silu_and_mul(torch::Tensor out, // [..., d] + torch::Tensor input) // [..., 2 * d] +{ + LAUNCH_ACTIVATION_GATE_KERNEL(silu_kernel, true); +} + +void gelu_and_mul(torch::Tensor& out, // [..., d] + torch::Tensor& input) // [..., 2 * d] +{ + LAUNCH_ACTIVATION_GATE_KERNEL(gelu_kernel, true); +} + +void gelu_tanh_and_mul(torch::Tensor& out, // [..., d] + torch::Tensor& input) // [..., 2 * d] +{ + LAUNCH_ACTIVATION_GATE_KERNEL(gelu_tanh_kernel, true); +} +} // namespace + +namespace xllm::kernel::cuda { + +void act_and_mul(torch::Tensor out, + torch::Tensor input, + const std::string& act_mode) { + if (act_mode != "silu" && act_mode != "gelu" && act_mode != "gelu_tanh" && + act_mode != "gelu_pytorch_tanh") { + LOG(FATAL) << "Unsupported act mode: " << act_mode + << ", only support silu, gelu, gelu_tanh, gelu_pytorch_tanh"; + } + + // flashinfer act_and_mul ops + // std::string uri = act_mode + "_and_mul"; + // FunctionFactory::get_instance().act_and_mul(uri).call( + // out, input, support_pdl()); + + if (act_mode == "silu") { + silu_and_mul(out, input); + } else if (act_mode == "gelu") { + gelu_and_mul(out, input); + } else if (act_mode == "gelu_tanh" || act_mode == "gelu_pytorch_tanh") { + // gelu_tanh or gelu_pytorch_tanh (mathematically equivalent) + gelu_tanh_and_mul(out, input); + } +} + +} // namespace xllm::kernel::cuda diff --git a/ex_engine/xllm_kernels/cuda/block_copy.cu b/ex_engine/xllm_kernels/cuda/block_copy.cu new file mode 100644 index 00000000..05651058 --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/block_copy.cu @@ -0,0 +1,209 @@ +/* Copyright 2025-2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "cuda_ops_api.h" +#include "utils.h" + +namespace xllm::kernel::cuda { +namespace { + +template +struct VecType; + +template <> +struct VecType { + using type = uint4; + static constexpr int32_t vec_width = 8; +}; + +template <> +struct VecType { + using type = uint4; + static constexpr int32_t vec_width = 8; +}; + +template <> +struct VecType { + using type = float4; + static constexpr int32_t vec_width = 4; +}; + +DEVICE_INLINE int32_t find_group_idx(const int32_t* __restrict__ cum_sum, + const int32_t num_groups, + const int32_t dst_idx) { + int32_t left = 0; + int32_t right = num_groups - 1; + while (left < right) { + const int32_t mid = left + ((right - left) >> 1); + const bool move_left = dst_idx < cum_sum[mid]; + right = move_left ? mid : right; + left = move_left ? left : mid + 1; + } + return left; +} + +template +__global__ void block_copy_kernel(const int64_t* __restrict__ key_cache_ptrs, + const int64_t* __restrict__ value_cache_ptrs, + const int32_t* __restrict__ src_block_indices, + const int32_t* __restrict__ dst_block_indices, + const int32_t* __restrict__ cum_sum, + const int32_t num_groups, + const int64_t numel_per_block) { + const int64_t layer_idx = static_cast(blockIdx.x); + const int32_t dst_linear_idx = static_cast(blockIdx.y); + const int64_t tile_idx = static_cast(blockIdx.z); + + scalar_t* __restrict__ key_cache = reinterpret_cast( + static_cast(key_cache_ptrs[layer_idx])); + scalar_t* __restrict__ value_cache = reinterpret_cast( + static_cast(value_cache_ptrs[layer_idx])); + + const int32_t group_idx = find_group_idx(cum_sum, num_groups, dst_linear_idx); + const int32_t src_block = src_block_indices[group_idx]; + const int32_t dst_block = dst_block_indices[dst_linear_idx]; + const int64_t src_offset = static_cast(src_block) * numel_per_block; + const int64_t dst_offset = static_cast(dst_block) * numel_per_block; + + if constexpr (kVectorized) { + using VecTypeT = typename VecType::type; + constexpr int32_t kVecWidth = VecType::vec_width; + const int64_t num_vecs_per_block = numel_per_block / kVecWidth; + const int64_t vec_idx = tile_idx * static_cast(blockDim.x) + + static_cast(threadIdx.x); + if (vec_idx >= num_vecs_per_block) { + return; + } + + const int64_t elem_offset = vec_idx * kVecWidth; + const auto* key_src_vec = + reinterpret_cast(key_cache + src_offset + elem_offset); + const auto* value_src_vec = reinterpret_cast( + value_cache + src_offset + elem_offset); + auto* key_dst_vec = + reinterpret_cast(key_cache + dst_offset + elem_offset); + auto* value_dst_vec = + reinterpret_cast(value_cache + dst_offset + elem_offset); + *key_dst_vec = *key_src_vec; + *value_dst_vec = *value_src_vec; + } else { + const int64_t elem_idx = tile_idx * static_cast(blockDim.x) + + static_cast(threadIdx.x); + if (elem_idx >= numel_per_block) { + return; + } + + key_cache[dst_offset + elem_idx] = key_cache[src_offset + elem_idx]; + value_cache[dst_offset + elem_idx] = value_cache[src_offset + elem_idx]; + } +} + +} // namespace + +void block_copy(torch::Tensor key_cache_ptrs, + torch::Tensor value_cache_ptrs, + torch::Tensor src_block_indices, + torch::Tensor dst_block_indices, + torch::Tensor cum_sum, + int64_t numel_per_block, + torch::ScalarType cache_dtype) { + if (src_block_indices.numel() == 0) { + return; + } + + CHECK(key_cache_ptrs.is_cuda()); + CHECK(value_cache_ptrs.is_cuda()); + CHECK(src_block_indices.is_cuda()); + CHECK(dst_block_indices.is_cuda()); + CHECK(cum_sum.is_cuda()); + CHECK_EQ(key_cache_ptrs.scalar_type(), torch::kInt64); + CHECK_EQ(value_cache_ptrs.scalar_type(), torch::kInt64); + CHECK_EQ(src_block_indices.scalar_type(), torch::kInt32); + CHECK_EQ(dst_block_indices.scalar_type(), torch::kInt32); + CHECK_EQ(cum_sum.scalar_type(), torch::kInt32); + CHECK_EQ(key_cache_ptrs.dim(), 1); + CHECK_EQ(value_cache_ptrs.dim(), 1); + CHECK_EQ(src_block_indices.dim(), 1); + CHECK_EQ(dst_block_indices.dim(), 1); + CHECK_EQ(cum_sum.dim(), 1); + CHECK(key_cache_ptrs.is_contiguous()); + CHECK(value_cache_ptrs.is_contiguous()); + CHECK(src_block_indices.is_contiguous()); + CHECK(dst_block_indices.is_contiguous()); + CHECK(cum_sum.is_contiguous()); + CHECK_EQ(key_cache_ptrs.size(0), value_cache_ptrs.size(0)); + CHECK_EQ(src_block_indices.size(0), cum_sum.size(0)); + CHECK_GT(numel_per_block, 0); + + const at::cuda::OptionalCUDAGuard device_guard(key_cache_ptrs.device()); + constexpr int32_t kThreadsPerBlock = 256; + const int32_t num_layers = static_cast(key_cache_ptrs.size(0)); + const int32_t num_groups = static_cast(src_block_indices.size(0)); + const int32_t num_dst_blocks = + static_cast(dst_block_indices.size(0)); + const cudaStream_t stream = + c10::cuda::getCurrentCUDAStream(key_cache_ptrs.get_device()); + + DISPATCH_FLOATING_TYPES(cache_dtype, "block_copy_kernel", [&] { + constexpr bool kHasVecType = std::is_same_v || + std::is_same_v || + std::is_same_v; + + if constexpr (kHasVecType) { + constexpr int32_t kVecWidth = VecType::vec_width; + if (numel_per_block % kVecWidth == 0) { + const int64_t tiles_per_block = + ceil_div(numel_per_block / kVecWidth, kThreadsPerBlock); + const dim3 grid(num_layers, num_dst_blocks, tiles_per_block); + block_copy_kernel + <<>>( + key_cache_ptrs.data_ptr(), + value_cache_ptrs.data_ptr(), + src_block_indices.data_ptr(), + dst_block_indices.data_ptr(), + cum_sum.data_ptr(), + num_groups, + numel_per_block); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return; + } + } + + const int64_t tiles_per_block = + ceil_div(numel_per_block, kThreadsPerBlock); + const dim3 grid(num_layers, num_dst_blocks, tiles_per_block); + block_copy_kernel<<>>( + key_cache_ptrs.data_ptr(), + value_cache_ptrs.data_ptr(), + src_block_indices.data_ptr(), + dst_block_indices.data_ptr(), + cum_sum.data_ptr(), + num_groups, + numel_per_block); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); +} + +} // namespace xllm::kernel::cuda diff --git a/ex_engine/xllm_kernels/cuda/headers/cuda_ops_api.h b/ex_engine/xllm_kernels/cuda/headers/cuda_ops_api.h new file mode 100644 index 00000000..95fed093 --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/headers/cuda_ops_api.h @@ -0,0 +1,306 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include +#include +#include + +#include +#include +#include + +#include "utils.h" + +namespace xllm::kernel::cuda { + +// TODO: add head_size parameter +void rotary_embedding(torch::Tensor& positions, + torch::Tensor& query, + std::optional key, + torch::Tensor& cos_sin_cache, + // int64_t head_size, + bool is_neox); + +// act_mode only support silu, gelu, gelu_tanh +void act_and_mul(torch::Tensor out, + torch::Tensor input, + const std::string& act_mode); + +void reshape_paged_cache( + torch::Tensor slot_ids, // [n_tokens] + torch::Tensor keys, // [n_tokens, n_kv_heads, head_dim] + torch::Tensor values, // [n_tokens, n_kv_heads, head_dim] + torch::Tensor key_cache, // [n_blocks, block_size, n_heads, head_dim] + torch::Tensor value_cache); + +void block_copy(torch::Tensor key_cache_ptrs, + torch::Tensor value_cache_ptrs, + torch::Tensor src_block_indices, + torch::Tensor dst_block_indices, + torch::Tensor cum_sum, + int64_t numel_per_block, + torch::ScalarType cache_dtype); +#if !defined(USE_DCU) +void batch_prefill(const std::string& uri, + ffi::Array plan_info, + torch::Tensor float_workspace_buffer, + torch::Tensor int_workspace_buffer, + torch::Tensor page_locked_int_workspace_buffer, + torch::Tensor query, + torch::Tensor key, + torch::Tensor value, + torch::Tensor q_cu_seq_lens, + torch::Tensor kv_cu_seq_lens, + int64_t window_left, + double sm_scale, + torch::Tensor output, + std::optional& output_lse, + const std::optional& mask = std::nullopt); + +// Wrapper function for batch_prefill that conditionally uses AttentionRunner +// for piecewise CUDA Graph capture +void batch_prefill_with_optional_piecewise_capture( + const std::string& uri, + ffi::Array plan_info, + torch::Tensor float_workspace_buffer, + torch::Tensor int_workspace_buffer, + torch::Tensor page_locked_int_workspace_buffer, + torch::Tensor query, + torch::Tensor key, + torch::Tensor value, + torch::Tensor q_cu_seq_lens, + torch::Tensor kv_cu_seq_lens, + int64_t window_left, + double sm_scale, + torch::Tensor output, + std::optional& output_lse); + +void batch_prefill_non_causal( + const std::string& uri, + ffi::Array plan_info, + torch::Tensor float_workspace_buffer, + torch::Tensor int_workspace_buffer, + torch::Tensor page_locked_int_workspace_buffer, + torch::Tensor query, + torch::Tensor key, + torch::Tensor value, + torch::Tensor q_cu_seq_lens, + torch::Tensor kv_cu_seq_lens, + int64_t window_left, + double sm_scale, + torch::Tensor output, + std::optional& output_lse, + const std::optional& mask = std::nullopt); + +void batch_chunked_prefill( + const std::string& uri, + ffi::Array plan_info, + torch::Tensor float_workspace_buffer, + torch::Tensor int_workspace_buffer, + torch::Tensor page_locked_int_workspace_buffer, + torch::Tensor query, + torch::Tensor k_cache, + torch::Tensor v_cache, + torch::Tensor paged_kv_indptr, + torch::Tensor paged_kv_indices, + torch::Tensor paged_kv_last_page_len, + int64_t window_left, + double sm_scale, + torch::Tensor output, + std::optional& output_lse, + std::optional qo_indptr = std::nullopt, + bool causal = true); + +void batch_decode(const std::string& uri, + ffi::Array plan_info, + torch::Tensor float_workspace_buffer, + torch::Tensor int_workspace_buffer, + torch::Tensor page_locked_int_workspace_buffer, + torch::Tensor query, + torch::Tensor k_cache, + torch::Tensor v_cache, + torch::Tensor paged_kv_indptr, + torch::Tensor paged_kv_indices, + torch::Tensor paged_kv_last_page_len, + int64_t window_left, + double sm_scale, + torch::Tensor output, + std::optional& output_lse, + bool use_tensor_core, + std::optional qo_indptr = std::nullopt); +#endif // !defined(USE_DCU) +void rms_norm(torch::Tensor output, + torch::Tensor input, + torch::Tensor weight, + double eps); + +void fused_add_rms_norm(torch::Tensor& input, // [..., hidden_size] + torch::Tensor& residual, // [..., hidden_size] + torch::Tensor& weight, // [hidden_size] + double epsilon); + +torch::Tensor matmul(torch::Tensor a, + torch::Tensor b, + std::optional bias); + +void cutlass_scaled_mm(torch::Tensor& c, + torch::Tensor const& a, + torch::Tensor const& b, + torch::Tensor const& a_scales, + torch::Tensor const& b_scales, + std::optional const& bias); + +// Static scaled FP8 quantization +// Quantizes input tensor to FP8 using a pre-computed scale factor +void static_scaled_fp8_quant(torch::Tensor& out, // [..., d] + torch::Tensor const& input, // [..., d] + torch::Tensor const& scale); // [1] + +// FP8 scaled quantize: quantizes input tensor to FP8 e4m3 format +// Returns: (quantized_output, scale) +std::tuple fp8_scaled_quantize( + const torch::Tensor& input, + const std::optional& output = std::nullopt, + const std::optional& scale = std::nullopt); + +// ============================================================================ +// Fused RMSNorm + Static FP8 Quantization +// ============================================================================ +// These functions combine RMSNorm and FP8 quantization to reduce memory +// bandwidth by avoiding the intermediate write-back to global memory. + +// Fused RMSNorm + Static FP8 Quantization (without residual) +// Combines RMSNorm normalization and FP8 quantization in a single kernel. +// This is optimal for the first layer where no residual connection exists. +void rms_norm_static_fp8_quant( + torch::Tensor& out, // [..., hidden_size], FP8 output + torch::Tensor& input, // [..., hidden_size], input tensor + torch::Tensor& weight, // [hidden_size], RMSNorm weight + torch::Tensor& scale, // [1], FP8 quantization scale + double epsilon); // RMSNorm epsilon + +// Fused Add + RMSNorm + Static FP8 Quantization (with residual) +// Combines residual addition, RMSNorm, and FP8 quantization in a single kernel. +// The residual tensor is updated in-place with the sum of input and residual. +void fused_add_rms_norm_static_fp8_quant( + torch::Tensor& out, // [..., hidden_size], FP8 output + torch::Tensor& input, // [..., hidden_size], input tensor + torch::Tensor& residual, // [..., hidden_size], residual (updated in-place) + torch::Tensor& weight, // [hidden_size], RMSNorm weight + torch::Tensor& scale, // [1], FP8 quantization scale + double epsilon); // RMSNorm epsilon + +// FP8 scaled matmul for W8A8 quantization using CUTLASS kernels +// Performs: c = (a @ b.T) with scales applied +torch::Tensor fp8_scaled_matmul( + const torch::Tensor& a, + const torch::Tensor& b, + const torch::Tensor& a_scale, + const torch::Tensor& b_scale, + torch::ScalarType output_dtype, + const std::optional& bias = std::nullopt, + const std::optional& output = std::nullopt); + +std::pair compute_topk_for_beam_search( + torch::Tensor combined_probs, + uint32_t batch_size, + uint32_t beam_size, + uint32_t top_k, + torch::Device device); + +std::pair compute_topk_general( + torch::Tensor input, + uint32_t batch_size, + uint32_t input_length, + uint32_t k, + torch::Device device); + +torch::Tensor air_log_softmax_last_dim(const torch::Tensor& input, + const torch::Tensor& temperatures); + +void fused_qk_norm_rope( + torch::Tensor& qkv, // Combined QKV tensor [num_tokens, + // (num_heads_q+num_heads_k+num_heads_v)*head_dim] + int64_t num_heads_q, // Number of query heads + int64_t num_heads_k, // Number of key heads + int64_t num_heads_v, // Number of value heads + int64_t head_dim, // Dimension per head + double eps, // Epsilon for RMS normalization + const torch::Tensor& q_weight, // RMSNorm weights for query [head_dim] + const torch::Tensor& k_weight, // RMSNorm weights for key [head_dim] + const torch::Tensor& + cos_sin_cache, // Cos/sin cache [max_position, rotary_dim] + bool interleaved, // Whether RoPE is applied in interleaved style + const torch::Tensor& position_ids // Position IDs for RoPE [num_tokens] +); + +std::tuple moe_fused_topk( + torch::Tensor& gating_output, + int64_t topk, + bool renormalize, + const std::optional& correction_bias, + const std::string& scoring_func); + +torch::Tensor random_sample(const torch::Tensor& probs); + +torch::Tensor cutlass_fused_moe( + const torch::Tensor& input, // [num_tokens, hidden] + const torch::Tensor& token_selected_experts, // [num_tokens, top_k] + const torch::Tensor& token_final_scales, // [num_tokens, top_k] + const torch::Tensor& + fc1_expert_weights, // [num_experts, inter_dim, hidden] + const torch::Tensor& + fc2_expert_weights, // [num_experts, hidden, inter_dim] + torch::ScalarType output_dtype, + const std::vector& quant_scales, + int32_t tp_size, + int32_t tp_rank, + int32_t ep_size, + int32_t ep_rank, + int32_t cluster_size, + int32_t cluster_rank, + const std::optional& fc1_expert_biases = std::nullopt, + const std::optional& fc2_expert_biases = std::nullopt, + const std::optional& input_sf = std::nullopt, + const std::optional& swiglu_alpha = std::nullopt, + const std::optional& swiglu_beta = std::nullopt, + const std::optional& swiglu_limit = std::nullopt, + const std::optional& output = std::nullopt, + bool enable_alltoall = false, + bool use_deepseek_fp8_block_scale = false, + bool use_w4_group_scaling = false, + bool use_mxfp8_act_scaling = false, + bool min_latency_mode = false, + bool use_packed_weights = false, + int32_t tune_max_num_tokens = 8192, + ActivationType activation_type = ActivationType::SWIGLU); + +// ---- moe_compute_index (moe_compute_index.cu) ---- +// Fused routing index: bincount + argsort replacement. +// Returns {src_dst, dst_src, expert_sizes}. +std::tuple moe_compute_index( + const torch::Tensor& expert_id, + int64_t num_experts); + +// ---- moe_combine_result (moe_combine.cu) ---- +// Fused combine: reorder + weighted sum in one pass. +torch::Tensor moe_combine_result(const torch::Tensor& gemm2, + const torch::Tensor& reduce_weight, + int64_t N, + int32_t topk); + +} // namespace xllm::kernel::cuda diff --git a/ex_engine/xllm_kernels/cuda/headers/device_utils.cuh b/ex_engine/xllm_kernels/cuda/headers/device_utils.cuh new file mode 100644 index 00000000..77e155b4 --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/headers/device_utils.cuh @@ -0,0 +1,116 @@ +/* Copyright 2025-2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#if defined(USE_DCU) +#include + +#include + +namespace cub = hipcub; +#else +#include +#if CUB_VERSION >= 200800 +#include +#endif +#endif + +namespace xllm::kernel::cuda { +#if !defined(USE_DCU) +using BFloat16Type = __nv_bfloat16; + +#define WARP_SIZE 32 +#define XLLM_KERNEL_ATTR(MAX_THREADS) +#else +using BFloat16Type = hip_bfloat16; + +#define WARP_SIZE 64 +#define XLLM_KERNEL_ATTR(MAX_THREADS) __launch_bounds__(MAX_THREADS, 1) +#endif +#define MAX(a, b) ((a) > (b) ? (a) : (b)) +#define MIN(a, b) ((a) < (b) ? (a) : (b)) + +// Aligned array type +template +class alignas(Alignment) AlignedArray { + T data[N]; +}; + +#define XLLM_SHFL_XOR_SYNC(mask, var, lane_mask) \ + __shfl_xor_sync((mask), (var), (lane_mask)) +#define XLLM_SHFL_XOR_SYNC_WIDTH(mask, var, lane_mask, width) \ + __shfl_xor_sync((mask), (var), (lane_mask), (width)) + +template +__device__ __forceinline__ T xllm_ldg(const T* ptr) { +#if defined(USE_DCU) + return *ptr; +#else + return __ldg(ptr); +#endif +} + +// Define reduction operators based on CUB version. +#if defined(USE_DCU) +using MaxReduceOp = hipcub::Max; +using MinReduceOp = hipcub::Min; +#elif CUB_VERSION >= 200800 +using MaxReduceOp = ::cuda::maximum<>; +using MinReduceOp = ::cuda::minimum<>; +#else +using MaxReduceOp = cub::Max; +using MinReduceOp = cub::Min; +#endif + +template +__device__ float convert_to_float(T x) { + if constexpr (std::is_same_v) { + return __half2float(x); +#if defined(USE_DCU) + } else if constexpr (std::is_same_v) { + return __bfloat162float(reinterpret_cast(x)); +#else + } else if constexpr (std::is_same_v) { + return __bfloat162float(x); +#endif + + } else if constexpr (std::is_same_v) { + return x; + } else { + return static_cast(x); + } +} + +// Constructs some constants needed to partition the work across threads at +// compile time. +template +struct TopkConstants { + static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T); + static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE) == 0 || + EXPERTS % (ELTS_PER_LDG * WARP_SIZE) == 0, + ""); + static constexpr int VECs_PER_THREAD = + MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE)); + static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG; + static constexpr int THREADS_PER_ROW = EXPERTS / VPT; + static constexpr int ROWS_PER_WARP = WARP_SIZE / THREADS_PER_ROW; +}; + +} // namespace xllm::kernel::cuda diff --git a/ex_engine/xllm_kernels/cuda/headers/fp8_quant_utils.cuh b/ex_engine/xllm_kernels/cuda/headers/fp8_quant_utils.cuh new file mode 100644 index 00000000..99b29948 --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/headers/fp8_quant_utils.cuh @@ -0,0 +1,239 @@ +/* Copyright 2025-2026 The xLLM Authors. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://github.com/jd-opensource/xllm/blob/main/LICENSE + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * ===========================================================================*/ + +#pragma once +// clang-format off +#include +#include +#include +// clang-format on +namespace xllm { +namespace kernel { +namespace cuda { + +// FP8 type max value definitions +template || + std::is_same_v>> +struct quant_type_max { + static constexpr T val() { return std::numeric_limits::max(); } +}; + +template +__host__ __device__ static constexpr T quant_type_max_v = + quant_type_max::val(); + +// Minimum scaling factor for quantization types +template || + std::is_same_v>> +struct min_scaling_factor { + __device__ __host__ static inline float val() { + return 1.0f / (quant_type_max_v * 512.0f); + } +}; + +template <> +struct min_scaling_factor { + __device__ __host__ static inline float val() { + return std::numeric_limits::epsilon(); + } +}; + +// Vectorization containers +template +struct __align__(vec_size * sizeof(scalar_t)) vec_n_t { + scalar_t val[vec_size]; +}; + +template +struct __align__(vec_size * sizeof(quant_type_t)) q8_n_t { + static_assert(std::is_same_v || + std::is_same_v); + quant_type_t val[vec_size]; +}; + +// Atomic max for float +__device__ __forceinline__ float atomicMaxFloat(float* addr, float value) { + float old; + old = (value >= 0) + ? __int_as_float(atomicMax((int*)addr, __float_as_int(value))) + : __uint_as_float( + atomicMin((unsigned int*)addr, __float_as_uint(value))); + return old; +} + +// FP8 conversion functions +namespace fp8 { + +#ifdef ENABLE_FP8 + +#include + +// float -> c10::Float8_e4m3fn conversion +template +__inline__ __device__ Tout +vec_conversion(const Tin& x, + const __nv_fp8_interpretation_t fp8_type = __NV_E4M3) { + return x; +} + +template <> +__inline__ __device__ c10::Float8_e4m3fn +vec_conversion( + const float& a, + const __nv_fp8_interpretation_t fp8_type) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return static_cast(a); +#else + return c10::Float8_e4m3fn(__nv_cvt_float_to_fp8(a, __NV_SATFINITE, fp8_type), + c10::Float8_e4m3fn::from_bits()); +#endif +} + +#endif // ENABLE_FP8 + +} // namespace fp8 + +// Scaled FP8 conversion with saturation +template +__device__ __forceinline__ fp8_type scaled_fp8_conversion(float const val, + float const scale) { + float x = 0.0f; + if constexpr (is_scale_inverted) { + x = val * scale; + } else { + x = val / scale; + } + + float r = + fmaxf(-quant_type_max_v, fminf(x, quant_type_max_v)); + +#ifdef ENABLE_FP8 + // Use hardware cvt instruction for fp8 on nvidia + return fp8::vec_conversion(r); +#else + return static_cast(r); +#endif +} + +// Vectorization utilities +template +struct DefaultVecOp { + ScaOp scalar_op; + + __device__ __forceinline__ void operator()( + vec_n_t& dst, + const vec_n_t& src) const { +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + scalar_op(dst.val[i], src.val[i]); + } + } +}; + +template +__device__ inline void vectorize_with_alignment( + const InT* in, + OutT* out, + int len, + int tid, + int stride, + VecOp&& vec_op, // vec_n_t -> vec_n_t + ScaOp&& scalar_op) { // InT -> OutT + static_assert(VEC_SIZE > 0 && (VEC_SIZE & (VEC_SIZE - 1)) == 0, + "VEC_SIZE must be a positive power-of-two"); + constexpr int WIDTH = VEC_SIZE * sizeof(InT); + uintptr_t addr = reinterpret_cast(in); + + // Fast path when the whole region is already aligned + bool can_vec = ((addr & (WIDTH - 1)) == 0) && ((len & (VEC_SIZE - 1)) == 0); + if (can_vec) { + int num_vec = len / VEC_SIZE; + + using vin_t = vec_n_t; + using vout_t = vec_n_t; + auto* v_in = reinterpret_cast(in); + auto* v_out = reinterpret_cast(out); + + for (int i = tid; i < num_vec; i += stride) { + vout_t tmp; + vin_t src = v_in[i]; + vec_op(tmp, src); + v_out[i] = tmp; + } + return; + } + + int misalignment_offset = addr & (WIDTH - 1); + int alignment_bytes = WIDTH - misalignment_offset; + int prefix_elems = alignment_bytes & (WIDTH - 1); + prefix_elems /= sizeof(InT); + prefix_elems = min(prefix_elems, len); + + // Prefix handling + for (int i = tid; i < prefix_elems; i += stride) { + scalar_op(out[i], in[i]); + } + + in += prefix_elems; + out += prefix_elems; + len -= prefix_elems; + + int num_vec = len / VEC_SIZE; + using vin_t = vec_n_t; + using vout_t = vec_n_t; + auto* v_in = reinterpret_cast(in); + auto* v_out = reinterpret_cast(out); + + // Vectorized main part + for (int i = tid; i < num_vec; i += stride) { + vout_t tmp; + vin_t src = v_in[i]; + vec_op(tmp, src); + v_out[i] = tmp; + } + + // Tail handling + int tail_start = num_vec * VEC_SIZE; + for (int i = tid + tail_start; i < len; i += stride) { + scalar_op(out[i], in[i]); + } +} + +template +__device__ __forceinline__ void vectorize_with_alignment(const InT* in, + OutT* out, + int len, + int tid, + int stride, + ScaOp&& scalar_op) { + using Vec = DefaultVecOp>; + vectorize_with_alignment(in, + out, + len, + tid, + stride, + Vec{scalar_op}, + std::forward(scalar_op)); +} + +} // namespace cuda +} // namespace kernel +} // namespace xllm diff --git a/ex_engine/xllm_kernels/cuda/headers/type_convert.cuh b/ex_engine/xllm_kernels/cuda/headers/type_convert.cuh new file mode 100644 index 00000000..835e1656 --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/headers/type_convert.cuh @@ -0,0 +1,231 @@ +/* Copyright 2025 The vLLM Authors and The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include +#include +#include + +// ref to: +// https://github.com/vllm-project/vllm/blob/main/csrc/type_convert.cuh + +/* Converter helpers for the conversion from torch types to HIP/CUDA types, + and the associated type conversions within HIP/CUDA. These helpers need + to be implemented for now because the relevant type conversion + operators/constructors are not consistently implemented by HIP/CUDA, so + a generic conversion via type casts cannot be implemented. + + Each helper should have the member static constexpr bool `exists`: + If false, the optimized kernel is not used for the corresponding torch type. + If true, the helper should be fully defined as shown in the examples below. + */ +namespace xllm::kernel::cuda { +template +class _typeConvert { + public: + static constexpr bool exists = false; +}; + +template <> +class _typeConvert { + public: + static constexpr bool exists = true; + using hip_type = float; + using packed_hip_type = float2; + using packed_hip_type4 = float4; // For 128-bit vectorization + + __device__ static __forceinline__ float convert(hip_type x) { return x; } + __device__ static __forceinline__ float2 convert(packed_hip_type x) { + return x; + } + __device__ static __forceinline__ float4 convert(packed_hip_type4 x) { + return x; + } +}; + +#if defined(USE_DCU) || (defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)) || \ + defined(USE_MACA) +// CUDA < 12.0 runs into issues with packed type conversion +template <> +class _typeConvert { + public: + static constexpr bool exists = true; + using hip_type = __half; + using packed_hip_type = __half2; + + __device__ static __forceinline__ float convert(hip_type x) { + return __half2float(x); + } + __device__ static __forceinline__ float2 convert(packed_hip_type x) { + return __half22float2(x); + } + __device__ static __forceinline__ hip_type convert(float x) { + return __float2half_rn(x); + } + __device__ static __forceinline__ packed_hip_type convert(float2 x) { + return __float22half2_rn(x); + } +}; +#endif // defined(USE_DCU) || CUDA_VERSION >= 12000 + +#if defined(USE_DCU) +template <> +class _typeConvert { + public: + static constexpr bool exists = true; + using hip_type = __hip_bfloat16; + using packed_hip_type = __hip_bfloat162; + + __device__ static __forceinline__ float convert(hip_type x) { + return __bfloat162float(x); + } + __device__ static __forceinline__ float2 convert(packed_hip_type x) { + return __bfloat1622float2(x); + } + __device__ static __forceinline__ hip_type convert(float x) { + return __float2bfloat16(x); + } + __device__ static __forceinline__ packed_hip_type convert(float2 x) { + return __float22bfloat162_rn(x); + } +}; +#elif defined(CUDA_VERSION) && (CUDA_VERSION >= 12000) && \ + defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) || \ + defined(USE_MACA) + +// CUDA_ARCH < 800 does not have BF16 support. +template <> +class _typeConvert { + public: + static constexpr bool exists = true; + using hip_type = __nv_bfloat16; + using packed_hip_type = __nv_bfloat162; + + __device__ static __forceinline__ float convert(hip_type x) { + return __bfloat162float(x); + } + __device__ static __forceinline__ float2 convert(packed_hip_type x) { + return __bfloat1622float2(x); + } + __device__ static __forceinline__ hip_type convert(float x) { + return __float2bfloat16(x); + } + __device__ static __forceinline__ packed_hip_type convert(float2 x) { + return __float22bfloat162_rn(x); + } +}; +#endif + +/* Vector helper to generate vectorized and packed FP16/BF16 ops + for appropriate specializations of fused_add_rms_norm_kernel. + Only functions that are necessary in that kernel are implemented. + Alignment to 16 bytes is required to use 128-bit global memory ops. + */ + +template +class alignas(16) _f16Vec { + public: + /* Not theoretically necessary that width is a power of 2 but should + almost always be the case for optimization purposes */ + static_assert(width > 0 && (width & (width - 1)) == 0, + "Width is not a positive power of 2!"); + using Converter = _typeConvert; + using T1 = typename Converter::hip_type; + using T2 = typename Converter::packed_hip_type; + T1 data[width]; + + __device__ _f16Vec& operator+=(const _f16Vec& other) { + if constexpr (width % 2 == 0) { +#pragma unroll + for (int i = 0; i < width; i += 2) { + if constexpr (std::is_same_v) { + data[i] += other.data[i]; + data[i + 1] += other.data[i + 1]; + } else { + T2 temp{data[i], data[i + 1]}; + temp += T2{other.data[i], other.data[i + 1]}; + data[i] = temp.x; + data[i + 1] = temp.y; + } + } + } else { +#pragma unroll + for (int i = 0; i < width; ++i) data[i] += other.data[i]; + } + return *this; + } + + __device__ _f16Vec& operator*=(const _f16Vec& other) { + if constexpr (width % 2 == 0) { +#pragma unroll + for (int i = 0; i < width; i += 2) { + if constexpr (std::is_same_v) { + data[i] *= other.data[i]; + data[i + 1] *= other.data[i + 1]; + } else { + T2 temp{data[i], data[i + 1]}; + temp *= T2{other.data[i], other.data[i + 1]}; + data[i] = temp.x; + data[i + 1] = temp.y; + } + } + } else { +#pragma unroll + for (int i = 0; i < width; ++i) data[i] *= other.data[i]; + } + return *this; + } + + __device__ _f16Vec& operator*=(const float scale) { + if constexpr (width % 2 == 0) { +#pragma unroll + for (int i = 0; i < width; i += 2) { + float2 temp_f = Converter::convert(T2{data[i], data[i + 1]}); + temp_f.x *= scale; + temp_f.y *= scale; + T2 temp = Converter::convert(temp_f); + data[i] = temp.x; + data[i + 1] = temp.y; + } + } else { +#pragma unroll + for (int i = 0; i < width; ++i) { + float temp = Converter::convert(data[i]) * scale; + data[i] = Converter::convert(temp); + } + } + return *this; + } + + __device__ float sum_squares() const { + float result = 0.0f; + if constexpr (width % 2 == 0) { +#pragma unroll + for (int i = 0; i < width; i += 2) { + float2 z = Converter::convert(T2{data[i], data[i + 1]}); + result += z.x * z.x + z.y * z.y; + } + } else { +#pragma unroll + for (int i = 0; i < width; ++i) { + float x = Converter::convert(data[i]); + result += x * x; + } + } + return result; + } +}; +} // namespace xllm::kernel::cuda diff --git a/ex_engine/xllm_kernels/cuda/headers/utils.h b/ex_engine/xllm_kernels/cuda/headers/utils.h new file mode 100644 index 00000000..020c6ab6 --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/headers/utils.h @@ -0,0 +1,163 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include +#if defined(USE_DCU) +#include +#else +#include +#endif +#include +#include +#if !defined(USE_DCU) +#include +#include +#include +#include +#include +#endif + +#include +#include +#include +#include + +#if defined(__CUDACC__) || defined(_NVHPC_CUDA) || defined(__HIPCC__) +#define HOST_DEVICE_INLINE __host__ __device__ __forceinline__ +#define DEVICE_INLINE __device__ __forceinline__ +#define HOST_INLINE __host__ __forceinline__ +#else +#define HOST_DEVICE_INLINE inline +#define DEVICE_INLINE inline +#define HOST_INLINE inline +#endif + +#if !defined(USE_DCU) +namespace ffi = tvm::ffi; +#endif + +namespace xllm::kernel::cuda { + +template +HOST_DEVICE_INLINE constexpr std::enable_if_t, T> +ceil_div(T a, T b) { + return (a + b - 1) / b; +} + +enum class ActivationType : int8_t { + GELU = 0, + RELU = 1, + SILU = 2, + SWIGLU = 3, + GEGLU = 4, + SWIGLU_BIAS = 5, + RELU2 = 6, + IDENTITY = 7, + INVALID_TYPE = 8 +}; + +// torch tensor is only on cpu +torch::Tensor get_cache_buffer(const int32_t seq_len, + const torch::Device& device); + +// NOLINTBEGIN(cppcoreguidelines-macro-usage) +#define DISPATCH_CASE_FLOATING_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) +#define DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH(TYPE, NAME, DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) +#define DISPATCH_CASE_HALF_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) +#define DISPATCH_HALF_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH(TYPE, NAME, DISPATCH_CASE_HALF_TYPES(__VA_ARGS__)) +// NOLINTEND(cppcoreguidelines-macro-usage) + +bool should_use_tensor_core(torch::ScalarType kv_cache_dtype, + int64_t num_attention_heads, + int64_t num_kv_heads); + +bool support_pdl(); + +std::string path_to_uri_so_lib(const std::string& uri); + +std::string determine_attention_backend(int64_t pos_encoding_mode, + bool use_fp16_qk_reduction, + bool use_custom_mask); + +std::string get_batch_prefill_uri(const std::string& backend, + torch::ScalarType dtype_q, + torch::ScalarType dtype_kv, + torch::ScalarType dtype_o, + torch::ScalarType dtype_idx, + int64_t head_dim_qk, + int64_t head_dim_vo, + int64_t pos_encoding_mode, + bool use_sliding_window, + bool use_logits_soft_cap, + bool use_fp16_qk_reduction); + +std::string get_batch_decode_uri(torch::ScalarType dtype_q, + torch::ScalarType dtype_kv, + torch::ScalarType dtype_o, + torch::ScalarType dtype_idx, + int64_t head_dim_qk, + int64_t head_dim_vo, + int64_t pos_encoding_mode, + bool use_sliding_window, + bool use_logits_soft_cap); + +std::tuple split_scale_param(const torch::Tensor& scale); + +#if !defined(USE_DCU) +DLDataType to_dl_data_type(torch::ScalarType scalar_type); + +// below are tvm-ffi related functions +ffi::Tensor to_ffi_tensor(const torch::Tensor& torch_tensor); + +ffi::Optional to_ffi_optional_tensor( + const std::optional& optional); + +ffi::Array to_ffi_array_tensors( + const std::vector& torch_tensors); + +ffi::Optional> to_ffi_optional_array_tensors( + const std::optional>& optional); + +ffi::Module get_module(const std::string& uri); + +ffi::Function get_function(const std::string& uri, + const std::string& func_name); + +inline void bind_tvmffi_stream_to_current_torch_stream( + const torch::Device& device) { + const auto cur = c10::cuda::getCurrentCUDAStream(device.index()); + // DLPack device type for CUDA is 2 (kDLCUDA). + void* original_stream = nullptr; + const int rc = TVMFFIEnvSetStream( + /*device_type=*/2, + /*device_id=*/device.index(), + reinterpret_cast(cur.stream()), + &original_stream); + if (rc != 0) { + LOG(WARNING) << "[tvmffi.stream] failed to set stream, rc=" << rc + << " dev=" << device.index(); + } +} +#endif // !defined(USE_DCU) +} // namespace xllm::kernel::cuda diff --git a/ex_engine/xllm_kernels/cuda/norm.cu b/ex_engine/xllm_kernels/cuda/norm.cu new file mode 100644 index 00000000..30e70084 --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/norm.cu @@ -0,0 +1,600 @@ +/* Copyright 2025 The vLLM Authors and The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include + +#include +#include + +#include "cuda_ops_api.h" +#include "device_utils.cuh" +#include "fp8_quant_utils.cuh" +#include "type_convert.cuh" + +// ref to: +// https://github.com/vllm-project/vllm/blob/main/csrc/layernorm_kernels.cu + +#if CUB_VERSION >= 200800 +#include +using CubAddOp = ::cuda::std::plus<>; +using CubMaxOp = ::cuda::maximum<>; +#else // if CUB_VERSION < 200800 +using CubAddOp = cub::Sum; +using CubMaxOp = cub::Max; +#endif // CUB_VERSION + +namespace { + +using namespace xllm::kernel::cuda; + +template +__global__ void XLLM_KERNEL_ATTR(1024) + rms_norm_kernel(scalar_t* __restrict__ out, // [..., hidden_size] + const scalar_t* __restrict__ input, // [..., hidden_size] + const int64_t input_stride, + const scalar_t* __restrict__ weight, // [hidden_size] + const float epsilon, + const int num_tokens, + const int hidden_size) { + __shared__ float s_variance; + float variance = 0.0f; + + for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { + const float x = static_cast(input[blockIdx.x * input_stride + idx]); + variance += x * x; + } + + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage reduceStore; + variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x); + + if (threadIdx.x == 0) { + s_variance = rsqrtf(variance / hidden_size + epsilon); + } + __syncthreads(); + + for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { + float x = static_cast(input[blockIdx.x * input_stride + idx]); + out[blockIdx.x * hidden_size + idx] = + (static_cast(x * s_variance)) * weight[idx]; + } +} + +/* Function specialization in the case of FP16/BF16 tensors. + Additional optimizations we can make in this case are + packed and vectorized operations, which help with the + memory latency bottleneck. */ +template +__global__ std::enable_if_t<(width > 0) && _typeConvert::exists> +XLLM_KERNEL_ATTR(1024) fused_add_rms_norm_kernel( + scalar_t* __restrict__ input, // [..., hidden_size] + const int64_t input_stride, + scalar_t* __restrict__ residual, // [..., hidden_size] + const scalar_t* __restrict__ weight, // [hidden_size] + const float epsilon, + const int num_tokens, + const int hidden_size) { + // Sanity checks on our vector struct and type-punned pointer arithmetic + static_assert(std::is_pod_v<_f16Vec>); + static_assert(sizeof(_f16Vec) == sizeof(scalar_t) * width); + + const int vec_hidden_size = hidden_size / width; + const int64_t vec_input_stride = input_stride / width; + __shared__ float s_variance; + float variance = 0.0f; + /* These and the argument pointers are all declared `restrict` as they are + not aliased in practice. Argument pointers should not be dereferenced + in this kernel as that would be undefined behavior */ + auto* __restrict__ input_v = + reinterpret_cast<_f16Vec*>(input); + auto* __restrict__ residual_v = + reinterpret_cast<_f16Vec*>(residual); + auto* __restrict__ weight_v = + reinterpret_cast*>(weight); + + for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) { + int id = blockIdx.x * vec_hidden_size + idx; + int64_t strided_id = blockIdx.x * vec_input_stride + idx; + _f16Vec temp = input_v[strided_id]; + temp += residual_v[id]; + variance += temp.sum_squares(); + residual_v[id] = temp; + } + + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage reduceStore; + variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x); + + if (threadIdx.x == 0) { + s_variance = rsqrtf(variance / hidden_size + epsilon); + } + __syncthreads(); + + for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) { + int id = blockIdx.x * vec_hidden_size + idx; + int64_t strided_id = blockIdx.x * vec_input_stride + idx; + _f16Vec temp = residual_v[id]; + temp *= s_variance; + temp *= weight_v[idx]; + input_v[strided_id] = temp; + } +} + +/* Generic fused_add_rms_norm_kernel + The width field is not used here but necessary for other specializations. + */ +template +__global__ std::enable_if_t<(width == 0) || !_typeConvert::exists> +XLLM_KERNEL_ATTR(1024) fused_add_rms_norm_kernel( + scalar_t* __restrict__ input, // [..., hidden_size] + const int64_t input_stride, + scalar_t* __restrict__ residual, // [..., hidden_size] + const scalar_t* __restrict__ weight, // [hidden_size] + const float epsilon, + const int num_tokens, + const int hidden_size) { + __shared__ float s_variance; + float variance = 0.0f; + + for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { + scalar_t z = input[blockIdx.x * input_stride + idx]; + z += residual[blockIdx.x * hidden_size + idx]; + float x = static_cast(z); + variance += x * x; + residual[blockIdx.x * hidden_size + idx] = z; + } + + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage reduceStore; + variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x); + + if (threadIdx.x == 0) { + s_variance = rsqrtf(variance / hidden_size + epsilon); + } + __syncthreads(); + + for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { + float x = static_cast(residual[blockIdx.x * hidden_size + idx]); + input[blockIdx.x * input_stride + idx] = + (static_cast(x * s_variance)) * weight[idx]; + } +} + +#define LAUNCH_FUSED_ADD_RMS_NORM(width) \ + DISPATCH_FLOATING_TYPES( \ + input.scalar_type(), "fused_add_rms_norm_kernel", [&] { \ + fused_add_rms_norm_kernel \ + <<>>(input.data_ptr(), \ + input_stride, \ + residual.data_ptr(), \ + weight.data_ptr(), \ + epsilon, \ + num_tokens, \ + hidden_size); \ + }); + +// ============================================================================ +// Fused RMSNorm + Static FP8 Quantization Kernels +// ============================================================================ +// These kernels combine RMSNorm and FP8 quantization to reduce memory +// bandwidth by avoiding the intermediate write-back to global memory. + +// Dispatch macro for FP8 types +#define DISPATCH_FP8_TYPES(TYPE, NAME, ...) \ + [&] { \ + const auto& the_type = TYPE; \ + switch (the_type) { \ + case at::ScalarType::Float8_e4m3fn: { \ + using fp8_t = c10::Float8_e4m3fn; \ + return __VA_ARGS__(); \ + } \ + default: \ + AT_ERROR(#NAME, \ + " not implemented for FP8 type '", \ + toString(the_type), \ + "'"); \ + } \ + }() + +/** + * Fused RMSNorm + Static FP8 Quantization kernel (without residual) + * Combines RMSNorm and FP8 quantization in a single kernel to reduce + * memory bandwidth by avoiding intermediate write-back. + * + * @tparam scalar_t Input data type (float, half, bfloat16) + * @tparam fp8_type Output FP8 type (c10::Float8_e4m3fn) + * @param out Output FP8 tensor [num_tokens, hidden_size] + * @param input Input tensor [num_tokens, hidden_size] + * @param input_stride Stride of input tensor in the token dimension + * @param weight RMSNorm weight tensor [hidden_size] + * @param scale FP8 quantization scale (scalar) + * @param epsilon RMSNorm epsilon + * @param num_tokens Number of tokens + * @param hidden_size Hidden dimension size + */ +template +__global__ void rms_norm_static_fp8_quant_kernel( + fp8_type* __restrict__ out, // [num_tokens, hidden_size] + const scalar_t* __restrict__ input, // [num_tokens, hidden_size] + const int64_t input_stride, + const scalar_t* __restrict__ weight, // [hidden_size] + const float* __restrict__ scale, // [1] + const float epsilon, + const int num_tokens, + const int hidden_size) { + __shared__ float s_variance; + float variance = 0.0f; + + const scalar_t* input_row = input + blockIdx.x * input_stride; + + // Step 1: Compute variance for RMSNorm + for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { + const float x = static_cast(input_row[idx]); + variance += x * x; + } + + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage reduceStore; + variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x); + + if (threadIdx.x == 0) { + s_variance = rsqrtf(variance / hidden_size + epsilon); + } + __syncthreads(); + + // Step 2: Precompute scale inverse to avoid division + const float scale_inv = 1.0f / (*scale); + + // Step 3: Fused RMSNorm + FP8 quantization + for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { + float x = static_cast(input_row[idx]); + float out_norm = (static_cast(x * s_variance)) * + static_cast(weight[idx]); + out[blockIdx.x * hidden_size + idx] = + xllm::kernel::cuda::scaled_fp8_conversion(out_norm, + scale_inv); + } +} + +/** + * Fused Add + RMSNorm + Static FP8 Quantization kernel (with residual) + * Optimized version with packed + vectorized operations for FP16/BF16. + * + * @tparam scalar_t Input data type (float, half, bfloat16) + * @tparam width Vector width for optimization (0, 8) + * @tparam fp8_type Output FP8 type (c10::Float8_e4m3fn) + */ +template +__global__ std::enable_if_t<(width > 0) && _typeConvert::exists> +fused_add_rms_norm_static_fp8_quant_kernel( + fp8_type* __restrict__ out, // [num_tokens, hidden_size] + scalar_t* __restrict__ input, // [num_tokens, hidden_size] + const int64_t input_stride, + scalar_t* __restrict__ residual, // [num_tokens, hidden_size] + const scalar_t* __restrict__ weight, // [hidden_size] + const float* __restrict__ scale, // [1] + const float epsilon, + const int num_tokens, + const int hidden_size) { + static_assert(std::is_pod_v<_f16Vec>); + static_assert(sizeof(_f16Vec) == sizeof(scalar_t) * width); + + const int vec_hidden_size = hidden_size / width; + const int64_t vec_input_stride = input_stride / width; + __shared__ float s_variance; + float variance = 0.0f; + + auto* __restrict__ input_v = + reinterpret_cast<_f16Vec*>(input); + auto* __restrict__ residual_v = + reinterpret_cast<_f16Vec*>(residual); + auto* __restrict__ weight_v = + reinterpret_cast*>(weight); + + // Step 1: Fused add and compute variance + for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) { + int id = blockIdx.x * vec_hidden_size + idx; + int64_t strided_id = blockIdx.x * vec_input_stride + idx; + _f16Vec temp = input_v[strided_id]; + temp += residual_v[id]; + variance += temp.sum_squares(); + residual_v[id] = temp; // Store updated residual + } + + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage reduceStore; + variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x); + + if (threadIdx.x == 0) { + s_variance = rsqrtf(variance / hidden_size + epsilon); + } + __syncthreads(); + + // Step 2: Precompute scale inverse + const float scale_inv = 1.0f / (*scale); + + // Step 3: Fused RMSNorm + FP8 quantization + for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) { + int id = blockIdx.x * vec_hidden_size + idx; + _f16Vec temp = residual_v[id]; + temp *= s_variance; + temp *= weight_v[idx]; + + // Convert each element to FP8 +#pragma unroll + for (int i = 0; i < width; ++i) { + float val = _typeConvert::convert(temp.data[i]); + out[id * width + i] = + xllm::kernel::cuda::scaled_fp8_conversion(val, + scale_inv); + } + } +} + +/** + * Generic fused add + RMSNorm + FP8 quant kernel (fallback for unaligned data) + */ +template +__global__ std::enable_if_t<(width == 0) || !_typeConvert::exists> +fused_add_rms_norm_static_fp8_quant_kernel( + fp8_type* __restrict__ out, // [num_tokens, hidden_size] + scalar_t* __restrict__ input, // [num_tokens, hidden_size] + const int64_t input_stride, + scalar_t* __restrict__ residual, // [num_tokens, hidden_size] + const scalar_t* __restrict__ weight, // [hidden_size] + const float* __restrict__ scale, // [1] + const float epsilon, + const int num_tokens, + const int hidden_size) { + __shared__ float s_variance; + float variance = 0.0f; + + // Step 1: Fused add and compute variance + for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { + scalar_t z = input[blockIdx.x * input_stride + idx]; + z += residual[blockIdx.x * hidden_size + idx]; + float x = static_cast(z); + variance += x * x; + residual[blockIdx.x * hidden_size + idx] = z; // Store updated residual + } + + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage reduceStore; + variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x); + + if (threadIdx.x == 0) { + s_variance = rsqrtf(variance / hidden_size + epsilon); + } + __syncthreads(); + + // Step 2: Precompute scale inverse + const float scale_inv = 1.0f / (*scale); + + // Step 3: Fused RMSNorm + FP8 quantization + for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { + float x = static_cast(residual[blockIdx.x * hidden_size + idx]); + float out_norm = (static_cast(x * s_variance)) * + static_cast(weight[idx]); + out[blockIdx.x * hidden_size + idx] = + xllm::kernel::cuda::scaled_fp8_conversion(out_norm, + scale_inv); + } +} + +#define LAUNCH_FUSED_ADD_RMS_NORM_STATIC_FP8_QUANT(width) \ + DISPATCH_FLOATING_TYPES( \ + input.scalar_type(), "fused_add_rms_norm_static_fp8_quant", [&] { \ + DISPATCH_FP8_TYPES( \ + out.scalar_type(), "fused_add_rms_norm_static_fp8_quant", [&] { \ + fused_add_rms_norm_static_fp8_quant_kernel \ + <<>>(out.data_ptr(), \ + input.data_ptr(), \ + input_stride, \ + residual.data_ptr(), \ + weight.data_ptr(), \ + scale.data_ptr(), \ + epsilon, \ + num_tokens, \ + hidden_size); \ + }); \ + }); + +} // namespace + +namespace xllm::kernel::cuda { + +// flashinfer rmsnorm ops +// void rmsnorm(torch::Tensor output, +// torch::Tensor input, +// torch::Tensor weight, +// double eps) { +// FunctionFactory::get_instance().rmsnorm_func("norm").call( +// output, input, weight, eps, support_pdl()); +// } + +void rms_norm(torch::Tensor output, // [..., hidden_size] + torch::Tensor input, // [..., hidden_size] + torch::Tensor weight, // [hidden_size] + double eps) { + CHECK(output.is_contiguous()); + CHECK(weight.is_contiguous()); + + // The kernel addresses tokens as `blockIdx.x * input_stride + idx`, which + // can only represent contiguous inputs or simple 2D strided rows. Flux q/k + // tensors reach this path as high-dimensional transposed views, so make that + // layout explicit before flattening tokens for the kernel. + if (input.dim() > 2 && !input.is_contiguous()) { + input = input.contiguous(); + } + CHECK(input.stride(-1) == 1); + + int hidden_size = input.size(-1); + int num_tokens = input.numel() / hidden_size; + int64_t input_stride = input.stride(-2); + + dim3 grid(num_tokens); + dim3 block(std::min(hidden_size, 1024)); + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_kernel", [&] { + rms_norm_kernel + <<>>(output.data_ptr(), + input.data_ptr(), + input_stride, + weight.data_ptr(), + eps, + num_tokens, + hidden_size); + }); +} + +void fused_add_rms_norm(torch::Tensor& input, // [..., hidden_size] + torch::Tensor& residual, // [..., hidden_size] + torch::Tensor& weight, // [hidden_size] + double epsilon) { + CHECK(weight.scalar_type() == input.scalar_type()); + CHECK(input.scalar_type() == residual.scalar_type()); + CHECK(residual.is_contiguous()); + CHECK(weight.is_contiguous()); + int hidden_size = input.size(-1); + int64_t input_stride = input.stride(-2); + int num_tokens = input.numel() / hidden_size; + + dim3 grid(num_tokens); + /* This kernel is memory-latency bound in many scenarios. + When num_tokens is large, a smaller block size allows + for increased block occupancy on CUs and better latency + hiding on global mem ops. */ + const int max_block_size = (num_tokens < 256) ? 1024 : 256; + dim3 block(std::min(hidden_size, max_block_size)); + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + /*If the tensor types are FP16/BF16, try to use the optimized kernel + with packed + vectorized ops. + Max optimization is achieved with a width-8 vector of FP16/BF16s + since we can load at most 128 bits at once in a global memory op. + However, this requires each tensor's data to be aligned to 16 + bytes. + */ + auto inp_ptr = reinterpret_cast(input.data_ptr()); + auto res_ptr = reinterpret_cast(residual.data_ptr()); + auto wt_ptr = reinterpret_cast(weight.data_ptr()); + constexpr int kVectorWidth = 8; + constexpr int kReqAlignmentBytes = + kVectorWidth * 2; // kVectorWidth * sizeof(bfloat16 or float16) (float32 + // falls back to non-vectorized version anyway) + bool ptrs_are_aligned = inp_ptr % kReqAlignmentBytes == 0 && + res_ptr % kReqAlignmentBytes == 0 && + wt_ptr % kReqAlignmentBytes == 0; + bool offsets_are_multiple_of_vector_width = + hidden_size % kVectorWidth == 0 && input_stride % kVectorWidth == 0; + if (ptrs_are_aligned && offsets_are_multiple_of_vector_width) { + LAUNCH_FUSED_ADD_RMS_NORM(8); + } else { + LAUNCH_FUSED_ADD_RMS_NORM(0); + } +} + +// ============================================================================ +// Fused RMSNorm + Static FP8 Quantization Host Functions +// ============================================================================ + +void rms_norm_static_fp8_quant(torch::Tensor& out, // [..., hidden_size], FP8 + torch::Tensor& input, // [..., hidden_size] + torch::Tensor& weight, // [hidden_size] + torch::Tensor& scale, // [1] + double epsilon) { + CHECK(out.is_contiguous()); + CHECK(input.stride(-1) == 1); + CHECK(weight.is_contiguous()); + CHECK(scale.is_contiguous()); + + int hidden_size = input.size(-1); + int64_t input_stride = input.stride(-2); + int num_tokens = input.numel() / hidden_size; + + // For large num_tokens, use smaller blocks to increase SM concurrency + const int max_block_size = (num_tokens < 256) ? 1024 : 256; + dim3 grid(num_tokens); + dim3 block(std::min(hidden_size, max_block_size)); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + DISPATCH_FLOATING_TYPES( + input.scalar_type(), "rms_norm_static_fp8_quant", [&] { + DISPATCH_FP8_TYPES(out.scalar_type(), "rms_norm_static_fp8_quant", [&] { + rms_norm_static_fp8_quant_kernel + <<>>(out.data_ptr(), + input.data_ptr(), + input_stride, + weight.data_ptr(), + scale.data_ptr(), + epsilon, + num_tokens, + hidden_size); + }); + }); +} + +void fused_add_rms_norm_static_fp8_quant( + torch::Tensor& out, // [..., hidden_size], FP8 + torch::Tensor& input, // [..., hidden_size] + torch::Tensor& residual, // [..., hidden_size] + torch::Tensor& weight, // [hidden_size] + torch::Tensor& scale, // [1] + double epsilon) { + CHECK(out.is_contiguous()); + CHECK(residual.is_contiguous()); + CHECK(weight.is_contiguous()); + CHECK(scale.is_contiguous()); + CHECK(residual.scalar_type() == input.scalar_type()); + CHECK(weight.scalar_type() == input.scalar_type()); + + int hidden_size = input.size(-1); + int64_t input_stride = input.stride(-2); + int num_tokens = input.numel() / hidden_size; + + dim3 grid(num_tokens); + const int max_block_size = (num_tokens < 256) ? 1024 : 256; + dim3 block(std::min(hidden_size, max_block_size)); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + // Check alignment for vectorized kernel + auto inp_ptr = reinterpret_cast(input.data_ptr()); + auto res_ptr = reinterpret_cast(residual.data_ptr()); + auto wt_ptr = reinterpret_cast(weight.data_ptr()); + constexpr int kVectorWidth = 8; + constexpr int kReqAlignmentBytes = kVectorWidth * 2; + + bool ptrs_are_aligned = inp_ptr % kReqAlignmentBytes == 0 && + res_ptr % kReqAlignmentBytes == 0 && + wt_ptr % kReqAlignmentBytes == 0; + bool offsets_are_multiple_of_vector_width = + hidden_size % kVectorWidth == 0 && input_stride % kVectorWidth == 0; + + if (ptrs_are_aligned && offsets_are_multiple_of_vector_width) { + LAUNCH_FUSED_ADD_RMS_NORM_STATIC_FP8_QUANT(8); + } else { + LAUNCH_FUSED_ADD_RMS_NORM_STATIC_FP8_QUANT(0); + } +} + +} // namespace xllm::kernel::cuda diff --git a/ex_engine/xllm_kernels/cuda/reshape_paged_cache.cu b/ex_engine/xllm_kernels/cuda/reshape_paged_cache.cu new file mode 100644 index 00000000..3fe3db91 --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/reshape_paged_cache.cu @@ -0,0 +1,101 @@ +/* Copyright 2025-2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include + +#include "cuda_ops_api.h" +#include "device_utils.cuh" + +namespace xllm::kernel::cuda { + +template +__global__ void XLLM_KERNEL_ATTR(1024) reshape_paged_cache_kernel( + const int* __restrict__ slot_ids, // [n_tokens] + const T* __restrict__ keys, // [n_tokens, n_heads, head_dim] + const T* __restrict__ values, // [n_tokens, n_heads, head_dim] + T* __restrict__ key_cache, + T* __restrict__ value_cache, + int64_t k_stride, + int64_t v_stride, + int64_t n_kv_heads, + int64_t head_dim, + int64_t block_size) { + // block/token index + const int64_t bid = blockIdx.x; + // which slot to write to + const int64_t slot_id = slot_ids[bid]; + if (slot_id < 0) { + return; + } + // block index + const int64_t block_idx = slot_id / block_size; + // offset within block + const int64_t block_offset = slot_id % block_size; + // base index for the block in cache + const int64_t block_base_idx = block_idx * block_size * n_kv_heads * head_dim; + // copy value one by one for the token + for (int64_t i = threadIdx.x; i < n_kv_heads * head_dim; i += blockDim.x) { + const int64_t k_src_idx = bid * k_stride + i; + const int64_t v_src_idx = bid * v_stride + i; + // cache: [n_blocks, block_size, n_heads, head_dim] + const int64_t head_base_idx = + block_base_idx + block_offset * n_kv_heads * head_dim; + // which head to write to + const int head_idx = i / head_dim; + // which dim within head to write to + const int head_offset = i % head_dim; + const int64_t dst_idx = head_base_idx + head_idx * head_dim + head_offset; + key_cache[dst_idx] = keys[k_src_idx]; + value_cache[dst_idx] = values[v_src_idx]; + } +} + +void reshape_paged_cache( + torch::Tensor slot_ids, // [n_tokens] + torch::Tensor keys, // [n_tokens, n_kv_heads, head_dim] + torch::Tensor values, // [n_tokens, n_kv_heads, head_dim] + torch::Tensor key_cache, // [n_blocks, block_size, n_heads, head_dim] + torch::Tensor value_cache) { + // keys and values should be continuous at n_kv_heads and head_dim dims + CHECK(keys.stride(-1) == 1 && keys.stride(-2) == keys.size(-1)); + CHECK(values.stride(-1) == 1 && values.stride(-2) == values.size(-1)); + const int64_t n_tokens = keys.size(-3); + const int64_t n_kv_heads = keys.size(-2); + const int64_t head_dim = keys.size(-1); + const int64_t block_size = key_cache.size(-3); + // it is possible that keys and values have different strides + const int64_t k_stride = keys.stride(-3); + const int64_t v_stride = values.stride(-3); + const int64_t n = n_kv_heads * head_dim; + dim3 grid(n_tokens); + dim3 block(std::min(n, 1024)); + DISPATCH_FLOATING_TYPES( + keys.scalar_type(), "reshape_paged_cache_kernel", [&] { + reshape_paged_cache_kernel + <<>>( + slot_ids.data_ptr(), + keys.data_ptr(), + values.data_ptr(), + key_cache.data_ptr(), + value_cache.data_ptr(), + k_stride, + v_stride, + n_kv_heads, + head_dim, + block_size); + }); +} + +} // namespace xllm::kernel::cuda diff --git a/ex_engine/xllm_kernels/cuda/rope.cu b/ex_engine/xllm_kernels/cuda/rope.cu new file mode 100644 index 00000000..585cdf99 --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/rope.cu @@ -0,0 +1,258 @@ +/* Copyright 2025 The vLLM Authors and The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ +#include +#include +#include + +#include "cuda_ops_api.h" +#include "device_utils.cuh" + +// ref to: +// https://github.com/vllm-project/vllm/blob/main/csrc/pos_encoding_kernels.cu + +namespace { + +template +inline __device__ void apply_token_rotary_embedding( + scalar_t* __restrict__ arr, + const scalar_t* __restrict__ cos_ptr, + const scalar_t* __restrict__ sin_ptr, + int rot_offset, + int embed_dim) { + int x_index, y_index; + scalar_t cos, sin; + if (IS_NEOX) { + // GPT-NeoX style rotary embedding. + x_index = rot_offset; + y_index = embed_dim + rot_offset; + cos = *(cos_ptr + x_index); + sin = *(sin_ptr + x_index); + } else { + // GPT-J style rotary embedding. + x_index = 2 * rot_offset; + y_index = 2 * rot_offset + 1; + cos = *(cos_ptr + x_index / 2); + sin = *(sin_ptr + x_index / 2); + } + + const scalar_t x = arr[x_index]; + const scalar_t y = arr[y_index]; + arr[x_index] = x * cos - y * sin; + arr[y_index] = y * cos + x * sin; +} + +template +inline __device__ void apply_rotary_embedding( + scalar_t* __restrict__ query, // [batch_size, seq_len, num_heads, + // head_size] or [num_tokens, num_heads, + // head_size] + scalar_t* __restrict__ key, // nullptr or + // [batch_size, seq_len, num_kv_heads, + // head_size] or [num_tokens, num_kv_heads, + // head_size] + const scalar_t* cache_ptr, + const int head_size, + const int num_heads, + const int num_kv_heads, + const int rot_dim, + const int token_idx, + const int64_t query_stride, + const int64_t key_stride, + const int64_t head_stride) { + const int embed_dim = rot_dim / 2; + const scalar_t* cos_ptr = cache_ptr; + const scalar_t* sin_ptr = cache_ptr + embed_dim; + + const int nq = num_heads * embed_dim; + for (int i = threadIdx.x; i < nq; i += blockDim.x) { + const int head_idx = i / embed_dim; + const int64_t token_head = + token_idx * query_stride + head_idx * head_stride; + const int rot_offset = i % embed_dim; + apply_token_rotary_embedding( + query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim); + } + + if (key != nullptr) { + const int nk = num_kv_heads * embed_dim; + for (int i = threadIdx.x; i < nk; i += blockDim.x) { + const int head_idx = i / embed_dim; + const int64_t token_head = + token_idx * key_stride + head_idx * head_stride; + const int rot_offset = i % embed_dim; + apply_token_rotary_embedding( + key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim); + } + } +} + +template +__global__ void XLLM_KERNEL_ATTR(512) rotary_embedding_kernel( + const int64_t* __restrict__ positions, // [batch_size, seq_len] or + // [num_tokens] + scalar_t* __restrict__ query, // [batch_size, seq_len, num_heads, + // head_size] or [num_tokens, num_heads, + // head_size] + scalar_t* __restrict__ key, // nullptr or + // [batch_size, seq_len, num_kv_heads, + // head_size] or [num_tokens, num_kv_heads, + // head_size] + const scalar_t* __restrict__ cos_sin_cache, // [max_position, 2, + // rot_dim // 2] + const int rot_dim, + const int64_t query_stride, + const int64_t key_stride, + const int64_t head_stride, + const int num_heads, + const int num_kv_heads, + const int head_size) { + // Each thread block is responsible for one token. + const int token_idx = blockIdx.x; + int64_t pos = positions[token_idx]; + const scalar_t* cache_ptr = cos_sin_cache + pos * rot_dim; + + apply_rotary_embedding(query, + key, + cache_ptr, + head_size, + num_heads, + num_kv_heads, + rot_dim, + token_idx, + query_stride, + key_stride, + head_stride); +} +} // namespace + +namespace xllm::kernel::cuda { + +// flashinfer rope ops +// void apply_rope_pos_ids_cos_sin_cache(torch::Tensor q, +// torch::Tensor k, +// torch::Tensor cos_sin_cache, +// torch::Tensor pos_ids, +// bool interleave) { +// const int64_t head_dim = cos_sin_cache.size(-1) / 2; +// q = q.view({q.size(0), -1, head_dim}); +// k = k.view({k.size(0), -1, head_dim}); + +// FunctionFactory::get_instance().rope_func("rope").call( +// q, k, q, k, cos_sin_cache, pos_ids, interleave); +// } + +void rotary_embedding( + torch::Tensor& positions, // [batch_size, seq_len] or [num_tokens] + torch::Tensor& query, // [batch_size, seq_len, num_heads * head_size] or + // [num_tokens, num_heads * head_size] or + // [batch_size, seq_len, num_heads, head_size] or + // [num_tokens, num_heads, head_size] + std::optional key, + // null or + // [batch_size, seq_len, num_kv_heads * head_size] or + // [num_tokens, num_kv_heads * head_size] or + // [batch_size, seq_len, num_heads, head_size] or + // [num_tokens, num_heads, head_size] + // int64_t head_size, + torch::Tensor& cos_sin_cache, // [max_position, rot_dim] + bool is_neox) { + // num_tokens = batch_size * seq_len + const int positions_ndim = positions.dim(); + const int query_ndim = query.dim(); + // For partial rotary models, e.g. MiniMax-M2 with head_dim=128 and + // rotary_dim=64, the cache width is the rotary dimension rather than the + // physical per-head stride. When query is already shaped as + // [*, num_heads, head_size], infer the real head_size from query itself. + int64_t head_size = (query_ndim == positions_ndim + 2) + ? query.size(-1) + : cos_sin_cache.size(-1); + int64_t num_tokens = positions.numel(); + + // Make sure num_tokens dim is consistent across positions, query, and key + CHECK(positions_ndim == 1 || positions_ndim == 2) + << "positions must have shape [num_tokens] or [batch_size, seq_len]"; + + if (positions_ndim == 1) { + CHECK(query.size(0) == positions.size(0) && + (!key.has_value() || key->size(0) == positions.size(0))) + << "query, key and positions must have the same number of tokens"; + } + if (positions_ndim == 2) { + CHECK(query.size(0) == positions.size(0) && + (!key.has_value() || key->size(0) == positions.size(0)) && + query.size(1) == positions.size(1) && + (!key.has_value() || key->size(1) == positions.size(1))) + << "query, key and positions must have the same batch_size and seq_len"; + } + + // Make sure head_size is valid for query and key + // hidden_size = num_heads * head_size + int query_hidden_size = query.numel() / num_tokens; + int key_hidden_size = key.has_value() ? key->numel() / num_tokens : 0; + CHECK(query_hidden_size % head_size == 0); + CHECK(key_hidden_size % head_size == 0); + + // Make sure query and key have consistent number of heads + int num_heads = query_hidden_size / head_size; + int num_kv_heads = key.has_value() ? key_hidden_size / head_size : num_heads; + CHECK(num_heads % num_kv_heads == 0); + + int rot_dim = cos_sin_cache.size(1); + int seq_dim_idx = positions_ndim - 1; + int64_t query_stride = query.stride(seq_dim_idx); + int64_t key_stride = key.has_value() ? key->stride(seq_dim_idx) : 0; + // Determine head stride: for [*, heads, head_size] use stride of last dim; + // for flat [*, heads*head_size], heads blocks are contiguous of size + // head_size + int64_t head_stride = + (query_ndim == positions_ndim + 2) ? query.stride(-2) : head_size; + + dim3 grid(num_tokens); + dim3 block(std::min(num_heads * rot_dim / 2, 512)); + const at::cuda::OptionalCUDAGuard device_guard(device_of(query)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + DISPATCH_FLOATING_TYPES( + query.scalar_type(), "apply_rope_pos_ids_cos_sin_cache", [&] { + if (is_neox) { + rotary_embedding_kernel<<>>( + positions.data_ptr(), + query.data_ptr(), + key.has_value() ? key->data_ptr() : nullptr, + cos_sin_cache.data_ptr(), + rot_dim, + query_stride, + key_stride, + head_stride, + num_heads, + num_kv_heads, + head_size); + } else { + rotary_embedding_kernel<<>>( + positions.data_ptr(), + query.data_ptr(), + key.has_value() ? key->data_ptr() : nullptr, + cos_sin_cache.data_ptr(), + rot_dim, + query_stride, + key_stride, + head_stride, + num_heads, + num_kv_heads, + head_size); + } + }); +} + +} // namespace xllm::kernel::cuda diff --git a/ex_engine/xllm_kernels/ilu/activation.cpp b/ex_engine/xllm_kernels/ilu/activation.cpp new file mode 100644 index 00000000..1ad364a4 --- /dev/null +++ b/ex_engine/xllm_kernels/ilu/activation.cpp @@ -0,0 +1,32 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "ilu_ops_api.h" + +using namespace ixformer; + +namespace xllm::kernel::ilu { + +void act_and_mul(torch::Tensor out, + torch::Tensor input, + const std::string& act_mode) { + if (act_mode == "silu") { + infer::silu_and_mul(input, out); + } else { + LOG(FATAL) << "Unsupported act mode: " << act_mode + << ", only support silu, gelu, gelu_tanh"; + } +} +} // namespace xllm::kernel::ilu diff --git a/ex_engine/xllm_kernels/ilu/attention.cpp b/ex_engine/xllm_kernels/ilu/attention.cpp new file mode 100644 index 00000000..ad3cd295 --- /dev/null +++ b/ex_engine/xllm_kernels/ilu/attention.cpp @@ -0,0 +1,163 @@ + +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "ilu_ops_api.h" +#include "ixinfer.h" +#include "utils.h" + +using namespace ixformer; + +namespace xllm::kernel::ilu { + +void reshape_paged_cache(torch::Tensor& key, + std::optional& value, + torch::Tensor& key_cache, + std::optional& value_cache, + torch::Tensor& slot_mapping) { + auto value_ = value.value_or(torch::Tensor()); + auto value_cache_ = value_cache.value_or(torch::Tensor()); + + int64_t key_token_stride = key.stride(0); + int64_t value_token_stride = 0; + if (value_.defined()) { + value_token_stride = value_.stride(0); + } + slot_mapping = slot_mapping.to(at::kLong); + infer::xllm_reshape_and_cache(key, + value_, + key_cache, + value_cache_, + slot_mapping, + key_token_stride, + value_token_stride); +} + +void batch_prefill(torch::Tensor& query, + const torch::Tensor& key, + const std::optional& value, + torch::Tensor& output, + std::optional& output_lse, + const std::optional& q_cu_seq_lens, + const std::optional& kv_cu_seq_lens, + const std::optional& alibi_slope, + const std::optional& attn_bias, + const std::optional& q_quant_scale, + const std::optional& k_quant_scale, + const std::optional& v_quant_scale, + const torch::Tensor& block_tables, + int64_t max_query_len, + int64_t max_seq_len, + float scale, + bool is_causal, + int64_t window_size_left, + int64_t window_size_right, + const std::string& compute_dtype, + bool return_lse) { + double softcap = 0.0; + bool sqrt_alibi = false; + auto q_cu_seq_lens_ = q_cu_seq_lens.value_or(torch::Tensor()); + auto kv_cu_seq_lens_ = kv_cu_seq_lens.value_or(torch::Tensor()); + auto q_quant_scale_ = q_quant_scale.value_or(torch::Tensor()); + auto k_quant_scale_ = k_quant_scale.value_or(torch::Tensor()); + auto v_quant_scale_ = v_quant_scale.value_or(torch::Tensor()); + auto block_tables_ = block_tables; + auto key_ = key; + auto value_ = value.value(); + infer::ixinfer_flash_attn_unpad_with_block_tables(query, + key_, + value_, + output, + block_tables_, + q_cu_seq_lens_, + kv_cu_seq_lens_, + max_query_len, + max_seq_len, + is_causal, + window_size_left, + window_size_right, + static_cast(scale), + softcap, + sqrt_alibi, + alibi_slope, + c10::nullopt, + output_lse); +} + +void batch_decode(torch::Tensor& query, + const torch::Tensor& k_cache, + torch::Tensor& output, + const torch::Tensor& block_table, + const torch::Tensor& seq_lens, + const std::optional& v_cache, + std::optional& output_lse, + const std::optional& q_quant_scale, + const std::optional& k_cache_quant_scale, + const std::optional& v_cache_quant_scale, + const std::optional& out_quant_scale, + const std::optional& alibi_slope, + const std::optional& mask, + const std::string& compute_dtype, + int64_t max_seq_len, + int64_t window_size_left, + int64_t window_size_right, + float scale, + bool return_lse, + bool is_causal, + int64_t kv_cache_quant_bit_size) { + if (query.dim() == 4) { + query = + query + .view({query.size(0) * query.size(1), query.size(2), query.size(3)}) + .contiguous(); + } + if (output.dim() == 4) { + output = output + .view({output.size(0) * output.size(1), + output.size(2), + output.size(3)}) + .contiguous(); + ; + } + auto v_cache_ = v_cache.value_or(torch::Tensor()); + int64_t num_kv_heads = k_cache.size(1); + int64_t page_block_size = k_cache.size(2); + double softcap = 0.0; + bool enable_cuda_graph = false; + bool use_sqrt_alibi = false; + auto block_table_ = block_table; + auto k_cache_ = k_cache; + auto seq_lens_ = seq_lens; + infer::xllm_paged_attention(output, + query, + k_cache_, + v_cache_, + num_kv_heads, + scale, + block_table_, + seq_lens_, + page_block_size, + max_seq_len, + alibi_slope, + is_causal, + (int32_t)window_size_left, + (int32_t)window_size_right, + softcap, + enable_cuda_graph, + use_sqrt_alibi, + c10::nullopt); +} + +} // namespace xllm::kernel::ilu \ No newline at end of file diff --git a/ex_engine/xllm_kernels/ilu/fused_moe.cpp b/ex_engine/xllm_kernels/ilu/fused_moe.cpp new file mode 100644 index 00000000..21c15d8c --- /dev/null +++ b/ex_engine/xllm_kernels/ilu/fused_moe.cpp @@ -0,0 +1,99 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include + +#include "ilu_ops_api.h" + +namespace xllm::kernel::ilu { + +std::tuple moe_active_topk( + const torch::Tensor& input, + int64_t topk, + int64_t num_expert_group, + int64_t topk_group, + bool normalize, + const std::optional& mask, + const std::string& normed_by, + const std::string& scoring_func, + double route_scale, + const std::optional& e_score_correction_bias) { + torch::Tensor input_ = input.to(torch::kFloat32); + auto reduce_weight = + torch::empty({input.size(0), topk}, + torch::dtype(torch::kFloat).device(input.device())); + auto topk_indices = + torch::empty({input.size(0), topk}, + torch::dtype(torch::kInt32).device(input.device())); + auto token_expert_indices = + torch::empty({input.size(0), topk}, + torch::dtype(torch::kInt32).device(input.device())); + + infer::topk_softmax( + reduce_weight, topk_indices, token_expert_indices, input_, false); + + auto tt = reduce_weight.sum(-1); + if (normalize) { + reduce_weight = reduce_weight / reduce_weight.sum(-1).unsqueeze(-1); + } + return std::make_tuple(reduce_weight, topk_indices); +} + +std::vector moe_gen_idx(torch::Tensor& expert_id, + int64_t expert_num) { + auto src_dst = expert_id.new_empty({expert_id.numel()}); + auto dst_src = torch::empty_like(src_dst); + auto expert_sizes_gpu = expert_id.new_empty({expert_num}); + auto expert_sizes_gpu_cumsum = expert_id.new_zeros({expert_id.numel() + 1}); + infer::moe_compute_token_index_api(expert_id, + src_dst, + dst_src, + expert_sizes_gpu, + /*expert_mask=*/std::nullopt, + /*expert_sizes_cpu*/ std::nullopt, + /*expert_sizes_gpu*/ std::nullopt, + 0, + expert_num, + expert_num); + + expert_sizes_gpu_cumsum = expert_sizes_gpu.cumsum(-1); + return {src_dst, dst_src, expert_sizes_gpu, expert_sizes_gpu_cumsum}; +} + +torch::Tensor moe_expand_input(const torch::Tensor& input, + const torch::Tensor& gather_index, + const torch::Tensor& combine_idx, + int64_t topk) { + int64_t dst_tokens = input.size(0) * topk; + auto output = input.new_empty({dst_tokens, input.size(1)}); + infer::moe_expand_input( + output, input, combine_idx, gather_index, dst_tokens, topk); + + return output; +} + +torch::Tensor moe_combine_result(torch::Tensor& input, torch::Tensor& weight) { + input = input.view({-1, weight.size(1), input.size(1)}); + auto output = input.new_empty({input.size(0), input.size(2)}); + infer::moe_output_reduce_sum(output, + input, + weight, + /*mask=*/std::nullopt, + /*extra_residual*/ std::nullopt, + /*scaling_factor=*/1.0); + return output; +} + +} // namespace xllm::kernel::ilu diff --git a/ex_engine/xllm_kernels/ilu/group_gemm.cpp b/ex_engine/xllm_kernels/ilu/group_gemm.cpp new file mode 100644 index 00000000..290299a0 --- /dev/null +++ b/ex_engine/xllm_kernels/ilu/group_gemm.cpp @@ -0,0 +1,39 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "ilu_ops_api.h" + +namespace xllm::kernel::ilu { + +torch::Tensor group_gemm(torch::Tensor& input, + torch::Tensor& weight, + torch::Tensor& tokens_per_experts, + const std::optional& dst_to_src, + torch::Tensor& output) { + infer::moe_w16a16_group_gemm( + output, + input, + weight, + tokens_per_experts, + dst_to_src, + /*bias=*/std::nullopt, + /*format=*/"TN", + /*persistent=*/0, + /*output_n=*/tokens_per_experts.sum().item()); + + return output; +} + +} // namespace xllm::kernel::ilu diff --git a/ex_engine/xllm_kernels/ilu/ilu_ops_api.h b/ex_engine/xllm_kernels/ilu/ilu_ops_api.h new file mode 100644 index 00000000..3dedd7da --- /dev/null +++ b/ex_engine/xllm_kernels/ilu/ilu_ops_api.h @@ -0,0 +1,153 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include +#include +#include +#include +#include + +#include + +#include "ATen/Tensor.h" +#include "ATen/cuda/CUDAEvent.h" +#include "c10/core/Device.h" +#include "c10/core/DeviceGuard.h" +#include "c10/core/GradMode.h" +#include "c10/core/InferenceMode.h" +#include "c10/core/MemoryFormat.h" +#include "c10/core/ScalarType.h" +#include "c10/core/TensorOptions.h" +#include "c10/cuda/CUDAFunctions.h" +#include "c10/cuda/CUDAGuard.h" +#include "c10/cuda/CUDAStream.h" +#include "ixformer.h" +#include "kernels/kernels.h" + +// #include "utils.h" +using namespace ixformer; + +namespace xllm::kernel::ilu { + +void apply_rope_pos_ids_cos_sin_cache(torch::Tensor& query, + torch::Tensor& key, + torch::Tensor& cos_sin_cache, + torch::Tensor& positions, + bool interleave); + +// act_mode only support silu, gelu, gelu_tanh +void act_and_mul(torch::Tensor out, + torch::Tensor input, + const std::string& act_mode); + +void reshape_paged_cache( + torch::Tensor& key, // (num_tokens, num_heads, head_size) + std::optional& value, // (num_tokens, num_heads, head_size) + torch::Tensor& key_cache, // (num_blocks, num_heads, block_size, head_size) + std::optional& + value_cache, // (num_blocks, num_heads, block_size, head_size) + torch::Tensor& slot_mapping); //(num_tokens) + +void batch_prefill(torch::Tensor& query, + const torch::Tensor& key, + const std::optional& value, + torch::Tensor& output, + std::optional& output_lse, + const std::optional& q_cu_seq_lens, + const std::optional& kv_cu_seq_lens, + const std::optional& alibi_slope, + const std::optional& attn_bias, + const std::optional& q_quant_scale, + const std::optional& k_quant_scale, + const std::optional& v_quant_scale, + const torch::Tensor& block_tables, + int64_t max_query_len, + int64_t max_seq_len, + float scale, + bool is_causal, + int64_t window_size_left, + int64_t window_size_right, + const std::string& compute_dtype, + bool return_lse); + +void batch_decode(torch::Tensor& query, + const torch::Tensor& k_cache, + torch::Tensor& output, + const torch::Tensor& block_table, + const torch::Tensor& seq_lens, + const std::optional& v_cache, + std::optional& output_lse, + const std::optional& q_quant_scale, + const std::optional& k_cache_quant_scale, + const std::optional& v_cache_quant_scale, + const std::optional& out_quant_scale, + const std::optional& alibi_slope, + const std::optional& mask, + const std::string& compute_dtype, + int64_t max_seq_len, + int64_t window_size_left, + int64_t window_size_right, + float scale, + bool return_lse, + bool is_causal, + int64_t kv_cache_quant_bit_size); + +void residual_layer_norm(torch::Tensor& input, + torch::Tensor& output, + std::optional& residual, + torch::Tensor& weight, + std::optional& bias, + std::optional& residual_out, + double eps); + +void rms_norm(torch::Tensor& output, + torch::Tensor& input, + torch::Tensor& weight, + double eps); + +torch::Tensor matmul(torch::Tensor a, + torch::Tensor b, + std::optional bias); + +std::tuple moe_active_topk( + const torch::Tensor& input, + int64_t topk, + int64_t num_expert_group, + int64_t topk_group, + bool normalize, + const std::optional& mask, + const std::string& normed_by, + const std::string& scoring_func, + double route_scale, + const std::optional& e_score_correction_bias); + +std::vector moe_gen_idx(torch::Tensor& expert_id, + int64_t expert_num); + +torch::Tensor moe_expand_input(const torch::Tensor& input, + const torch::Tensor& gather_index, + const torch::Tensor& combine_idx, + int64_t topk); + +torch::Tensor group_gemm(torch::Tensor& input, + torch::Tensor& weight, + torch::Tensor& tokens_per_experts, + const std::optional& dst_to_src, + torch::Tensor& output); + +torch::Tensor moe_combine_result(torch::Tensor& input, torch::Tensor& weight); +} // namespace xllm::kernel::ilu diff --git a/ex_engine/xllm_kernels/ilu/ixformer.h b/ex_engine/xllm_kernels/ilu/ixformer.h new file mode 100644 index 00000000..83bad88e --- /dev/null +++ b/ex_engine/xllm_kernels/ilu/ixformer.h @@ -0,0 +1,147 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ +#include + +#include "ATen/Tensor.h" +#include "utils.h" + +namespace ixformer::infer { +torch::Tensor ixinfer_flash_attn_unpad_with_block_tables( + torch::Tensor& query, + torch::Tensor& key_cache, + torch::Tensor& value_cache, + torch::Tensor& out, + torch::Tensor& block_tables, + torch::Tensor& cu_seq_q, + torch::Tensor& cu_seq_k, + int64_t max_seq_q, + int64_t max_seq_k, + bool is_causal, + int64_t window_left, + int64_t window_right, + double scale, + double softcap, + bool sqrt_alibi, + const std::optional& alibi_slopes, + const std::optional& sinks, + std::optional& lse); + +void silu_and_mul(torch::Tensor& input, torch::Tensor& output); + +torch::Tensor xllm_paged_attention( + torch::Tensor& out, + torch::Tensor& query, + torch::Tensor& key_cache, + torch::Tensor& value_cache, + int64_t num_kv_heads, + double scale, + torch::Tensor& block_tables, + torch::Tensor& context_lens, + int64_t block_size, + int64_t max_context_len, + const std::optional& alibi_slopes, + bool causal, + int32_t window_left, + int32_t window_right, + double softcap, + bool enable_cuda_graph, + bool use_sqrt_alibi, + const std::optional& sinks); + +torch::Tensor ixformer_linear(torch::Tensor& input, + torch::Tensor& weight, + int64_t act_type, + const std::optional& bias, + const std::optional& out, + const std::optional persistent); + +torch::Tensor ixformer_linear_ex(torch::Tensor& input, + torch::Tensor& weight, + const c10::optional& bias, + const c10::optional& out); + +void xllm_reshape_and_cache(torch::Tensor& key, + torch::Tensor& value, + torch::Tensor& key_cache, + torch::Tensor& value_cache, + torch::Tensor& slot_mapping, + int64_t key_token_stride, + int64_t value_token_stride); + +void xllm_rotary_embedding(torch::Tensor& positions, + torch::Tensor& query, + torch::Tensor& key, + int64_t head_size, + torch::Tensor& cos_sin_cache, + bool is_neox); + +void residual_rms_norm(torch::Tensor& input, + torch::Tensor& residual, + torch::Tensor& weight, + torch::Tensor& output, + torch::Tensor& residual_output, + const std::optional& fused_bias, + double alpha, + double eps, + bool is_post); + +void rms_norm(torch::Tensor& input, + torch::Tensor& weight, + torch::Tensor& output, + const std::optional& fused_bias, + double eps); + +void topk_softmax(torch::Tensor& topk_weights, + torch::Tensor& topk_indices, + torch::Tensor& token_expert_indices, + torch::Tensor& gating_output, + bool renormalize); + +void moe_compute_token_index_api( + torch::Tensor& topk_ids, + torch::Tensor& src_dst, + torch::Tensor& dst_src, + torch::Tensor& expert_sizes_gpu, + const c10::optional& expert_mask, + const c10::optional& expert_sizes_cpu, + const c10::optional& expand_tokens_gpu, + int64_t start_expert_id, + int64_t end_expert_id, + int64_t num_experts); + +void moe_expand_input(torch::Tensor outputs, + torch::Tensor inputs, + torch::Tensor dst_to_src, + const c10::optional& src_to_dst, + int64_t dst_tokens, + int64_t expand_factor); + +void moe_w16a16_group_gemm(torch::Tensor output, + torch::Tensor inputs, + torch::Tensor weights, + torch::Tensor tokens_per_experts, + const c10::optional& dst_to_src, + const c10::optional& bias, + std::string format, + int64_t persistent, + int64_t output_n); + +void moe_output_reduce_sum(torch::Tensor outputs, + torch::Tensor inputs, + const c10::optional& mul_weight, + const c10::optional& mask, + const c10::optional& extra_residual, + double scaling_factor); +} // namespace ixformer::infer diff --git a/ex_engine/xllm_kernels/ilu/matmul.cpp b/ex_engine/xllm_kernels/ilu/matmul.cpp new file mode 100644 index 00000000..f90c0d47 --- /dev/null +++ b/ex_engine/xllm_kernels/ilu/matmul.cpp @@ -0,0 +1,73 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "ilu_ops_api.h" +#include "util/env_var.h" + +namespace xllm::kernel::ilu { + +bool gemv_conditions(const torch::Tensor& input, + const torch::Tensor& weight, + const torch::Tensor& bias, + int64_t gemv_max_batch) { + // gemv input:[m,k] weight:[n,k] + // 1. m <= gemv_max_batch + // 2. k % 32 == 0 && n % 2 == 0 + // 3. bias is None + + torch::Tensor input_view = input.view({-1, input.size(-1)}); + torch::Tensor weight_view = weight.view({-1, weight.size(-1)}); + + int64_t m = input_view.size(0); + int64_t k = input_view.size(1); + int64_t n = weight_view.size(0); + + if (bias.defined() == false && m <= gemv_max_batch && k % 32 == 0 && + n % 2 == 0) { + return true; + } + return false; +} + +torch::Tensor matmul(torch::Tensor a, + torch::Tensor b, + std::optional bias) { + int64_t act_type = -1; + bool persistent = false; + std::vector output_shape = a.sizes().vec(); + if (!output_shape.empty()) { + output_shape[output_shape.size() - 1] = b.size(0); + } + torch::Tensor output = a.new_empty(output_shape); + + bool use_gemv = true; + const int64_t gemv_max_batch = 1; + const bool disable_infer_gemm_ex = + xllm::util::get_bool_env("DISABLE_INFER_GEMM_EX", false); + + use_gemv = + use_gemv && + gemv_conditions(a, b, bias.value_or(at::Tensor()), gemv_max_batch) && + !disable_infer_gemm_ex && (act_type == -1); + + if (use_gemv) { + output = infer::ixformer_linear_ex(a, b, bias, output); + } else { + output = infer::ixformer_linear(a, b, act_type, bias, output, persistent); + } + return output; +} + +} // namespace xllm::kernel::ilu diff --git a/ex_engine/xllm_kernels/ilu/norm.cpp b/ex_engine/xllm_kernels/ilu/norm.cpp new file mode 100644 index 00000000..e451bc36 --- /dev/null +++ b/ex_engine/xllm_kernels/ilu/norm.cpp @@ -0,0 +1,51 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "ilu_ops_api.h" +#include "utils.h" + +using namespace ixformer; + +namespace xllm::kernel::ilu { + +void residual_layer_norm(torch::Tensor& input, + torch::Tensor& output, + std::optional& residual, + torch::Tensor& weight, + std::optional& bias, + std::optional& residual_out, + double eps) { + auto residual_ = residual.value_or(torch::zeros_like(input)); + torch::Tensor residual_out_ = residual_out.value_or(torch::zeros_like(input)); + infer::residual_rms_norm(input, + residual_, + weight, + output, + residual_out_, + bias, + /*alpha=*/1.0, + eps, + false); +} + +void rms_norm(torch::Tensor& output, + torch::Tensor& input, + torch::Tensor& weight, + double eps) { + std::optional fused_bias = std::nullopt; + infer::rms_norm(input, weight, output, fused_bias, eps); +} + +} // namespace xllm::kernel::ilu \ No newline at end of file diff --git a/ex_engine/xllm_kernels/ilu/rope.cpp b/ex_engine/xllm_kernels/ilu/rope.cpp new file mode 100644 index 00000000..45af7656 --- /dev/null +++ b/ex_engine/xllm_kernels/ilu/rope.cpp @@ -0,0 +1,31 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "ilu_ops_api.h" +#include "utils.h" + +namespace xllm::kernel::ilu { + +void apply_rope_pos_ids_cos_sin_cache(torch::Tensor& query, + torch::Tensor& key, + torch::Tensor& cos_sin_cache, + torch::Tensor& positions, + bool interleave) { + const int64_t head_size = cos_sin_cache.size(-1); + infer::xllm_rotary_embedding( + positions, query, key, head_size, cos_sin_cache, !interleave); +} + +} // namespace xllm::kernel::ilu diff --git a/ex_engine/xllm_kernels/ilu/utils.h b/ex_engine/xllm_kernels/ilu/utils.h new file mode 100644 index 00000000..9fd15298 --- /dev/null +++ b/ex_engine/xllm_kernels/ilu/utils.h @@ -0,0 +1,63 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ +#pragma once +namespace xllm::kernel::ilu { +#undef check_tensor_contiguous +#define check_tensor_contiguous(x, type) \ + TORCH_CHECK(x.scalar_type() == type); \ + TORCH_CHECK(x.is_cuda()); \ + TORCH_CHECK(x.is_contiguous()); + +#undef check_tensor_half_bf_float +#define check_tensor_half_bf_float(x) \ + TORCH_CHECK(x.scalar_type() == at::ScalarType::Half || \ + x.scalar_type() == at::ScalarType::Float || \ + x.scalar_type() == at::ScalarType::BFloat16); \ + TORCH_CHECK(x.is_cuda()); + +// from torchCheckMsgImpl +inline const char* ixformer_check_msg_impl(const char* msg) { return msg; } +// // If there is just 1 user-provided C-string argument, use it. + +#define IXFORMER_CHECK_MSG(cond, type, ...) \ + (ixformer_check_msg_impl( \ + "Expected " #cond \ + " to be true, but got false. " \ + "(Could this error message be improved? If so, " \ + "please report an enhancement request to ixformer.)", \ + ##__VA_ARGS__)) + +#define IXFORMER_CHECK(cond, ...) \ + { \ + if (!(cond)) { \ + std::cerr << __FILE__ << " (" << __LINE__ << ")" \ + << "-" << __FUNCTION__ << " : " \ + << IXFORMER_CHECK_MSG(cond, "", ##__VA_ARGS__) << std::endl; \ + throw std::runtime_error("IXFORMER_CHECK ERROR"); \ + } \ + } + +#undef CUINFER_CHECK +#define CUINFER_CHECK(func) \ + do { \ + cuinferStatus_t status = (func); \ + if (status != CUINFER_STATUS_SUCCESS) { \ + std::cerr << "Error in file " << __FILE__ << " on line " << __LINE__ \ + << ": " << cuinferGetErrorString(status) << std::endl; \ + throw std::runtime_error("CUINFER_CHECK ERROR"); \ + } \ + } while (0) + +} // namespace xllm::kernel::ilu \ No newline at end of file diff --git a/upstream_ref/xllm/docs/zh/getting_started/quick_start_GLM5.md b/upstream_ref/xllm/docs/zh/getting_started/quick_start_GLM5.md index 8b815071..1c2f402b 100644 --- a/upstream_ref/xllm/docs/zh/getting_started/quick_start_GLM5.md +++ b/upstream_ref/xllm/docs/zh/getting_started/quick_start_GLM5.md @@ -1,583 +1,583 @@ -# 使用 xLLM 在 Ascend A3设备 推理 GLM-5.0-W8A8 基座模型 - -+ 源码地址：https://github.com/jd-opensource/xllm - -+ 国内可用: https://gitcode.com/xLLM-AI/xllm - -+ 权重下载: [modelscope-GLM-5-W8A8](https://www.modelscope.cn/models/Eco-Tech/GLM-5-W8A8-xLLM/files) - -## 1.拉取镜像环境 - -首先下载xLLM提供的镜像： - -```bash -# A2 x86 -docker pull quay.io/jd_xllm/xllm-ai:xllm-dev-a2-x86-20260429 -# A2 arm -docker pull quay.io/jd_xllm/xllm-ai:xllm-dev-a2-arm-20260429 -# A3 arm -docker pull quay.io/jd_xllm/xllm-ai:xllm-dev-a3-arm-20260429 -``` - -**注意**: A2 机器性能未进行压测。 - -然后创建对应的容器 - -```bash -sudo docker run -it --ipc=host -u 0 --privileged --name mydocker --network=host \ - -v /var/queue_schedule:/var/queue_schedule \ - -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \ - -v /usr/local/Ascend/add-ons/:/usr/local/Ascend/add-ons/ \ - -v /usr/local/sbin/npu-smi:/usr/local/sbin/npu-smi \ - -v /var/log/npu/conf/slog/slog.conf:/var/log/npu/conf/slog/slog.conf \ - -v /var/log/npu/slog/:/var/log/npu/slog \ - -v ~/.ssh:/root/.ssh \ - -v /var/log/npu/profiling/:/var/log/npu/profiling \ - -v /var/log/npu/dump/:/var/log/npu/dump \ - -v /runtime/:/runtime/ -v /etc/hccn.conf:/etc/hccn.conf \ - -v /export/home:/export/home \ - -v /home/:/home/ \ - -w /export/home \ - quay.io/jd_xllm/xllm-ai:xllm-dev-hb-rc2-x86 -``` - -## 2.拉取源码并编译 - -下载官方仓库与模块依赖： - -```bash -git clone https://github.com/jd-opensource/xllm -cd xllm -git checkout preview/glm-5 -git submodule init -git submodule update -``` - -下载安装依赖: - -```bash -pip install --upgrade pre-commit -yum install numactl -``` - -执行编译，在`build/`下生成可执行文件`build/xllm/core/server/xllm`： - -```bash -python setup.py build -``` - -## 3.启动模型 - -### 若机器为重启后初次拉起服务，需先执行以下脚本对device进行初始化 - -#若不执行且npu未初始化可能导致xllm进程拉起失败 - -```bash -python -c "import torch_npu -for i in range(16):torch_npu.npu.set_device(i)" -``` - -### 环境变量 - -```bash -##### 1， 配置依赖路径相关环境变量 -# export PYTHON_INCLUDE_PATH="$(python3 -c 'from sysconfig import get_paths; print(get_paths()["include"])')" -# export PYTHON_LIB_PATH="$(python3 -c 'from sysconfig import get_paths; print(get_paths()["include"])')" -# export PYTORCH_NPU_INSTALL_PATH=/usr/local/libtorch_npu/ -# export PYTORCH_INSTALL_PATH="$(python3 -c 'import torch, os; print(os.path.dirname(os.path.abspath(torch.__file__)))')" -# export LIBTORCH_ROOT="$(python3 -c 'import torch, os; print(os.path.dirname(os.path.abspath(torch.__file__)))')" - -# export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/latest/opp/vendors/xllm/op_api/lib/:$LD_LIBRARY_PATH -# export LD_LIBRARY_PATH=/usr/local/libtorch_npu/lib:$LD_LIBRARY_PATH -export LD_PRELOAD=/usr/lib64/libjemalloc.so.2:$LD_PRELOAD - -# source /usr/local/Ascend/ascend-toolkit/set_env.sh -# source /usr/local/Ascend/nnal/atb/set_env.sh - -##### 2， 配置日志相关环境变量 -rm -rf /root/ascend/log/ -rm -rf core.* - -##### 3. 配置性能、通信相关环境变量 -export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True -export NPU_MEMORY_FRACTION=0.96 -export ATB_WORKSPACE_MEM_ALLOC_ALG_TYPE=3 -export ATB_WORKSPACE_MEM_ALLOC_GLOBAL=1 - -export OMP_NUM_THREADS=12 -export ALLOW_INTERNAL_FORMAT=1 - -export ATB_LAYER_INTERNAL_TENSOR_REUSE=1 -export ATB_LLM_ENABLE_AUTO_TRANSPOSE=0 -export ATB_CONVERT_NCHW_TO_AND=1 -export ATB_LAUNCH_KERNEL_WITH_TILING=1 -export ATB_OPERATION_EXECUTE_ASYNC=2 -export ATB_CONTEXT_WORKSPACE_SIZE=0 -export INF_NAN_MODE_ENABLE=1 -export HCCL_EXEC_TIMEOUT=300 -export HCCL_CONNECT_TIMEOUT=300 -export HCCL_OP_EXPANSION_MODE="AIV" -export HCCL_IF_BASE_PORT=2864 -``` - -## 启动命令 - GLM-5 （W8A8权重可单机拉起） - -```bash -BATCH_SIZE=256 -#推理最大batch数量 -XLLM_PATH="./myxllm/xllm/build/xllm/core/server/xllm" -#推理入口文件路径（上一步中编译产物） -MODEL_PATH=/path/to/GLM-5-W8A8/ -#模型路径（此处为int8量化的Glm-5） -DRAFT_MODEL_PATH=/path/to/GLM-5-W8A8/GLM-5-W8A8-MTP/ -#Glm-5 导出的mtp权重 - -MASTER_NODE_ADDR="11.87.49.110:10015" -LOCAL_HOST="11.87.49.110" -# Service Port -START_PORT=18994 -START_DEVICE=0 -LOG_DIR="logs" -NNODES=16 - -for (( i=0; i<$NNODES; i++ )) -do - PORT=$((START_PORT + i)) - DEVICE=$((START_DEVICE + i)) - LOG_FILE="$LOG_DIR/node_$i.log" - nohup numactl -C $((DEVICE*40))-$((DEVICE*40+39)) $XLLM_PATH \ - --model $MODEL_PATH \ - --port $PORT \ - --devices="npu:$DEVICE" \ - --master_node_addr=$MASTER_NODE_ADDR \ - --nnodes=$NNODES \ - --node_rank=$i \ - --max_memory_utilization=0.85 \ - --max_tokens_per_batch=8192 \ - --max_seqs_per_batch=32 \ - --block_size=128 \ - --enable_prefix_cache=false \ - --enable_chunked_prefill=true \ - --communication_backend="hccl" \ - --enable_schedule_overlap=true \ - --enable_graph=true \ - --enable_graph_mode_decode_no_padding=true \ - --draft_model=$DRAFT_MODEL_PATH \ - --draft_devices="npu:$DEVICE" \ - --num_speculative_tokens=1 \ - --ep_size=8 \ - --dp_size=1 \ - > $LOG_FILE 2>&1 & -done - -# numactl -C xxxxx 亲和性绑核(NUMA亲和性查询命令： npu-smi info -t topo) -#--max_memory_utilization 单卡最大显存占用比例 -#--max_tokens_per_batch 单batch最大token数 （主要限制prefill） -#--max_seqs_per_batch 单batch最大请求数 （主要限制decoe） -#--communication_backend 通信backend 可选(hccl / lccl) 此处建议hccl -#--enable_schedule_overlap 开启异步调度 -#--enable_prefix_cache 开启prefix_cache -#--enable_chunked_prefill 开启chunked_prefill -#--enable_graph 开启aclgraph -#--draft_model mtp - mtp权重路径 -#--draft_devices mtp - mtp推理设备(与主模型同一) -#--num_speculative_tokens mtp - 预测token数 -``` - -日志出现"Brpc Server Started"表示服务成功拉起。 - -## 其他可选环境变量 - -```bash -#开启确定性计算 -export LCCL_DETERMINISTIC=1 -export HCCL_DETERMINISTIC=true -export ATB_MATMUL_SHUFFLE_K_ENABLE=0 - -# #开启动态profiling模式 -# export PROFILING_MODE=dynamic -# \rm -rf ~/dynamic_profiling_socket_* -``` - -## 启动命令 - 双机拉起样例 - -### Node0 (master) - -```bash -MASTER_NODE_ADDR="11.87.49.110:19990" -LOCAL_HOST="11.87.49.110" -START_PORT=15890 -START_DEVICE=0 -LOG_DIR="logs" -NNODES=32 -LOCAL_NODES=16 -export HCCL_IF_BASE_PORT=48439 -unset HCCL_OP_EXPANSION_MODE - -for (( i=0; i<$LOCAL_NODES; i++ ))do - PORT=$((START_PORT + i)) - DEVICE=$((START_DEVICE + i)); LOG_FILE="$LOG_DIR/node_$i.log" - nohup numactl -C $((DEVICE*40))-$((DEVICE*40+39)) $XLLM_PATH \ --model $MODEL_PATH \ - --host $LOCAL_HOST \ - --port $PORT \ - --devices="npu:$DEVICE" \ - --master_node_addr=$MASTER_NODE_ADDR \ - --nnodes=$NNODES \ - --node_rank=$i \ - --max_memory_utilization=0.85 \ - --max_tokens_per_batch=8192 \ - --max_seqs_per_batch=4 \ - --block_size=128 \ - --enable_prefix_cache=false \ - --enable_chunked_prefill=true \ - --communication_backend="hccl" \ - --enable_schedule_overlap=true \ - --enable_graph=true \ - --enable_graph_mode_decode_no_padding=true \ - --ep_size=16 \ - --dp_size=1 \ - --rank_tablefile=/yourPath/ranktable.json \ - > $LOG_FILE 2>&1 & -done -``` - -#### Node1 (worker) - -```bash -MASTER_NODE_ADDR="11.87.49.110:19990" -LOCAL_HOST="11.87.49.111" -START_PORT=15890 -START_DEVICE=0 -LOG_DIR="logs" -NNODES=32 -LOCAL_NODES=16 -export HCCL_IF_BASE_PORT=48439 -unset HCCL_OP_EXPANSION_MODE - -for (( i=0; i<$LOCAL_NODES; i++ ))do - PORT=$((START_PORT + i)) - DEVICE=$((START_DEVICE + i)); LOG_FILE="$LOG_DIR/node_$i.log" - nohup numactl -C $((DEVICE*40))-$((DEVICE*40+39)) $XLLM_PATH \ --model $MODEL_PATH \ - --host $LOCAL_HOST \ - --port $PORT \ - --devices="npu:$DEVICE" \ - --master_node_addr=$MASTER_NODE_ADDR \ - --nnodes=$NNODES \ - --node_rank=$((i + LOCAL_NODES)) \ - --max_memory_utilization=0.85 \ - --max_tokens_per_batch=8192 \ - --max_seqs_per_batch=4 \ - --block_size=128 \ - --enable_prefix_cache=false \ - --enable_chunked_prefill=true \ - --communication_backend="hccl" \ - --enable_schedule_overlap=true \ - --enable_graph=true \ - --enable_graph_mode_decode_no_padding=true \ - --ep_size=16 \ - --dp_size=1 \ - --rank_tablefile=/yourPath/ranktable.json \ - > $LOG_FILE 2>&1 & -done -``` - -#### ranktable样例 - - ranktable配置指导：https://www.hiascend.com/document/detail/zh/canncommercial/83RC1/hccl/hcclug/hcclug_000014.html - -```json -{ - "version": "1.0", - "server_count": "2", - "server_list": [ - { - "server_id": "11.87.49.110", - "device": [ - { - "device_id": "0", - "device_ip": "11.86.23.210", - "rank_id": "0" - }, - ... - { - "device_id": "7", - "device_ip": "11.86.23.217", - "rank_id": "7" - } - ], - "host_nic_ip": "reserve" - }, - { - "server_id": "11.87.49.111", - "device": [ - { - "device_id": "0", - "device_ip": "11.87.63.202", - "rank_id": "8" - }, - ... - { - "device_id": "7", - "device_ip": "11.87.63.209", - "rank_id": "15" - } - ], - "host_nic_ip": "reserve" - } - ], - "status": "completed" -} -``` - - -## device NUMA亲和性查看 - -命令： - -```bash -npu-smi info -t topo -``` - -前述命令中 - -```bash -numactl -C $((DEVICE*12))-$((DEVICE*12+11)) -``` - -表示该进程绑在对应亲和的核上，可根据机器具体情况修改绑定的核id - -## EX3.Glm-5 权重量化 - -### 安装msmodelslim - -```bash -git clone https://gitcode.com/shenxiaolong/msmodelslim.git -cd msmodelslim -bash install.sh -``` - -### 修改tokenizer_config.json - -```bash - "extra_special_tokens" - 改成 "additional_special_tokens" - - "tokenizer_class": "TokenizersBackend" - 改成 "tokenizer_class": "PreTrainedTokenizer" -``` - -### 基于GLM-5-BF16 权重量化W8A8权重 - -```bash -### 预处理mtp相关权重 -python example/GLM5/extract_mtp.py --model-dir ${model_path} - -#指定transformers版本 -pip install transformers==4.48.2 - -#量化执行（生成量化权重） -msmodelslim quant --model_path ${model_path} --save_path ${save_path} --model_type DeepSeek-V3.2 --quant_type w8a8 --trust_remote_code True - -#拷贝chat_template文件 -cp ${model_path}/chat_template.jinja ${save_path} - -#量化mtp权重导出（用于xllm推理） -python example/GLM5/export_mtp.py --input-dir ${int8_save_path} --output-dir ${mtp_save_path} -``` - -## PD分离 - -### etcd\xllm-service 安装 - -#### PD分离部署 - -`xllm`支持PD分离部署，这需要与另一个开源库[xllm service](https://github.com/jd-opensource/xllm-service)配套使用。 - -##### xLLM Service依赖 - -首先，我们下载安装`xllm service`，与安装编译`xllm`类似： - -```bash -git clone https://github.com/jd-opensource/xllm-service -cd xllm_service -git submodule init -git submodule update -``` - -##### etcd安装 - -`xllm_service`依赖[etcd](https://github.com/etcd-io/etcd)，使用etcd官方提供的[安装脚本](https://github.com/etcd-io/etcd/releases)进行安装，其脚本提供的默认安装路径是`/tmp/etcd-download-test/etcd`，我们可以手动修改其脚本中的安装路径，也可以运行完脚本之后手动迁移： - -```bash -mv /tmp/etcd-download-test/etcd /path/to/your/etcd -``` - -##### xLLM Service编译 - -先应用patch: - -```bash -sh prepare.sh -``` - -再执行编译: - -```bash -mkdir -p build -cd build -cmake .. -make -j 8 -cd .. -``` - -!!! warning "可能的错误" - 这里能会遇到关于`boost-locale`和`boost-interprocess`的安装错误：`vcpkg-src/packages/boost-locale_x64-linux/include: No such file or directory`,`/vcpkg-src/packages/boost-interprocess_x64-linux/include: No such file or directory` - 我们使用`vcpkg`重新安装这些包: - ```bash - /path/to/vcpkg remove boost-locale boost-interprocess - /path/to/vcpkg install boost-locale:x64-linux - /path/to/vcpkg install boost-interprocess:x64-linux - ``` - -### PD分离运行 - -启动etcd: - -```bash -./etcd-download-test/etcd --listen-peer-urls 'http://localhost:2390' --listen-client-urls 'http://localhost:2389' --advertise-client-urls 'http://localhost:2391' -``` - -跨机配置时，etcd参考如下： - -```bash -/tmp/etcd-download-test/etcd --listen-peer-urls 'http://0.0.0.0:3390' --listen-client-urls 'http://0.0.0.0:3389' --advertise-client-urls 'http://11.87.191.82:3389' -``` - -启动xllm service: - -```bash -ENABLE_DECODE_RESPONSE_TO_SERVICE=true ./xllm_master_serving --etcd_addr="127.0.0.1:12389" --http_server_port 28888 --rpc_server_port 28889 --tokenizer_path=/export/home/models/GLM-5-W8A8/ -``` - -跨机配置时，启动xllm service: - -```bash -ENABLE_DECODE_RESPONSE_TO_SERVICE=true ../xllm-service/build/xllm_service/xllm_master_serving --etcd_addr="11.87.191.82:3389" --http_server_port 38888 --rpc_server_port 38889 --tokenizer_path=/export/home/models/GLM-5-W8A8/ -``` -- 启动Prefill实例 -```bash - BATCH_SIZE=256 - #推理最大batch数量 - XLLM_PATH="./myxllm/xllm/build/xllm/core/server/xllm" - #推理入口文件路径（上一步中编译产物） - MODEL_PATH=/export/home/models/GLM-5-w8a8/ - #模型路径（此处为int量化的Glm-5） - DRAFT_MODEL_PATH=/export/home/models/GLM-5-MTP/ - - MASTER_NODE_ADDR="11.87.49.110:10015" - LOCAL_HOST="11.87.49.110" - # Service Port - START_PORT=18994 - START_DEVICE=0 - LOG_DIR="logs" - NNODES=16 - - for (( i=0; i<$NNODES; i++ )) - do - PORT=$((START_PORT + i)) - DEVICE=$((START_DEVICE + i)) - LOG_FILE="$LOG_DIR/node_$i.log" - nohup numactl -C $((i*40))-$((i*40+39)) $XLLM_PATH \ - --model $MODEL_PATH --model_id glmmoe \ - --host $LOCAL_HOST \ - --port $PORT \ - --devices="npu:$DEVICE" \ - --master_node_addr=$MASTER_NODE_ADDR \ - --nnodes=$NNODES \ - --node_rank=$i \ - --max_memory_utilization=0.86 \ - --max_tokens_per_batch=5000 \ - --max_seqs_per_batch=$BATCH_SIZE \ - --communication_backend=hccl \ - --enable_schedule_overlap=true \ - --enable_prefix_cache=false \ - --enable_chunked_prefill=false \ - --enable_graph=true \ - --draft_model $DRAFT_MODEL_PATH \ - --draft_devices="npu:$DEVICE" \ - --num_speculative_tokens 1 \ - --enable_disagg_pd=true \ - --instance_role=PREFILL \ - --etcd_addr=$LOCAL_HOST:3389 \ - --transfer_listen_port=$((36100 + i)) \ - --disagg_pd_port=8877 \ - > $LOG_FILE 2>&1 & - done - - #--etcd_addr=$LOCAL_HOST:3389 参考etcd中advertise-client-urls的配置 - #--instance_role=DECODE PD配置，DECODE\PREFILL - ``` - -- 启动Decode实例 - - ```bash - BATCH_SIZE=256 - #推理最大batch数量 - XLLM_PATH="./myxllm/xllm/build/xllm/core/server/xllm" - #推理入口文件路径（上一步中编译产物） - MODEL_PATH=/export/home/models/GLM-5-w8a8/ - #模型路径（此处为int量化的Glm-5） - DRAFT_MODEL_PATH=/export/home/models/GLM-5-MTP/ - - MASTER_NODE_ADDR="11.87.49.110:10015" - LOCAL_HOST="11.87.49.110" - # Service Port - START_PORT=18994 - START_DEVICE=0 - LOG_DIR="logs" - NNODES=16 - - for (( i=0; i<$NNODES; i++ )) - do - PORT=$((START_PORT + i)) - DEVICE=$((START_DEVICE + i)) - LOG_FILE="$LOG_DIR/node_$i.log" - nohup numactl -C $((i*40))-$((i*40+39)) $XLLM_PATH \ - --model $MODEL_PATH --model_id glmmoe \ - --host $LOCAL_HOST \ - --port $PORT \ - --devices="npu:$DEVICE" \ - --master_node_addr=$MASTER_NODE_ADDR \ - --nnodes=$NNODES \ - --node_rank=$i \ - --max_memory_utilization=0.86 \ - --max_tokens_per_batch=5000 \ - --max_seqs_per_batch=$BATCH_SIZE \ - --communication_backend=hccl \ - --enable_schedule_overlap=true \ - --enable_prefix_cache=false \ - --enable_chunked_prefill=false \ - --enable_graph=true \ - --draft_model $DRAFT_MODEL_PATH \ - --draft_devices="npu:$DEVICE" \ - --num_speculative_tokens 1 \ - --enable_disagg_pd=true \ - --instance_role=DECODE \ - --etcd_addr=$LOCAL_HOST:3389 \ - --transfer_listen_port=$((36100 + i)) \ - --disagg_pd_port=8877 \ - > $LOG_FILE 2>&1 & - done - - #--etcd_addr=$LOCAL_HOST:3389 参考etcd中advertise-client-urls的配置 - #--instance_role=DECODE PD配置，DECODE\PREFILL - ``` - - 需要注意： - -- PD分离需要读取`/etc/hccn.conf`文件，确保将物理机上的该文件映射到了容器中 - -- `etcd_addr`需与`xllm_service`的`etcd_addr`相同 - 测试命令和上面类似，注意`curl http://localhost:{PORT}/v1/chat/completions ...`的`PORT`选择为启动xLLM service的`http_server_port`。 - -- 多机部署P或者Q时(例如部署两个P)，需要增加--rank_tablefile来完成通信。 +# 使用 xLLM 在 Ascend A3设备 推理 GLM-5.0-W8A8 基座模型 + ++ 源码地址：https://github.com/jd-opensource/xllm + ++ 国内可用: https://gitcode.com/xLLM-AI/xllm + ++ 权重下载: [modelscope-GLM-5-W8A8](https://www.modelscope.cn/models/Eco-Tech/GLM-5-W8A8-xLLM/files) + +## 1.拉取镜像环境 + +首先下载xLLM提供的镜像： + +```bash +# A2 x86 +docker pull quay.io/jd_xllm/xllm-ai:xllm-dev-a2-x86-20260429 +# A2 arm +docker pull quay.io/jd_xllm/xllm-ai:xllm-dev-a2-arm-20260429 +# A3 arm +docker pull quay.io/jd_xllm/xllm-ai:xllm-dev-a3-arm-20260429 +``` + +**注意**: A2 机器性能未进行压测。 + +然后创建对应的容器 + +```bash +sudo docker run -it --ipc=host -u 0 --privileged --name mydocker --network=host \ + -v /var/queue_schedule:/var/queue_schedule \ + -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \ + -v /usr/local/Ascend/add-ons/:/usr/local/Ascend/add-ons/ \ + -v /usr/local/sbin/npu-smi:/usr/local/sbin/npu-smi \ + -v /var/log/npu/conf/slog/slog.conf:/var/log/npu/conf/slog/slog.conf \ + -v /var/log/npu/slog/:/var/log/npu/slog \ + -v ~/.ssh:/root/.ssh \ + -v /var/log/npu/profiling/:/var/log/npu/profiling \ + -v /var/log/npu/dump/:/var/log/npu/dump \ + -v /runtime/:/runtime/ -v /etc/hccn.conf:/etc/hccn.conf \ + -v /export/home:/export/home \ + -v /home/:/home/ \ + -w /export/home \ + quay.io/jd_xllm/xllm-ai:xllm-dev-hb-rc2-x86 +``` + +## 2.拉取源码并编译 + +下载官方仓库与模块依赖： + +```bash +git clone https://github.com/jd-opensource/xllm +cd xllm +git checkout preview/glm-5 +git submodule init +git submodule update +``` + +下载安装依赖: + +```bash +pip install --upgrade pre-commit +yum install numactl +``` + +执行编译，在`build/`下生成可执行文件`build/xllm/core/server/xllm`： + +```bash +python setup.py build +``` + +## 3.启动模型 + +### 若机器为重启后初次拉起服务，需先执行以下脚本对device进行初始化 + +#若不执行且npu未初始化可能导致xllm进程拉起失败 + +```bash +python -c "import torch_npu +for i in range(16):torch_npu.npu.set_device(i)" +``` + +### 环境变量 + +```bash +##### 1， 配置依赖路径相关环境变量 +# export PYTHON_INCLUDE_PATH="$(python3 -c 'from sysconfig import get_paths; print(get_paths()["include"])')" +# export PYTHON_LIB_PATH="$(python3 -c 'from sysconfig import get_paths; print(get_paths()["include"])')" +# export PYTORCH_NPU_INSTALL_PATH=/usr/local/libtorch_npu/ +# export PYTORCH_INSTALL_PATH="$(python3 -c 'import torch, os; print(os.path.dirname(os.path.abspath(torch.__file__)))')" +# export LIBTORCH_ROOT="$(python3 -c 'import torch, os; print(os.path.dirname(os.path.abspath(torch.__file__)))')" + +# export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/latest/opp/vendors/xllm/op_api/lib/:$LD_LIBRARY_PATH +# export LD_LIBRARY_PATH=/usr/local/libtorch_npu/lib:$LD_LIBRARY_PATH +export LD_PRELOAD=/usr/lib64/libjemalloc.so.2:$LD_PRELOAD + +# source /usr/local/Ascend/ascend-toolkit/set_env.sh +# source /usr/local/Ascend/nnal/atb/set_env.sh + +##### 2， 配置日志相关环境变量 +rm -rf /root/ascend/log/ +rm -rf core.* + +##### 3. 配置性能、通信相关环境变量 +export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True +export NPU_MEMORY_FRACTION=0.96 +export ATB_WORKSPACE_MEM_ALLOC_ALG_TYPE=3 +export ATB_WORKSPACE_MEM_ALLOC_GLOBAL=1 + +export OMP_NUM_THREADS=12 +export ALLOW_INTERNAL_FORMAT=1 + +export ATB_LAYER_INTERNAL_TENSOR_REUSE=1 +export ATB_LLM_ENABLE_AUTO_TRANSPOSE=0 +export ATB_CONVERT_NCHW_TO_AND=1 +export ATB_LAUNCH_KERNEL_WITH_TILING=1 +export ATB_OPERATION_EXECUTE_ASYNC=2 +export ATB_CONTEXT_WORKSPACE_SIZE=0 +export INF_NAN_MODE_ENABLE=1 +export HCCL_EXEC_TIMEOUT=300 +export HCCL_CONNECT_TIMEOUT=300 +export HCCL_OP_EXPANSION_MODE="AIV" +export HCCL_IF_BASE_PORT=2864 +``` + +## 启动命令 - GLM-5 （W8A8权重可单机拉起） + +```bash +BATCH_SIZE=256 +#推理最大batch数量 +XLLM_PATH="./myxllm/xllm/build/xllm/core/server/xllm" +#推理入口文件路径（上一步中编译产物） +MODEL_PATH=/path/to/GLM-5-W8A8/ +#模型路径（此处为int8量化的Glm-5） +DRAFT_MODEL_PATH=/path/to/GLM-5-W8A8/GLM-5-W8A8-MTP/ +#Glm-5 导出的mtp权重 + +MASTER_NODE_ADDR="11.87.49.110:10015" +LOCAL_HOST="11.87.49.110" +# Service Port +START_PORT=18994 +START_DEVICE=0 +LOG_DIR="logs" +NNODES=16 + +for (( i=0; i<$NNODES; i++ )) +do + PORT=$((START_PORT + i)) + DEVICE=$((START_DEVICE + i)) + LOG_FILE="$LOG_DIR/node_$i.log" + nohup numactl -C $((DEVICE*40))-$((DEVICE*40+39)) $XLLM_PATH \ + --model $MODEL_PATH \ + --port $PORT \ + --devices="npu:$DEVICE" \ + --master_node_addr=$MASTER_NODE_ADDR \ + --nnodes=$NNODES \ + --node_rank=$i \ + --max_memory_utilization=0.85 \ + --max_tokens_per_batch=8192 \ + --max_seqs_per_batch=32 \ + --block_size=128 \ + --enable_prefix_cache=false \ + --enable_chunked_prefill=true \ + --communication_backend="hccl" \ + --enable_schedule_overlap=true \ + --enable_graph=true \ + --enable_graph_mode_decode_no_padding=true \ + --draft_model=$DRAFT_MODEL_PATH \ + --draft_devices="npu:$DEVICE" \ + --num_speculative_tokens=1 \ + --ep_size=8 \ + --dp_size=1 \ + > $LOG_FILE 2>&1 & +done + +# numactl -C xxxxx 亲和性绑核(NUMA亲和性查询命令： npu-smi info -t topo) +#--max_memory_utilization 单卡最大显存占用比例 +#--max_tokens_per_batch 单batch最大token数 （主要限制prefill） +#--max_seqs_per_batch 单batch最大请求数 （主要限制decoe） +#--communication_backend 通信backend 可选(hccl / lccl) 此处建议hccl +#--enable_schedule_overlap 开启异步调度 +#--enable_prefix_cache 开启prefix_cache +#--enable_chunked_prefill 开启chunked_prefill +#--enable_graph 开启aclgraph +#--draft_model mtp - mtp权重路径 +#--draft_devices mtp - mtp推理设备(与主模型同一) +#--num_speculative_tokens mtp - 预测token数 +``` + +日志出现"Brpc Server Started"表示服务成功拉起。 + +## 其他可选环境变量 + +```bash +#开启确定性计算 +export LCCL_DETERMINISTIC=1 +export HCCL_DETERMINISTIC=true +export ATB_MATMUL_SHUFFLE_K_ENABLE=0 + +# #开启动态profiling模式 +# export PROFILING_MODE=dynamic +# \rm -rf ~/dynamic_profiling_socket_* +``` + +## 启动命令 - 双机拉起样例 + +### Node0 (master) + +```bash +MASTER_NODE_ADDR="11.87.49.110:19990" +LOCAL_HOST="11.87.49.110" +START_PORT=15890 +START_DEVICE=0 +LOG_DIR="logs" +NNODES=32 +LOCAL_NODES=16 +export HCCL_IF_BASE_PORT=48439 +unset HCCL_OP_EXPANSION_MODE + +for (( i=0; i<$LOCAL_NODES; i++ ))do + PORT=$((START_PORT + i)) + DEVICE=$((START_DEVICE + i)); LOG_FILE="$LOG_DIR/node_$i.log" + nohup numactl -C $((DEVICE*40))-$((DEVICE*40+39)) $XLLM_PATH \ --model $MODEL_PATH \ + --host $LOCAL_HOST \ + --port $PORT \ + --devices="npu:$DEVICE" \ + --master_node_addr=$MASTER_NODE_ADDR \ + --nnodes=$NNODES \ + --node_rank=$i \ + --max_memory_utilization=0.85 \ + --max_tokens_per_batch=8192 \ + --max_seqs_per_batch=4 \ + --block_size=128 \ + --enable_prefix_cache=false \ + --enable_chunked_prefill=true \ + --communication_backend="hccl" \ + --enable_schedule_overlap=true \ + --enable_graph=true \ + --enable_graph_mode_decode_no_padding=true \ + --ep_size=16 \ + --dp_size=1 \ + --rank_tablefile=/yourPath/ranktable.json \ + > $LOG_FILE 2>&1 & +done +``` + +#### Node1 (worker) + +```bash +MASTER_NODE_ADDR="11.87.49.110:19990" +LOCAL_HOST="11.87.49.111" +START_PORT=15890 +START_DEVICE=0 +LOG_DIR="logs" +NNODES=32 +LOCAL_NODES=16 +export HCCL_IF_BASE_PORT=48439 +unset HCCL_OP_EXPANSION_MODE + +for (( i=0; i<$LOCAL_NODES; i++ ))do + PORT=$((START_PORT + i)) + DEVICE=$((START_DEVICE + i)); LOG_FILE="$LOG_DIR/node_$i.log" + nohup numactl -C $((DEVICE*40))-$((DEVICE*40+39)) $XLLM_PATH \ --model $MODEL_PATH \ + --host $LOCAL_HOST \ + --port $PORT \ + --devices="npu:$DEVICE" \ + --master_node_addr=$MASTER_NODE_ADDR \ + --nnodes=$NNODES \ + --node_rank=$((i + LOCAL_NODES)) \ + --max_memory_utilization=0.85 \ + --max_tokens_per_batch=8192 \ + --max_seqs_per_batch=4 \ + --block_size=128 \ + --enable_prefix_cache=false \ + --enable_chunked_prefill=true \ + --communication_backend="hccl" \ + --enable_schedule_overlap=true \ + --enable_graph=true \ + --enable_graph_mode_decode_no_padding=true \ + --ep_size=16 \ + --dp_size=1 \ + --rank_tablefile=/yourPath/ranktable.json \ + > $LOG_FILE 2>&1 & +done +``` + +#### ranktable样例 + + ranktable配置指导：https://www.hiascend.com/document/detail/zh/canncommercial/83RC1/hccl/hcclug/hcclug_000014.html + +```json +{ + "version": "1.0", + "server_count": "2", + "server_list": [ + { + "server_id": "11.87.49.110", + "device": [ + { + "device_id": "0", + "device_ip": "11.86.23.210", + "rank_id": "0" + }, + ... + { + "device_id": "7", + "device_ip": "11.86.23.217", + "rank_id": "7" + } + ], + "host_nic_ip": "reserve" + }, + { + "server_id": "11.87.49.111", + "device": [ + { + "device_id": "0", + "device_ip": "11.87.63.202", + "rank_id": "8" + }, + ... + { + "device_id": "7", + "device_ip": "11.87.63.209", + "rank_id": "15" + } + ], + "host_nic_ip": "reserve" + } + ], + "status": "completed" +} +``` + + +## device NUMA亲和性查看 + +命令： + +```bash +npu-smi info -t topo +``` + +前述命令中 + +```bash +numactl -C $((DEVICE*12))-$((DEVICE*12+11)) +``` + +表示该进程绑在对应亲和的核上，可根据机器具体情况修改绑定的核id + +## EX3.Glm-5 权重量化 + +### 安装msmodelslim + +```bash +git clone https://gitcode.com/shenxiaolong/msmodelslim.git +cd msmodelslim +bash install.sh +``` + +### 修改tokenizer_config.json + +```bash + "extra_special_tokens" + 改成 "additional_special_tokens" + + "tokenizer_class": "TokenizersBackend" + 改成 "tokenizer_class": "PreTrainedTokenizer" +``` + +### 基于GLM-5-BF16 权重量化W8A8权重 + +```bash +### 预处理mtp相关权重 +python example/GLM5/extract_mtp.py --model-dir ${model_path} + +#指定transformers版本 +pip install transformers==4.48.2 + +#量化执行（生成量化权重） +msmodelslim quant --model_path ${model_path} --save_path ${save_path} --model_type DeepSeek-V3.2 --quant_type w8a8 --trust_remote_code True + +#拷贝chat_template文件 +cp ${model_path}/chat_template.jinja ${save_path} + +#量化mtp权重导出（用于xllm推理） +python example/GLM5/export_mtp.py --input-dir ${int8_save_path} --output-dir ${mtp_save_path} +``` + +## PD分离 + +### etcd\xllm-service 安装 + +#### PD分离部署 + +`xllm`支持PD分离部署，这需要与另一个开源库[xllm service](https://github.com/jd-opensource/xllm-service)配套使用。 + +##### xLLM Service依赖 + +首先，我们下载安装`xllm service`，与安装编译`xllm`类似： + +```bash +git clone https://github.com/jd-opensource/xllm-service +cd xllm_service +git submodule init +git submodule update +``` + +##### etcd安装 + +`xllm_service`依赖[etcd](https://github.com/etcd-io/etcd)，使用etcd官方提供的[安装脚本](https://github.com/etcd-io/etcd/releases)进行安装，其脚本提供的默认安装路径是`/tmp/etcd-download-test/etcd`，我们可以手动修改其脚本中的安装路径，也可以运行完脚本之后手动迁移： + +```bash +mv /tmp/etcd-download-test/etcd /path/to/your/etcd +``` + +##### xLLM Service编译 + +先应用patch: + +```bash +sh prepare.sh +``` + +再执行编译: + +```bash +mkdir -p build +cd build +cmake .. +make -j 8 +cd .. +``` + +!!! warning "可能的错误" + 这里能会遇到关于`boost-locale`和`boost-interprocess`的安装错误：`vcpkg-src/packages/boost-locale_x64-linux/include: No such file or directory`,`/vcpkg-src/packages/boost-interprocess_x64-linux/include: No such file or directory` + 我们使用`vcpkg`重新安装这些包: + ```bash + /path/to/vcpkg remove boost-locale boost-interprocess + /path/to/vcpkg install boost-locale:x64-linux + /path/to/vcpkg install boost-interprocess:x64-linux + ``` + +### PD分离运行 + +启动etcd: + +```bash +./etcd-download-test/etcd --listen-peer-urls 'http://localhost:2390' --listen-client-urls 'http://localhost:2389' --advertise-client-urls 'http://localhost:2391' +``` + +跨机配置时，etcd参考如下： + +```bash +/tmp/etcd-download-test/etcd --listen-peer-urls 'http://0.0.0.0:3390' --listen-client-urls 'http://0.0.0.0:3389' --advertise-client-urls 'http://11.87.191.82:3389' +``` + +启动xllm service: + +```bash +ENABLE_DECODE_RESPONSE_TO_SERVICE=true ./xllm_master_serving --etcd_addr="127.0.0.1:12389" --http_server_port 28888 --rpc_server_port 28889 --tokenizer_path=/export/home/models/GLM-5-W8A8/ +``` + +跨机配置时，启动xllm service: + +```bash +ENABLE_DECODE_RESPONSE_TO_SERVICE=true ../xllm-service/build/xllm_service/xllm_master_serving --etcd_addr="11.87.191.82:3389" --http_server_port 38888 --rpc_server_port 38889 --tokenizer_path=/export/home/models/GLM-5-W8A8/ +``` +- 启动Prefill实例 +```bash + BATCH_SIZE=256 + #推理最大batch数量 + XLLM_PATH="./myxllm/xllm/build/xllm/core/server/xllm" + #推理入口文件路径（上一步中编译产物） + MODEL_PATH=/export/home/models/GLM-5-w8a8/ + #模型路径（此处为int量化的Glm-5） + DRAFT_MODEL_PATH=/export/home/models/GLM-5-MTP/ + + MASTER_NODE_ADDR="11.87.49.110:10015" + LOCAL_HOST="11.87.49.110" + # Service Port + START_PORT=18994 + START_DEVICE=0 + LOG_DIR="logs" + NNODES=16 + + for (( i=0; i<$NNODES; i++ )) + do + PORT=$((START_PORT + i)) + DEVICE=$((START_DEVICE + i)) + LOG_FILE="$LOG_DIR/node_$i.log" + nohup numactl -C $((i*40))-$((i*40+39)) $XLLM_PATH \ + --model $MODEL_PATH --model_id glmmoe \ + --host $LOCAL_HOST \ + --port $PORT \ + --devices="npu:$DEVICE" \ + --master_node_addr=$MASTER_NODE_ADDR \ + --nnodes=$NNODES \ + --node_rank=$i \ + --max_memory_utilization=0.86 \ + --max_tokens_per_batch=5000 \ + --max_seqs_per_batch=$BATCH_SIZE \ + --communication_backend=hccl \ + --enable_schedule_overlap=true \ + --enable_prefix_cache=false \ + --enable_chunked_prefill=false \ + --enable_graph=true \ + --draft_model $DRAFT_MODEL_PATH \ + --draft_devices="npu:$DEVICE" \ + --num_speculative_tokens 1 \ + --enable_disagg_pd=true \ + --instance_role=PREFILL \ + --etcd_addr=$LOCAL_HOST:3389 \ + --transfer_listen_port=$((36100 + i)) \ + --disagg_pd_port=8877 \ + > $LOG_FILE 2>&1 & + done + + #--etcd_addr=$LOCAL_HOST:3389 参考etcd中advertise-client-urls的配置 + #--instance_role=DECODE PD配置，DECODE\PREFILL + ``` + +- 启动Decode实例 + + ```bash + BATCH_SIZE=256 + #推理最大batch数量 + XLLM_PATH="./myxllm/xllm/build/xllm/core/server/xllm" + #推理入口文件路径（上一步中编译产物） + MODEL_PATH=/export/home/models/GLM-5-w8a8/ + #模型路径（此处为int量化的Glm-5） + DRAFT_MODEL_PATH=/export/home/models/GLM-5-MTP/ + + MASTER_NODE_ADDR="11.87.49.110:10015" + LOCAL_HOST="11.87.49.110" + # Service Port + START_PORT=18994 + START_DEVICE=0 + LOG_DIR="logs" + NNODES=16 + + for (( i=0; i<$NNODES; i++ )) + do + PORT=$((START_PORT + i)) + DEVICE=$((START_DEVICE + i)) + LOG_FILE="$LOG_DIR/node_$i.log" + nohup numactl -C $((i*40))-$((i*40+39)) $XLLM_PATH \ + --model $MODEL_PATH --model_id glmmoe \ + --host $LOCAL_HOST \ + --port $PORT \ + --devices="npu:$DEVICE" \ + --master_node_addr=$MASTER_NODE_ADDR \ + --nnodes=$NNODES \ + --node_rank=$i \ + --max_memory_utilization=0.86 \ + --max_tokens_per_batch=5000 \ + --max_seqs_per_batch=$BATCH_SIZE \ + --communication_backend=hccl \ + --enable_schedule_overlap=true \ + --enable_prefix_cache=false \ + --enable_chunked_prefill=false \ + --enable_graph=true \ + --draft_model $DRAFT_MODEL_PATH \ + --draft_devices="npu:$DEVICE" \ + --num_speculative_tokens 1 \ + --enable_disagg_pd=true \ + --instance_role=DECODE \ + --etcd_addr=$LOCAL_HOST:3389 \ + --transfer_listen_port=$((36100 + i)) \ + --disagg_pd_port=8877 \ + > $LOG_FILE 2>&1 & + done + + #--etcd_addr=$LOCAL_HOST:3389 参考etcd中advertise-client-urls的配置 + #--instance_role=DECODE PD配置，DECODE\PREFILL + ``` + + 需要注意： + +- PD分离需要读取`/etc/hccn.conf`文件，确保将物理机上的该文件映射到了容器中 + +- `etcd_addr`需与`xllm_service`的`etcd_addr`相同 + 测试命令和上面类似，注意`curl http://localhost:{PORT}/v1/chat/completions ...`的`PORT`选择为启动xLLM service的`http_server_port`。 + +- 多机部署P或者Q时(例如部署两个P)，需要增加--rank_tablefile来完成通信。 diff --git a/vllm/model_executor/layers/quantization/utils/__init__.py b/vllm/model_executor/layers/quantization/utils/__init__.py index e60f0c79..6d18fa3b 100644 --- a/vllm/model_executor/layers/quantization/utils/__init__.py +++ b/vllm/model_executor/layers/quantization/utils/__init__.py @@ -1,3 +1,3 @@ -from .layer_utils import replace_parameter, update_tensor_inplace - -__all__ = ['update_tensor_inplace', 'replace_parameter'] +from .layer_utils import replace_parameter, update_tensor_inplace + +__all__ = ['update_tensor_inplace', 'replace_parameter'] diff --git a/vllm_overrides/core/block/block_table.py b/vllm_overrides/core/block/block_table.py index 40bfd231..b553dc74 100644 --- a/vllm_overrides/core/block/block_table.py +++ b/vllm_overrides/core/block/block_table.py @@ -1,43 +1,43 @@ -import math -from typing import List, Optional - -from vllm.core.block.common import BlockList +import math +from typing import List, Optional + +from vllm.core.block.common import BlockList from vllm.core.block.interfaces import Block, DeviceAwareBlockAllocator -from vllm.utils import Device, cdiv, chunk_list - - -class BlockTable: - """A class to manage blocks for a specific sequence. - - The BlockTable maps a sequence of tokens to a list of blocks, where each - block represents a contiguous memory allocation for a portion of the - sequence. The blocks are managed by a DeviceAwareBlockAllocator, which is - responsible for allocating and freeing memory for the blocks. - - Args: - block_size (int): The maximum number of tokens that can be stored in a - single block. - block_allocator (DeviceAwareBlockAllocator): The block allocator used to - manage memory for the blocks. - _blocks (Optional[List[Block]], optional): An optional list of existing - blocks to initialize the BlockTable with. If not provided, an empty - BlockTable is created. - max_block_sliding_window (Optional[int], optional): The number of - blocks to keep around for each sequance. If None, all blocks - are kept (eg., when sliding window is not used). - It should at least fit the sliding window size of the model. - - Attributes: - _block_size (int): The maximum number of tokens that can be stored in a - single block. - _allocator (DeviceAwareBlockAllocator): The block allocator used to - manage memory for the blocks. - _blocks (Optional[List[Block]]): The list of blocks managed by this - BlockTable. - _num_full_slots (int): The number of tokens currently stored in the - blocks. - """ - +from vllm.utils import Device, cdiv, chunk_list + + +class BlockTable: + """A class to manage blocks for a specific sequence. + + The BlockTable maps a sequence of tokens to a list of blocks, where each + block represents a contiguous memory allocation for a portion of the + sequence. The blocks are managed by a DeviceAwareBlockAllocator, which is + responsible for allocating and freeing memory for the blocks. + + Args: + block_size (int): The maximum number of tokens that can be stored in a + single block. + block_allocator (DeviceAwareBlockAllocator): The block allocator used to + manage memory for the blocks. + _blocks (Optional[List[Block]], optional): An optional list of existing + blocks to initialize the BlockTable with. If not provided, an empty + BlockTable is created. + max_block_sliding_window (Optional[int], optional): The number of + blocks to keep around for each sequance. If None, all blocks + are kept (eg., when sliding window is not used). + It should at least fit the sliding window size of the model. + + Attributes: + _block_size (int): The maximum number of tokens that can be stored in a + single block. + _allocator (DeviceAwareBlockAllocator): The block allocator used to + manage memory for the blocks. + _blocks (Optional[List[Block]]): The list of blocks managed by this + BlockTable. + _num_full_slots (int): The number of tokens currently stored in the + blocks. + """ + def __init__( self, block_size: int, @@ -52,55 +52,55 @@ class BlockTable: if _blocks is None: _blocks = [] self._blocks: BlockList = BlockList(_blocks) - - self._max_block_sliding_window = max_block_sliding_window - self._num_full_slots = self._get_num_token_ids() - - @staticmethod - def get_num_required_blocks(token_ids: List[int], - block_size: int, - num_lookahead_slots: int = 0) -> int: - """Calculates the minimum number of blocks required to store a given - sequence of token IDs along with any look-ahead slots that may be - required (like in multi-step + chunked-prefill). - - This assumes worst-case scenario, where every block requires a new - allocation (e.g. ignoring prefix caching). - - Args: - token_ids (List[int]): The sequence of token IDs to be stored. - block_size (int): The maximum number of tokens that can be stored in - a single block. - num_lookahead_slots (int): look-ahead slots that the sequence may - require. - - Returns: - int: The minimum number of blocks required to store the given - sequence of token IDs along with any required look-ahead slots. - """ - return cdiv(len(token_ids) + num_lookahead_slots, block_size) - + + self._max_block_sliding_window = max_block_sliding_window + self._num_full_slots = self._get_num_token_ids() + + @staticmethod + def get_num_required_blocks(token_ids: List[int], + block_size: int, + num_lookahead_slots: int = 0) -> int: + """Calculates the minimum number of blocks required to store a given + sequence of token IDs along with any look-ahead slots that may be + required (like in multi-step + chunked-prefill). + + This assumes worst-case scenario, where every block requires a new + allocation (e.g. ignoring prefix caching). + + Args: + token_ids (List[int]): The sequence of token IDs to be stored. + block_size (int): The maximum number of tokens that can be stored in + a single block. + num_lookahead_slots (int): look-ahead slots that the sequence may + require. + + Returns: + int: The minimum number of blocks required to store the given + sequence of token IDs along with any required look-ahead slots. + """ + return cdiv(len(token_ids) + num_lookahead_slots, block_size) + def allocate(self, token_ids: List[int], device: Device = Device.GPU) -> None: - """Allocates memory blocks for storing the given sequence of token IDs. - - This method allocates the required number of blocks to store the given - sequence of token IDs. - - Args: - token_ids (List[int]): The sequence of token IDs to be stored. - device (Device, optional): The device on which the blocks should be - allocated. Defaults to Device.GPU. - """ + """Allocates memory blocks for storing the given sequence of token IDs. + + This method allocates the required number of blocks to store the given + sequence of token IDs. + + Args: + token_ids (List[int]): The sequence of token IDs to be stored. + device (Device, optional): The device on which the blocks should be + allocated. Defaults to Device.GPU. + """ assert not self._is_allocated assert token_ids blocks = self._allocate_blocks_for_token_ids(prev_block=None, token_ids=token_ids, device=device) - self.update(blocks) - self._num_full_slots = len(token_ids) - + self.update(blocks) + self._num_full_slots = len(token_ids) + def update(self, blocks: List[Block]) -> None: """Resets the table to the newly provided blocks (with their corresponding block ids) @@ -115,106 +115,106 @@ class BlockTable: if block_hash is not None: content_hashes.append(block_hash) return content_hashes - - def append_token_ids(self, - token_ids: List[int], - num_lookahead_slots: int = 0, - num_computed_slots: Optional[int] = None) -> None: - """Appends a sequence of token IDs to the existing blocks in the - BlockTable. - - This method appends the given sequence of token IDs to the existing - blocks in the BlockTable. If there is not enough space in the existing - blocks, new blocks are allocated using the `ensure_num_empty_slots` - method to accommodate the additional tokens. - - The token IDs are divided into chunks of size `block_size` (except for - the first chunk, which may be smaller), and each chunk is appended to a - separate block. - - Args: - token_ids (List[int]): The sequence of token IDs to be appended. - num_computed_slots (Optional[int]): The number of KV cache slots - that are already filled (computed). - When sliding window is enabled, this is used to compute how many - blocks to drop at the front of the sequence. - Without sliding window, None can be passed. - Without chunked prefill, it should be the same as - _num_full_slots. - """ - assert self._is_allocated, "no blocks have been allocated" - assert len(self._blocks) > 0 - - # Drop blocks that are no longer needed due to sliding window - if self._max_block_sliding_window is not None: - null_block = self._allocator.allocate_or_get_null_block() - assert num_computed_slots is not None - end_block_idx = (num_computed_slots // - self._block_size) - self._max_block_sliding_window - for idx in range(0, end_block_idx): - b = self._blocks[idx] - if b is not null_block: - self._allocator.free(b) - self._blocks[idx] = null_block - - # Ensure there are enough empty slots for the new tokens plus - # lookahead slots - self.ensure_num_empty_slots(num_empty_slots=len(token_ids) + - num_lookahead_slots) - - # Update the blocks with the new tokens - first_block_idx = self._num_full_slots // self._block_size - token_blocks = self._chunk_token_blocks_for_append(token_ids) - - for i, token_block in enumerate(token_blocks): - self._blocks.append_token_ids(first_block_idx + i, token_block) - - self._num_full_slots += len(token_ids) - - def ensure_num_empty_slots(self, num_empty_slots: int) -> None: - """Ensures that the BlockTable has at least the specified number of - empty slots available. - - This method checks if the BlockTable has enough empty slots (i.e., - available space) to accommodate the requested number of tokens. If not, - it allocates additional blocks on the GPU to ensure that the required - number of empty slots is available. - - Args: - num_empty_slots (int): The minimum number of empty slots required. - """ - # Currently the block table only supports - # appending tokens to GPU blocks. - device = Device.GPU - assert self._is_allocated - - if self._num_empty_slots >= num_empty_slots: - return - - slots_to_allocate = num_empty_slots - self._num_empty_slots - blocks_to_allocate = cdiv(slots_to_allocate, self._block_size) - - for _ in range(blocks_to_allocate): - assert len(self._blocks) > 0 - self._blocks.append( - self._allocator.allocate_mutable_block( - prev_block=self._blocks[-1], device=device)) - - def fork(self) -> "BlockTable": - """Creates a new BlockTable instance with a copy of the blocks from the - current instance. - - This method creates a new BlockTable instance with the same block size, - block allocator, and a copy of the blocks from the current instance. The - new BlockTable has its own independent set of blocks, but shares the - same underlying memory allocation with the original BlockTable. - - Returns: - BlockTable: A new BlockTable instance with a copy of the blocks from - the current instance. - """ - assert self._is_allocated - assert len(self._blocks) > 0 + + def append_token_ids(self, + token_ids: List[int], + num_lookahead_slots: int = 0, + num_computed_slots: Optional[int] = None) -> None: + """Appends a sequence of token IDs to the existing blocks in the + BlockTable. + + This method appends the given sequence of token IDs to the existing + blocks in the BlockTable. If there is not enough space in the existing + blocks, new blocks are allocated using the `ensure_num_empty_slots` + method to accommodate the additional tokens. + + The token IDs are divided into chunks of size `block_size` (except for + the first chunk, which may be smaller), and each chunk is appended to a + separate block. + + Args: + token_ids (List[int]): The sequence of token IDs to be appended. + num_computed_slots (Optional[int]): The number of KV cache slots + that are already filled (computed). + When sliding window is enabled, this is used to compute how many + blocks to drop at the front of the sequence. + Without sliding window, None can be passed. + Without chunked prefill, it should be the same as + _num_full_slots. + """ + assert self._is_allocated, "no blocks have been allocated" + assert len(self._blocks) > 0 + + # Drop blocks that are no longer needed due to sliding window + if self._max_block_sliding_window is not None: + null_block = self._allocator.allocate_or_get_null_block() + assert num_computed_slots is not None + end_block_idx = (num_computed_slots // + self._block_size) - self._max_block_sliding_window + for idx in range(0, end_block_idx): + b = self._blocks[idx] + if b is not null_block: + self._allocator.free(b) + self._blocks[idx] = null_block + + # Ensure there are enough empty slots for the new tokens plus + # lookahead slots + self.ensure_num_empty_slots(num_empty_slots=len(token_ids) + + num_lookahead_slots) + + # Update the blocks with the new tokens + first_block_idx = self._num_full_slots // self._block_size + token_blocks = self._chunk_token_blocks_for_append(token_ids) + + for i, token_block in enumerate(token_blocks): + self._blocks.append_token_ids(first_block_idx + i, token_block) + + self._num_full_slots += len(token_ids) + + def ensure_num_empty_slots(self, num_empty_slots: int) -> None: + """Ensures that the BlockTable has at least the specified number of + empty slots available. + + This method checks if the BlockTable has enough empty slots (i.e., + available space) to accommodate the requested number of tokens. If not, + it allocates additional blocks on the GPU to ensure that the required + number of empty slots is available. + + Args: + num_empty_slots (int): The minimum number of empty slots required. + """ + # Currently the block table only supports + # appending tokens to GPU blocks. + device = Device.GPU + assert self._is_allocated + + if self._num_empty_slots >= num_empty_slots: + return + + slots_to_allocate = num_empty_slots - self._num_empty_slots + blocks_to_allocate = cdiv(slots_to_allocate, self._block_size) + + for _ in range(blocks_to_allocate): + assert len(self._blocks) > 0 + self._blocks.append( + self._allocator.allocate_mutable_block( + prev_block=self._blocks[-1], device=device)) + + def fork(self) -> "BlockTable": + """Creates a new BlockTable instance with a copy of the blocks from the + current instance. + + This method creates a new BlockTable instance with the same block size, + block allocator, and a copy of the blocks from the current instance. The + new BlockTable has its own independent set of blocks, but shares the + same underlying memory allocation with the original BlockTable. + + Returns: + BlockTable: A new BlockTable instance with a copy of the blocks from + the current instance. + """ + assert self._is_allocated + assert len(self._blocks) > 0 forked_blocks = self._allocator.fork(self._blocks[-1]) return BlockTable( block_size=self._block_size, @@ -223,84 +223,84 @@ class BlockTable: max_block_sliding_window=self._max_block_sliding_window, cache_namespace=self._cache_namespace, ) - - def free(self) -> None: - """Frees the memory occupied by the blocks in the BlockTable. - - This method iterates over all the blocks in the `_blocks` list and calls - the `free` method of the `_allocator` object to release the memory - occupied by each block. After freeing all the blocks, the `_blocks` list - is set to `None`. - """ - for block in self.blocks: - self._allocator.free(block) - self._blocks.reset() - - @property - def physical_block_ids(self) -> List[int]: - """Returns a list of physical block indices for the blocks in the - BlockTable. - - This property returns a list of integers, where each integer represents - the physical block index of a corresponding block in the `_blocks` list. - The physical block index is a unique identifier for the memory location - occupied by the block. - - Returns: - List[int]: A list of physical block indices for the blocks in the - BlockTable. - """ - return self._blocks.ids() - - def get_unseen_token_ids(self, sequence_token_ids: List[int]) -> List[int]: - """Get the number of "unseen" tokens in the sequence. - - Unseen tokens are tokens in the sequence corresponding to this block - table, but are not yet appended to this block table. - - Args: - sequence_token_ids (List[int]): The list of token ids in the - sequence. - - Returns: - List[int]: The postfix of sequence_token_ids that has not yet been - appended to the block table. - """ - - # Since the block table is append-only, the unseen token ids are the - # ones after the appended ones. - return sequence_token_ids[self.num_full_slots:] - + + def free(self) -> None: + """Frees the memory occupied by the blocks in the BlockTable. + + This method iterates over all the blocks in the `_blocks` list and calls + the `free` method of the `_allocator` object to release the memory + occupied by each block. After freeing all the blocks, the `_blocks` list + is set to `None`. + """ + for block in self.blocks: + self._allocator.free(block) + self._blocks.reset() + + @property + def physical_block_ids(self) -> List[int]: + """Returns a list of physical block indices for the blocks in the + BlockTable. + + This property returns a list of integers, where each integer represents + the physical block index of a corresponding block in the `_blocks` list. + The physical block index is a unique identifier for the memory location + occupied by the block. + + Returns: + List[int]: A list of physical block indices for the blocks in the + BlockTable. + """ + return self._blocks.ids() + + def get_unseen_token_ids(self, sequence_token_ids: List[int]) -> List[int]: + """Get the number of "unseen" tokens in the sequence. + + Unseen tokens are tokens in the sequence corresponding to this block + table, but are not yet appended to this block table. + + Args: + sequence_token_ids (List[int]): The list of token ids in the + sequence. + + Returns: + List[int]: The postfix of sequence_token_ids that has not yet been + appended to the block table. + """ + + # Since the block table is append-only, the unseen token ids are the + # ones after the appended ones. + return sequence_token_ids[self.num_full_slots:] + def _allocate_blocks_for_token_ids(self, prev_block: Optional[Block], token_ids: List[int], device: Device) -> List[Block]: blocks: List[Block] = [] block_token_ids = [] - tail_token_ids = [] - for cur_token_ids in chunk_list(token_ids, self._block_size): - if len(cur_token_ids) == self._block_size: - block_token_ids.append(cur_token_ids) - else: - tail_token_ids.append(cur_token_ids) - + tail_token_ids = [] + for cur_token_ids in chunk_list(token_ids, self._block_size): + if len(cur_token_ids) == self._block_size: + block_token_ids.append(cur_token_ids) + else: + tail_token_ids.append(cur_token_ids) + if block_token_ids: blocks.extend(self._allocate_immutable_blocks( prev_block=prev_block, block_token_ids=block_token_ids, device=device)) prev_block = blocks[-1] - - if tail_token_ids: - assert len(tail_token_ids) == 1 - cur_token_ids = tail_token_ids[0] - + + if tail_token_ids: + assert len(tail_token_ids) == 1 + cur_token_ids = tail_token_ids[0] + block = self._allocate_mutable_block(prev_block=prev_block, device=device) - block.append_token_ids(cur_token_ids) - - blocks.append(block) - + block.append_token_ids(cur_token_ids) + + blocks.append(block) + return blocks def _allocate_mutable_block(self, prev_block: Optional[Block], @@ -372,85 +372,85 @@ class BlockTable: prev_block, block_token_ids=block_token_ids, device=device) - - def _get_all_token_ids(self) -> List[int]: - # NOTE: This function is O(seq_len); use sparingly. - token_ids: List[int] = [] - - if not self._is_allocated: - return token_ids - - for block in self.blocks: - token_ids.extend(block.token_ids) - - return token_ids - - def _get_num_token_ids(self) -> int: - res = 0 - for block in self.blocks: - res += len(block.token_ids) - - return res - - @property - def _is_allocated(self) -> bool: - return len(self._blocks) > 0 - - @property - def blocks(self) -> List[Block]: - return self._blocks.list() - - @property - def _num_empty_slots(self) -> int: - assert self._is_allocated - return len(self._blocks) * self._block_size - self._num_full_slots - - @property - def num_full_slots(self) -> int: - """Returns the total number of tokens currently stored in the - BlockTable. - - Returns: - int: The total number of tokens currently stored in the BlockTable. - """ - return self._num_full_slots - - def get_num_blocks_touched_by_append_slots( - self, token_ids: List[int], num_lookahead_slots: int) -> int: - """Determine how many blocks will be "touched" by appending the token - ids. - - This is required for the scheduler to determine whether a sequence can - continue generation, or if it must be preempted. - """ - # Math below is equivalent to: - # all_token_ids = token_ids + [-1] * num_lookahead_slots - # token_blocks = self._chunk_token_blocks_for_append(all_token_ids) - # return len(token_blocks) - - num_token_ids = len(token_ids) + num_lookahead_slots - first_chunk_size = self._block_size - (self._num_full_slots % - self._block_size) - num_token_blocks = (1 + math.ceil( - (num_token_ids - first_chunk_size) / self._block_size)) - return num_token_blocks - - def _chunk_token_blocks_for_append( - self, token_ids: List[int]) -> List[List[int]]: - """Split the token ids into block-sized chunks so they can be easily - appended to blocks. The first such "token block" may have less token ids - than the block size, since the last allocated block may be partially - full. - - If no token ids are provided, then no chunks are returned. - """ - - if not token_ids: - return [] - - first_chunk_size = self._block_size - (self._num_full_slots % - self._block_size) - token_blocks = [token_ids[:first_chunk_size]] - token_blocks.extend( - chunk_list(token_ids[first_chunk_size:], self._block_size)) - return token_blocks + + def _get_all_token_ids(self) -> List[int]: + # NOTE: This function is O(seq_len); use sparingly. + token_ids: List[int] = [] + + if not self._is_allocated: + return token_ids + + for block in self.blocks: + token_ids.extend(block.token_ids) + + return token_ids + + def _get_num_token_ids(self) -> int: + res = 0 + for block in self.blocks: + res += len(block.token_ids) + + return res + + @property + def _is_allocated(self) -> bool: + return len(self._blocks) > 0 + + @property + def blocks(self) -> List[Block]: + return self._blocks.list() + + @property + def _num_empty_slots(self) -> int: + assert self._is_allocated + return len(self._blocks) * self._block_size - self._num_full_slots + + @property + def num_full_slots(self) -> int: + """Returns the total number of tokens currently stored in the + BlockTable. + + Returns: + int: The total number of tokens currently stored in the BlockTable. + """ + return self._num_full_slots + + def get_num_blocks_touched_by_append_slots( + self, token_ids: List[int], num_lookahead_slots: int) -> int: + """Determine how many blocks will be "touched" by appending the token + ids. + + This is required for the scheduler to determine whether a sequence can + continue generation, or if it must be preempted. + """ + # Math below is equivalent to: + # all_token_ids = token_ids + [-1] * num_lookahead_slots + # token_blocks = self._chunk_token_blocks_for_append(all_token_ids) + # return len(token_blocks) + + num_token_ids = len(token_ids) + num_lookahead_slots + first_chunk_size = self._block_size - (self._num_full_slots % + self._block_size) + num_token_blocks = (1 + math.ceil( + (num_token_ids - first_chunk_size) / self._block_size)) + return num_token_blocks + + def _chunk_token_blocks_for_append( + self, token_ids: List[int]) -> List[List[int]]: + """Split the token ids into block-sized chunks so they can be easily + appended to blocks. The first such "token block" may have less token ids + than the block size, since the last allocated block may be partially + full. + + If no token ids are provided, then no chunks are returned. + """ + + if not token_ids: + return [] + + first_chunk_size = self._block_size - (self._num_full_slots % + self._block_size) + token_blocks = [token_ids[:first_chunk_size]] + token_blocks.extend( + chunk_list(token_ids[first_chunk_size:], self._block_size)) + return token_blocks diff --git a/vllm_overrides/core/block/cpu_gpu_block_allocator.py b/vllm_overrides/core/block/cpu_gpu_block_allocator.py index 9628289c..09a4bdd2 100644 --- a/vllm_overrides/core/block/cpu_gpu_block_allocator.py +++ b/vllm_overrides/core/block/cpu_gpu_block_allocator.py @@ -4,56 +4,56 @@ from vllm.core.block.cpu_kv_content_cache import (CpuKvContentCache, cpu_kv_offload_enabled) from vllm.core.block.interfaces import (Block, BlockAllocator, BlockId, DeviceAwareBlockAllocator) -from vllm.core.block.naive_block import NaiveBlock, NaiveBlockAllocator -from vllm.core.block.prefix_caching_block import PrefixCachingBlockAllocator -from vllm.utils import Device - - -class CpuGpuBlockAllocator(DeviceAwareBlockAllocator): - """A block allocator that can allocate blocks on both CPU and GPU memory. - - This class implements the `DeviceAwareBlockAllocator` interface and provides - functionality for allocating and managing blocks of memory on both CPU and - GPU devices. - - The `CpuGpuBlockAllocator` maintains separate memory pools for CPU and GPU - blocks, and allows for allocation, deallocation, forking, and swapping of - blocks across these memory pools. - """ - - @staticmethod - def create( - allocator_type: str, - num_gpu_blocks: int, - num_cpu_blocks: int, - block_size: int, - ) -> DeviceAwareBlockAllocator: - """Creates a CpuGpuBlockAllocator instance with the specified - configuration. - - This static method creates and returns a CpuGpuBlockAllocator instance - based on the provided parameters. It initializes the CPU and GPU block - allocators with the specified number of blocks, block size, and - allocator type. - - Args: - allocator_type (str): The type of block allocator to use for CPU - and GPU blocks. Currently supported values are "naive" and - "prefix_caching". - num_gpu_blocks (int): The number of blocks to allocate for GPU - memory. - num_cpu_blocks (int): The number of blocks to allocate for CPU - memory. - block_size (int): The size of each block in number of tokens. - - Returns: - DeviceAwareBlockAllocator: A CpuGpuBlockAllocator instance with the - specified configuration. - - Notes: - - The block IDs are assigned contiguously, with GPU block IDs coming - before CPU block IDs. - """ +from vllm.core.block.naive_block import NaiveBlock, NaiveBlockAllocator +from vllm.core.block.prefix_caching_block import PrefixCachingBlockAllocator +from vllm.utils import Device + + +class CpuGpuBlockAllocator(DeviceAwareBlockAllocator): + """A block allocator that can allocate blocks on both CPU and GPU memory. + + This class implements the `DeviceAwareBlockAllocator` interface and provides + functionality for allocating and managing blocks of memory on both CPU and + GPU devices. + + The `CpuGpuBlockAllocator` maintains separate memory pools for CPU and GPU + blocks, and allows for allocation, deallocation, forking, and swapping of + blocks across these memory pools. + """ + + @staticmethod + def create( + allocator_type: str, + num_gpu_blocks: int, + num_cpu_blocks: int, + block_size: int, + ) -> DeviceAwareBlockAllocator: + """Creates a CpuGpuBlockAllocator instance with the specified + configuration. + + This static method creates and returns a CpuGpuBlockAllocator instance + based on the provided parameters. It initializes the CPU and GPU block + allocators with the specified number of blocks, block size, and + allocator type. + + Args: + allocator_type (str): The type of block allocator to use for CPU + and GPU blocks. Currently supported values are "naive" and + "prefix_caching". + num_gpu_blocks (int): The number of blocks to allocate for GPU + memory. + num_cpu_blocks (int): The number of blocks to allocate for CPU + memory. + block_size (int): The size of each block in number of tokens. + + Returns: + DeviceAwareBlockAllocator: A CpuGpuBlockAllocator instance with the + specified configuration. + + Notes: + - The block IDs are assigned contiguously, with GPU block IDs coming + before CPU block IDs. + """ content_offload = cpu_kv_offload_enabled() if content_offload and allocator_type != "prefix_caching": raise RuntimeError( @@ -63,38 +63,38 @@ class CpuGpuBlockAllocator(DeviceAwareBlockAllocator): "BI100_CPU_KV_OFFLOAD=1 requires at least one CPU KV block") block_ids = list(range(num_gpu_blocks + num_cpu_blocks)) - gpu_block_ids = block_ids[:num_gpu_blocks] - cpu_block_ids = block_ids[num_gpu_blocks:] - - if allocator_type == "naive": - gpu_allocator: BlockAllocator = NaiveBlockAllocator( - create_block=NaiveBlock, # type: ignore - num_blocks=num_gpu_blocks, - block_size=block_size, - block_ids=gpu_block_ids, - ) - - cpu_allocator: BlockAllocator = NaiveBlockAllocator( - create_block=NaiveBlock, # type: ignore - num_blocks=num_cpu_blocks, - block_size=block_size, - block_ids=cpu_block_ids, - ) - elif allocator_type == "prefix_caching": - gpu_allocator = PrefixCachingBlockAllocator( - num_blocks=num_gpu_blocks, - block_size=block_size, - block_ids=gpu_block_ids, - ) - - cpu_allocator = PrefixCachingBlockAllocator( - num_blocks=num_cpu_blocks, - block_size=block_size, - block_ids=cpu_block_ids, - ) - else: - raise ValueError(f"Unknown allocator type {allocator_type=}") - + gpu_block_ids = block_ids[:num_gpu_blocks] + cpu_block_ids = block_ids[num_gpu_blocks:] + + if allocator_type == "naive": + gpu_allocator: BlockAllocator = NaiveBlockAllocator( + create_block=NaiveBlock, # type: ignore + num_blocks=num_gpu_blocks, + block_size=block_size, + block_ids=gpu_block_ids, + ) + + cpu_allocator: BlockAllocator = NaiveBlockAllocator( + create_block=NaiveBlock, # type: ignore + num_blocks=num_cpu_blocks, + block_size=block_size, + block_ids=cpu_block_ids, + ) + elif allocator_type == "prefix_caching": + gpu_allocator = PrefixCachingBlockAllocator( + num_blocks=num_gpu_blocks, + block_size=block_size, + block_ids=gpu_block_ids, + ) + + cpu_allocator = PrefixCachingBlockAllocator( + num_blocks=num_cpu_blocks, + block_size=block_size, + block_ids=cpu_block_ids, + ) + else: + raise ValueError(f"Unknown allocator type {allocator_type=}") + return CpuGpuBlockAllocator( cpu_block_allocator=cpu_allocator, gpu_block_allocator=gpu_allocator, @@ -105,21 +105,21 @@ class CpuGpuBlockAllocator(DeviceAwareBlockAllocator): def __init__(self, cpu_block_allocator: BlockAllocator, gpu_block_allocator: BlockAllocator, cpu_content_cache: Optional[CpuKvContentCache] = None): - assert not ( - cpu_block_allocator.all_block_ids - & gpu_block_allocator.all_block_ids - ), "cpu and gpu block allocators can't have intersection of block ids" - - self._allocators = { - Device.CPU: cpu_block_allocator, - Device.GPU: gpu_block_allocator, - } - + assert not ( + cpu_block_allocator.all_block_ids + & gpu_block_allocator.all_block_ids + ), "cpu and gpu block allocators can't have intersection of block ids" + + self._allocators = { + Device.CPU: cpu_block_allocator, + Device.GPU: gpu_block_allocator, + } + self._swap_mapping: Dict[int, int] = {} self._null_block: Optional[Block] = None self._cpu_content_cache = cpu_content_cache - - self._block_ids_to_allocator: Dict[int, BlockAllocator] = {} + + self._block_ids_to_allocator: Dict[int, BlockAllocator] = {} for _, allocator in self._allocators.items(): for block_id in allocator.all_block_ids: self._block_ids_to_allocator[block_id] = allocator @@ -164,236 +164,236 @@ class CpuGpuBlockAllocator(DeviceAwareBlockAllocator): assert self._cpu_content_cache is not None gpu_slot = self.get_physical_block_id(Device.GPU, gpu_block_id) return self._cpu_content_cache.stage_store(content_hash, gpu_slot) - - def allocate_or_get_null_block(self) -> Block: - if self._null_block is None: - self._null_block = NullBlock( - self.allocate_mutable_block(None, Device.GPU)) - return self._null_block - - def allocate_mutable_block(self, prev_block: Optional[Block], - device: Device) -> Block: - """Allocates a new mutable block on the specified device. - - Args: - prev_block (Optional[Block]): The previous block to in the sequence. - Used for prefix hashing. - device (Device): The device on which to allocate the new block. - - Returns: - Block: The newly allocated mutable block. - """ - return self._allocators[device].allocate_mutable_block(prev_block) - - def allocate_immutable_blocks(self, prev_block: Optional[Block], - block_token_ids: List[List[int]], - device: Device) -> List[Block]: - """Allocates a new group of immutable blocks with the provided block - token IDs on the specified device. - - Args: - prev_block (Optional[Block]): The previous block in the sequence. - Used for prefix hashing. - block_token_ids (List[int]): The list of block token IDs to be - stored in the new blocks. - device (Device): The device on which to allocate the new block. - - Returns: - List[Block]: The newly allocated list of immutable blocks - containing the provided block token IDs. - """ - return self._allocators[device].allocate_immutable_blocks( - prev_block, block_token_ids) - - def allocate_immutable_block(self, prev_block: Optional[Block], - token_ids: List[int], - device: Device) -> Block: - """Allocates a new immutable block with the provided token IDs on the - specified device. - - Args: - prev_block (Optional[Block]): The previous block in the sequence. - Used for prefix hashing. - token_ids (List[int]): The list of token IDs to be stored in the new - block. - device (Device): The device on which to allocate the new block. - - Returns: - Block: The newly allocated immutable block containing the provided - token IDs. - """ - return self._allocators[device].allocate_immutable_block( - prev_block, token_ids) - - def free(self, block: Block) -> None: - """Frees the memory occupied by the given block. - - Args: - block (Block): The block to be freed. - """ - # Null block should never be freed - if isinstance(block, NullBlock): - return - block_id = block.block_id - assert block_id is not None - allocator = self._block_ids_to_allocator[block_id] - allocator.free(block) - - def fork(self, last_block: Block) -> List[Block]: - """Creates a new sequence of blocks that shares the same underlying - memory as the original sequence. - - Args: - last_block (Block): The last block in the original sequence. - - Returns: - List[Block]: A new list of blocks that shares the same memory as the - original sequence. - """ - # do not attempt to fork the null block - assert not isinstance(last_block, NullBlock) - block_id = last_block.block_id - assert block_id is not None - allocator = self._block_ids_to_allocator[block_id] - return allocator.fork(last_block) - - def get_num_free_blocks(self, device: Device) -> int: - """Returns the number of free blocks available on the specified device. - - Args: - device (Device): The device for which to query the number of free - blocks. AssertionError is raised if None is passed. - - Returns: - int: The number of free blocks available on the specified device. - """ - return self._allocators[device].get_num_free_blocks() - - def get_num_total_blocks(self, device: Device) -> int: - return self._allocators[device].get_num_total_blocks() - - def get_physical_block_id(self, device: Device, absolute_id: int) -> int: - """Returns the zero-offset block id on certain device given the - absolute block id. - - Args: - device (Device): The device for which to query relative block id. - absolute_id (int): The absolute block id for the block in - whole allocator. - - Returns: - int: The zero-offset block id on certain device. - """ - return self._allocators[device].get_physical_block_id(absolute_id) - + + def allocate_or_get_null_block(self) -> Block: + if self._null_block is None: + self._null_block = NullBlock( + self.allocate_mutable_block(None, Device.GPU)) + return self._null_block + + def allocate_mutable_block(self, prev_block: Optional[Block], + device: Device) -> Block: + """Allocates a new mutable block on the specified device. + + Args: + prev_block (Optional[Block]): The previous block to in the sequence. + Used for prefix hashing. + device (Device): The device on which to allocate the new block. + + Returns: + Block: The newly allocated mutable block. + """ + return self._allocators[device].allocate_mutable_block(prev_block) + + def allocate_immutable_blocks(self, prev_block: Optional[Block], + block_token_ids: List[List[int]], + device: Device) -> List[Block]: + """Allocates a new group of immutable blocks with the provided block + token IDs on the specified device. + + Args: + prev_block (Optional[Block]): The previous block in the sequence. + Used for prefix hashing. + block_token_ids (List[int]): The list of block token IDs to be + stored in the new blocks. + device (Device): The device on which to allocate the new block. + + Returns: + List[Block]: The newly allocated list of immutable blocks + containing the provided block token IDs. + """ + return self._allocators[device].allocate_immutable_blocks( + prev_block, block_token_ids) + + def allocate_immutable_block(self, prev_block: Optional[Block], + token_ids: List[int], + device: Device) -> Block: + """Allocates a new immutable block with the provided token IDs on the + specified device. + + Args: + prev_block (Optional[Block]): The previous block in the sequence. + Used for prefix hashing. + token_ids (List[int]): The list of token IDs to be stored in the new + block. + device (Device): The device on which to allocate the new block. + + Returns: + Block: The newly allocated immutable block containing the provided + token IDs. + """ + return self._allocators[device].allocate_immutable_block( + prev_block, token_ids) + + def free(self, block: Block) -> None: + """Frees the memory occupied by the given block. + + Args: + block (Block): The block to be freed. + """ + # Null block should never be freed + if isinstance(block, NullBlock): + return + block_id = block.block_id + assert block_id is not None + allocator = self._block_ids_to_allocator[block_id] + allocator.free(block) + + def fork(self, last_block: Block) -> List[Block]: + """Creates a new sequence of blocks that shares the same underlying + memory as the original sequence. + + Args: + last_block (Block): The last block in the original sequence. + + Returns: + List[Block]: A new list of blocks that shares the same memory as the + original sequence. + """ + # do not attempt to fork the null block + assert not isinstance(last_block, NullBlock) + block_id = last_block.block_id + assert block_id is not None + allocator = self._block_ids_to_allocator[block_id] + return allocator.fork(last_block) + + def get_num_free_blocks(self, device: Device) -> int: + """Returns the number of free blocks available on the specified device. + + Args: + device (Device): The device for which to query the number of free + blocks. AssertionError is raised if None is passed. + + Returns: + int: The number of free blocks available on the specified device. + """ + return self._allocators[device].get_num_free_blocks() + + def get_num_total_blocks(self, device: Device) -> int: + return self._allocators[device].get_num_total_blocks() + + def get_physical_block_id(self, device: Device, absolute_id: int) -> int: + """Returns the zero-offset block id on certain device given the + absolute block id. + + Args: + device (Device): The device for which to query relative block id. + absolute_id (int): The absolute block id for the block in + whole allocator. + + Returns: + int: The zero-offset block id on certain device. + """ + return self._allocators[device].get_physical_block_id(absolute_id) + def swap(self, blocks: List[Block], src_device: Device, dst_device: Device) -> Dict[int, int]: - """Execute the swap for the given blocks from source_device - on to dest_device, save the current swap mapping and append - them to the accumulated `self._swap_mapping` for each - scheduling move. - - Args: - blocks: List of blocks to be swapped. - src_device (Device): Device to swap the 'blocks' from. - dst_device (Device): Device to swap the 'blocks' to. - - Returns: - Dict[int, int]: Swap mapping from source_device - on to dest_device. - """ + """Execute the swap for the given blocks from source_device + on to dest_device, save the current swap mapping and append + them to the accumulated `self._swap_mapping` for each + scheduling move. + + Args: + blocks: List of blocks to be swapped. + src_device (Device): Device to swap the 'blocks' from. + dst_device (Device): Device to swap the 'blocks' to. + + Returns: + Dict[int, int]: Swap mapping from source_device + on to dest_device. + """ if self.content_offload_enabled: raise RuntimeError( "request-level preemption swap cannot share CPU slots with " "BI100_CPU_KV_OFFLOAD") src_block_ids = [block.block_id for block in blocks] - self._allocators[src_device].swap_out(blocks) - self._allocators[dst_device].swap_in(blocks) - dst_block_ids = [block.block_id for block in blocks] - - current_swap_mapping: Dict[int, int] = {} - for src_block_id, dst_block_id in zip(src_block_ids, dst_block_ids): - if src_block_id is not None and dst_block_id is not None: - self._swap_mapping[src_block_id] = dst_block_id - current_swap_mapping[src_block_id] = dst_block_id - return current_swap_mapping - - def get_num_full_blocks_touched(self, blocks: List[Block], - device: Device) -> int: - """Returns the number of full blocks that will be touched by - swapping in/out the given blocks on to the 'device'. - - Args: - blocks: List of blocks to be swapped. - device (Device): Device to swap the 'blocks' on. - - Returns: - int: the number of full blocks that will be touched by - swapping in/out the given blocks on to the 'device'. - Non full blocks are ignored when deciding the number - of blocks to touch. - """ - return self._allocators[device].get_num_full_blocks_touched(blocks) - - def clear_copy_on_writes(self) -> List[Tuple[int, int]]: - """Clears the copy-on-write (CoW) state and returns the mapping of - source to destination block IDs. - - Returns: - List[Tuple[int, int]]: A list mapping source block IDs to - destination block IDs. - """ - # CoW only supported on GPU - device = Device.GPU - return self._allocators[device].clear_copy_on_writes() - - def mark_blocks_as_accessed(self, block_ids: List[int], - now: float) -> None: - """Mark blocks as accessed, only use for prefix caching.""" - # Prefix caching only supported on GPU. - device = Device.GPU - return self._allocators[device].mark_blocks_as_accessed(block_ids, now) - - def mark_blocks_as_computed(self, block_ids: List[int]) -> None: - """Mark blocks as accessed, only use for prefix caching.""" - # Prefix caching only supported on GPU. - device = Device.GPU - return self._allocators[device].mark_blocks_as_computed(block_ids) - - def get_computed_block_ids(self, prev_computed_block_ids: List[int], - block_ids: List[int], - skip_last_block_id: bool) -> List[int]: - # Prefix caching only supported on GPU. - device = Device.GPU - return self._allocators[device].get_computed_block_ids( - prev_computed_block_ids, block_ids, skip_last_block_id) - - def get_common_computed_block_ids( - self, computed_seq_block_ids: List[List[int]]) -> List[int]: - # Prefix caching only supported on GPU. - device = Device.GPU - return self._allocators[device].get_common_computed_block_ids( - computed_seq_block_ids) - - @property - def all_block_ids(self) -> FrozenSet[int]: - return frozenset(self._block_ids_to_allocator.keys()) - - def get_prefix_cache_hit_rate(self, device: Device) -> float: - """Prefix cache hit rate. -1 means not supported or disabled.""" - assert device in self._allocators - return self._allocators[device].get_prefix_cache_hit_rate() - + self._allocators[src_device].swap_out(blocks) + self._allocators[dst_device].swap_in(blocks) + dst_block_ids = [block.block_id for block in blocks] + + current_swap_mapping: Dict[int, int] = {} + for src_block_id, dst_block_id in zip(src_block_ids, dst_block_ids): + if src_block_id is not None and dst_block_id is not None: + self._swap_mapping[src_block_id] = dst_block_id + current_swap_mapping[src_block_id] = dst_block_id + return current_swap_mapping + + def get_num_full_blocks_touched(self, blocks: List[Block], + device: Device) -> int: + """Returns the number of full blocks that will be touched by + swapping in/out the given blocks on to the 'device'. + + Args: + blocks: List of blocks to be swapped. + device (Device): Device to swap the 'blocks' on. + + Returns: + int: the number of full blocks that will be touched by + swapping in/out the given blocks on to the 'device'. + Non full blocks are ignored when deciding the number + of blocks to touch. + """ + return self._allocators[device].get_num_full_blocks_touched(blocks) + + def clear_copy_on_writes(self) -> List[Tuple[int, int]]: + """Clears the copy-on-write (CoW) state and returns the mapping of + source to destination block IDs. + + Returns: + List[Tuple[int, int]]: A list mapping source block IDs to + destination block IDs. + """ + # CoW only supported on GPU + device = Device.GPU + return self._allocators[device].clear_copy_on_writes() + + def mark_blocks_as_accessed(self, block_ids: List[int], + now: float) -> None: + """Mark blocks as accessed, only use for prefix caching.""" + # Prefix caching only supported on GPU. + device = Device.GPU + return self._allocators[device].mark_blocks_as_accessed(block_ids, now) + + def mark_blocks_as_computed(self, block_ids: List[int]) -> None: + """Mark blocks as accessed, only use for prefix caching.""" + # Prefix caching only supported on GPU. + device = Device.GPU + return self._allocators[device].mark_blocks_as_computed(block_ids) + + def get_computed_block_ids(self, prev_computed_block_ids: List[int], + block_ids: List[int], + skip_last_block_id: bool) -> List[int]: + # Prefix caching only supported on GPU. + device = Device.GPU + return self._allocators[device].get_computed_block_ids( + prev_computed_block_ids, block_ids, skip_last_block_id) + + def get_common_computed_block_ids( + self, computed_seq_block_ids: List[List[int]]) -> List[int]: + # Prefix caching only supported on GPU. + device = Device.GPU + return self._allocators[device].get_common_computed_block_ids( + computed_seq_block_ids) + + @property + def all_block_ids(self) -> FrozenSet[int]: + return frozenset(self._block_ids_to_allocator.keys()) + + def get_prefix_cache_hit_rate(self, device: Device) -> float: + """Prefix cache hit rate. -1 means not supported or disabled.""" + assert device in self._allocators + return self._allocators[device].get_prefix_cache_hit_rate() + def get_and_reset_swaps(self) -> List[Tuple[int, int]]: - """Returns and clears the mapping of source to destination block IDs. - Will be called after every swapping operations for now, and after every - schedule when BlockManagerV2 become default. Currently not useful. - - Returns: - List[Tuple[int, int]]: A mapping of source to destination block IDs. - """ - mapping = self._swap_mapping.copy() + """Returns and clears the mapping of source to destination block IDs. + Will be called after every swapping operations for now, and after every + schedule when BlockManagerV2 become default. Currently not useful. + + Returns: + List[Tuple[int, int]]: A mapping of source to destination block IDs. + """ + mapping = self._swap_mapping.copy() self._swap_mapping.clear() return list(mapping.items()) @@ -407,69 +407,69 @@ class CpuGpuBlockAllocator(DeviceAwareBlockAllocator): def begin_prefix_cache_step(self) -> None: if self._cpu_content_cache is not None: self._cpu_content_cache.begin_step() - - -class NullBlock(Block): - """ - Null blocks are used as a placeholders for KV cache blocks that have - been dropped due to sliding window. - This implementation just wraps an ordinary block and prevents it from - being modified. It also allows for testing if a block is NullBlock - via isinstance(). - """ - - def __init__(self, proxy: Block): - super().__init__() - self._proxy = proxy - - def append_token_ids(self, token_ids: List[BlockId]): - raise ValueError("null block should not be modified") - - @property - def block_id(self): - return self._proxy.block_id - - @block_id.setter - def block_id(self, value: Optional[BlockId]): - raise ValueError("null block should not be modified") - - @property - def token_ids(self) -> List[BlockId]: - return self._proxy.token_ids - - @property - def num_tokens_total(self) -> int: - raise NotImplementedError( - "num_tokens_total is not used for null block") - - @property - def num_empty_slots(self) -> BlockId: - return self._proxy.num_empty_slots - - @property - def is_full(self): - return self._proxy.is_full - - @property - def prev_block(self): - return self._proxy.prev_block - - @property - def computed(self): - return self._proxy.computed - - @computed.setter - def computed(self, value): - self._proxy.computed = value - - @property - def last_accessed(self) -> float: - return self._proxy.last_accessed - - @last_accessed.setter - def last_accessed(self, last_accessed_ts: float): - self._proxy.last_accessed = last_accessed_ts - - @property - def content_hash(self): - return self._proxy.content_hash + + +class NullBlock(Block): + """ + Null blocks are used as a placeholders for KV cache blocks that have + been dropped due to sliding window. + This implementation just wraps an ordinary block and prevents it from + being modified. It also allows for testing if a block is NullBlock + via isinstance(). + """ + + def __init__(self, proxy: Block): + super().__init__() + self._proxy = proxy + + def append_token_ids(self, token_ids: List[BlockId]): + raise ValueError("null block should not be modified") + + @property + def block_id(self): + return self._proxy.block_id + + @block_id.setter + def block_id(self, value: Optional[BlockId]): + raise ValueError("null block should not be modified") + + @property + def token_ids(self) -> List[BlockId]: + return self._proxy.token_ids + + @property + def num_tokens_total(self) -> int: + raise NotImplementedError( + "num_tokens_total is not used for null block") + + @property + def num_empty_slots(self) -> BlockId: + return self._proxy.num_empty_slots + + @property + def is_full(self): + return self._proxy.is_full + + @property + def prev_block(self): + return self._proxy.prev_block + + @property + def computed(self): + return self._proxy.computed + + @computed.setter + def computed(self, value): + self._proxy.computed = value + + @property + def last_accessed(self) -> float: + return self._proxy.last_accessed + + @last_accessed.setter + def last_accessed(self, last_accessed_ts: float): + self._proxy.last_accessed = last_accessed_ts + + @property + def content_hash(self): + return self._proxy.content_hash diff --git a/vllm_overrides/core/block/prefix_caching_block.py b/vllm_overrides/core/block/prefix_caching_block.py index 772fa71f..c2c38e38 100644 --- a/vllm_overrides/core/block/prefix_caching_block.py +++ b/vllm_overrides/core/block/prefix_caching_block.py @@ -4,105 +4,105 @@ import struct from os.path import commonprefix from typing import (Callable, Dict, FrozenSet, Iterable, List, Optional, Set, Tuple) - -from vllm.core.block.common import (CacheMetricData, CopyOnWriteTracker, - get_all_blocks_recursively) -from vllm.core.block.interfaces import Block, BlockAllocator, BlockId, Device -from vllm.core.block.naive_block import (BlockPool, NaiveBlock, - NaiveBlockAllocator) + +from vllm.core.block.common import (CacheMetricData, CopyOnWriteTracker, + get_all_blocks_recursively) +from vllm.core.block.interfaces import Block, BlockAllocator, BlockId, Device +from vllm.core.block.naive_block import (BlockPool, NaiveBlock, + NaiveBlockAllocator) from vllm.core.evictor_v2 import (EvictionPolicy, Evictor, eviction_policy_from_env, make_evictor) - + PrefixHash = bytes - -# By default, we init our block access time as _DEFAULT_LAST_ACCESSED_TIME -# so that if we find one block is still hold _DEFAULT_LAST_ACCESSED_TIME, -# then we know this block hasn't been accessed yet. -_DEFAULT_LAST_ACCESSED_TIME = -1 - - -class BlockTracker: - """Used to track the status of a block inside the prefix caching allocator - """ - __slots__ = ("active", "last_accessed", "computed") - - def reset(self): - self.last_accessed: float = _DEFAULT_LAST_ACCESSED_TIME - self.computed: bool = False - - def __init__(self): - self.active: bool = False - self.reset() - - def enable(self): - assert not self.active - self.active = True - self.reset() - - def disable(self): - assert self.active - self.active = False - self.reset() - - -class PrefixCachingBlockAllocator(BlockAllocator): - """A block allocator that implements prefix caching. - - The PrefixCachingBlockAllocator maintains a cache of blocks based on their - content hash. It reuses blocks with the same content hash to avoid redundant - memory allocation. The allocator also supports copy-on-write operations. - - Args: - num_blocks (int): The total number of blocks to manage. - block_size (int): The size of each block in tokens. - block_ids(Optional[Iterable[int]], optional): An optional iterable of - block IDs. If not provided, block IDs will be assigned sequentially - from 0 to num_blocks - 1. - """ - - def __init__( - self, - num_blocks: int, - block_size: int, - block_ids: Optional[Iterable[int]] = None, + +# By default, we init our block access time as _DEFAULT_LAST_ACCESSED_TIME +# so that if we find one block is still hold _DEFAULT_LAST_ACCESSED_TIME, +# then we know this block hasn't been accessed yet. +_DEFAULT_LAST_ACCESSED_TIME = -1 + + +class BlockTracker: + """Used to track the status of a block inside the prefix caching allocator + """ + __slots__ = ("active", "last_accessed", "computed") + + def reset(self): + self.last_accessed: float = _DEFAULT_LAST_ACCESSED_TIME + self.computed: bool = False + + def __init__(self): + self.active: bool = False + self.reset() + + def enable(self): + assert not self.active + self.active = True + self.reset() + + def disable(self): + assert self.active + self.active = False + self.reset() + + +class PrefixCachingBlockAllocator(BlockAllocator): + """A block allocator that implements prefix caching. + + The PrefixCachingBlockAllocator maintains a cache of blocks based on their + content hash. It reuses blocks with the same content hash to avoid redundant + memory allocation. The allocator also supports copy-on-write operations. + + Args: + num_blocks (int): The total number of blocks to manage. + block_size (int): The size of each block in tokens. + block_ids(Optional[Iterable[int]], optional): An optional iterable of + block IDs. If not provided, block IDs will be assigned sequentially + from 0 to num_blocks - 1. + """ + + def __init__( + self, + num_blocks: int, + block_size: int, + block_ids: Optional[Iterable[int]] = None, eviction_policy: Optional[EvictionPolicy] = None, - ): - if block_ids is None: - block_ids = range(num_blocks) - - self._block_size = block_size - + ): + if block_ids is None: + block_ids = range(num_blocks) + + self._block_size = block_size + # A mapping of prefix hash to block index. All blocks which have a # prefix hash will be in this dict, even if they have refcount 0. self._cached_blocks: Dict[PrefixHash, BlockId] = {} self._cache_namespace: Optional[bytes] = None - - # A list of immutable block IDs that have been touched by scheduler - # and should be marked as computed after an entire batch of sequences - # are scheduled. - self._touched_blocks: Set[BlockId] = set() - - # Used to track status of each physical block id - self._block_tracker: Dict[BlockId, BlockTracker] = {} - for block_id in block_ids: - self._block_tracker[block_id] = BlockTracker() - - # Pre-allocate "num_blocks * extra_factor" block objects. - # The "* extra_factor" is a buffer to allow more block objects - # than physical blocks - extra_factor = 4 - self._block_pool = BlockPool(self._block_size, self._create_block, - self, num_blocks * extra_factor) - - # An allocator for blocks that do not have prefix hashes. - self._hashless_allocator = NaiveBlockAllocator( - create_block=self._create_block, # type: ignore - num_blocks=num_blocks, - block_size=block_size, - block_ids=block_ids, - block_pool=self._block_pool, # Share block pool here - ) - + + # A list of immutable block IDs that have been touched by scheduler + # and should be marked as computed after an entire batch of sequences + # are scheduled. + self._touched_blocks: Set[BlockId] = set() + + # Used to track status of each physical block id + self._block_tracker: Dict[BlockId, BlockTracker] = {} + for block_id in block_ids: + self._block_tracker[block_id] = BlockTracker() + + # Pre-allocate "num_blocks * extra_factor" block objects. + # The "* extra_factor" is a buffer to allow more block objects + # than physical blocks + extra_factor = 4 + self._block_pool = BlockPool(self._block_size, self._create_block, + self, num_blocks * extra_factor) + + # An allocator for blocks that do not have prefix hashes. + self._hashless_allocator = NaiveBlockAllocator( + create_block=self._create_block, # type: ignore + num_blocks=num_blocks, + block_size=block_size, + block_ids=block_ids, + block_pool=self._block_pool, # Share block pool here + ) + if eviction_policy is None: eviction_policy = eviction_policy_from_env() self.eviction_policy = eviction_policy @@ -110,15 +110,15 @@ class PrefixCachingBlockAllocator(BlockAllocator): # Evitor used to maintain how we want to handle those computed blocks # if we find memory pressure is high. self.evictor: Evictor = make_evictor(eviction_policy) - - # We share the refcounter between allocators. This allows us to promote - # blocks originally allocated in the hashless allocator to immutable - # blocks. - self._refcounter = self._hashless_allocator.refcounter - - self._cow_tracker = CopyOnWriteTracker( - refcounter=self._refcounter.as_readonly()) - + + # We share the refcounter between allocators. This allows us to promote + # blocks originally allocated in the hashless allocator to immutable + # blocks. + self._refcounter = self._hashless_allocator.refcounter + + self._cow_tracker = CopyOnWriteTracker( + refcounter=self._refcounter.as_readonly()) + self.metric_data = CacheMetricData() self._external_cache_claim: Optional[ @@ -153,8 +153,8 @@ class PrefixCachingBlockAllocator(BlockAllocator): self._external_cache_load = load self._external_cache_cancel = cancel self._external_cache_store = store - - # Implements Block.Factory. + + # Implements Block.Factory. def _create_block( self, prev_block: Optional[Block], @@ -334,48 +334,48 @@ class PrefixCachingBlockAllocator(BlockAllocator): prev_block: Optional[Block], token_ids: List[int], device: Optional[Device] = None) -> Block: - """Allocates an immutable block with the given token IDs, reusing cached - blocks if possible. - - Args: - prev_block (Optional[Block]): The previous block in the sequence. - token_ids (List[int]): The token IDs to be stored in the block. - - Returns: - Block: The allocated immutable block. - """ + """Allocates an immutable block with the given token IDs, reusing cached + blocks if possible. + + Args: + prev_block (Optional[Block]): The previous block in the sequence. + token_ids (List[int]): The token IDs to be stored in the block. + + Returns: + Block: The allocated immutable block. + """ return self.allocate_immutable_block_with_cache_namespace( prev_block=prev_block, token_ids=token_ids, cache_namespace=b"", device=device) - - def allocate_immutable_blocks( - self, - prev_block: Optional[Block], - block_token_ids: List[List[int]], - device: Optional[Device] = None) -> List[Block]: - blocks = [] - for token_ids in block_token_ids: - prev_block = self.allocate_immutable_block(prev_block=prev_block, - token_ids=token_ids, - device=device) - blocks.append(prev_block) - return blocks - - def allocate_mutable_block(self, - prev_block: Optional[Block], - device: Optional[Device] = None) -> Block: - """Allocates a mutable block. If there are no free blocks, this will - evict unused cached blocks. - - Args: - prev_block (Block): The previous block in the sequence. - None is not allowed unlike it is super class. - - Returns: - Block: The allocated mutable block. - """ + + def allocate_immutable_blocks( + self, + prev_block: Optional[Block], + block_token_ids: List[List[int]], + device: Optional[Device] = None) -> List[Block]: + blocks = [] + for token_ids in block_token_ids: + prev_block = self.allocate_immutable_block(prev_block=prev_block, + token_ids=token_ids, + device=device) + blocks.append(prev_block) + return blocks + + def allocate_mutable_block(self, + prev_block: Optional[Block], + device: Optional[Device] = None) -> Block: + """Allocates a mutable block. If there are no free blocks, this will + evict unused cached blocks. + + Args: + prev_block (Block): The previous block in the sequence. + None is not allowed unlike it is super class. + + Returns: + Block: The allocated mutable block. + """ assert device is None assert_prefix_caching_block_or_none(prev_block) @@ -387,105 +387,105 @@ class PrefixCachingBlockAllocator(BlockAllocator): assert not block.computed assert block.content_hash is None return block - - def _incr_refcount_cached_block(self, block: Block) -> None: - # Set this block to be "computed" since it is pointing to a - # cached block id (which was already computed) - block.computed = True - - block_id = block.block_id - assert block_id is not None - - refcount = self._refcounter.incr(block_id) - if refcount == 1: - # In case a cached block was evicted, restore its tracking - if block_id in self.evictor: - self.evictor.remove(block_id) - - self._track_block_id(block_id, computed=True) - - def _decr_refcount_cached_block(self, block: Block) -> None: - # Ensure this is immutable/cached block - assert block.content_hash is not None - - block_id = block.block_id - assert block_id is not None - - refcount = self._refcounter.decr(block_id) - if refcount > 0: - block.block_id = None - return - else: - assert refcount == 0 - - # No longer used - assert block.content_hash in self._cached_blocks - - # Add the cached block to the evictor - # (This keeps the cached block around so it can be reused) - self.evictor.add(block_id, block.content_hash, block.num_tokens_total, - self._block_tracker[block_id].last_accessed) - - # Stop tracking the block - self._untrack_block_id(block_id) - - block.block_id = None - - def _decr_refcount_hashless_block(self, block: Block) -> None: - block_id = block.block_id - assert block_id is not None - - # We may have a fork case where block is shared, - # in which case, we cannot remove it from tracking - refcount = self._refcounter.get(block_id) - if refcount == 1: - self._untrack_block_id(block_id) - - # Decrement refcount of the block_id, but do not free the block object - # itself (will be handled by the caller) - self._hashless_allocator.free(block, keep_block_object=True) - - def _allocate_block_id(self) -> BlockId: - """First tries to allocate a block id from the hashless allocator, - and if there are no blocks, then tries to evict an unused cached block. - """ - hashless_block_id = self._maybe_allocate_hashless_block_id() - if hashless_block_id is not None: - return hashless_block_id - - evicted_block_id = self._maybe_allocate_evicted_block_id() - if evicted_block_id is not None: - return evicted_block_id - - # No block available in hashless allocator, nor in unused cache blocks. - raise BlockAllocator.NoFreeBlocksError() - - def _maybe_allocate_hashless_block_id(self) -> Optional[BlockId]: - try: - # Allocate mutable block and extract its block_id - block = self._hashless_allocator.allocate_mutable_block( - prev_block=None) - block_id = block.block_id - self._block_pool.free_block(block) - - self._track_block_id(block_id, computed=False) - return block_id - except BlockAllocator.NoFreeBlocksError: - return None - - def _maybe_allocate_evicted_block_id(self) -> Optional[BlockId]: - if self.evictor.num_blocks == 0: - return None - - # Here we get an evicted block, which is only added - # into evictor if its ref counter is 0 - # and since its content would be changed, we need - # to remove it from _cached_blocks's tracking list - block_id, content_hash_to_evict = self.evictor.evict() - - # Sanity checks - assert content_hash_to_evict in self._cached_blocks - _block_id = self._cached_blocks[content_hash_to_evict] + + def _incr_refcount_cached_block(self, block: Block) -> None: + # Set this block to be "computed" since it is pointing to a + # cached block id (which was already computed) + block.computed = True + + block_id = block.block_id + assert block_id is not None + + refcount = self._refcounter.incr(block_id) + if refcount == 1: + # In case a cached block was evicted, restore its tracking + if block_id in self.evictor: + self.evictor.remove(block_id) + + self._track_block_id(block_id, computed=True) + + def _decr_refcount_cached_block(self, block: Block) -> None: + # Ensure this is immutable/cached block + assert block.content_hash is not None + + block_id = block.block_id + assert block_id is not None + + refcount = self._refcounter.decr(block_id) + if refcount > 0: + block.block_id = None + return + else: + assert refcount == 0 + + # No longer used + assert block.content_hash in self._cached_blocks + + # Add the cached block to the evictor + # (This keeps the cached block around so it can be reused) + self.evictor.add(block_id, block.content_hash, block.num_tokens_total, + self._block_tracker[block_id].last_accessed) + + # Stop tracking the block + self._untrack_block_id(block_id) + + block.block_id = None + + def _decr_refcount_hashless_block(self, block: Block) -> None: + block_id = block.block_id + assert block_id is not None + + # We may have a fork case where block is shared, + # in which case, we cannot remove it from tracking + refcount = self._refcounter.get(block_id) + if refcount == 1: + self._untrack_block_id(block_id) + + # Decrement refcount of the block_id, but do not free the block object + # itself (will be handled by the caller) + self._hashless_allocator.free(block, keep_block_object=True) + + def _allocate_block_id(self) -> BlockId: + """First tries to allocate a block id from the hashless allocator, + and if there are no blocks, then tries to evict an unused cached block. + """ + hashless_block_id = self._maybe_allocate_hashless_block_id() + if hashless_block_id is not None: + return hashless_block_id + + evicted_block_id = self._maybe_allocate_evicted_block_id() + if evicted_block_id is not None: + return evicted_block_id + + # No block available in hashless allocator, nor in unused cache blocks. + raise BlockAllocator.NoFreeBlocksError() + + def _maybe_allocate_hashless_block_id(self) -> Optional[BlockId]: + try: + # Allocate mutable block and extract its block_id + block = self._hashless_allocator.allocate_mutable_block( + prev_block=None) + block_id = block.block_id + self._block_pool.free_block(block) + + self._track_block_id(block_id, computed=False) + return block_id + except BlockAllocator.NoFreeBlocksError: + return None + + def _maybe_allocate_evicted_block_id(self) -> Optional[BlockId]: + if self.evictor.num_blocks == 0: + return None + + # Here we get an evicted block, which is only added + # into evictor if its ref counter is 0 + # and since its content would be changed, we need + # to remove it from _cached_blocks's tracking list + block_id, content_hash_to_evict = self.evictor.evict() + + # Sanity checks + assert content_hash_to_evict in self._cached_blocks + _block_id = self._cached_blocks[content_hash_to_evict] assert self._refcounter.get(_block_id) == 0 assert _block_id == block_id @@ -493,242 +493,242 @@ class PrefixCachingBlockAllocator(BlockAllocator): self._external_cache_store(content_hash_to_evict, block_id) self._cached_blocks.pop(content_hash_to_evict) - - self._refcounter.incr(block_id) - self._track_block_id(block_id, computed=False) - - return block_id - - def _free_block_id(self, block: Block) -> None: - """Decrements the refcount of the block. The block may be in two - possible states: (1) immutable/cached or (2) mutable/hashless. - In the first case, the refcount is decremented directly and the block - may be possibly added to the evictor. In other case, hashless - allocator free(..) with keep_block_object=True is called to only free - the block id (since the block object may be reused by the caller) - """ - block_id = block.block_id - assert block_id is not None, "Freeing unallocated block is undefined" - - if block.content_hash is not None: - # Immutable: This type of block is always cached, and we want to - # keep it in the evictor for future reuse - self._decr_refcount_cached_block(block) - else: - # Mutable: This type of block is not cached, so we release it - # directly to the hashless allocator - self._decr_refcount_hashless_block(block) - - assert block.block_id is None - - def free(self, block: Block, keep_block_object: bool = False) -> None: - """Release the block (look at free_block_id(..) docs) - """ - # Release the physical block index - self._free_block_id(block) - - # Release the block object to the pool - if not keep_block_object: - self._block_pool.free_block(block) - - def fork(self, last_block: Block) -> List[Block]: - """Creates a new sequence of blocks that shares the same underlying - memory as the original sequence. - - Args: - last_block (Block): The last block in the original sequence. - - Returns: - List[Block]: The new sequence of blocks that shares the same memory - as the original sequence. - """ - source_blocks = get_all_blocks_recursively(last_block) - - forked_blocks: List[Block] = [] - prev_block = None - for block in source_blocks: - block_id = block.block_id - assert block_id is not None - - refcount = self._refcounter.incr(block_id) - assert refcount != 1, "can't fork free'd block_id = {}".format( - block_id) - + + self._refcounter.incr(block_id) + self._track_block_id(block_id, computed=False) + + return block_id + + def _free_block_id(self, block: Block) -> None: + """Decrements the refcount of the block. The block may be in two + possible states: (1) immutable/cached or (2) mutable/hashless. + In the first case, the refcount is decremented directly and the block + may be possibly added to the evictor. In other case, hashless + allocator free(..) with keep_block_object=True is called to only free + the block id (since the block object may be reused by the caller) + """ + block_id = block.block_id + assert block_id is not None, "Freeing unallocated block is undefined" + + if block.content_hash is not None: + # Immutable: This type of block is always cached, and we want to + # keep it in the evictor for future reuse + self._decr_refcount_cached_block(block) + else: + # Mutable: This type of block is not cached, so we release it + # directly to the hashless allocator + self._decr_refcount_hashless_block(block) + + assert block.block_id is None + + def free(self, block: Block, keep_block_object: bool = False) -> None: + """Release the block (look at free_block_id(..) docs) + """ + # Release the physical block index + self._free_block_id(block) + + # Release the block object to the pool + if not keep_block_object: + self._block_pool.free_block(block) + + def fork(self, last_block: Block) -> List[Block]: + """Creates a new sequence of blocks that shares the same underlying + memory as the original sequence. + + Args: + last_block (Block): The last block in the original sequence. + + Returns: + List[Block]: The new sequence of blocks that shares the same memory + as the original sequence. + """ + source_blocks = get_all_blocks_recursively(last_block) + + forked_blocks: List[Block] = [] + prev_block = None + for block in source_blocks: + block_id = block.block_id + assert block_id is not None + + refcount = self._refcounter.incr(block_id) + assert refcount != 1, "can't fork free'd block_id = {}".format( + block_id) + forked_block = self._init_block( prev_block=prev_block, token_ids=block.token_ids, block_size=self._block_size, physical_block_id=block_id, cache_namespace=block.cache_namespace) - - forked_blocks.append(forked_block) - prev_block = forked_blocks[-1] - - return forked_blocks - - def get_num_free_blocks(self, device: Optional[Device] = None) -> int: - assert device is None - # The number of free blocks is the number of hashless free blocks - # plus the number of blocks evictor could free from its list. - return self._hashless_allocator.get_num_free_blocks( - ) + self.evictor.num_blocks - - def get_num_total_blocks(self) -> int: - return self._hashless_allocator.get_num_total_blocks() - - def get_physical_block_id(self, absolute_id: int) -> int: - """Returns the zero-offset block id on certain block allocator - given the absolute block id. - - Args: - absolute_id (int): The absolute block id for the block - in whole allocator. - - Returns: - int: The rzero-offset block id on certain device. - """ - return sorted(self.all_block_ids).index(absolute_id) - - @property - def all_block_ids(self) -> FrozenSet[int]: - return self._hashless_allocator.all_block_ids - - def get_prefix_cache_hit_rate(self) -> float: - return self.metric_data.get_hit_rate() - - def is_block_cached(self, block: Block) -> bool: - assert block.content_hash is not None - return block.content_hash in self._cached_blocks - - def promote_to_immutable_block(self, block: Block) -> BlockId: - """Once a mutable block is full, it can be promoted to an immutable - block. This means that its content can be referenced by future blocks - having the same prefix. - - Note that if we already have a cached block with the same content, we - will replace the newly-promoted block's mapping with the existing cached - block id. - - Args: - block: The mutable block to be promoted. - - Returns: - BlockId: Either the original block index, or the block index of - the previously cached block matching the same content. - """ - # Ensure block can be promoted - assert block.content_hash is not None - assert block.block_id is not None - assert self._refcounter.get(block.block_id) > 0 - - if block.content_hash not in self._cached_blocks: - # No cached content hash => Set this block as cached. - # Note that this block cannot be marked as computed yet - # because other sequences in the same batch cannot reuse - # this block. - self._cached_blocks[block.content_hash] = block.block_id - # Mark this block as touched so that it can be marked as - # computed after the entire batch of sequences are scheduled. - self._touched_blocks.add(block.block_id) - return block.block_id - - # Reuse the cached content hash - self._decr_refcount_hashless_block(block) - block.block_id = self._cached_blocks[block.content_hash] - - # Increment refcount of the cached block and (possibly) restore - # it from the evictor. - # Note that in this case, the block is marked as computed - self._incr_refcount_cached_block(block) - - return block.block_id - - def cow_block_if_not_appendable(self, block: Block) -> BlockId: - """Performs a copy-on-write operation on the given block if it is not - appendable. - - Args: - block (Block): The block to check for copy-on-write. - - Returns: - BlockId: The block index of the new block if a copy-on-write - operation was performed, or the original block index if - no copy-on-write was necessary. - """ - src_block_id = block.block_id - assert src_block_id is not None - - if self._cow_tracker.is_appendable(block): - return src_block_id - - self._free_block_id(block) - trg_block_id = self._allocate_block_id() - - self._cow_tracker.record_cow(src_block_id, trg_block_id) - - return trg_block_id - - def clear_copy_on_writes(self) -> List[Tuple[BlockId, BlockId]]: - """Returns the copy-on-write source->destination mapping and clears it. - - Returns: - List[Tuple[BlockId, BlockId]]: A list mapping source - block indices to destination block indices. - """ - return self._cow_tracker.clear_cows() - - def mark_blocks_as_accessed(self, block_ids: List[int], - now: float) -> None: - """Mark blocks as accessed, used in prefix caching. - - If the block is added into evictor, we need to update corresponding - info in evictor's metadata. - """ - - for block_id in block_ids: - if self._block_tracker[block_id].active: - self._block_tracker[block_id].last_accessed = now - elif block_id in self.evictor: - self.evictor.update(block_id, now) - else: - raise ValueError( - "Mark block as accessed which is not belonged to GPU") - - def mark_blocks_as_computed(self, block_ids: List[int]) -> None: - # Mark all touched blocks as computed. - for block_id in self._touched_blocks: - self._block_tracker[block_id].computed = True - self._touched_blocks.clear() - - def _track_block_id(self, block_id: Optional[BlockId], - computed: bool) -> None: - assert block_id is not None - self._block_tracker[block_id].enable() - self._block_tracker[block_id].computed = computed - - def _untrack_block_id(self, block_id: Optional[BlockId]) -> None: - assert block_id is not None - self._block_tracker[block_id].disable() - - def block_is_computed(self, block_id: int) -> bool: - if self._block_tracker[block_id].active: - return self._block_tracker[block_id].computed - else: - return block_id in self.evictor - - def get_computed_block_ids(self, - prev_computed_block_ids: List[int], - block_ids: List[int], - skip_last_block_id: bool = True) -> List[int]: - prev_prefix_size = len(prev_computed_block_ids) - cur_size = len(block_ids) - if skip_last_block_id: - cur_size -= 1 - - # Sanity checks - assert cur_size >= 0 - assert prev_prefix_size <= cur_size - + + forked_blocks.append(forked_block) + prev_block = forked_blocks[-1] + + return forked_blocks + + def get_num_free_blocks(self, device: Optional[Device] = None) -> int: + assert device is None + # The number of free blocks is the number of hashless free blocks + # plus the number of blocks evictor could free from its list. + return self._hashless_allocator.get_num_free_blocks( + ) + self.evictor.num_blocks + + def get_num_total_blocks(self) -> int: + return self._hashless_allocator.get_num_total_blocks() + + def get_physical_block_id(self, absolute_id: int) -> int: + """Returns the zero-offset block id on certain block allocator + given the absolute block id. + + Args: + absolute_id (int): The absolute block id for the block + in whole allocator. + + Returns: + int: The rzero-offset block id on certain device. + """ + return sorted(self.all_block_ids).index(absolute_id) + + @property + def all_block_ids(self) -> FrozenSet[int]: + return self._hashless_allocator.all_block_ids + + def get_prefix_cache_hit_rate(self) -> float: + return self.metric_data.get_hit_rate() + + def is_block_cached(self, block: Block) -> bool: + assert block.content_hash is not None + return block.content_hash in self._cached_blocks + + def promote_to_immutable_block(self, block: Block) -> BlockId: + """Once a mutable block is full, it can be promoted to an immutable + block. This means that its content can be referenced by future blocks + having the same prefix. + + Note that if we already have a cached block with the same content, we + will replace the newly-promoted block's mapping with the existing cached + block id. + + Args: + block: The mutable block to be promoted. + + Returns: + BlockId: Either the original block index, or the block index of + the previously cached block matching the same content. + """ + # Ensure block can be promoted + assert block.content_hash is not None + assert block.block_id is not None + assert self._refcounter.get(block.block_id) > 0 + + if block.content_hash not in self._cached_blocks: + # No cached content hash => Set this block as cached. + # Note that this block cannot be marked as computed yet + # because other sequences in the same batch cannot reuse + # this block. + self._cached_blocks[block.content_hash] = block.block_id + # Mark this block as touched so that it can be marked as + # computed after the entire batch of sequences are scheduled. + self._touched_blocks.add(block.block_id) + return block.block_id + + # Reuse the cached content hash + self._decr_refcount_hashless_block(block) + block.block_id = self._cached_blocks[block.content_hash] + + # Increment refcount of the cached block and (possibly) restore + # it from the evictor. + # Note that in this case, the block is marked as computed + self._incr_refcount_cached_block(block) + + return block.block_id + + def cow_block_if_not_appendable(self, block: Block) -> BlockId: + """Performs a copy-on-write operation on the given block if it is not + appendable. + + Args: + block (Block): The block to check for copy-on-write. + + Returns: + BlockId: The block index of the new block if a copy-on-write + operation was performed, or the original block index if + no copy-on-write was necessary. + """ + src_block_id = block.block_id + assert src_block_id is not None + + if self._cow_tracker.is_appendable(block): + return src_block_id + + self._free_block_id(block) + trg_block_id = self._allocate_block_id() + + self._cow_tracker.record_cow(src_block_id, trg_block_id) + + return trg_block_id + + def clear_copy_on_writes(self) -> List[Tuple[BlockId, BlockId]]: + """Returns the copy-on-write source->destination mapping and clears it. + + Returns: + List[Tuple[BlockId, BlockId]]: A list mapping source + block indices to destination block indices. + """ + return self._cow_tracker.clear_cows() + + def mark_blocks_as_accessed(self, block_ids: List[int], + now: float) -> None: + """Mark blocks as accessed, used in prefix caching. + + If the block is added into evictor, we need to update corresponding + info in evictor's metadata. + """ + + for block_id in block_ids: + if self._block_tracker[block_id].active: + self._block_tracker[block_id].last_accessed = now + elif block_id in self.evictor: + self.evictor.update(block_id, now) + else: + raise ValueError( + "Mark block as accessed which is not belonged to GPU") + + def mark_blocks_as_computed(self, block_ids: List[int]) -> None: + # Mark all touched blocks as computed. + for block_id in self._touched_blocks: + self._block_tracker[block_id].computed = True + self._touched_blocks.clear() + + def _track_block_id(self, block_id: Optional[BlockId], + computed: bool) -> None: + assert block_id is not None + self._block_tracker[block_id].enable() + self._block_tracker[block_id].computed = computed + + def _untrack_block_id(self, block_id: Optional[BlockId]) -> None: + assert block_id is not None + self._block_tracker[block_id].disable() + + def block_is_computed(self, block_id: int) -> bool: + if self._block_tracker[block_id].active: + return self._block_tracker[block_id].computed + else: + return block_id in self.evictor + + def get_computed_block_ids(self, + prev_computed_block_ids: List[int], + block_ids: List[int], + skip_last_block_id: bool = True) -> List[int]: + prev_prefix_size = len(prev_computed_block_ids) + cur_size = len(block_ids) + if skip_last_block_id: + cur_size -= 1 + + # Sanity checks + assert cur_size >= 0 + assert prev_prefix_size <= cur_size + ret = prev_computed_block_ids for i in range(prev_prefix_size, cur_size): block_id = block_ids[i] @@ -736,73 +736,73 @@ class PrefixCachingBlockAllocator(BlockAllocator): break ret.append(block_id) return ret - - def get_common_computed_block_ids( - self, computed_seq_block_ids: List[List[int]]) -> List[int]: - """Return the block ids that are common for a given sequence group. - - Only those blocks that are immutable and already be marked - compyted would be taken consideration. - """ - - # NOTE We exclude the last block to avoid the case where the entire - # prompt is cached. This would cause erroneous behavior in model - # runner. - - # It returns a list of int although type annotation says list of string. - if len(computed_seq_block_ids) == 1: - return computed_seq_block_ids[0] - - return commonprefix([ - ids for ids in computed_seq_block_ids # type: ignore - if ids - ]) - - def get_num_full_blocks_touched(self, blocks: List[Block]) -> int: - """Returns the number of full blocks that will be touched by - swapping in/out. - - Args: - blocks: List of blocks to be swapped. - Returns: - int: the number of full blocks that will be touched by - swapping in/out the given blocks. Non full blocks are ignored - when deciding the number of blocks to touch. - """ - num_touched_blocks: int = 0 - for block in blocks: - # If the block has a match in the cache and the cached - # block is not referenced, then we still count it as a - # touched block - if block.is_full and (not self.is_block_cached(block) or \ - (block.content_hash is not None and \ - self._cached_blocks[block.content_hash] in \ - self.evictor)): - num_touched_blocks += 1 - return num_touched_blocks - - def swap_out(self, blocks: List[Block]) -> None: - """Execute the swap out actions. Basically just free the - given blocks. - - Args: - blocks: List of blocks to be swapped out. - """ - for block in blocks: - self._free_block_id(block) - - def swap_in(self, blocks: List[Block]) -> None: - """Execute the swap in actions. Change the block id from - old allocator to current allocator for each block to finish - the block table update. - - Args: - blocks: List of blocks to be swapped in. - """ - for block in blocks: - # Here we allocate either immutable or mutable block and then - # extract its block_id. Note that the block object is released - # and the block_id is assigned to "block" to allow reusing the + + def get_common_computed_block_ids( + self, computed_seq_block_ids: List[List[int]]) -> List[int]: + """Return the block ids that are common for a given sequence group. + + Only those blocks that are immutable and already be marked + compyted would be taken consideration. + """ + + # NOTE We exclude the last block to avoid the case where the entire + # prompt is cached. This would cause erroneous behavior in model + # runner. + + # It returns a list of int although type annotation says list of string. + if len(computed_seq_block_ids) == 1: + return computed_seq_block_ids[0] + + return commonprefix([ + ids for ids in computed_seq_block_ids # type: ignore + if ids + ]) + + def get_num_full_blocks_touched(self, blocks: List[Block]) -> int: + """Returns the number of full blocks that will be touched by + swapping in/out. + + Args: + blocks: List of blocks to be swapped. + Returns: + int: the number of full blocks that will be touched by + swapping in/out the given blocks. Non full blocks are ignored + when deciding the number of blocks to touch. + """ + num_touched_blocks: int = 0 + for block in blocks: + # If the block has a match in the cache and the cached + # block is not referenced, then we still count it as a + # touched block + if block.is_full and (not self.is_block_cached(block) or \ + (block.content_hash is not None and \ + self._cached_blocks[block.content_hash] in \ + self.evictor)): + num_touched_blocks += 1 + return num_touched_blocks + + def swap_out(self, blocks: List[Block]) -> None: + """Execute the swap out actions. Basically just free the + given blocks. + + Args: + blocks: List of blocks to be swapped out. + """ + for block in blocks: + self._free_block_id(block) + + def swap_in(self, blocks: List[Block]) -> None: + """Execute the swap in actions. Change the block id from + old allocator to current allocator for each block to finish + the block table update. + + Args: + blocks: List of blocks to be swapped in. + """ + for block in blocks: + # Here we allocate either immutable or mutable block and then + # extract its block_id. Note that the block object is released + # and the block_id is assigned to "block" to allow reusing the # existing "block" object if block.is_full: tmp_block = ( @@ -816,33 +816,33 @@ class PrefixCachingBlockAllocator(BlockAllocator): prev_block=block.prev_block, cache_namespace=block.cache_namespace)) tmp_block.append_token_ids(block.token_ids) - - block_id = tmp_block.block_id - self._block_pool.free_block(tmp_block) - - block.block_id = block_id # Assign block_id - - -class PrefixCachingBlock(Block): - """A block implementation that supports prefix caching. - - The PrefixCachingBlock class represents a block of token IDs with prefix - caching capabilities. It wraps a NaiveBlock internally and provides - additional functionality for content hashing and promoting immutable blocks - with the prefix caching allocator. - - Args: - prev_block (Optional[PrefixCachingBlock]): The previous block in the - sequence. - token_ids (List[int]): The initial token IDs to be stored in the block. - block_size (int): The maximum number of token IDs that can be stored in - the block. - allocator (BlockAllocator): The prefix - caching block allocator associated with this block. - block_id (Optional[int], optional): The physical block index - of this block. Defaults to None. - """ - + + block_id = tmp_block.block_id + self._block_pool.free_block(tmp_block) + + block.block_id = block_id # Assign block_id + + +class PrefixCachingBlock(Block): + """A block implementation that supports prefix caching. + + The PrefixCachingBlock class represents a block of token IDs with prefix + caching capabilities. It wraps a NaiveBlock internally and provides + additional functionality for content hashing and promoting immutable blocks + with the prefix caching allocator. + + Args: + prev_block (Optional[PrefixCachingBlock]): The previous block in the + sequence. + token_ids (List[int]): The initial token IDs to be stored in the block. + block_size (int): The maximum number of token IDs that can be stored in + the block. + allocator (BlockAllocator): The prefix + caching block allocator associated with this block. + block_id (Optional[int], optional): The physical block index + of this block. Defaults to None. + """ + def __init__( self, prev_block: Optional[Block], @@ -853,12 +853,12 @@ class PrefixCachingBlock(Block): computed: bool = False, cache_namespace: Optional[bytes] = None, ): - assert isinstance(allocator, PrefixCachingBlockAllocator), ( - "Currently this class is only tested with " - "PrefixCachingBlockAllocator. Got instead allocator = {}".format( - allocator)) - assert_prefix_caching_block_or_none(prev_block) - + assert isinstance(allocator, PrefixCachingBlockAllocator), ( + "Currently this class is only tested with " + "PrefixCachingBlockAllocator. Got instead allocator = {}".format( + allocator)) + assert_prefix_caching_block_or_none(prev_block) + self._prev_block = prev_block self._cached_content_hash: Optional[bytes] = None self._cache_namespace = cache_namespace or b"" @@ -868,16 +868,16 @@ class PrefixCachingBlock(Block): self._allocator = allocator self._last_accessed: float = _DEFAULT_LAST_ACCESSED_TIME self._computed = computed - - # On the first time, we create the block object, and next we only - # reinitialize it - if hasattr(self, "_block"): - self._block.__init__( # type: ignore[has-type] - prev_block=prev_block, - token_ids=token_ids, - block_size=block_size, - block_id=block_id, - allocator=self._allocator) + + # On the first time, we create the block object, and next we only + # reinitialize it + if hasattr(self, "_block"): + self._block.__init__( # type: ignore[has-type] + prev_block=prev_block, + token_ids=token_ids, + block_size=block_size, + block_id=block_id, + allocator=self._allocator) else: self._block = NaiveBlock(prev_block=prev_block, token_ids=token_ids, @@ -886,94 +886,94 @@ class PrefixCachingBlock(Block): allocator=self._allocator) self._update_num_tokens_total() - - def _update_num_tokens_total(self): - """Incrementally computes the number of tokens that there is - till the current block (included) - """ - res = 0 - - # Add all previous blocks - if self._prev_block is not None: - res += self._prev_block.num_tokens_total - - # Add current block - res += len(self.token_ids) - - self._cached_num_tokens_total = res - - @property - def computed(self) -> bool: - return self._computed - - @computed.setter - def computed(self, value) -> None: - self._computed = value - - @property - def last_accessed(self) -> float: - return self._last_accessed - - @last_accessed.setter - def last_accessed(self, last_accessed_ts: float): - self._last_accessed = last_accessed_ts - - def append_token_ids(self, token_ids: List[int]) -> None: - """Appends the given token IDs to the block and registers the block as - immutable if the block becomes full. - - Args: - token_ids (List[int]): The token IDs to be appended to the block. - """ - # Ensure this is mutable block (not promoted) - assert self.content_hash is None - assert not self.computed - - if len(token_ids) == 0: - return - - # Ensure there are input tokens - assert token_ids, "Got token_ids = {}".format(token_ids) - - # Naive block handles CoW. - self._block.append_token_ids(token_ids) - self._update_num_tokens_total() - - # If the content hash is present, then the block can be made immutable. - # Register ourselves with the allocator, potentially replacing the - # physical block index. - if self.content_hash is not None: - self.block_id = self._allocator.promote_to_immutable_block(self) - - @property - def block_id(self) -> Optional[int]: - return self._block.block_id - - @block_id.setter - def block_id(self, value) -> None: - self._block.block_id = value - - @property - def is_full(self) -> bool: - return self._block.is_full - - @property - def num_empty_slots(self) -> int: - return self._block.num_empty_slots - - @property - def num_tokens_total(self) -> int: - return self._cached_num_tokens_total - - @property - def block_size(self) -> int: - return self._block.block_size - - @property - def token_ids(self) -> List[int]: - return self._block.token_ids - - @property + + def _update_num_tokens_total(self): + """Incrementally computes the number of tokens that there is + till the current block (included) + """ + res = 0 + + # Add all previous blocks + if self._prev_block is not None: + res += self._prev_block.num_tokens_total + + # Add current block + res += len(self.token_ids) + + self._cached_num_tokens_total = res + + @property + def computed(self) -> bool: + return self._computed + + @computed.setter + def computed(self, value) -> None: + self._computed = value + + @property + def last_accessed(self) -> float: + return self._last_accessed + + @last_accessed.setter + def last_accessed(self, last_accessed_ts: float): + self._last_accessed = last_accessed_ts + + def append_token_ids(self, token_ids: List[int]) -> None: + """Appends the given token IDs to the block and registers the block as + immutable if the block becomes full. + + Args: + token_ids (List[int]): The token IDs to be appended to the block. + """ + # Ensure this is mutable block (not promoted) + assert self.content_hash is None + assert not self.computed + + if len(token_ids) == 0: + return + + # Ensure there are input tokens + assert token_ids, "Got token_ids = {}".format(token_ids) + + # Naive block handles CoW. + self._block.append_token_ids(token_ids) + self._update_num_tokens_total() + + # If the content hash is present, then the block can be made immutable. + # Register ourselves with the allocator, potentially replacing the + # physical block index. + if self.content_hash is not None: + self.block_id = self._allocator.promote_to_immutable_block(self) + + @property + def block_id(self) -> Optional[int]: + return self._block.block_id + + @block_id.setter + def block_id(self, value) -> None: + self._block.block_id = value + + @property + def is_full(self) -> bool: + return self._block.is_full + + @property + def num_empty_slots(self) -> int: + return self._block.num_empty_slots + + @property + def num_tokens_total(self) -> int: + return self._cached_num_tokens_total + + @property + def block_size(self) -> int: + return self._block.block_size + + @property + def token_ids(self) -> List[int]: + return self._block.token_ids + + @property def prev_block(self) -> Optional[Block]: return self._prev_block @@ -983,31 +983,31 @@ class PrefixCachingBlock(Block): @property def content_hash(self) -> Optional[bytes]: - """Return the content-based hash of the current block, or None if it is - not yet defined. - - For the content-based hash to be defined, the current block must be - full. - """ - # If the hash is already computed, return it. - if self._cached_content_hash is not None: - return self._cached_content_hash - - # We cannot compute a hash for the current block because it is not full. - if not self.is_full: - return None - + """Return the content-based hash of the current block, or None if it is + not yet defined. + + For the content-based hash to be defined, the current block must be + full. + """ + # If the hash is already computed, return it. + if self._cached_content_hash is not None: + return self._cached_content_hash + + # We cannot compute a hash for the current block because it is not full. + if not self.is_full: + return None + is_first_block = self._prev_block is None prev_block_hash = ( None if is_first_block else self._prev_block.content_hash # type: ignore ) - - # Previous block exists but does not yet have a hash. - # Return no hash in this case. - if prev_block_hash is None and not is_first_block: - return None - + + # Previous block exists but does not yet have a hash. + # Return no hash in this case. + if prev_block_hash is None and not is_first_block: + return None + self._cached_content_hash = PrefixCachingBlock.hash_block_tokens( is_first_block, prev_block_hash, @@ -1022,21 +1022,21 @@ class PrefixCachingBlock(Block): cur_block_token_ids: List[int], cache_namespace: Optional[bytes] = None, ) -> bytes: - """Computes a hash value corresponding to the contents of a block and - the contents of the preceding block(s). The hash value is used for - prefix caching. - - NOTE: Content-based hashing does not yet support LoRA. - - Parameters: - - is_first_block (bool): A flag indicating if the block is the first in - the sequence. - - prev_block_hash (Optional[int]): The hash of the previous block. None - if this is the first block. - - cur_block_token_ids (List[int]): A list of token ids in the current - block. The current block is assumed to be full. - - Returns: + """Computes a hash value corresponding to the contents of a block and + the contents of the preceding block(s). The hash value is used for + prefix caching. + + NOTE: Content-based hashing does not yet support LoRA. + + Parameters: + - is_first_block (bool): A flag indicating if the block is the first in + the sequence. + - prev_block_hash (Optional[int]): The hash of the previous block. None + if this is the first block. + - cur_block_token_ids (List[int]): A list of token ids in the current + block. The current block is assumed to be full. + + Returns: - bytes: The computed hash value for the block. """ assert (prev_block_hash is None) == is_first_block @@ -1053,131 +1053,131 @@ class PrefixCachingBlock(Block): f"!{len(cur_block_token_ids)}q", *(int(token_id) for token_id in cur_block_token_ids))) return digest.digest() - - -class ComputedBlocksTracker: - """Handles caching of per-sequence computed block ids. - When a sequence appears for the first time, it traverses all of the - blocks and detects the prefix of blocks that is computed. On the - subsequent times, it only traverses the new blocks that were added - and updates the already recorded prefix of blocks with the newly - computed blocks. - - To avoid redundant traversals, the algorithm also detects when there - is a "gap" in the computed prefix. For example, if we have blocks = - [1,2,3,4,5], and we have detected [1,2,3] as the computed prefix, then - we won't try to add more computed blocks to [1,2,3] in this sequence - iteration, and will add more computed blocks only after the sequence is - freed and reused again. - - Note that currently, for a given sequence, we also skip the last - block id for caching purposes, to avoid caching of a full sequence - """ - - def __init__(self, allocator): - self._allocator = allocator - self._cached_computed_seq_blocks: Dict[int, Tuple[List[int], - bool]] = {} - - def add_seq(self, seq_id: int) -> None: - """Start tracking seq_id - """ - assert seq_id not in self._cached_computed_seq_blocks - self._cached_computed_seq_blocks[seq_id] = ([], False) - - def remove_seq(self, seq_id: int) -> None: - """Stop tracking seq_id - """ - assert seq_id in self._cached_computed_seq_blocks - del self._cached_computed_seq_blocks[seq_id] - - def get_cached_computed_blocks_and_update( - self, seq_id: int, block_ids: List[int]) -> List[int]: - """ Look at the class documentation for details - """ - # Ensure seq_id is already tracked - assert seq_id in self._cached_computed_seq_blocks - - # Get cached data (may be empty on the first time) - prev_computed_block_ids, has_gap = self._cached_computed_seq_blocks[ - seq_id] - - if has_gap: - # When gap is detected, we do not add more computed blocks at this - # sequence iteration - return prev_computed_block_ids - - # We do not consider the last block id for caching purposes. - num_cur_blocks = len(block_ids) - 1 - assert num_cur_blocks >= 0 - - if len(prev_computed_block_ids) >= num_cur_blocks: - # Cache HIT - assert len(prev_computed_block_ids) == num_cur_blocks - return prev_computed_block_ids - - # If here, then we may possibly add more computed blocks. As a result, - # traverse the additional blocks after prev_computed_block_ids to - # detect more computed blocks and add them. - - # Incremental init for seq_id => Look only at the new blocks - computed_block_ids = self._allocator.get_computed_block_ids( # noqa: E501 - prev_computed_block_ids, - block_ids, - skip_last_block_id= - True, # We skip last block id to avoid caching of full seq - ) - - # Detect if there is a "gap" - has_gap = len(computed_block_ids) < num_cur_blocks - - # Record - self._cached_computed_seq_blocks[seq_id] = (computed_block_ids, - has_gap) - - return computed_block_ids - - -class LastAccessBlocksTracker: - """Manages the last access time of the tracked sequences, in order to allow - an efficient update of allocator's block last access times - """ - - def __init__(self, allocator): - self._allocator = allocator - self._seq_last_access: Dict[int, Optional[float]] = {} - - def add_seq(self, seq_id: int) -> None: - """Start tracking seq_id - """ - assert seq_id not in self._seq_last_access - self._seq_last_access[seq_id] = None - - def remove_seq(self, seq_id: int) -> None: - """Stop tracking seq_id - """ - assert seq_id in self._seq_last_access - del self._seq_last_access[seq_id] - - def update_last_access(self, seq_id: int, time: float) -> None: - assert seq_id in self._seq_last_access - self._seq_last_access[seq_id] = time - - def update_seq_blocks_last_access(self, seq_id: int, - block_ids: List[int]) -> None: - assert seq_id in self._seq_last_access - - ts = self._seq_last_access[seq_id] - - if ts is None: - # No last access was recorded, no need to update. - return - - self._allocator.mark_blocks_as_accessed(block_ids, ts) - - -def assert_prefix_caching_block_or_none(block: Optional[Block]): - if block is None: - return - assert isinstance(block, - PrefixCachingBlock), "Got block = {}".format(block) + + +class ComputedBlocksTracker: + """Handles caching of per-sequence computed block ids. + When a sequence appears for the first time, it traverses all of the + blocks and detects the prefix of blocks that is computed. On the + subsequent times, it only traverses the new blocks that were added + and updates the already recorded prefix of blocks with the newly + computed blocks. + + To avoid redundant traversals, the algorithm also detects when there + is a "gap" in the computed prefix. For example, if we have blocks = + [1,2,3,4,5], and we have detected [1,2,3] as the computed prefix, then + we won't try to add more computed blocks to [1,2,3] in this sequence + iteration, and will add more computed blocks only after the sequence is + freed and reused again. + + Note that currently, for a given sequence, we also skip the last + block id for caching purposes, to avoid caching of a full sequence + """ + + def __init__(self, allocator): + self._allocator = allocator + self._cached_computed_seq_blocks: Dict[int, Tuple[List[int], + bool]] = {} + + def add_seq(self, seq_id: int) -> None: + """Start tracking seq_id + """ + assert seq_id not in self._cached_computed_seq_blocks + self._cached_computed_seq_blocks[seq_id] = ([], False) + + def remove_seq(self, seq_id: int) -> None: + """Stop tracking seq_id + """ + assert seq_id in self._cached_computed_seq_blocks + del self._cached_computed_seq_blocks[seq_id] + + def get_cached_computed_blocks_and_update( + self, seq_id: int, block_ids: List[int]) -> List[int]: + """ Look at the class documentation for details + """ + # Ensure seq_id is already tracked + assert seq_id in self._cached_computed_seq_blocks + + # Get cached data (may be empty on the first time) + prev_computed_block_ids, has_gap = self._cached_computed_seq_blocks[ + seq_id] + + if has_gap: + # When gap is detected, we do not add more computed blocks at this + # sequence iteration + return prev_computed_block_ids + + # We do not consider the last block id for caching purposes. + num_cur_blocks = len(block_ids) - 1 + assert num_cur_blocks >= 0 + + if len(prev_computed_block_ids) >= num_cur_blocks: + # Cache HIT + assert len(prev_computed_block_ids) == num_cur_blocks + return prev_computed_block_ids + + # If here, then we may possibly add more computed blocks. As a result, + # traverse the additional blocks after prev_computed_block_ids to + # detect more computed blocks and add them. + + # Incremental init for seq_id => Look only at the new blocks + computed_block_ids = self._allocator.get_computed_block_ids( # noqa: E501 + prev_computed_block_ids, + block_ids, + skip_last_block_id= + True, # We skip last block id to avoid caching of full seq + ) + + # Detect if there is a "gap" + has_gap = len(computed_block_ids) < num_cur_blocks + + # Record + self._cached_computed_seq_blocks[seq_id] = (computed_block_ids, + has_gap) + + return computed_block_ids + + +class LastAccessBlocksTracker: + """Manages the last access time of the tracked sequences, in order to allow + an efficient update of allocator's block last access times + """ + + def __init__(self, allocator): + self._allocator = allocator + self._seq_last_access: Dict[int, Optional[float]] = {} + + def add_seq(self, seq_id: int) -> None: + """Start tracking seq_id + """ + assert seq_id not in self._seq_last_access + self._seq_last_access[seq_id] = None + + def remove_seq(self, seq_id: int) -> None: + """Stop tracking seq_id + """ + assert seq_id in self._seq_last_access + del self._seq_last_access[seq_id] + + def update_last_access(self, seq_id: int, time: float) -> None: + assert seq_id in self._seq_last_access + self._seq_last_access[seq_id] = time + + def update_seq_blocks_last_access(self, seq_id: int, + block_ids: List[int]) -> None: + assert seq_id in self._seq_last_access + + ts = self._seq_last_access[seq_id] + + if ts is None: + # No last access was recorded, no need to update. + return + + self._allocator.mark_blocks_as_accessed(block_ids, ts) + + +def assert_prefix_caching_block_or_none(block: Optional[Block]): + if block is None: + return + assert isinstance(block, + PrefixCachingBlock), "Got block = {}".format(block) diff --git a/vllm_overrides/core/block_manager_v2.py b/vllm_overrides/core/block_manager_v2.py index 321a60d2..5f9254f0 100644 --- a/vllm_overrides/core/block_manager_v2.py +++ b/vllm_overrides/core/block_manager_v2.py @@ -32,86 +32,86 @@ logger = init_logger(__name__) class BlockSpaceManagerV2(BlockSpaceManager): - """BlockSpaceManager which manages the allocation of KV cache. - - It owns responsibility for allocation, swapping, allocating memory for - autoregressively-generated tokens, and other advanced features such as - prefix caching, forking/copy-on-write, and sliding-window memory allocation. - - This class implements the design described in - https://github.com/vllm-project/vllm/pull/3492. - - Lookahead slots - The block manager has the notion of a "lookahead slot". These are slots - in the KV cache that are allocated for a sequence. Unlike the other - allocated slots, the content of these slots is undefined -- the worker - may use the memory allocations in any way. - - In practice, a worker could use these lookahead slots to run multiple - forward passes for a single scheduler invocation. Each successive - forward pass would write KV activations to the corresponding lookahead - slot. This allows low inter-token latency use-cases, where the overhead - of continuous batching scheduling is amortized over >1 generated tokens. - - Speculative decoding uses lookahead slots to store KV activations of - proposal tokens. - - See https://github.com/vllm-project/vllm/pull/3250 for more information - on lookahead scheduling. - - Args: - block_size (int): The size of each memory block. - num_gpu_blocks (int): The number of memory blocks allocated on GPU. - num_cpu_blocks (int): The number of memory blocks allocated on CPU. - watermark (float, optional): The threshold used for memory swapping. - Defaults to 0.01. - sliding_window (Optional[int], optional): The size of the sliding - window. Defaults to None. - enable_caching (bool, optional): Flag indicating whether caching is - enabled. Defaults to False. - """ - + """BlockSpaceManager which manages the allocation of KV cache. + + It owns responsibility for allocation, swapping, allocating memory for + autoregressively-generated tokens, and other advanced features such as + prefix caching, forking/copy-on-write, and sliding-window memory allocation. + + This class implements the design described in + https://github.com/vllm-project/vllm/pull/3492. + + Lookahead slots + The block manager has the notion of a "lookahead slot". These are slots + in the KV cache that are allocated for a sequence. Unlike the other + allocated slots, the content of these slots is undefined -- the worker + may use the memory allocations in any way. + + In practice, a worker could use these lookahead slots to run multiple + forward passes for a single scheduler invocation. Each successive + forward pass would write KV activations to the corresponding lookahead + slot. This allows low inter-token latency use-cases, where the overhead + of continuous batching scheduling is amortized over >1 generated tokens. + + Speculative decoding uses lookahead slots to store KV activations of + proposal tokens. + + See https://github.com/vllm-project/vllm/pull/3250 for more information + on lookahead scheduling. + + Args: + block_size (int): The size of each memory block. + num_gpu_blocks (int): The number of memory blocks allocated on GPU. + num_cpu_blocks (int): The number of memory blocks allocated on CPU. + watermark (float, optional): The threshold used for memory swapping. + Defaults to 0.01. + sliding_window (Optional[int], optional): The size of the sliding + window. Defaults to None. + enable_caching (bool, optional): Flag indicating whether caching is + enabled. Defaults to False. + """ + def __init__( self, block_size: int, num_gpu_blocks: int, num_cpu_blocks: int, - watermark: float = 0.01, - sliding_window: Optional[int] = None, - enable_caching: bool = False, - ) -> None: - self.block_size = block_size - self.num_total_gpu_blocks = num_gpu_blocks - self.num_total_cpu_blocks = num_cpu_blocks - - self.sliding_window = sliding_window - # max_block_sliding_window is the max number of blocks that need to be - # allocated - self.max_block_sliding_window = None - if sliding_window is not None: - # +1 here because // rounds down - num_blocks = sliding_window // block_size + 1 - # +1 here because the last block may not be full, - # and so the sequence stretches one more block at the beginning - # For example, if sliding_window is 3 and block_size is 4, - # we may need 2 blocks when the second block only holds 1 token. - self.max_block_sliding_window = num_blocks + 1 - - self.watermark = watermark - assert watermark >= 0.0 - - self.enable_caching = enable_caching - - self.watermark_blocks = int(watermark * num_gpu_blocks) - - self.block_allocator = CpuGpuBlockAllocator.create( - allocator_type="prefix_caching" if enable_caching else "naive", - num_gpu_blocks=num_gpu_blocks, - num_cpu_blocks=num_cpu_blocks, - block_size=block_size, - ) - - self.block_tables: Dict[SeqId, BlockTable] = {} + watermark: float = 0.01, + sliding_window: Optional[int] = None, + enable_caching: bool = False, + ) -> None: + self.block_size = block_size + self.num_total_gpu_blocks = num_gpu_blocks + self.num_total_cpu_blocks = num_cpu_blocks + + self.sliding_window = sliding_window + # max_block_sliding_window is the max number of blocks that need to be + # allocated + self.max_block_sliding_window = None + if sliding_window is not None: + # +1 here because // rounds down + num_blocks = sliding_window // block_size + 1 + # +1 here because the last block may not be full, + # and so the sequence stretches one more block at the beginning + # For example, if sliding_window is 3 and block_size is 4, + # we may need 2 blocks when the second block only holds 1 token. + self.max_block_sliding_window = num_blocks + 1 + + self.watermark = watermark + assert watermark >= 0.0 + + self.enable_caching = enable_caching + + self.watermark_blocks = int(watermark * num_gpu_blocks) + + self.block_allocator = CpuGpuBlockAllocator.create( + allocator_type="prefix_caching" if enable_caching else "naive", + num_gpu_blocks=num_gpu_blocks, + num_cpu_blocks=num_cpu_blocks, + block_size=block_size, + ) + + self.block_tables: Dict[SeqId, BlockTable] = {} self.cross_block_tables: Dict[EncoderSeqId, BlockTable] = {} self._warned_mm_namespace_requests = set[str]() self._request_local_namespace: Dict[str, bytes] = {} @@ -121,46 +121,46 @@ class BlockSpaceManagerV2(BlockSpaceManager): self.block_allocator) self._last_access_blocks_tracker = LastAccessBlocksTracker( self.block_allocator) - - def can_allocate(self, - seq_group: SequenceGroup, - num_lookahead_slots: int = 0) -> AllocStatus: - # FIXME(woosuk): Here we assume that all sequences in the group share - # the same prompt. This may not be true for preempted sequences. - - check_no_caching_or_swa_for_blockmgr_encdec(self, seq_group) - - seq = seq_group.get_seqs(status=SequenceStatus.WAITING)[0] - num_required_blocks = BlockTable.get_num_required_blocks( - seq.get_token_ids(), - block_size=self.block_size, - num_lookahead_slots=num_lookahead_slots, - ) - - if seq_group.is_encoder_decoder(): - encoder_seq = seq_group.get_encoder_seq() - assert encoder_seq is not None - num_required_blocks += BlockTable.get_num_required_blocks( - encoder_seq.get_token_ids(), - block_size=self.block_size, - ) - - if self.max_block_sliding_window is not None: - num_required_blocks = min(num_required_blocks, - self.max_block_sliding_window) - - num_free_gpu_blocks = self.block_allocator.get_num_free_blocks( - device=Device.GPU) - - # Use watermark to avoid frequent cache eviction. - if (self.num_total_gpu_blocks - num_required_blocks < - self.watermark_blocks): - return AllocStatus.NEVER - if num_free_gpu_blocks - num_required_blocks >= self.watermark_blocks: - return AllocStatus.OK - else: - return AllocStatus.LATER - + + def can_allocate(self, + seq_group: SequenceGroup, + num_lookahead_slots: int = 0) -> AllocStatus: + # FIXME(woosuk): Here we assume that all sequences in the group share + # the same prompt. This may not be true for preempted sequences. + + check_no_caching_or_swa_for_blockmgr_encdec(self, seq_group) + + seq = seq_group.get_seqs(status=SequenceStatus.WAITING)[0] + num_required_blocks = BlockTable.get_num_required_blocks( + seq.get_token_ids(), + block_size=self.block_size, + num_lookahead_slots=num_lookahead_slots, + ) + + if seq_group.is_encoder_decoder(): + encoder_seq = seq_group.get_encoder_seq() + assert encoder_seq is not None + num_required_blocks += BlockTable.get_num_required_blocks( + encoder_seq.get_token_ids(), + block_size=self.block_size, + ) + + if self.max_block_sliding_window is not None: + num_required_blocks = min(num_required_blocks, + self.max_block_sliding_window) + + num_free_gpu_blocks = self.block_allocator.get_num_free_blocks( + device=Device.GPU) + + # Use watermark to avoid frequent cache eviction. + if (self.num_total_gpu_blocks - num_required_blocks < + self.watermark_blocks): + return AllocStatus.NEVER + if num_free_gpu_blocks - num_required_blocks >= self.watermark_blocks: + return AllocStatus.OK + else: + return AllocStatus.LATER + def _allocate_sequence( self, seq: Sequence, @@ -181,10 +181,10 @@ class BlockSpaceManagerV2(BlockSpaceManager): def allocate(self, seq_group: SequenceGroup) -> None: # Allocate self-attention block tables for decoder sequences - waiting_seqs = seq_group.get_seqs(status=SequenceStatus.WAITING) - assert not (set(seq.seq_id for seq in waiting_seqs) - & self.block_tables.keys()), "block table already exists" - + waiting_seqs = seq_group.get_seqs(status=SequenceStatus.WAITING) + assert not (set(seq.seq_id for seq in waiting_seqs) + & self.block_tables.keys()), "block table already exists" + # NOTE: Here we assume that all sequences in the group have the same # prompt. seq = waiting_seqs[0] @@ -199,31 +199,31 @@ class BlockSpaceManagerV2(BlockSpaceManager): cache_namespace=cache_namespace, ) self.block_tables[seq.seq_id] = block_table - - # Track seq + + # Track seq self._computed_blocks_tracker.add_seq(seq.seq_id) self._last_access_blocks_tracker.add_seq(seq.seq_id) - - # Assign the block table for each sequence. - for seq in waiting_seqs[1:]: - self.block_tables[seq.seq_id] = block_table.fork() - - # Track seq - self._computed_blocks_tracker.add_seq(seq.seq_id) - self._last_access_blocks_tracker.add_seq(seq.seq_id) - - # Allocate cross-attention block table for encoder sequence - # - # NOTE: Here we assume that all sequences in the group have the same - # encoder prompt. + + # Assign the block table for each sequence. + for seq in waiting_seqs[1:]: + self.block_tables[seq.seq_id] = block_table.fork() + + # Track seq + self._computed_blocks_tracker.add_seq(seq.seq_id) + self._last_access_blocks_tracker.add_seq(seq.seq_id) + + # Allocate cross-attention block table for encoder sequence + # + # NOTE: Here we assume that all sequences in the group have the same + # encoder prompt. request_id = seq_group.request_id assert (request_id not in self.cross_block_tables), \ "block table already exists" - - check_no_caching_or_swa_for_blockmgr_encdec(self, seq_group) - + + check_no_caching_or_swa_for_blockmgr_encdec(self, seq_group) + if seq_group.is_encoder_decoder(): encoder_seq = seq_group.get_encoder_seq() assert encoder_seq is not None @@ -279,129 +279,129 @@ class BlockSpaceManagerV2(BlockSpaceManager): else: digest.update(b"text|") return digest.digest() - - def can_append_slots(self, seq_group: SequenceGroup, - num_lookahead_slots: int) -> bool: - """Determine if there is enough space in the GPU KV cache to continue - generation of the specified sequence group. - - We use a worst-case heuristic: assume each touched block will require a - new allocation (either via CoW or new block). We can append slots if the - number of touched blocks is less than the number of free blocks. - - "Lookahead slots" are slots that are allocated in addition to the slots - for known tokens. The contents of the lookahead slots are not defined. - This is used by speculative decoding when speculating future tokens. - """ - - num_touched_blocks = 0 - for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING): - block_table = self.block_tables[seq.seq_id] - - num_touched_blocks += ( - block_table.get_num_blocks_touched_by_append_slots( - token_ids=block_table.get_unseen_token_ids( - seq.get_token_ids()), - num_lookahead_slots=num_lookahead_slots, - )) - - num_free_gpu_blocks = self.block_allocator.get_num_free_blocks( - Device.GPU) - return num_touched_blocks <= num_free_gpu_blocks - - def append_slots( - self, - seq: Sequence, - num_lookahead_slots: int, - ) -> List[Tuple[int, int]]: - - block_table = self.block_tables[seq.seq_id] - - block_table.append_token_ids( - token_ids=block_table.get_unseen_token_ids(seq.get_token_ids()), - num_lookahead_slots=num_lookahead_slots, - num_computed_slots=seq.data.get_num_computed_tokens(), - ) - # Return any new copy-on-writes. - new_cows = self.block_allocator.clear_copy_on_writes() - return new_cows - - def free(self, seq: Sequence) -> None: - seq_id = seq.seq_id - - if seq_id not in self.block_tables: - # Already freed or haven't been scheduled yet. - return - - # Update seq block ids with the latest access time - self._last_access_blocks_tracker.update_seq_blocks_last_access( - seq_id, self.block_tables[seq.seq_id].physical_block_ids) - - # Untrack seq - self._last_access_blocks_tracker.remove_seq(seq_id) - self._computed_blocks_tracker.remove_seq(seq_id) - - # Free table/blocks - self.block_tables[seq_id].free() - del self.block_tables[seq_id] - - def free_cross(self, seq_group: SequenceGroup) -> None: - request_id = seq_group.request_id - if request_id not in self.cross_block_tables: - # Already freed or hasn't been scheduled yet. - return - self.cross_block_tables[request_id].free() - del self.cross_block_tables[request_id] - + + def can_append_slots(self, seq_group: SequenceGroup, + num_lookahead_slots: int) -> bool: + """Determine if there is enough space in the GPU KV cache to continue + generation of the specified sequence group. + + We use a worst-case heuristic: assume each touched block will require a + new allocation (either via CoW or new block). We can append slots if the + number of touched blocks is less than the number of free blocks. + + "Lookahead slots" are slots that are allocated in addition to the slots + for known tokens. The contents of the lookahead slots are not defined. + This is used by speculative decoding when speculating future tokens. + """ + + num_touched_blocks = 0 + for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING): + block_table = self.block_tables[seq.seq_id] + + num_touched_blocks += ( + block_table.get_num_blocks_touched_by_append_slots( + token_ids=block_table.get_unseen_token_ids( + seq.get_token_ids()), + num_lookahead_slots=num_lookahead_slots, + )) + + num_free_gpu_blocks = self.block_allocator.get_num_free_blocks( + Device.GPU) + return num_touched_blocks <= num_free_gpu_blocks + + def append_slots( + self, + seq: Sequence, + num_lookahead_slots: int, + ) -> List[Tuple[int, int]]: + + block_table = self.block_tables[seq.seq_id] + + block_table.append_token_ids( + token_ids=block_table.get_unseen_token_ids(seq.get_token_ids()), + num_lookahead_slots=num_lookahead_slots, + num_computed_slots=seq.data.get_num_computed_tokens(), + ) + # Return any new copy-on-writes. + new_cows = self.block_allocator.clear_copy_on_writes() + return new_cows + + def free(self, seq: Sequence) -> None: + seq_id = seq.seq_id + + if seq_id not in self.block_tables: + # Already freed or haven't been scheduled yet. + return + + # Update seq block ids with the latest access time + self._last_access_blocks_tracker.update_seq_blocks_last_access( + seq_id, self.block_tables[seq.seq_id].physical_block_ids) + + # Untrack seq + self._last_access_blocks_tracker.remove_seq(seq_id) + self._computed_blocks_tracker.remove_seq(seq_id) + + # Free table/blocks + self.block_tables[seq_id].free() + del self.block_tables[seq_id] + + def free_cross(self, seq_group: SequenceGroup) -> None: + request_id = seq_group.request_id + if request_id not in self.cross_block_tables: + # Already freed or hasn't been scheduled yet. + return + self.cross_block_tables[request_id].free() + del self.cross_block_tables[request_id] + def get_block_table(self, seq: Sequence) -> List[int]: block_ids = self.block_tables[seq.seq_id].physical_block_ids return block_ids # type: ignore def get_cross_block_table(self, seq_group: SequenceGroup) -> List[int]: - request_id = seq_group.request_id - assert request_id in self.cross_block_tables - block_ids = self.cross_block_tables[request_id].physical_block_ids - assert all(b is not None for b in block_ids) - return block_ids # type: ignore - - def access_all_blocks_in_seq(self, seq: Sequence, now: float): - if self.enable_caching: - # Record the latest access time for the sequence. The actual update - # of the block ids is deferred to the sequence free(..) call, since - # only during freeing of block ids, the blocks are actually added to - # the evictor (which is when the most updated time is required) - # (This avoids expensive calls to mark_blocks_as_accessed(..)) - self._last_access_blocks_tracker.update_last_access( - seq.seq_id, now) - - def mark_blocks_as_computed(self, seq_group: SequenceGroup, - token_chunk_size: int): - # If prefix caching is enabled, mark immutable blocks as computed - # right after they have been scheduled (for prefill). This assumes - # the scheduler is synchronous so blocks are actually computed when - # scheduling the next batch. - self.block_allocator.mark_blocks_as_computed([]) - + request_id = seq_group.request_id + assert request_id in self.cross_block_tables + block_ids = self.cross_block_tables[request_id].physical_block_ids + assert all(b is not None for b in block_ids) + return block_ids # type: ignore + + def access_all_blocks_in_seq(self, seq: Sequence, now: float): + if self.enable_caching: + # Record the latest access time for the sequence. The actual update + # of the block ids is deferred to the sequence free(..) call, since + # only during freeing of block ids, the blocks are actually added to + # the evictor (which is when the most updated time is required) + # (This avoids expensive calls to mark_blocks_as_accessed(..)) + self._last_access_blocks_tracker.update_last_access( + seq.seq_id, now) + + def mark_blocks_as_computed(self, seq_group: SequenceGroup, + token_chunk_size: int): + # If prefix caching is enabled, mark immutable blocks as computed + # right after they have been scheduled (for prefill). This assumes + # the scheduler is synchronous so blocks are actually computed when + # scheduling the next batch. + self.block_allocator.mark_blocks_as_computed([]) + def get_common_computed_block_ids( self, seqs: List[Sequence]) -> GenericSequence[int]: - """Determine which blocks for which we skip prefill. - - With prefix caching we can skip prefill for previously-generated blocks. - Currently, the attention implementation only supports skipping cached - blocks if they are a contiguous prefix of cached blocks. - - This method determines which blocks can be safely skipped for all - sequences in the sequence group. - """ - computed_seq_block_ids = [] - for seq in seqs: - computed_seq_block_ids.append( - self._computed_blocks_tracker. - get_cached_computed_blocks_and_update( - seq.seq_id, - self.block_tables[seq.seq_id].physical_block_ids)) - - # NOTE(sang): This assumes seq_block_ids doesn't contain any None. + """Determine which blocks for which we skip prefill. + + With prefix caching we can skip prefill for previously-generated blocks. + Currently, the attention implementation only supports skipping cached + blocks if they are a contiguous prefix of cached blocks. + + This method determines which blocks can be safely skipped for all + sequences in the sequence group. + """ + computed_seq_block_ids = [] + for seq in seqs: + computed_seq_block_ids.append( + self._computed_blocks_tracker. + get_cached_computed_blocks_and_update( + seq.seq_id, + self.block_tables[seq.seq_id].physical_block_ids)) + + # NOTE(sang): This assumes seq_block_ids doesn't contain any None. return self.block_allocator.get_common_computed_block_ids( computed_seq_block_ids) # type: ignore @@ -583,187 +583,187 @@ class BlockSpaceManagerV2(BlockSpaceManager): raise TypeError(f"Unsupported multimodal namespace value type {type(value)}") - - def fork(self, parent_seq: Sequence, child_seq: Sequence) -> None: - if parent_seq.seq_id not in self.block_tables: - # Parent sequence has either been freed or never existed. - return - src_block_table = self.block_tables[parent_seq.seq_id] - self.block_tables[child_seq.seq_id] = src_block_table.fork() - - # Track child seq - self._computed_blocks_tracker.add_seq(child_seq.seq_id) - self._last_access_blocks_tracker.add_seq(child_seq.seq_id) - + + def fork(self, parent_seq: Sequence, child_seq: Sequence) -> None: + if parent_seq.seq_id not in self.block_tables: + # Parent sequence has either been freed or never existed. + return + src_block_table = self.block_tables[parent_seq.seq_id] + self.block_tables[child_seq.seq_id] = src_block_table.fork() + + # Track child seq + self._computed_blocks_tracker.add_seq(child_seq.seq_id) + self._last_access_blocks_tracker.add_seq(child_seq.seq_id) + def can_swap_in(self, seq_group: SequenceGroup, num_lookahead_slots: int) -> AllocStatus: - """Returns the AllocStatus for the given sequence_group - with num_lookahead_slots. - - Args: - sequence_group (SequenceGroup): The sequence group to swap in. - num_lookahead_slots (int): Number of lookahead slots used in - speculative decoding, default to 0. - - Returns: - AllocStatus: The AllocStatus for the given sequence group. - """ + """Returns the AllocStatus for the given sequence_group + with num_lookahead_slots. + + Args: + sequence_group (SequenceGroup): The sequence group to swap in. + num_lookahead_slots (int): Number of lookahead slots used in + speculative decoding, default to 0. + + Returns: + AllocStatus: The AllocStatus for the given sequence group. + """ if self.block_allocator.content_offload_enabled: return AllocStatus.NEVER return self._can_swap(seq_group, Device.GPU, SequenceStatus.SWAPPED, num_lookahead_slots) - - def swap_in(self, seq_group: SequenceGroup) -> List[Tuple[int, int]]: - """Returns the block id mapping (from CPU to GPU) generated by - swapping in the given seq_group with num_lookahead_slots. - - Args: - seq_group (SequenceGroup): The sequence group to swap in. - - Returns: - List[Tuple[int, int]]: The mapping of swapping block from CPU - to GPU. - """ - physical_block_id_mapping = [] - for seq in seq_group.get_seqs(status=SequenceStatus.SWAPPED): - blocks = self.block_tables[seq.seq_id].blocks - if len(blocks) == 0: - continue - - seq_swap_mapping = self.block_allocator.swap(blocks=blocks, - src_device=Device.CPU, - dst_device=Device.GPU) - - # Refresh the block ids of the table (post-swap) - self.block_tables[seq.seq_id].update(blocks) - - seq_physical_block_id_mapping = { - self.block_allocator.get_physical_block_id( - Device.CPU, cpu_block_id): - self.block_allocator.get_physical_block_id( - Device.GPU, gpu_block_id) - for cpu_block_id, gpu_block_id in seq_swap_mapping.items() - } - - physical_block_id_mapping.extend( - list(seq_physical_block_id_mapping.items())) - - return physical_block_id_mapping - + + def swap_in(self, seq_group: SequenceGroup) -> List[Tuple[int, int]]: + """Returns the block id mapping (from CPU to GPU) generated by + swapping in the given seq_group with num_lookahead_slots. + + Args: + seq_group (SequenceGroup): The sequence group to swap in. + + Returns: + List[Tuple[int, int]]: The mapping of swapping block from CPU + to GPU. + """ + physical_block_id_mapping = [] + for seq in seq_group.get_seqs(status=SequenceStatus.SWAPPED): + blocks = self.block_tables[seq.seq_id].blocks + if len(blocks) == 0: + continue + + seq_swap_mapping = self.block_allocator.swap(blocks=blocks, + src_device=Device.CPU, + dst_device=Device.GPU) + + # Refresh the block ids of the table (post-swap) + self.block_tables[seq.seq_id].update(blocks) + + seq_physical_block_id_mapping = { + self.block_allocator.get_physical_block_id( + Device.CPU, cpu_block_id): + self.block_allocator.get_physical_block_id( + Device.GPU, gpu_block_id) + for cpu_block_id, gpu_block_id in seq_swap_mapping.items() + } + + physical_block_id_mapping.extend( + list(seq_physical_block_id_mapping.items())) + + return physical_block_id_mapping + def can_swap_out(self, seq_group: SequenceGroup) -> bool: - """Returns whether we can swap out the given sequence_group - with num_lookahead_slots. - - Args: - seq_group (SequenceGroup): The sequence group to swap in. - num_lookahead_slots (int): Number of lookahead slots used in - speculative decoding, default to 0. - - Returns: - bool: Whether it's possible to swap out current sequence group. - """ + """Returns whether we can swap out the given sequence_group + with num_lookahead_slots. + + Args: + seq_group (SequenceGroup): The sequence group to swap in. + num_lookahead_slots (int): Number of lookahead slots used in + speculative decoding, default to 0. + + Returns: + bool: Whether it's possible to swap out current sequence group. + """ if self.block_allocator.content_offload_enabled: return False alloc_status = self._can_swap(seq_group, Device.CPU, SequenceStatus.RUNNING) - return alloc_status == AllocStatus.OK - - def swap_out(self, seq_group: SequenceGroup) -> List[Tuple[int, int]]: - """Returns the block id mapping (from GPU to CPU) generated by - swapping out the given sequence_group with num_lookahead_slots. - - Args: - sequence_group (SequenceGroup): The sequence group to swap in. - - Returns: - List[Tuple[int, int]]: The mapping of swapping block from - GPU to CPU. - """ - physical_block_id_mapping = [] - for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING): - blocks = self.block_tables[seq.seq_id].blocks - if len(blocks) == 0: - continue - - seq_swap_mapping = self.block_allocator.swap(blocks=blocks, - src_device=Device.GPU, - dst_device=Device.CPU) - - # Refresh the block ids of the table (post-swap) - self.block_tables[seq.seq_id].update(blocks) - - seq_physical_block_id_mapping = { - self.block_allocator.get_physical_block_id( - Device.GPU, gpu_block_id): - self.block_allocator.get_physical_block_id( - Device.CPU, cpu_block_id) - for gpu_block_id, cpu_block_id in seq_swap_mapping.items() - } - - physical_block_id_mapping.extend( - list(seq_physical_block_id_mapping.items())) - - return physical_block_id_mapping - - def get_num_free_gpu_blocks(self) -> int: - return self.block_allocator.get_num_free_blocks(Device.GPU) - - def get_num_free_cpu_blocks(self) -> int: - return self.block_allocator.get_num_free_blocks(Device.CPU) - - def get_prefix_cache_hit_rate(self, device: Device) -> float: - return self.block_allocator.get_prefix_cache_hit_rate(device) - - def _can_swap(self, - seq_group: SequenceGroup, - device: Device, - status: SequenceStatus, - num_lookahead_slots: int = 0) -> AllocStatus: - """Returns the AllocStatus for swapping in/out the given sequence_group - on to the 'device'. - - Args: - sequence_group (SequenceGroup): The sequence group to swap in. - device (Device): device to swap the 'seq_group' on. - status (SequenceStatus): The status of sequence which is needed - for action. RUNNING for swap out and SWAPPED for swap in - num_lookahead_slots (int): Number of lookahead slots used in - speculative decoding, default to 0. - - Returns: - AllocStatus: The AllocStatus for swapping in/out the given - sequence_group on to the 'device'. - """ - # First determine the number of blocks that will be touched by this - # swap. Then verify if there are available blocks in the device - # to perform the swap. - num_blocks_touched = 0 - blocks: List[Block] = [] - for seq in seq_group.get_seqs(status=status): - block_table = self.block_tables[seq.seq_id] - if block_table.blocks is not None: - # Compute the number blocks to touch for the tokens to be - # appended. This does NOT include the full blocks that need - # to be touched for the swap. - num_blocks_touched += \ - block_table.get_num_blocks_touched_by_append_slots( - block_table.get_unseen_token_ids(seq.get_token_ids()), - num_lookahead_slots=num_lookahead_slots) - blocks.extend(block_table.blocks) - # Compute the number of full blocks to touch and add it to the - # existing count of blocks to touch. - num_blocks_touched += self.block_allocator.get_num_full_blocks_touched( - blocks, device=device) - - watermark_blocks = 0 - if device == Device.GPU: - watermark_blocks = self.watermark_blocks - - if self.block_allocator.get_num_total_blocks( - device) < num_blocks_touched: - return AllocStatus.NEVER - elif self.block_allocator.get_num_free_blocks( - device) - num_blocks_touched >= watermark_blocks: - return AllocStatus.OK - else: - return AllocStatus.LATER + return alloc_status == AllocStatus.OK + + def swap_out(self, seq_group: SequenceGroup) -> List[Tuple[int, int]]: + """Returns the block id mapping (from GPU to CPU) generated by + swapping out the given sequence_group with num_lookahead_slots. + + Args: + sequence_group (SequenceGroup): The sequence group to swap in. + + Returns: + List[Tuple[int, int]]: The mapping of swapping block from + GPU to CPU. + """ + physical_block_id_mapping = [] + for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING): + blocks = self.block_tables[seq.seq_id].blocks + if len(blocks) == 0: + continue + + seq_swap_mapping = self.block_allocator.swap(blocks=blocks, + src_device=Device.GPU, + dst_device=Device.CPU) + + # Refresh the block ids of the table (post-swap) + self.block_tables[seq.seq_id].update(blocks) + + seq_physical_block_id_mapping = { + self.block_allocator.get_physical_block_id( + Device.GPU, gpu_block_id): + self.block_allocator.get_physical_block_id( + Device.CPU, cpu_block_id) + for gpu_block_id, cpu_block_id in seq_swap_mapping.items() + } + + physical_block_id_mapping.extend( + list(seq_physical_block_id_mapping.items())) + + return physical_block_id_mapping + + def get_num_free_gpu_blocks(self) -> int: + return self.block_allocator.get_num_free_blocks(Device.GPU) + + def get_num_free_cpu_blocks(self) -> int: + return self.block_allocator.get_num_free_blocks(Device.CPU) + + def get_prefix_cache_hit_rate(self, device: Device) -> float: + return self.block_allocator.get_prefix_cache_hit_rate(device) + + def _can_swap(self, + seq_group: SequenceGroup, + device: Device, + status: SequenceStatus, + num_lookahead_slots: int = 0) -> AllocStatus: + """Returns the AllocStatus for swapping in/out the given sequence_group + on to the 'device'. + + Args: + sequence_group (SequenceGroup): The sequence group to swap in. + device (Device): device to swap the 'seq_group' on. + status (SequenceStatus): The status of sequence which is needed + for action. RUNNING for swap out and SWAPPED for swap in + num_lookahead_slots (int): Number of lookahead slots used in + speculative decoding, default to 0. + + Returns: + AllocStatus: The AllocStatus for swapping in/out the given + sequence_group on to the 'device'. + """ + # First determine the number of blocks that will be touched by this + # swap. Then verify if there are available blocks in the device + # to perform the swap. + num_blocks_touched = 0 + blocks: List[Block] = [] + for seq in seq_group.get_seqs(status=status): + block_table = self.block_tables[seq.seq_id] + if block_table.blocks is not None: + # Compute the number blocks to touch for the tokens to be + # appended. This does NOT include the full blocks that need + # to be touched for the swap. + num_blocks_touched += \ + block_table.get_num_blocks_touched_by_append_slots( + block_table.get_unseen_token_ids(seq.get_token_ids()), + num_lookahead_slots=num_lookahead_slots) + blocks.extend(block_table.blocks) + # Compute the number of full blocks to touch and add it to the + # existing count of blocks to touch. + num_blocks_touched += self.block_allocator.get_num_full_blocks_touched( + blocks, device=device) + + watermark_blocks = 0 + if device == Device.GPU: + watermark_blocks = self.watermark_blocks + + if self.block_allocator.get_num_total_blocks( + device) < num_blocks_touched: + return AllocStatus.NEVER + elif self.block_allocator.get_num_free_blocks( + device) - num_blocks_touched >= watermark_blocks: + return AllocStatus.OK + else: + return AllocStatus.LATER diff --git a/vllm_overrides/core/evictor_v2.py b/vllm_overrides/core/evictor_v2.py index 393c8426..3ce89711 100644 --- a/vllm_overrides/core/evictor_v2.py +++ b/vllm_overrides/core/evictor_v2.py @@ -7,128 +7,128 @@ from typing import Dict, List, OrderedDict, Tuple ContentHash = bytes - - + + class EvictionPolicy(enum.Enum): - """Enum for eviction policy used by make_evictor to instantiate the correct - Evictor subclass. + """Enum for eviction policy used by make_evictor to instantiate the correct + Evictor subclass. """ LRU = enum.auto() FREQUENCY_AWARE = enum.auto() - - -class Evictor(ABC): - """The Evictor subclasses should be used by the BlockAllocator class to - handle eviction of freed PhysicalTokenBlocks. - """ - - @abstractmethod - def __init__(self): - pass - - @abstractmethod - def __contains__(self, block_id: int) -> bool: - pass - - @abstractmethod + + +class Evictor(ABC): + """The Evictor subclasses should be used by the BlockAllocator class to + handle eviction of freed PhysicalTokenBlocks. + """ + + @abstractmethod + def __init__(self): + pass + + @abstractmethod + def __contains__(self, block_id: int) -> bool: + pass + + @abstractmethod def evict(self) -> Tuple[int, ContentHash]: - """Runs the eviction algorithm and returns the evicted block's - content hash along with physical block id along with physical block id - """ - pass - - @abstractmethod + """Runs the eviction algorithm and returns the evicted block's + content hash along with physical block id along with physical block id + """ + pass + + @abstractmethod def add(self, block_id: int, content_hash: ContentHash, num_hashed_tokens: int, last_accessed: float): - """Adds block to the evictor, making it a candidate for eviction""" - pass - - @abstractmethod - def update(self, block_id: int, last_accessed: float): - """Update corresponding block's access time in metadata""" - pass - - @abstractmethod - def remove(self, block_id: int): - """Remove a given block id from the cache.""" - pass - - @property - @abstractmethod - def num_blocks(self) -> int: - pass - - -class BlockMetaData(): - """Data structure for storing key data describe cached block, so that - evitor could use to make its decision which one to choose for eviction - - Here we use physical block id as the dict key, as there maybe several - blocks with the same content hash, but their physical id is unique. - """ - + """Adds block to the evictor, making it a candidate for eviction""" + pass + + @abstractmethod + def update(self, block_id: int, last_accessed: float): + """Update corresponding block's access time in metadata""" + pass + + @abstractmethod + def remove(self, block_id: int): + """Remove a given block id from the cache.""" + pass + + @property + @abstractmethod + def num_blocks(self) -> int: + pass + + +class BlockMetaData(): + """Data structure for storing key data describe cached block, so that + evitor could use to make its decision which one to choose for eviction + + Here we use physical block id as the dict key, as there maybe several + blocks with the same content hash, but their physical id is unique. + """ + def __init__(self, content_hash: ContentHash, num_hashed_tokens: int, last_accessed: float): - self.content_hash = content_hash - self.num_hashed_tokens = num_hashed_tokens - self.last_accessed = last_accessed - - -class LRUEvictor(Evictor): - """Evicts in a least-recently-used order using the last_accessed timestamp - that's recorded in the PhysicalTokenBlock. If there are multiple blocks with - the same last_accessed time, then the one with the largest num_hashed_tokens - will be evicted. If two blocks each have the lowest last_accessed time and - highest num_hashed_tokens value, then one will be chose arbitrarily - """ - - def __init__(self): - self.free_table: OrderedDict[int, BlockMetaData] = OrderedDict() - - def __contains__(self, block_id: int) -> bool: - return block_id in self.free_table - + self.content_hash = content_hash + self.num_hashed_tokens = num_hashed_tokens + self.last_accessed = last_accessed + + +class LRUEvictor(Evictor): + """Evicts in a least-recently-used order using the last_accessed timestamp + that's recorded in the PhysicalTokenBlock. If there are multiple blocks with + the same last_accessed time, then the one with the largest num_hashed_tokens + will be evicted. If two blocks each have the lowest last_accessed time and + highest num_hashed_tokens value, then one will be chose arbitrarily + """ + + def __init__(self): + self.free_table: OrderedDict[int, BlockMetaData] = OrderedDict() + + def __contains__(self, block_id: int) -> bool: + return block_id in self.free_table + def evict(self) -> Tuple[int, ContentHash]: - if len(self.free_table) == 0: - raise ValueError("No usable cache memory left") - - evicted_block, evicted_block_id = None, None - # The blocks with the lowest timestamps should be placed consecutively - # at the start of OrderedDict. Loop through all these blocks to - # find the one with maximum number of hashed tokens. - for _id, block in self.free_table.items(): - if evicted_block is None: - evicted_block, evicted_block_id = block, _id - continue - if evicted_block.last_accessed < block.last_accessed: - break - if evicted_block.num_hashed_tokens < block.num_hashed_tokens: - evicted_block, evicted_block_id = block, _id - - assert evicted_block is not None - assert evicted_block_id is not None - self.free_table.pop(evicted_block_id) - - return evicted_block_id, evicted_block.content_hash - + if len(self.free_table) == 0: + raise ValueError("No usable cache memory left") + + evicted_block, evicted_block_id = None, None + # The blocks with the lowest timestamps should be placed consecutively + # at the start of OrderedDict. Loop through all these blocks to + # find the one with maximum number of hashed tokens. + for _id, block in self.free_table.items(): + if evicted_block is None: + evicted_block, evicted_block_id = block, _id + continue + if evicted_block.last_accessed < block.last_accessed: + break + if evicted_block.num_hashed_tokens < block.num_hashed_tokens: + evicted_block, evicted_block_id = block, _id + + assert evicted_block is not None + assert evicted_block_id is not None + self.free_table.pop(evicted_block_id) + + return evicted_block_id, evicted_block.content_hash + def add(self, block_id: int, content_hash: ContentHash, num_hashed_tokens: int, last_accessed: float): - self.free_table[block_id] = BlockMetaData(content_hash, - num_hashed_tokens, - last_accessed) - - def update(self, block_id: int, last_accessed: float): - self.free_table[block_id].last_accessed = last_accessed - - def remove(self, block_id: int): - if block_id not in self.free_table: - raise ValueError( - "Attempting to remove block that's not in the evictor") - self.free_table.pop(block_id) - - @property + self.free_table[block_id] = BlockMetaData(content_hash, + num_hashed_tokens, + last_accessed) + + def update(self, block_id: int, last_accessed: float): + self.free_table[block_id].last_accessed = last_accessed + + def remove(self, block_id: int): + if block_id not in self.free_table: + raise ValueError( + "Attempting to remove block that's not in the evictor") + self.free_table.pop(block_id) + + @property def num_blocks(self) -> int: return len(self.free_table) @@ -261,8 +261,8 @@ def eviction_policy_from_env( raise ValueError( "BI100_KV_EVICTION_POLICY must be one of: frequency, lru") return policies[value] - - + + def make_evictor(eviction_policy: EvictionPolicy) -> Evictor: if eviction_policy == EvictionPolicy.LRU: return LRUEvictor() diff --git a/vllm_overrides/model_executor/layers/sampler.py b/vllm_overrides/model_executor/layers/sampler.py index ef6db443..f47d35c4 100644 --- a/vllm_overrides/model_executor/layers/sampler.py +++ b/vllm_overrides/model_executor/layers/sampler.py @@ -1,1057 +1,1057 @@ -"""A layer that samples the next tokens from the model's outputs.""" -import itertools -import warnings -from dataclasses import dataclass -from importlib.util import find_spec -from math import inf -from typing import Dict, List, Optional, Tuple, Union - -import msgspec -import torch -import torch.nn as nn - -import vllm.envs as envs -from vllm.model_executor.sampling_metadata import (SamplingMetadata, - SamplingTensors, - SequenceGroupToSample) -from vllm.sampling_params import SamplingType -from vllm.sequence import (VLLM_INVALID_TOKEN_ID, - CompletionSequenceGroupOutput, Logprob, - PromptLogprobs, SampleLogprobs, SequenceOutput) -from vllm.spec_decode.metrics import SpecDecodeWorkerMetrics - -if envs.VLLM_USE_FLASHINFER_SAMPLER and find_spec("flashinfer"): - import flashinfer.sampling - # yapf: disable - from flashinfer.sampling import ( - top_k_top_p_sampling_from_probs as flashinfer_top_k_top_p_sampling) - - # yapf: enable -else: - flashinfer_top_k_top_p_sampling = None - -# (num_token_ids, num_parent_ids) per sequence group. -SampleResultType = List[Tuple[List[int], List[int]]] - -# Types of temporary data structures used for -# computing sample_result -SampleMetadataType = Dict[SamplingType, Tuple[List[int], - List[SequenceGroupToSample]]] -MultinomialSamplesType = Dict[SamplingType, torch.Tensor] -SampleResultsDictType = Dict[int, Tuple[List[int], List[int]]] - - -# Encapsulates temporary data structures for computing -# sample_result. -# -# * For multi-step scheduling: must be returned -# by `Sampler.forward()` and used later to compute the pythonized -# sample_result -# -# * For single-step scheduling: consumed immediately -# inside `Sampler.forward()` to compute pythonized sample_result. -@dataclass -class SampleResultArgsType: - sample_metadata: SampleMetadataType - multinomial_samples: MultinomialSamplesType - sample_results_dict: SampleResultsDictType - sampling_metadata: SamplingMetadata - greedy_samples: Optional[torch.Tensor] - beam_search_logprobs: Optional[torch.Tensor] - - -# Union of non-deferred (single-step scheduling) -# vs deferred (multi-step scheduling) -# sample result types -MaybeDeferredSampleResultType = Union[SampleResultType, SampleResultArgsType] - -# Abbreviation of the _sample() return type -SampleReturnType = Tuple[MaybeDeferredSampleResultType, Optional[torch.Tensor]] - - -class SamplerOutput( - msgspec.Struct, - omit_defaults=True, # type: ignore[call-arg] - array_like=True): # type: ignore[call-arg] - """For each sequence group, we generate a list of SequenceOutput object, - each of which contains one possible candidate for the next token. - - This data structure implements methods, so it can be used like a list, but - also has optional fields for device tensors. - """ - - outputs: List[CompletionSequenceGroupOutput] - - # On-device tensor containing probabilities of each token. - sampled_token_probs: Optional[torch.Tensor] = None - - # On-device tensor containing the logprobs of each token. - logprobs: Optional["torch.Tensor"] = None - - # Holds either (1) the pythonized sampler result (single-step scheduling) - # or (2) what will be arguments for later deferred pythonization of the - # sampler result (muliti-step scheduling) - deferred_sample_results_args: Optional[SampleResultArgsType] = None - - # On-device tensor containing the sampled token ids. - sampled_token_ids: Optional[torch.Tensor] = None - # CPU tensor containing the sampled token ids. Used during multi-step to - # return the sampled token ids from last rank to AsyncLLMEngine to be - # 'broadcasted' to all other PP ranks for next step. - sampled_token_ids_cpu: Optional[torch.Tensor] = None - - # Spec decode metrics populated by workers. - spec_decode_worker_metrics: Optional[SpecDecodeWorkerMetrics] = None - - # Optional last hidden states from the model. - hidden_states: Optional[torch.Tensor] = None - - # Optional prefill hidden states from the model - # (used for models like EAGLE). - prefill_hidden_states: Optional[torch.Tensor] = None - - # Time taken in the forward pass for this across all workers - model_forward_time: Optional[float] = None - - # Time taken in the model execute function. This will include model forward, - # block/sync across workers, cpu-gpu sync time and sampling time. - model_execute_time: Optional[float] = None - - def __getitem__(self, idx: int): - return self.outputs[idx] - - def __setitem__(self, idx: int, value): - self.outputs[idx] = value - - def __len__(self): - return len(self.outputs) - - def __eq__(self, other: object): - return isinstance(other, - self.__class__) and self.outputs == other.outputs - - def __repr__(self) -> str: - """Show the shape of a tensor instead of its values to reduce noise. - """ - sampled_token_probs_repr = ("None" if self.sampled_token_probs is None - else self.sampled_token_probs.shape) - sampled_token_ids_repr = ("None" if self.sampled_token_ids is None else - self.sampled_token_ids.shape) - return ( - f"SamplerOutput(outputs={self.outputs}, " - f"sampled_token_probs={sampled_token_probs_repr}, " - f"sampled_token_ids={sampled_token_ids_repr}, " - f"spec_decode_worker_metrics={self.spec_decode_worker_metrics})") - - -class Sampler(nn.Module): - """Samples the next tokens from the model's outputs. - - This layer does the following: - 1. Discard the hidden states that are not used for sampling (i.e., all - tokens except the final one in each prompt). - 2. Compute the logits for the next tokens. - 3. Apply presence, frequency and repetition penalties. - 4. Apply temperature scaling. - 5. Apply top-p and top-k truncation. - 6. Sample the next tokens. - Here, each sequence group within the batch can have different sampling - parameters (e.g., sampling method, temperature, top-p, top-k, etc.). - - The structure of the logits tensor is coupled with the seq_groups in - sampling_metadata. Typically, each sequence in each seq_group has one row in - logits for the next token to be sampled; however, for a seq_group with a - prompt request with the prompt_logprobs sampling parameter, there are rows - in logits for each token in the input prompt. - """ - - def __init__(self): - super().__init__() - - # Whether or not the SamplerOutput should have on-device tensors - # containing the sampled token ids and probabilities. This is used by - # speculative decoding. - self.include_gpu_probs_tensor = False - self.should_modify_greedy_probs_inplace = False - - def _init_sampling_tensors( - self, - logits: torch.Tensor, - sampling_metadata: SamplingMetadata, - ): - """The goal here is to reuse sampling tensors between similar decode - runs. This is possible because sampling logic does not change between - decodes of the same sequences. - """ - _, vocab_size = logits.shape - - # First free any existing stored sampling tensors. - # This is necessary because some sampling tensors may - # have pinned memory. - self._sampling_tensors = None - - # Initialize new sampling tensors - (sampling_tensors, do_penalties, do_top_p_top_k, - do_min_p) = SamplingTensors.from_sampling_metadata( - sampling_metadata, vocab_size, logits.device, logits.dtype) - - self._sampling_tensors = sampling_tensors - self._do_penalties = do_penalties - self._do_top_p_top_k = do_top_p_top_k - self._do_min_p = do_min_p - - def forward( - self, - logits: torch.Tensor, - sampling_metadata: SamplingMetadata, - ) -> Optional[SamplerOutput]: - """ - Single-step scheduling: - * Perform GPU-side sampling computation & compute - GPU-side logprobs tensor - * Pythonize sampling result & logprobs tensor - - Multi-step scheduling: - * Perform GPU-side sampling computation & compute - GPU-side logprobs tensor - * Defer Pythonization of sampling result & logprobs - tensor - * Encapsulate arguments required for deferred Pythonization - in the :class:`SamplerOutput` structure - - Args: - logits: (num_tokens, vocab_size). - sampling_metadata: Metadata for sampling. - """ - assert logits is not None - _, vocab_size = logits.shape - - # Prepare sampling tensors with pinned memory to avoid blocking. - if not sampling_metadata.reuse_sampling_tensors: - self._init_sampling_tensors(logits, sampling_metadata) - elif self._do_penalties: - # In this case, the sampling tensors logic depends on - # "output_tokens" of a sequence. As a result, we cannot - # reuse sampling tensors, since "output_tokens" changes - # between decode runs. - self._init_sampling_tensors(logits, sampling_metadata) - - assert self._sampling_tensors is not None - sampling_tensors = self._sampling_tensors - do_penalties = self._do_penalties - do_top_p_top_k = self._do_top_p_top_k - do_min_p = self._do_min_p - - logits = _apply_min_tokens_penalty(logits, sampling_metadata) - - # Apply presence and frequency penalties. - if do_penalties: - logits = _apply_penalties(logits, sampling_tensors.prompt_tokens, - sampling_tensors.output_tokens, - sampling_tensors.presence_penalties, - sampling_tensors.frequency_penalties, - sampling_tensors.repetition_penalties) - - # Use float32 to apply temperature scaling. - # Use in-place division to avoid creating a new tensor. - logits = logits.to(torch.float) - logits.div_(sampling_tensors.temperatures.unsqueeze(dim=1)) - - if do_top_p_top_k and flashinfer_top_k_top_p_sampling is None: - logits = _apply_top_k_top_p(logits, sampling_tensors.top_ps, - sampling_tensors.top_ks) - - if do_min_p: - logits = _apply_min_p(logits, sampling_tensors.min_ps) - - # We use float32 for probabilities and log probabilities. - # Compute the probabilities. - probs = torch.softmax(logits, dim=-1, dtype=torch.float) - # Compute the log probabilities. - logprobs = torch.log_softmax(logits, dim=-1, dtype=torch.float) - - # Sample the next tokens. - maybe_deferred_sample_results, maybe_sampled_tokens_tensor = _sample( - probs, - logprobs, - sampling_metadata, - sampling_tensors, - include_gpu_probs_tensor=self.include_gpu_probs_tensor, - modify_greedy_probs=self._should_modify_greedy_probs_inplace, - ) - - if self.include_gpu_probs_tensor: - # Since we will defer sampler result Pythonization, - # preserve GPU-side tensors in support of later - # deferred pythonization of logprobs - assert maybe_sampled_tokens_tensor is not None - on_device_tensors = (probs, logprobs, maybe_sampled_tokens_tensor) - else: - # Since Pythonization has already happened, don't preserve - # GPU-side tensors. - on_device_tensors = None - - # Get the logprobs query results. - prompt_logprobs = None - sample_logprobs = None - if not sampling_metadata.skip_sampler_cpu_output: - # Pythonize logprobs now (GPU -> CPU); do not defer. - assert not isinstance(maybe_deferred_sample_results, - SampleResultArgsType) - prompt_logprobs, sample_logprobs = get_logprobs( - logprobs, sampling_metadata, maybe_deferred_sample_results) - - return _build_sampler_output( - maybe_deferred_sample_results, - sampling_metadata, - prompt_logprobs, - sample_logprobs, - on_device_tensors=on_device_tensors, - skip_sampler_cpu_output=sampling_metadata.skip_sampler_cpu_output) - - @property - def _should_modify_greedy_probs_inplace(self) -> bool: - """Whether or not the sampler should modify the probability distribution - of greedily-sampled tokens such that multinomial sampling would sample - the greedily-sampled token. - - In other words, if True then we set the probability of the greedily- - sampled token to 1. - - This is used by speculative decoding, which requires that the sampling - method be encoded into the probability distribution. - """ - return self.should_modify_greedy_probs_inplace - - -def _get_bin_counts_and_mask( - tokens: torch.Tensor, - vocab_size: int, - num_seqs: int, -) -> Tuple[torch.Tensor, torch.Tensor]: - # Compute the bin counts for the tokens. - # vocab_size + 1 for padding. - bin_counts = torch.zeros((num_seqs, vocab_size + 1), - dtype=torch.long, - device=tokens.device) - bin_counts.scatter_add_(1, tokens, torch.ones_like(tokens)) - bin_counts = bin_counts[:, :vocab_size] - mask = bin_counts > 0 - - return bin_counts, mask - - -def _apply_min_tokens_penalty( - logits: torch.Tensor, - sampling_metadata: SamplingMetadata, -) -> torch.Tensor: - """Apply min_tokens penalty which sets stop tokens to -inf if min_tokens - have not been generated yet - """ - # list of indices in logits that will be set to -inf - logits_to_penalize: List[Tuple[int, int]] = [] - logits_applied = 0 - for seq_group in sampling_metadata.seq_groups: - seq_ids = seq_group.seq_ids - sampling_params = seq_group.sampling_params - - sample_indices = seq_group.sample_indices - logits_applied += len(sample_indices) + len( - seq_group.prompt_logprob_indices) - if not seq_group.do_sample: - continue - - start_idx = sample_indices[0] - min_tokens = sampling_params.min_tokens - token_ids_to_penalize = sampling_params.all_stop_token_ids - if min_tokens > 0 and token_ids_to_penalize: - seqs_to_penalize: List[int] = [] - for j, seq_id in enumerate(seq_ids): - seq_data = seq_group.seq_data[seq_id] - if len(seq_data.output_token_ids_array) < min_tokens: - seqs_to_penalize.append(j) - - if seqs_to_penalize: - # convert to the index into logits - seqs_to_penalize = [start_idx + j for j in seqs_to_penalize] - # itertools.product pairs each seq index with every token id - logits_to_penalize.extend( - itertools.product(seqs_to_penalize, token_ids_to_penalize)) - - if logits_to_penalize: - # use zip and * to group indices along each dimension - # eg. [ (1,2), (1,3), (5,6) ] -> ( (1,1,5), (2,3,6) ) - logits[tuple(zip(*logits_to_penalize))] = -float("inf") - - # verifies that no rows in logits were missed unexpectedly - assert logits_applied == logits.shape[0] - return logits - - -def _apply_penalties(logits: torch.Tensor, prompt_tokens_tensor: torch.Tensor, - output_tokens_tensor: torch.Tensor, - presence_penalties: torch.Tensor, - frequency_penalties: torch.Tensor, - repetition_penalties: torch.Tensor) -> torch.Tensor: - num_seqs, vocab_size = logits.shape - _, prompt_mask = _get_bin_counts_and_mask(prompt_tokens_tensor, vocab_size, - num_seqs) - output_bin_counts, output_mask = _get_bin_counts_and_mask( - output_tokens_tensor, vocab_size, num_seqs) - - repetition_penalties = repetition_penalties[:, None].repeat(1, vocab_size) - repetition_penalties[~(prompt_mask | output_mask)] = 1.0 - logits = torch.where(logits > 0, logits / repetition_penalties, - logits * repetition_penalties) - - # We follow the definition in OpenAI API. - # Refer to https://platform.openai.com/docs/api-reference/parameter-details - logits -= frequency_penalties.unsqueeze_(dim=1) * output_bin_counts - logits -= presence_penalties.unsqueeze_(dim=1) * output_mask - return logits - - -def _apply_top_k_top_p( - logits: torch.Tensor, - p: torch.Tensor, - k: torch.Tensor, -) -> torch.Tensor: - logits_sort, logits_idx = logits.sort(dim=-1, descending=False) - - # Apply top-k. - top_k_mask = logits_sort.size(1) - k.to(torch.long) - # Get all the top_k values. - top_k_mask = logits_sort.gather(1, top_k_mask.unsqueeze(dim=1)) - top_k_mask = logits_sort < top_k_mask - logits_sort.masked_fill_(top_k_mask, -float("inf")) - - # Apply top-p. - probs_sort = logits_sort.softmax(dim=-1) - probs_sum = probs_sort.cumsum(dim=-1) - top_p_mask = probs_sum <= 1 - p.unsqueeze(dim=1) - # at least one - top_p_mask[:, -1] = False - logits_sort.masked_fill_(top_p_mask, -float("inf")) - - # Re-sort the probabilities. - logits = torch.empty_like(logits_sort).scatter_(dim=-1, - index=logits_idx, - src=logits_sort) - return logits - - -def _apply_min_p( - logits: torch.Tensor, - min_p: torch.Tensor, -) -> torch.Tensor: - """ - Adapted from - https://github.com/oobabooga/text-generation-webui/blob/3146124ec01f02c8fb1650a6517cf1b60b537aaf/modules/sampler_hijack.py#L16C17-L16C17 - """ - probs = torch.softmax(logits, dim=-1) - top_probs, _ = probs.max(dim=-1, keepdim=True) - scaled_min_p = min_p.unsqueeze_(dim=1) * top_probs - tokens_to_remove = probs < scaled_min_p - logits = logits.masked_fill_(tokens_to_remove, -float("inf")) - - return logits - - -def _greedy_sample( - selected_seq_groups: List[SequenceGroupToSample], - samples: torch.Tensor, -) -> SampleResultType: - """Run greedy sampling on a given samples. - - Args: - selected_seq_groups: A list of sequence groups batched. - samples: (num_selected_samples,) A tensor of samples. The length of - samples could be smaller than selected_seq_groups if - seq_group.do_sample is False. - Returns: - Tuple of (next_token_ids, parent_ids). The length of returned list is - same as the length of selected_seq_groups. If the corresponding - seq_group has do_sample=False, tuple contains ([], []) - """ - samples_lst = samples.tolist() - sample_idx = 0 - results: SampleResultType = [] - for seq_group in selected_seq_groups: - if not seq_group.do_sample: - results.append(([], [])) - continue - - seq_ids = seq_group.seq_ids - num_parent_seqs = len(seq_ids) - assert num_parent_seqs == 1, ( - "Greedy sampling should have only one seq.") - parent_ids = list(range(num_parent_seqs)) - next_token_ids = [samples_lst[sample_idx]] - results.append((next_token_ids, parent_ids)) - sample_idx += num_parent_seqs - return results - - -def _random_sample( - selected_seq_groups: List[SequenceGroupToSample], - random_samples: torch.Tensor, -) -> SampleResultType: - """Run random sampling on a given samples. - - Args: - selected_seq_groups: A list of sequence groups batched. - random_samples: (num_selected_samples,) A tensor of samples. The - length of samples could be smaller than selected_seq_groups if - seq_group.do_sample is False. - Returns: - Tuple of (next_token_ids, parent_ids). The length of returned list is - same as the length of selected_seq_groups. If the corresponding - seq_group has do_sample=False, tuple contains ([], []) - """ - # Find the maximum n value of the prompt phase requests. - random_samples = random_samples.cpu() - sample_idx = 0 - results: SampleResultType = [] - for seq_group in selected_seq_groups: - if not seq_group.do_sample: - results.append(([], [])) - continue - - seq_ids = seq_group.seq_ids - sampling_params = seq_group.sampling_params - is_prompt = seq_group.is_prompt - num_parent_seqs = len(seq_ids) - if is_prompt: - # Prompt phase. - parent_ids = [0] * sampling_params.n - next_token_ids = random_samples[ - sample_idx, :sampling_params.n].tolist() - else: - # Generation phase. - parent_ids = list(range(num_parent_seqs)) - next_token_ids = random_samples[sample_idx:sample_idx + - num_parent_seqs, 0].tolist() - results.append((next_token_ids, parent_ids)) - sample_idx += num_parent_seqs - return results - - -def _beam_search_sample( - selected_seq_groups: List[SequenceGroupToSample], - logprobs: torch.Tensor, -) -> SampleResultType: - """Run beam sampling on a given samples. - - Args: - selected_seq_groups: A list of sequence groups batched. - logprobs: (num_selected_samples, vocab_size,) A tensor of logprob - on selected sample indices. - Returns: - Tuple of (next_token_ids, parent_ids). The length of returned list is - same as the length of selected_seq_groups. If the corresponding - seq_group has do_sample=False, tuple contains ([], []) - """ - # We sample 2 * beam_width candidates to make sure that with high - # probability we can get `beam_width` candidates in addition to - # the finished sequences for the next iteration. See - # https://github.com/tensorflow/tensor2tensor/blob/bafdc1b67730430d38d6ab802cbd51f9d053ba2e/tensor2tensor/utils/beam_search.py#L557-L563 - # for details. See also HF reference: - # https://github.com/huggingface/transformers/blob/a4dd53d88e4852f023332d284ff07a01afcd5681/src/transformers/generation/utils.py#L3063-L3065 - # - # NOTE: Beam search is not vectorized, so its speed can be slower than - # other sampling methods. - sample_idx = 0 - results: SampleResultType = [] - for seq_group in selected_seq_groups: - if not seq_group.do_sample: - results.append(([], [])) - continue - - is_prompt = seq_group.is_prompt - seq_ids, sampling_params = seq_group.seq_ids, seq_group.sampling_params - num_parent_seqs = len(seq_ids) - beam_width = sampling_params.n - seq_group_logprobs = logprobs[sample_idx:sample_idx + num_parent_seqs] - if is_prompt: - # Prompt phase. - assert num_parent_seqs == 1, ( - "Prompt input should have only one seq.") - parent_ids = [0] * (2 * beam_width) - _, next_token_ids = torch.topk(seq_group_logprobs[0], - 2 * beam_width) - next_token_ids = next_token_ids.tolist() - else: - # Generation phase. - cumulative_logprobs: List[float] = [ - seq_group.seq_data[seq_id].cumulative_logprob - for seq_id in seq_ids - ] - cumulative_logprobs_tensor = torch.tensor( - cumulative_logprobs, - dtype=torch.float, - device=seq_group_logprobs.device) - seq_group_logprobs = (seq_group_logprobs + - cumulative_logprobs_tensor.unsqueeze(dim=1)) - _, topk_ids = torch.topk(seq_group_logprobs.flatten(), - 2 * beam_width) - topk_ids = topk_ids.tolist() - vocab_size = seq_group_logprobs.size(-1) - parent_ids = [i // vocab_size for i in topk_ids] - next_token_ids = [i % vocab_size for i in topk_ids] - results.append((next_token_ids, parent_ids)) - sample_idx += num_parent_seqs - assert sample_idx == logprobs.size(0) - return results - - -# torch.multinomial forces a GPU<->CPU sync. -# Therefore, we use an optimized implementation instead. -# Note that we always sample with replacement. -# probs will be modified in place, but this is fine, as we pass -# in a copy already. -def _multinomial( - probs: torch.Tensor, - num_samples: int, - seq_groups: Optional[List[SequenceGroupToSample]] = None, -) -> torch.Tensor: - if num_samples > 1: - probs = probs.repeat_interleave(num_samples, dim=0) - q = torch.empty_like(probs) - if seq_groups is None: - q.exponential_() - else: - sample_idx = 0 - for seq_group in seq_groups: - seq_ids = seq_group.seq_ids - stride = len(seq_ids) * num_samples - assert seq_group.generator is not None - q[sample_idx:sample_idx + - stride].exponential_(generator=seq_group.generator) - sample_idx += stride - return probs.div_(q).argmax(dim=1).view(-1, num_samples) - - -def _top_k_top_p_multinomial_with_flashinfer( - probs: torch.Tensor, top_ks: torch.Tensor, top_ps: torch.Tensor, - num_samples: int, seq_groups: Optional[List[SequenceGroupToSample]]): - max_top_k_round = 32 - if num_samples > 1: - probs = probs.repeat_interleave(num_samples, dim=0) - top_ks = top_ks.repeat_interleave(num_samples) - top_ps = top_ps.repeat_interleave(num_samples) - batch_size = probs.shape[0] - uniform_samples = torch.empty((max_top_k_round, batch_size), - device=probs.device) - if seq_groups is None: - uniform_samples.uniform_() - else: - sample_idx = 0 - for seq_group in seq_groups: - seq_ids = seq_group.seq_ids - stride = len(seq_ids) * num_samples - assert seq_group.generator is not None - uniform_samples[:, sample_idx:sample_idx + - stride].uniform_(generator=seq_group.generator) - sample_idx += stride - batch_next_token_ids, success = flashinfer_top_k_top_p_sampling( - probs, - uniform_samples, - top_ks, - top_ps, - ) - if not success.all(): - warnings.warn("FlashInfer rejection sampling failed, fallback.", - stacklevel=1) - probs = flashinfer.sampling.top_k_renorm_prob(probs, top_ks) - probs = flashinfer.sampling.top_p_renorm_prob(probs, top_ps) - batch_next_token_ids = flashinfer.sampling.sampling_from_probs( - probs, uniform_samples[0]) - return batch_next_token_ids.view(-1, num_samples) - - -def get_pythonized_sample_results( - sample_result_args: SampleResultArgsType) -> SampleResultType: - '''This function consumes GPU-side sampler results and computes - Pythonized CPU-side sampler results (GPU -> CPU sync.) - - Single-step scheduling: this function is invoked at sampling-time - for immediate Pythonization. - - Multi-step scheduling: Pythonization is deferred until after multiple - GPU-side steps have been completed. - - Args: - sample_result_args: GPU-side inputs to the Pythonization process - - Returns: - Pythonized sampler results - ''' - - ( - sample_metadata, - sampling_metadata, - greedy_samples, - multinomial_samples, - beam_search_logprobs, - sample_results_dict, - ) = ( - sample_result_args.sample_metadata, - sample_result_args.sampling_metadata, - sample_result_args.greedy_samples, - sample_result_args.multinomial_samples, - sample_result_args.beam_search_logprobs, - sample_result_args.sample_results_dict, - ) - - for sampling_type in SamplingType: - if sampling_type not in sample_metadata: - continue - (seq_group_id, seq_groups) = sample_metadata[sampling_type] - if sampling_type == SamplingType.GREEDY: - sample_results = _greedy_sample(seq_groups, greedy_samples) - elif sampling_type in (SamplingType.RANDOM, SamplingType.RANDOM_SEED): - sample_results = _random_sample(seq_groups, - multinomial_samples[sampling_type]) - elif sampling_type == SamplingType.BEAM: - sample_results = _beam_search_sample(seq_groups, - beam_search_logprobs) - sample_results_dict.update(zip(seq_group_id, sample_results)) - - return [ - sample_results_dict.get(i, ([], [])) - for i in range(len(sampling_metadata.seq_groups)) - ] - - -def _sample_with_torch( - probs: torch.Tensor, - logprobs: torch.Tensor, - sampling_metadata: SamplingMetadata, - sampling_tensors: SamplingTensors, - include_gpu_probs_tensor: bool, - modify_greedy_probs: bool, -) -> SampleReturnType: - '''Torch-oriented _sample() implementation. - - Single-step scheduling: - * Perform GPU-side sampling computation - * Immediately Pythonize sampling result - - Multi-step scheduling: - * Perform GPU-side sampling computation - * Defer Pythonization & preserve GPU-side - tensors required for Pythonization - ''' - - categorized_seq_group_ids: Dict[SamplingType, - List[int]] = {t: [] - for t in SamplingType} - categorized_sample_indices = sampling_metadata.categorized_sample_indices - for i, seq_group in enumerate(sampling_metadata.seq_groups): - sampling_params = seq_group.sampling_params - sampling_type = sampling_params.sampling_type - categorized_seq_group_ids[sampling_type].append(i) - - sample_results_dict: SampleResultsDictType = {} - sample_metadata: SampleMetadataType = {} - multinomial_samples: MultinomialSamplesType = {} - greedy_samples: Optional[torch.Tensor] = None - beam_search_logprobs: Optional[torch.Tensor] = None - - # Create output tensor for sampled token ids. - if include_gpu_probs_tensor: - sampled_token_ids_tensor = torch.full((logprobs.shape[0], 1), - VLLM_INVALID_TOKEN_ID, - dtype=torch.long, - device=logprobs.device) - else: - sampled_token_ids_tensor = None - - # Counterintiutively, having two loops here is actually faster. - # The first loop can run without waiting on GPU<->CPU sync. - for sampling_type in SamplingType: - sample_indices = categorized_sample_indices[sampling_type] - num_tokens = len(sample_indices) - if num_tokens == 0: - continue - - seq_group_id = categorized_seq_group_ids[sampling_type] - seq_groups = [sampling_metadata.seq_groups[i] for i in seq_group_id] - sample_metadata[sampling_type] = (seq_group_id, seq_groups) - long_sample_indices = sample_indices.long() - if sampling_type == SamplingType.GREEDY: - greedy_samples = torch.argmax(logprobs[long_sample_indices], - dim=-1) - - if sampled_token_ids_tensor is not None: - # Store sampled tokens in output tensor. - sampled_token_ids_tensor[ - long_sample_indices] = greedy_samples.unsqueeze(-1) - - if modify_greedy_probs: - # If required, modify the probabilities such that sampling from - # the modified distribution would always sample the argmax - # token id. - _modify_greedy_probs_inplace(logprobs, probs, - long_sample_indices, - greedy_samples) - - elif sampling_type in (SamplingType.RANDOM, SamplingType.RANDOM_SEED): - max_n_in_batch = 1 - for seq_group in seq_groups: - if seq_group.is_prompt: - sampling_params = seq_group.sampling_params - max_n_in_batch = max(max_n_in_batch, sampling_params.n) - seq_groups_arg = (None if sampling_type == SamplingType.RANDOM else - seq_groups) - - if flashinfer_top_k_top_p_sampling is not None: - multinomial_samples[ - sampling_type] = _top_k_top_p_multinomial_with_flashinfer( - probs[long_sample_indices], - sampling_tensors.top_ks[long_sample_indices], - sampling_tensors.top_ps[long_sample_indices], - max_n_in_batch, - seq_groups_arg, - ) - else: - multinomial_samples[sampling_type] = _multinomial( - probs[long_sample_indices], - max_n_in_batch, - seq_groups=seq_groups_arg) - - if sampled_token_ids_tensor is not None: - # Store sampled tokens in output tensor. - sampled_token_ids_tensor[long_sample_indices] = \ - multinomial_samples[sampling_type].to(torch.long) - - elif sampling_type == SamplingType.BEAM: - beam_search_logprobs = logprobs[sample_indices] - else: - raise ValueError(f"Unsupported sampling type: {sampling_type}") - - # Encapsulate arguments for computing Pythonized sampler - # results, whether deferred or otherwise. - maybe_deferred_args = SampleResultArgsType( - sampling_metadata=sampling_metadata, - sample_metadata=sample_metadata, - multinomial_samples=multinomial_samples, - greedy_samples=greedy_samples, - beam_search_logprobs=beam_search_logprobs, - sample_results_dict=sample_results_dict) - - if not sampling_metadata.skip_sampler_cpu_output: - # GPU<->CPU sync happens here. - # This also converts the sampler output to a Python object. - # Return Pythonized sampler result & sampled token ids - return get_pythonized_sample_results( - maybe_deferred_args), sampled_token_ids_tensor - else: - # Defer sampler result Pythonization; return deferred - # Pythonization args & sampled token ids - return ( - maybe_deferred_args, - sampled_token_ids_tensor, - ) - - -def _sample( - probs: torch.Tensor, - logprobs: torch.Tensor, - sampling_metadata: SamplingMetadata, - sampling_tensors: SamplingTensors, - include_gpu_probs_tensor: bool, - modify_greedy_probs: bool, -) -> SampleReturnType: - """ - Args: - probs: (num_query_tokens_in_batch, num_vocab) - logprobs: (num_query_tokens_in_batch, num_vocab) - sampling_metadata: The metadata for a batch for sampling. - sampling_tensors: Tensors that include sampling related metadata. - - Returns: - (next_token_ids, parent_seq_ids) for each seq group in a batch. - If sampling is skipped, it returns ([], []) - sampled_token_ids_tensor: A tensor of sampled token ids. - """ - return _sample_with_torch( - probs, - logprobs, - sampling_metadata, - sampling_tensors, - include_gpu_probs_tensor=include_gpu_probs_tensor, - modify_greedy_probs=modify_greedy_probs, - ) - - -def _get_ranks(x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: - """ - This function calculates the ranks of the chosen tokens in a logprob tensor. - - Args: - x (torch.Tensor): 2D logprob tensor of shape (N, M) - where N is the no. of tokens and M is the vocab dim. - indices (torch.Tensor): List of chosen token indices. - - Returns: - torch.Tensor: 1D tensor of shape (N,) where N is the no. of tokens. - Each element in the returned tensor represents the rank - of the chosen token in the input logprob tensor. - """ - vals = x[torch.arange(0, len(x), device=x.device, dtype=indices.dtype), - indices] - result = (x > vals[:, None]) - del vals - return result.sum(1).add_(1) - - -def get_logprobs( - logprobs: torch.Tensor, - sampling_metadata: SamplingMetadata, - sample_results: SampleResultType, -) -> Tuple[List[Optional[PromptLogprobs]], List[SampleLogprobs]]: - """Return sample logprobs and prompt logprobs. - - The logic consists of 3 parts. - - Select indices to compute logprob from, ranks of token ids, and - the top k token ids from logprobs. - - Compute prompt logprobs if required. - - Compute sample logprobs if required. - - Args: - logprobs: (num_query_tokens_across_batch, num_vocab). Each query token's - logprob per vocab. Sequence groups' query tokens are batched in a - single flattened tensor. For example, assuming there are N - seq groups, it is sorted by prefill tokens for seq_group_1 (if - prompt logprob is enabled), decode tokens for seq_group_1 (if - sampling is required), prefill tokens for seq_group_2, ... - sampling_metadata: The sampling metadata. - sample_results: (num_seq_groups) The tuple of (next_token_ids, - parent_ids) for each sequence group. When beam search is enabled, - sample_results can contain different number of seq_ids from - sampling_metadata.seq_groups. It is because beam search creates - 2 * BEAM_WIDTH number of samples (whereas there are only up to - BEAM_WIDTH number of seq_ids). - - Returns: - A tuple of prompt and sample logprobs per sequence group in a batch. - """ - # The index of query token to calculate logprobs. It includes both - # prompt and sample logprob indices. - query_indices: List[int] = [] - # The next token ids to get the logprob value from. - next_token_ids: List[int] = [] - # The largest requested number of logprobs. We find logprobs as many as the - # largest num logprobs in this API. If every logprobs is None, it will be - # set to -1. - largest_num_logprobs = -1 - - # Select indices to compute logprob from, ranks of token ids, and the top - # k token ids from logprobs. - for (seq_group, sample_result) in zip(sampling_metadata.seq_groups, - sample_results): - sampling_params = seq_group.sampling_params - - # Update indices and tokens for prompt logprobs. - if (seq_group.is_prompt - and sampling_params.prompt_logprobs is not None): - largest_num_logprobs = max(largest_num_logprobs, - sampling_params.prompt_logprobs) - next_prompt_tokens = _get_next_prompt_tokens(seq_group) - query_indices.extend(seq_group.prompt_logprob_indices) - next_token_ids.extend(next_prompt_tokens) - - # Update indices and next tokenes for sample logprob. - if seq_group.do_sample: - token_ids, parent_seq_ids = sample_result - # NOTE: We cannot directly use sample_indices because - # sample_indices only contain parent seq_ids of a previous step. - # The current step may have different number of seq_ids, and - # we can obtain it from `sample_result[1]`. - query_idx = seq_group.sample_indices[0] - query_indices.extend( - [query_idx + parent_id for parent_id in parent_seq_ids]) - next_token_ids.extend(token_ids) - - if sampling_params.logprobs is not None: - largest_num_logprobs = max(largest_num_logprobs, - sampling_params.logprobs) - - assert len(next_token_ids) == len(query_indices) - +"""A layer that samples the next tokens from the model's outputs.""" +import itertools +import warnings +from dataclasses import dataclass +from importlib.util import find_spec +from math import inf +from typing import Dict, List, Optional, Tuple, Union + +import msgspec +import torch +import torch.nn as nn + +import vllm.envs as envs +from vllm.model_executor.sampling_metadata import (SamplingMetadata, + SamplingTensors, + SequenceGroupToSample) +from vllm.sampling_params import SamplingType +from vllm.sequence import (VLLM_INVALID_TOKEN_ID, + CompletionSequenceGroupOutput, Logprob, + PromptLogprobs, SampleLogprobs, SequenceOutput) +from vllm.spec_decode.metrics import SpecDecodeWorkerMetrics + +if envs.VLLM_USE_FLASHINFER_SAMPLER and find_spec("flashinfer"): + import flashinfer.sampling + # yapf: disable + from flashinfer.sampling import ( + top_k_top_p_sampling_from_probs as flashinfer_top_k_top_p_sampling) + + # yapf: enable +else: + flashinfer_top_k_top_p_sampling = None + +# (num_token_ids, num_parent_ids) per sequence group. +SampleResultType = List[Tuple[List[int], List[int]]] + +# Types of temporary data structures used for +# computing sample_result +SampleMetadataType = Dict[SamplingType, Tuple[List[int], + List[SequenceGroupToSample]]] +MultinomialSamplesType = Dict[SamplingType, torch.Tensor] +SampleResultsDictType = Dict[int, Tuple[List[int], List[int]]] + + +# Encapsulates temporary data structures for computing +# sample_result. +# +# * For multi-step scheduling: must be returned +# by `Sampler.forward()` and used later to compute the pythonized +# sample_result +# +# * For single-step scheduling: consumed immediately +# inside `Sampler.forward()` to compute pythonized sample_result. +@dataclass +class SampleResultArgsType: + sample_metadata: SampleMetadataType + multinomial_samples: MultinomialSamplesType + sample_results_dict: SampleResultsDictType + sampling_metadata: SamplingMetadata + greedy_samples: Optional[torch.Tensor] + beam_search_logprobs: Optional[torch.Tensor] + + +# Union of non-deferred (single-step scheduling) +# vs deferred (multi-step scheduling) +# sample result types +MaybeDeferredSampleResultType = Union[SampleResultType, SampleResultArgsType] + +# Abbreviation of the _sample() return type +SampleReturnType = Tuple[MaybeDeferredSampleResultType, Optional[torch.Tensor]] + + +class SamplerOutput( + msgspec.Struct, + omit_defaults=True, # type: ignore[call-arg] + array_like=True): # type: ignore[call-arg] + """For each sequence group, we generate a list of SequenceOutput object, + each of which contains one possible candidate for the next token. + + This data structure implements methods, so it can be used like a list, but + also has optional fields for device tensors. + """ + + outputs: List[CompletionSequenceGroupOutput] + + # On-device tensor containing probabilities of each token. + sampled_token_probs: Optional[torch.Tensor] = None + + # On-device tensor containing the logprobs of each token. + logprobs: Optional["torch.Tensor"] = None + + # Holds either (1) the pythonized sampler result (single-step scheduling) + # or (2) what will be arguments for later deferred pythonization of the + # sampler result (muliti-step scheduling) + deferred_sample_results_args: Optional[SampleResultArgsType] = None + + # On-device tensor containing the sampled token ids. + sampled_token_ids: Optional[torch.Tensor] = None + # CPU tensor containing the sampled token ids. Used during multi-step to + # return the sampled token ids from last rank to AsyncLLMEngine to be + # 'broadcasted' to all other PP ranks for next step. + sampled_token_ids_cpu: Optional[torch.Tensor] = None + + # Spec decode metrics populated by workers. + spec_decode_worker_metrics: Optional[SpecDecodeWorkerMetrics] = None + + # Optional last hidden states from the model. + hidden_states: Optional[torch.Tensor] = None + + # Optional prefill hidden states from the model + # (used for models like EAGLE). + prefill_hidden_states: Optional[torch.Tensor] = None + + # Time taken in the forward pass for this across all workers + model_forward_time: Optional[float] = None + + # Time taken in the model execute function. This will include model forward, + # block/sync across workers, cpu-gpu sync time and sampling time. + model_execute_time: Optional[float] = None + + def __getitem__(self, idx: int): + return self.outputs[idx] + + def __setitem__(self, idx: int, value): + self.outputs[idx] = value + + def __len__(self): + return len(self.outputs) + + def __eq__(self, other: object): + return isinstance(other, + self.__class__) and self.outputs == other.outputs + + def __repr__(self) -> str: + """Show the shape of a tensor instead of its values to reduce noise. + """ + sampled_token_probs_repr = ("None" if self.sampled_token_probs is None + else self.sampled_token_probs.shape) + sampled_token_ids_repr = ("None" if self.sampled_token_ids is None else + self.sampled_token_ids.shape) + return ( + f"SamplerOutput(outputs={self.outputs}, " + f"sampled_token_probs={sampled_token_probs_repr}, " + f"sampled_token_ids={sampled_token_ids_repr}, " + f"spec_decode_worker_metrics={self.spec_decode_worker_metrics})") + + +class Sampler(nn.Module): + """Samples the next tokens from the model's outputs. + + This layer does the following: + 1. Discard the hidden states that are not used for sampling (i.e., all + tokens except the final one in each prompt). + 2. Compute the logits for the next tokens. + 3. Apply presence, frequency and repetition penalties. + 4. Apply temperature scaling. + 5. Apply top-p and top-k truncation. + 6. Sample the next tokens. + Here, each sequence group within the batch can have different sampling + parameters (e.g., sampling method, temperature, top-p, top-k, etc.). + + The structure of the logits tensor is coupled with the seq_groups in + sampling_metadata. Typically, each sequence in each seq_group has one row in + logits for the next token to be sampled; however, for a seq_group with a + prompt request with the prompt_logprobs sampling parameter, there are rows + in logits for each token in the input prompt. + """ + + def __init__(self): + super().__init__() + + # Whether or not the SamplerOutput should have on-device tensors + # containing the sampled token ids and probabilities. This is used by + # speculative decoding. + self.include_gpu_probs_tensor = False + self.should_modify_greedy_probs_inplace = False + + def _init_sampling_tensors( + self, + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + ): + """The goal here is to reuse sampling tensors between similar decode + runs. This is possible because sampling logic does not change between + decodes of the same sequences. + """ + _, vocab_size = logits.shape + + # First free any existing stored sampling tensors. + # This is necessary because some sampling tensors may + # have pinned memory. + self._sampling_tensors = None + + # Initialize new sampling tensors + (sampling_tensors, do_penalties, do_top_p_top_k, + do_min_p) = SamplingTensors.from_sampling_metadata( + sampling_metadata, vocab_size, logits.device, logits.dtype) + + self._sampling_tensors = sampling_tensors + self._do_penalties = do_penalties + self._do_top_p_top_k = do_top_p_top_k + self._do_min_p = do_min_p + + def forward( + self, + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + ) -> Optional[SamplerOutput]: + """ + Single-step scheduling: + * Perform GPU-side sampling computation & compute + GPU-side logprobs tensor + * Pythonize sampling result & logprobs tensor + + Multi-step scheduling: + * Perform GPU-side sampling computation & compute + GPU-side logprobs tensor + * Defer Pythonization of sampling result & logprobs + tensor + * Encapsulate arguments required for deferred Pythonization + in the :class:`SamplerOutput` structure + + Args: + logits: (num_tokens, vocab_size). + sampling_metadata: Metadata for sampling. + """ + assert logits is not None + _, vocab_size = logits.shape + + # Prepare sampling tensors with pinned memory to avoid blocking. + if not sampling_metadata.reuse_sampling_tensors: + self._init_sampling_tensors(logits, sampling_metadata) + elif self._do_penalties: + # In this case, the sampling tensors logic depends on + # "output_tokens" of a sequence. As a result, we cannot + # reuse sampling tensors, since "output_tokens" changes + # between decode runs. + self._init_sampling_tensors(logits, sampling_metadata) + + assert self._sampling_tensors is not None + sampling_tensors = self._sampling_tensors + do_penalties = self._do_penalties + do_top_p_top_k = self._do_top_p_top_k + do_min_p = self._do_min_p + + logits = _apply_min_tokens_penalty(logits, sampling_metadata) + + # Apply presence and frequency penalties. + if do_penalties: + logits = _apply_penalties(logits, sampling_tensors.prompt_tokens, + sampling_tensors.output_tokens, + sampling_tensors.presence_penalties, + sampling_tensors.frequency_penalties, + sampling_tensors.repetition_penalties) + + # Use float32 to apply temperature scaling. + # Use in-place division to avoid creating a new tensor. + logits = logits.to(torch.float) + logits.div_(sampling_tensors.temperatures.unsqueeze(dim=1)) + + if do_top_p_top_k and flashinfer_top_k_top_p_sampling is None: + logits = _apply_top_k_top_p(logits, sampling_tensors.top_ps, + sampling_tensors.top_ks) + + if do_min_p: + logits = _apply_min_p(logits, sampling_tensors.min_ps) + + # We use float32 for probabilities and log probabilities. + # Compute the probabilities. + probs = torch.softmax(logits, dim=-1, dtype=torch.float) + # Compute the log probabilities. + logprobs = torch.log_softmax(logits, dim=-1, dtype=torch.float) + + # Sample the next tokens. + maybe_deferred_sample_results, maybe_sampled_tokens_tensor = _sample( + probs, + logprobs, + sampling_metadata, + sampling_tensors, + include_gpu_probs_tensor=self.include_gpu_probs_tensor, + modify_greedy_probs=self._should_modify_greedy_probs_inplace, + ) + + if self.include_gpu_probs_tensor: + # Since we will defer sampler result Pythonization, + # preserve GPU-side tensors in support of later + # deferred pythonization of logprobs + assert maybe_sampled_tokens_tensor is not None + on_device_tensors = (probs, logprobs, maybe_sampled_tokens_tensor) + else: + # Since Pythonization has already happened, don't preserve + # GPU-side tensors. + on_device_tensors = None + + # Get the logprobs query results. + prompt_logprobs = None + sample_logprobs = None + if not sampling_metadata.skip_sampler_cpu_output: + # Pythonize logprobs now (GPU -> CPU); do not defer. + assert not isinstance(maybe_deferred_sample_results, + SampleResultArgsType) + prompt_logprobs, sample_logprobs = get_logprobs( + logprobs, sampling_metadata, maybe_deferred_sample_results) + + return _build_sampler_output( + maybe_deferred_sample_results, + sampling_metadata, + prompt_logprobs, + sample_logprobs, + on_device_tensors=on_device_tensors, + skip_sampler_cpu_output=sampling_metadata.skip_sampler_cpu_output) + + @property + def _should_modify_greedy_probs_inplace(self) -> bool: + """Whether or not the sampler should modify the probability distribution + of greedily-sampled tokens such that multinomial sampling would sample + the greedily-sampled token. + + In other words, if True then we set the probability of the greedily- + sampled token to 1. + + This is used by speculative decoding, which requires that the sampling + method be encoded into the probability distribution. + """ + return self.should_modify_greedy_probs_inplace + + +def _get_bin_counts_and_mask( + tokens: torch.Tensor, + vocab_size: int, + num_seqs: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + # Compute the bin counts for the tokens. + # vocab_size + 1 for padding. + bin_counts = torch.zeros((num_seqs, vocab_size + 1), + dtype=torch.long, + device=tokens.device) + bin_counts.scatter_add_(1, tokens, torch.ones_like(tokens)) + bin_counts = bin_counts[:, :vocab_size] + mask = bin_counts > 0 + + return bin_counts, mask + + +def _apply_min_tokens_penalty( + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, +) -> torch.Tensor: + """Apply min_tokens penalty which sets stop tokens to -inf if min_tokens + have not been generated yet + """ + # list of indices in logits that will be set to -inf + logits_to_penalize: List[Tuple[int, int]] = [] + logits_applied = 0 + for seq_group in sampling_metadata.seq_groups: + seq_ids = seq_group.seq_ids + sampling_params = seq_group.sampling_params + + sample_indices = seq_group.sample_indices + logits_applied += len(sample_indices) + len( + seq_group.prompt_logprob_indices) + if not seq_group.do_sample: + continue + + start_idx = sample_indices[0] + min_tokens = sampling_params.min_tokens + token_ids_to_penalize = sampling_params.all_stop_token_ids + if min_tokens > 0 and token_ids_to_penalize: + seqs_to_penalize: List[int] = [] + for j, seq_id in enumerate(seq_ids): + seq_data = seq_group.seq_data[seq_id] + if len(seq_data.output_token_ids_array) < min_tokens: + seqs_to_penalize.append(j) + + if seqs_to_penalize: + # convert to the index into logits + seqs_to_penalize = [start_idx + j for j in seqs_to_penalize] + # itertools.product pairs each seq index with every token id + logits_to_penalize.extend( + itertools.product(seqs_to_penalize, token_ids_to_penalize)) + + if logits_to_penalize: + # use zip and * to group indices along each dimension + # eg. [ (1,2), (1,3), (5,6) ] -> ( (1,1,5), (2,3,6) ) + logits[tuple(zip(*logits_to_penalize))] = -float("inf") + + # verifies that no rows in logits were missed unexpectedly + assert logits_applied == logits.shape[0] + return logits + + +def _apply_penalties(logits: torch.Tensor, prompt_tokens_tensor: torch.Tensor, + output_tokens_tensor: torch.Tensor, + presence_penalties: torch.Tensor, + frequency_penalties: torch.Tensor, + repetition_penalties: torch.Tensor) -> torch.Tensor: + num_seqs, vocab_size = logits.shape + _, prompt_mask = _get_bin_counts_and_mask(prompt_tokens_tensor, vocab_size, + num_seqs) + output_bin_counts, output_mask = _get_bin_counts_and_mask( + output_tokens_tensor, vocab_size, num_seqs) + + repetition_penalties = repetition_penalties[:, None].repeat(1, vocab_size) + repetition_penalties[~(prompt_mask | output_mask)] = 1.0 + logits = torch.where(logits > 0, logits / repetition_penalties, + logits * repetition_penalties) + + # We follow the definition in OpenAI API. + # Refer to https://platform.openai.com/docs/api-reference/parameter-details + logits -= frequency_penalties.unsqueeze_(dim=1) * output_bin_counts + logits -= presence_penalties.unsqueeze_(dim=1) * output_mask + return logits + + +def _apply_top_k_top_p( + logits: torch.Tensor, + p: torch.Tensor, + k: torch.Tensor, +) -> torch.Tensor: + logits_sort, logits_idx = logits.sort(dim=-1, descending=False) + + # Apply top-k. + top_k_mask = logits_sort.size(1) - k.to(torch.long) + # Get all the top_k values. + top_k_mask = logits_sort.gather(1, top_k_mask.unsqueeze(dim=1)) + top_k_mask = logits_sort < top_k_mask + logits_sort.masked_fill_(top_k_mask, -float("inf")) + + # Apply top-p. + probs_sort = logits_sort.softmax(dim=-1) + probs_sum = probs_sort.cumsum(dim=-1) + top_p_mask = probs_sum <= 1 - p.unsqueeze(dim=1) + # at least one + top_p_mask[:, -1] = False + logits_sort.masked_fill_(top_p_mask, -float("inf")) + + # Re-sort the probabilities. + logits = torch.empty_like(logits_sort).scatter_(dim=-1, + index=logits_idx, + src=logits_sort) + return logits + + +def _apply_min_p( + logits: torch.Tensor, + min_p: torch.Tensor, +) -> torch.Tensor: + """ + Adapted from + https://github.com/oobabooga/text-generation-webui/blob/3146124ec01f02c8fb1650a6517cf1b60b537aaf/modules/sampler_hijack.py#L16C17-L16C17 + """ + probs = torch.softmax(logits, dim=-1) + top_probs, _ = probs.max(dim=-1, keepdim=True) + scaled_min_p = min_p.unsqueeze_(dim=1) * top_probs + tokens_to_remove = probs < scaled_min_p + logits = logits.masked_fill_(tokens_to_remove, -float("inf")) + + return logits + + +def _greedy_sample( + selected_seq_groups: List[SequenceGroupToSample], + samples: torch.Tensor, +) -> SampleResultType: + """Run greedy sampling on a given samples. + + Args: + selected_seq_groups: A list of sequence groups batched. + samples: (num_selected_samples,) A tensor of samples. The length of + samples could be smaller than selected_seq_groups if + seq_group.do_sample is False. + Returns: + Tuple of (next_token_ids, parent_ids). The length of returned list is + same as the length of selected_seq_groups. If the corresponding + seq_group has do_sample=False, tuple contains ([], []) + """ + samples_lst = samples.tolist() + sample_idx = 0 + results: SampleResultType = [] + for seq_group in selected_seq_groups: + if not seq_group.do_sample: + results.append(([], [])) + continue + + seq_ids = seq_group.seq_ids + num_parent_seqs = len(seq_ids) + assert num_parent_seqs == 1, ( + "Greedy sampling should have only one seq.") + parent_ids = list(range(num_parent_seqs)) + next_token_ids = [samples_lst[sample_idx]] + results.append((next_token_ids, parent_ids)) + sample_idx += num_parent_seqs + return results + + +def _random_sample( + selected_seq_groups: List[SequenceGroupToSample], + random_samples: torch.Tensor, +) -> SampleResultType: + """Run random sampling on a given samples. + + Args: + selected_seq_groups: A list of sequence groups batched. + random_samples: (num_selected_samples,) A tensor of samples. The + length of samples could be smaller than selected_seq_groups if + seq_group.do_sample is False. + Returns: + Tuple of (next_token_ids, parent_ids). The length of returned list is + same as the length of selected_seq_groups. If the corresponding + seq_group has do_sample=False, tuple contains ([], []) + """ + # Find the maximum n value of the prompt phase requests. + random_samples = random_samples.cpu() + sample_idx = 0 + results: SampleResultType = [] + for seq_group in selected_seq_groups: + if not seq_group.do_sample: + results.append(([], [])) + continue + + seq_ids = seq_group.seq_ids + sampling_params = seq_group.sampling_params + is_prompt = seq_group.is_prompt + num_parent_seqs = len(seq_ids) + if is_prompt: + # Prompt phase. + parent_ids = [0] * sampling_params.n + next_token_ids = random_samples[ + sample_idx, :sampling_params.n].tolist() + else: + # Generation phase. + parent_ids = list(range(num_parent_seqs)) + next_token_ids = random_samples[sample_idx:sample_idx + + num_parent_seqs, 0].tolist() + results.append((next_token_ids, parent_ids)) + sample_idx += num_parent_seqs + return results + + +def _beam_search_sample( + selected_seq_groups: List[SequenceGroupToSample], + logprobs: torch.Tensor, +) -> SampleResultType: + """Run beam sampling on a given samples. + + Args: + selected_seq_groups: A list of sequence groups batched. + logprobs: (num_selected_samples, vocab_size,) A tensor of logprob + on selected sample indices. + Returns: + Tuple of (next_token_ids, parent_ids). The length of returned list is + same as the length of selected_seq_groups. If the corresponding + seq_group has do_sample=False, tuple contains ([], []) + """ + # We sample 2 * beam_width candidates to make sure that with high + # probability we can get `beam_width` candidates in addition to + # the finished sequences for the next iteration. See + # https://github.com/tensorflow/tensor2tensor/blob/bafdc1b67730430d38d6ab802cbd51f9d053ba2e/tensor2tensor/utils/beam_search.py#L557-L563 + # for details. See also HF reference: + # https://github.com/huggingface/transformers/blob/a4dd53d88e4852f023332d284ff07a01afcd5681/src/transformers/generation/utils.py#L3063-L3065 + # + # NOTE: Beam search is not vectorized, so its speed can be slower than + # other sampling methods. + sample_idx = 0 + results: SampleResultType = [] + for seq_group in selected_seq_groups: + if not seq_group.do_sample: + results.append(([], [])) + continue + + is_prompt = seq_group.is_prompt + seq_ids, sampling_params = seq_group.seq_ids, seq_group.sampling_params + num_parent_seqs = len(seq_ids) + beam_width = sampling_params.n + seq_group_logprobs = logprobs[sample_idx:sample_idx + num_parent_seqs] + if is_prompt: + # Prompt phase. + assert num_parent_seqs == 1, ( + "Prompt input should have only one seq.") + parent_ids = [0] * (2 * beam_width) + _, next_token_ids = torch.topk(seq_group_logprobs[0], + 2 * beam_width) + next_token_ids = next_token_ids.tolist() + else: + # Generation phase. + cumulative_logprobs: List[float] = [ + seq_group.seq_data[seq_id].cumulative_logprob + for seq_id in seq_ids + ] + cumulative_logprobs_tensor = torch.tensor( + cumulative_logprobs, + dtype=torch.float, + device=seq_group_logprobs.device) + seq_group_logprobs = (seq_group_logprobs + + cumulative_logprobs_tensor.unsqueeze(dim=1)) + _, topk_ids = torch.topk(seq_group_logprobs.flatten(), + 2 * beam_width) + topk_ids = topk_ids.tolist() + vocab_size = seq_group_logprobs.size(-1) + parent_ids = [i // vocab_size for i in topk_ids] + next_token_ids = [i % vocab_size for i in topk_ids] + results.append((next_token_ids, parent_ids)) + sample_idx += num_parent_seqs + assert sample_idx == logprobs.size(0) + return results + + +# torch.multinomial forces a GPU<->CPU sync. +# Therefore, we use an optimized implementation instead. +# Note that we always sample with replacement. +# probs will be modified in place, but this is fine, as we pass +# in a copy already. +def _multinomial( + probs: torch.Tensor, + num_samples: int, + seq_groups: Optional[List[SequenceGroupToSample]] = None, +) -> torch.Tensor: + if num_samples > 1: + probs = probs.repeat_interleave(num_samples, dim=0) + q = torch.empty_like(probs) + if seq_groups is None: + q.exponential_() + else: + sample_idx = 0 + for seq_group in seq_groups: + seq_ids = seq_group.seq_ids + stride = len(seq_ids) * num_samples + assert seq_group.generator is not None + q[sample_idx:sample_idx + + stride].exponential_(generator=seq_group.generator) + sample_idx += stride + return probs.div_(q).argmax(dim=1).view(-1, num_samples) + + +def _top_k_top_p_multinomial_with_flashinfer( + probs: torch.Tensor, top_ks: torch.Tensor, top_ps: torch.Tensor, + num_samples: int, seq_groups: Optional[List[SequenceGroupToSample]]): + max_top_k_round = 32 + if num_samples > 1: + probs = probs.repeat_interleave(num_samples, dim=0) + top_ks = top_ks.repeat_interleave(num_samples) + top_ps = top_ps.repeat_interleave(num_samples) + batch_size = probs.shape[0] + uniform_samples = torch.empty((max_top_k_round, batch_size), + device=probs.device) + if seq_groups is None: + uniform_samples.uniform_() + else: + sample_idx = 0 + for seq_group in seq_groups: + seq_ids = seq_group.seq_ids + stride = len(seq_ids) * num_samples + assert seq_group.generator is not None + uniform_samples[:, sample_idx:sample_idx + + stride].uniform_(generator=seq_group.generator) + sample_idx += stride + batch_next_token_ids, success = flashinfer_top_k_top_p_sampling( + probs, + uniform_samples, + top_ks, + top_ps, + ) + if not success.all(): + warnings.warn("FlashInfer rejection sampling failed, fallback.", + stacklevel=1) + probs = flashinfer.sampling.top_k_renorm_prob(probs, top_ks) + probs = flashinfer.sampling.top_p_renorm_prob(probs, top_ps) + batch_next_token_ids = flashinfer.sampling.sampling_from_probs( + probs, uniform_samples[0]) + return batch_next_token_ids.view(-1, num_samples) + + +def get_pythonized_sample_results( + sample_result_args: SampleResultArgsType) -> SampleResultType: + '''This function consumes GPU-side sampler results and computes + Pythonized CPU-side sampler results (GPU -> CPU sync.) + + Single-step scheduling: this function is invoked at sampling-time + for immediate Pythonization. + + Multi-step scheduling: Pythonization is deferred until after multiple + GPU-side steps have been completed. + + Args: + sample_result_args: GPU-side inputs to the Pythonization process + + Returns: + Pythonized sampler results + ''' + + ( + sample_metadata, + sampling_metadata, + greedy_samples, + multinomial_samples, + beam_search_logprobs, + sample_results_dict, + ) = ( + sample_result_args.sample_metadata, + sample_result_args.sampling_metadata, + sample_result_args.greedy_samples, + sample_result_args.multinomial_samples, + sample_result_args.beam_search_logprobs, + sample_result_args.sample_results_dict, + ) + + for sampling_type in SamplingType: + if sampling_type not in sample_metadata: + continue + (seq_group_id, seq_groups) = sample_metadata[sampling_type] + if sampling_type == SamplingType.GREEDY: + sample_results = _greedy_sample(seq_groups, greedy_samples) + elif sampling_type in (SamplingType.RANDOM, SamplingType.RANDOM_SEED): + sample_results = _random_sample(seq_groups, + multinomial_samples[sampling_type]) + elif sampling_type == SamplingType.BEAM: + sample_results = _beam_search_sample(seq_groups, + beam_search_logprobs) + sample_results_dict.update(zip(seq_group_id, sample_results)) + + return [ + sample_results_dict.get(i, ([], [])) + for i in range(len(sampling_metadata.seq_groups)) + ] + + +def _sample_with_torch( + probs: torch.Tensor, + logprobs: torch.Tensor, + sampling_metadata: SamplingMetadata, + sampling_tensors: SamplingTensors, + include_gpu_probs_tensor: bool, + modify_greedy_probs: bool, +) -> SampleReturnType: + '''Torch-oriented _sample() implementation. + + Single-step scheduling: + * Perform GPU-side sampling computation + * Immediately Pythonize sampling result + + Multi-step scheduling: + * Perform GPU-side sampling computation + * Defer Pythonization & preserve GPU-side + tensors required for Pythonization + ''' + + categorized_seq_group_ids: Dict[SamplingType, + List[int]] = {t: [] + for t in SamplingType} + categorized_sample_indices = sampling_metadata.categorized_sample_indices + for i, seq_group in enumerate(sampling_metadata.seq_groups): + sampling_params = seq_group.sampling_params + sampling_type = sampling_params.sampling_type + categorized_seq_group_ids[sampling_type].append(i) + + sample_results_dict: SampleResultsDictType = {} + sample_metadata: SampleMetadataType = {} + multinomial_samples: MultinomialSamplesType = {} + greedy_samples: Optional[torch.Tensor] = None + beam_search_logprobs: Optional[torch.Tensor] = None + + # Create output tensor for sampled token ids. + if include_gpu_probs_tensor: + sampled_token_ids_tensor = torch.full((logprobs.shape[0], 1), + VLLM_INVALID_TOKEN_ID, + dtype=torch.long, + device=logprobs.device) + else: + sampled_token_ids_tensor = None + + # Counterintiutively, having two loops here is actually faster. + # The first loop can run without waiting on GPU<->CPU sync. + for sampling_type in SamplingType: + sample_indices = categorized_sample_indices[sampling_type] + num_tokens = len(sample_indices) + if num_tokens == 0: + continue + + seq_group_id = categorized_seq_group_ids[sampling_type] + seq_groups = [sampling_metadata.seq_groups[i] for i in seq_group_id] + sample_metadata[sampling_type] = (seq_group_id, seq_groups) + long_sample_indices = sample_indices.long() + if sampling_type == SamplingType.GREEDY: + greedy_samples = torch.argmax(logprobs[long_sample_indices], + dim=-1) + + if sampled_token_ids_tensor is not None: + # Store sampled tokens in output tensor. + sampled_token_ids_tensor[ + long_sample_indices] = greedy_samples.unsqueeze(-1) + + if modify_greedy_probs: + # If required, modify the probabilities such that sampling from + # the modified distribution would always sample the argmax + # token id. + _modify_greedy_probs_inplace(logprobs, probs, + long_sample_indices, + greedy_samples) + + elif sampling_type in (SamplingType.RANDOM, SamplingType.RANDOM_SEED): + max_n_in_batch = 1 + for seq_group in seq_groups: + if seq_group.is_prompt: + sampling_params = seq_group.sampling_params + max_n_in_batch = max(max_n_in_batch, sampling_params.n) + seq_groups_arg = (None if sampling_type == SamplingType.RANDOM else + seq_groups) + + if flashinfer_top_k_top_p_sampling is not None: + multinomial_samples[ + sampling_type] = _top_k_top_p_multinomial_with_flashinfer( + probs[long_sample_indices], + sampling_tensors.top_ks[long_sample_indices], + sampling_tensors.top_ps[long_sample_indices], + max_n_in_batch, + seq_groups_arg, + ) + else: + multinomial_samples[sampling_type] = _multinomial( + probs[long_sample_indices], + max_n_in_batch, + seq_groups=seq_groups_arg) + + if sampled_token_ids_tensor is not None: + # Store sampled tokens in output tensor. + sampled_token_ids_tensor[long_sample_indices] = \ + multinomial_samples[sampling_type].to(torch.long) + + elif sampling_type == SamplingType.BEAM: + beam_search_logprobs = logprobs[sample_indices] + else: + raise ValueError(f"Unsupported sampling type: {sampling_type}") + + # Encapsulate arguments for computing Pythonized sampler + # results, whether deferred or otherwise. + maybe_deferred_args = SampleResultArgsType( + sampling_metadata=sampling_metadata, + sample_metadata=sample_metadata, + multinomial_samples=multinomial_samples, + greedy_samples=greedy_samples, + beam_search_logprobs=beam_search_logprobs, + sample_results_dict=sample_results_dict) + + if not sampling_metadata.skip_sampler_cpu_output: + # GPU<->CPU sync happens here. + # This also converts the sampler output to a Python object. + # Return Pythonized sampler result & sampled token ids + return get_pythonized_sample_results( + maybe_deferred_args), sampled_token_ids_tensor + else: + # Defer sampler result Pythonization; return deferred + # Pythonization args & sampled token ids + return ( + maybe_deferred_args, + sampled_token_ids_tensor, + ) + + +def _sample( + probs: torch.Tensor, + logprobs: torch.Tensor, + sampling_metadata: SamplingMetadata, + sampling_tensors: SamplingTensors, + include_gpu_probs_tensor: bool, + modify_greedy_probs: bool, +) -> SampleReturnType: + """ + Args: + probs: (num_query_tokens_in_batch, num_vocab) + logprobs: (num_query_tokens_in_batch, num_vocab) + sampling_metadata: The metadata for a batch for sampling. + sampling_tensors: Tensors that include sampling related metadata. + + Returns: + (next_token_ids, parent_seq_ids) for each seq group in a batch. + If sampling is skipped, it returns ([], []) + sampled_token_ids_tensor: A tensor of sampled token ids. + """ + return _sample_with_torch( + probs, + logprobs, + sampling_metadata, + sampling_tensors, + include_gpu_probs_tensor=include_gpu_probs_tensor, + modify_greedy_probs=modify_greedy_probs, + ) + + +def _get_ranks(x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: + """ + This function calculates the ranks of the chosen tokens in a logprob tensor. + + Args: + x (torch.Tensor): 2D logprob tensor of shape (N, M) + where N is the no. of tokens and M is the vocab dim. + indices (torch.Tensor): List of chosen token indices. + + Returns: + torch.Tensor: 1D tensor of shape (N,) where N is the no. of tokens. + Each element in the returned tensor represents the rank + of the chosen token in the input logprob tensor. + """ + vals = x[torch.arange(0, len(x), device=x.device, dtype=indices.dtype), + indices] + result = (x > vals[:, None]) + del vals + return result.sum(1).add_(1) + + +def get_logprobs( + logprobs: torch.Tensor, + sampling_metadata: SamplingMetadata, + sample_results: SampleResultType, +) -> Tuple[List[Optional[PromptLogprobs]], List[SampleLogprobs]]: + """Return sample logprobs and prompt logprobs. + + The logic consists of 3 parts. + - Select indices to compute logprob from, ranks of token ids, and + the top k token ids from logprobs. + - Compute prompt logprobs if required. + - Compute sample logprobs if required. + + Args: + logprobs: (num_query_tokens_across_batch, num_vocab). Each query token's + logprob per vocab. Sequence groups' query tokens are batched in a + single flattened tensor. For example, assuming there are N + seq groups, it is sorted by prefill tokens for seq_group_1 (if + prompt logprob is enabled), decode tokens for seq_group_1 (if + sampling is required), prefill tokens for seq_group_2, ... + sampling_metadata: The sampling metadata. + sample_results: (num_seq_groups) The tuple of (next_token_ids, + parent_ids) for each sequence group. When beam search is enabled, + sample_results can contain different number of seq_ids from + sampling_metadata.seq_groups. It is because beam search creates + 2 * BEAM_WIDTH number of samples (whereas there are only up to + BEAM_WIDTH number of seq_ids). + + Returns: + A tuple of prompt and sample logprobs per sequence group in a batch. + """ + # The index of query token to calculate logprobs. It includes both + # prompt and sample logprob indices. + query_indices: List[int] = [] + # The next token ids to get the logprob value from. + next_token_ids: List[int] = [] + # The largest requested number of logprobs. We find logprobs as many as the + # largest num logprobs in this API. If every logprobs is None, it will be + # set to -1. + largest_num_logprobs = -1 + + # Select indices to compute logprob from, ranks of token ids, and the top + # k token ids from logprobs. + for (seq_group, sample_result) in zip(sampling_metadata.seq_groups, + sample_results): + sampling_params = seq_group.sampling_params + + # Update indices and tokens for prompt logprobs. + if (seq_group.is_prompt + and sampling_params.prompt_logprobs is not None): + largest_num_logprobs = max(largest_num_logprobs, + sampling_params.prompt_logprobs) + next_prompt_tokens = _get_next_prompt_tokens(seq_group) + query_indices.extend(seq_group.prompt_logprob_indices) + next_token_ids.extend(next_prompt_tokens) + + # Update indices and next tokenes for sample logprob. + if seq_group.do_sample: + token_ids, parent_seq_ids = sample_result + # NOTE: We cannot directly use sample_indices because + # sample_indices only contain parent seq_ids of a previous step. + # The current step may have different number of seq_ids, and + # we can obtain it from `sample_result[1]`. + query_idx = seq_group.sample_indices[0] + query_indices.extend( + [query_idx + parent_id for parent_id in parent_seq_ids]) + next_token_ids.extend(token_ids) + + if sampling_params.logprobs is not None: + largest_num_logprobs = max(largest_num_logprobs, + sampling_params.logprobs) + + assert len(next_token_ids) == len(query_indices) + selected_logprobs, ranks = None, None - top_logprobs, top_token_ids = None, None - - # If largest_num_logprobs == -1, i.e. no logprobs are requested, we can - # skip the whole logprob calculation. + top_logprobs, top_token_ids = None, None + + # If largest_num_logprobs == -1, i.e. no logprobs are requested, we can + # skip the whole logprob calculation. if query_indices and largest_num_logprobs >= 0: - query_indices_gpu = torch.tensor(query_indices, device=logprobs.device) - next_token_ids_gpu = torch.tensor(next_token_ids, - device=logprobs.device) - - # (num_selected_query_tokens, num_logprobs). Note that query_indices can - # contain duplicates if beam search is enabled. - selected_logprobs = logprobs[[ - query_indices_gpu, - next_token_ids_gpu, - ]] - ranks = _get_ranks( - logprobs[query_indices_gpu], - next_token_ids_gpu, - ) - assert selected_logprobs.shape[0] == ranks.shape[0] - - # We need to compute top k only if there exists logprobs > 0. - if largest_num_logprobs > 0: - # Logprobs of topk tokens for a batch of sequence groups. - # (num_query_tokens_across_batch). - top_logprobs, top_token_ids = torch.topk(logprobs, - largest_num_logprobs, - dim=-1) - top_logprobs = top_logprobs.to('cpu') - top_token_ids = top_token_ids.to('cpu') - - selected_logprobs = selected_logprobs.to('cpu') - ranks = ranks.to('cpu') - - # Find prompt/sample logprobs. - prompt_logprobs_per_seq_group: List[Optional[PromptLogprobs]] = [] - sample_logprobs_per_seq_group: List[SampleLogprobs] = [] - top_logprob_idx = 0 - selected_logprobs_idx = 0 - - for seq_group, sample_result in zip(sampling_metadata.seq_groups, - sample_results): - (prompt_logprobs, top_logprob_idx, - selected_logprobs_idx) = _get_prompt_logprob_if_needed( - seq_group, selected_logprobs, ranks, top_token_ids, top_logprobs, - selected_logprobs_idx, top_logprob_idx) - prompt_logprobs_per_seq_group.append(prompt_logprobs) - - (sampled_logprobs, top_logprob_idx, - selected_logprobs_idx) = _get_sampled_logprob_if_needed( - seq_group, sample_result, selected_logprobs, ranks, top_token_ids, - top_logprobs, selected_logprobs_idx, top_logprob_idx) - sample_logprobs_per_seq_group.append(sampled_logprobs) - - return prompt_logprobs_per_seq_group, sample_logprobs_per_seq_group - - -def _get_prompt_logprob_if_needed( - seq_group: SequenceGroupToSample, - selected_logprobs: torch.Tensor, - ranks: torch.Tensor, - top_token_ids: torch.Tensor, - top_logprobs: torch.Tensor, - selected_logprobs_idx: int, - top_logprob_idx: int, -): - """Compute the prompt logprob from a sequence group if needed.""" - sampling_params = seq_group.sampling_params - is_prompt = seq_group.is_prompt - - # Find prompt logprobs + query_indices_gpu = torch.tensor(query_indices, device=logprobs.device) + next_token_ids_gpu = torch.tensor(next_token_ids, + device=logprobs.device) + + # (num_selected_query_tokens, num_logprobs). Note that query_indices can + # contain duplicates if beam search is enabled. + selected_logprobs = logprobs[[ + query_indices_gpu, + next_token_ids_gpu, + ]] + ranks = _get_ranks( + logprobs[query_indices_gpu], + next_token_ids_gpu, + ) + assert selected_logprobs.shape[0] == ranks.shape[0] + + # We need to compute top k only if there exists logprobs > 0. + if largest_num_logprobs > 0: + # Logprobs of topk tokens for a batch of sequence groups. + # (num_query_tokens_across_batch). + top_logprobs, top_token_ids = torch.topk(logprobs, + largest_num_logprobs, + dim=-1) + top_logprobs = top_logprobs.to('cpu') + top_token_ids = top_token_ids.to('cpu') + + selected_logprobs = selected_logprobs.to('cpu') + ranks = ranks.to('cpu') + + # Find prompt/sample logprobs. + prompt_logprobs_per_seq_group: List[Optional[PromptLogprobs]] = [] + sample_logprobs_per_seq_group: List[SampleLogprobs] = [] + top_logprob_idx = 0 + selected_logprobs_idx = 0 + + for seq_group, sample_result in zip(sampling_metadata.seq_groups, + sample_results): + (prompt_logprobs, top_logprob_idx, + selected_logprobs_idx) = _get_prompt_logprob_if_needed( + seq_group, selected_logprobs, ranks, top_token_ids, top_logprobs, + selected_logprobs_idx, top_logprob_idx) + prompt_logprobs_per_seq_group.append(prompt_logprobs) + + (sampled_logprobs, top_logprob_idx, + selected_logprobs_idx) = _get_sampled_logprob_if_needed( + seq_group, sample_result, selected_logprobs, ranks, top_token_ids, + top_logprobs, selected_logprobs_idx, top_logprob_idx) + sample_logprobs_per_seq_group.append(sampled_logprobs) + + return prompt_logprobs_per_seq_group, sample_logprobs_per_seq_group + + +def _get_prompt_logprob_if_needed( + seq_group: SequenceGroupToSample, + selected_logprobs: torch.Tensor, + ranks: torch.Tensor, + top_token_ids: torch.Tensor, + top_logprobs: torch.Tensor, + selected_logprobs_idx: int, + top_logprob_idx: int, +): + """Compute the prompt logprob from a sequence group if needed.""" + sampling_params = seq_group.sampling_params + is_prompt = seq_group.is_prompt + + # Find prompt logprobs prompt_logprobs: Optional[PromptLogprobs] = None if is_prompt and sampling_params.prompt_logprobs is not None: query_len = seq_group.query_len @@ -1093,243 +1093,243 @@ def _get_prompt_logprob_if_needed( for idx, (token_id, output_index) in enumerate(zip( next_prompt_tokens, seq_group.prompt_logprob_output_indices)): - # Calculate the prompt logprob of the real prompt tokens. - # {token_id: (logprob, rank_from_vocab)} - prompt_logprobs_dict: Dict[int, Tuple[float, int]] = { - token_id: (selected_logprob_items[idx], rank_items[idx]) - } - + # Calculate the prompt logprob of the real prompt tokens. + # {token_id: (logprob, rank_from_vocab)} + prompt_logprobs_dict: Dict[int, Tuple[float, int]] = { + token_id: (selected_logprob_items[idx], rank_items[idx]) + } + # Add top K prompt logprobs along with its rank. if num_logprobs > 0: assert top_token_ids is not None assert top_logprobs is not None top_ids = top_token_ids[ - top_logprob_idx, :num_logprobs].tolist() - top_probs = top_logprobs[ - top_logprob_idx, :num_logprobs].tolist() - # Top K is already sorted by rank, so we can use 1 ~ - # num_logprobs + 1 for rank. - top_ranks = range(1, num_logprobs + 1) - prompt_logprobs_dict.update({ - top_id: (top_prob, rank) - for top_id, top_prob, rank in zip(top_ids, top_probs, - top_ranks) - }) + top_logprob_idx, :num_logprobs].tolist() + top_probs = top_logprobs[ + top_logprob_idx, :num_logprobs].tolist() + # Top K is already sorted by rank, so we can use 1 ~ + # num_logprobs + 1 for rank. + top_ranks = range(1, num_logprobs + 1) + prompt_logprobs_dict.update({ + top_id: (top_prob, rank) + for top_id, top_prob, rank in zip(top_ids, top_probs, + top_ranks) + }) prompt_logprobs[output_index] = { token_id: Logprob(*logprob_and_rank) for token_id, logprob_and_rank in prompt_logprobs_dict.items() } - # + 1 to go to the next prompt token. - top_logprob_idx += 1 - - # + len(next_prompt_tokens) to go to the next prompt. + # + 1 to go to the next prompt token. + top_logprob_idx += 1 + + # + len(next_prompt_tokens) to go to the next prompt. selected_logprobs_idx += len(next_prompt_tokens) - return prompt_logprobs, top_logprob_idx, selected_logprobs_idx - - -def _get_sampled_logprob_if_needed( - seq_group: SequenceGroupToSample, - sample_result: Tuple[List[int], List[int]], - selected_logprobs: torch.Tensor, - ranks: torch.Tensor, - top_token_ids: torch.Tensor, - top_logprobs: torch.Tensor, - selected_logprobs_idx: int, - top_logprob_idx: int, -): - """Compute the sample logprob if needed.""" - seq_ids = seq_group.seq_ids - num_logprobs = seq_group.sampling_params.logprobs - sampled_logprobs: SampleLogprobs = [] - next_token_ids, parent_seq_ids = sample_result - - if seq_group.do_sample: - assert len(next_token_ids) > 0 - if num_logprobs is None: - for next_token_id in next_token_ids: - # Use a dummy logprob - sampled_logprobs.append({next_token_id: Logprob(inf)}) - else: - # Pre-select items from tensor. tolist() is faster than repetitive - # `.item()` calls. - selected_logprob_items = selected_logprobs[ - selected_logprobs_idx:selected_logprobs_idx + - len(next_token_ids)].tolist() - rank_items = ranks[selected_logprobs_idx:selected_logprobs_idx + - len(next_token_ids)].tolist() - for idx, (next_token_id, parent_id) in enumerate( - zip(next_token_ids, parent_seq_ids)): - # Get the logprob of a sampled token. - sampled_logprobs_dict = { - next_token_id: - (selected_logprob_items[idx], rank_items[idx]) - } - if num_logprobs is not None and num_logprobs > 0: - # Get top K logprobs. - top_ids = top_token_ids[top_logprob_idx + - parent_id, :num_logprobs].tolist() - top_probs = top_logprobs[ - top_logprob_idx + parent_id, :num_logprobs].tolist() - # Top K is already sorted by rank, so we can use 1 ~ - # num_logprobs + 1 for rank. - top_ranks = range(1, num_logprobs + 1) - sampled_logprobs_dict.update({ - top_id: (top_prob, rank) - for top_id, top_prob, rank in zip( - top_ids, top_probs, top_ranks) - }) - - sampled_logprobs.append({ - token_id: Logprob(*logprob_and_rank) - for token_id, logprob_and_rank in - sampled_logprobs_dict.items() - }) - - # NOTE: This part of code is not intuitive. `selected_logprobs` include - # logprobs for the current step, which has len(next_token_ids) tokens - # per sequence group. `logprobs` includes logprobs from the previous - # steps, which has len(seq_ids) tokens per sequence group. - - # Iterate to the next sequence group in a batch. - selected_logprobs_idx += len(next_token_ids) - # Iterate to the next sequence group in a batch. - top_logprob_idx += len(seq_ids) - return sampled_logprobs, top_logprob_idx, selected_logprobs_idx - - -def _modify_greedy_probs_inplace(logprobs: torch.Tensor, probs: torch.Tensor, - sample_indices: torch.Tensor, - greedy_samples: torch.Tensor) -> None: - """Modify the probability distributions of the greedily-sampled tokens such - that each sampled token has a "probability" of 1.0. This is required by - speculative decoding, which depends on the sampling method being encoded - within the probability distribution for correctness. - - # Why do we only need to do this for greedy sampling? - - vLLM's sampler performs the following steps for greedy or multinomial - (random) sampling: - 1. Get logits from model. - 2. Modify logits according to per-sequence sampling parameters. - - Multiply by temperature, top-k and top-p masking, penalize tokens - according to their frequency, etc. - 3. Sample a token. - - Random sampling simply samples from the modified probability - distribution. - - Greedy sampling performs `argmax` to obtain the token with the - highest likelihood. - - Ignoring greedy sampling for a moment, we find that the computed probability - distribution has the following property: we can sample from it independently - and find that the token sampled by the Sampler has a frequency corresponding - to how often we see it in our sampling. In other words, for tokens sampled - with vLLM's random SamplingType, the computed probability distribution - encodes the sampling methodology completely. - - Greedy sampling does not normally have this property. vLLM modifies logits - according to sampling params, then performs `argmax`, then returns the - sampled token and the computed probability distribution. If we sample from - the distribution, we'll find the likelihood of the greedily-sampled token - is not always 1.0. - - Since lossless speculative decoding requires that the sampling methodology - be encoded within the probability distribution, we are motivated to modify - the probability distribution such that the sampled token has probability 1 - when speculative decoding is used. - - NOTE: Alternatively, we could use an extremely low temperature to achieve - greedy sampling using multinomial computation and unite the codepaths. This - has implications on the overall design of the sampler, e.g. how to record - accurate logprobs for the user, so this improvement is deferred to later. - """ - # NOTE: logprobs are not modified so they can be returned to the user. - probs[sample_indices, :] = 0 - probs[sample_indices, greedy_samples] = 1.0 - - -def _build_sampler_output( - maybe_deferred_sample_results: MaybeDeferredSampleResultType, - sampling_metadata: SamplingMetadata, - prompt_logprobs: Optional[List[Optional[PromptLogprobs]]], - sample_logprobs: Optional[List[SampleLogprobs]], - on_device_tensors: Optional[Tuple[torch.Tensor, torch.Tensor, - torch.Tensor]], - skip_sampler_cpu_output: bool = False, -) -> SamplerOutput: - """Construct Python objects with the output of sampling. - - Args: - on_device_tensors: Tuple containing on-device tensors with the - probabilities used in sampling and the sampled token ids. This - allows post-processing without copies to CPU/serialization, e.g. in - speculative decoding rejection sampling. - """ - sampler_output: List[CompletionSequenceGroupOutput] = [] - - if skip_sampler_cpu_output: - assert isinstance(maybe_deferred_sample_results, SampleResultArgsType) - deferred_sample_results_args = maybe_deferred_sample_results - else: - assert prompt_logprobs is not None - assert sample_logprobs is not None - assert not isinstance(maybe_deferred_sample_results, - SampleResultArgsType) - deferred_sample_results_args = None - - for (seq_group, sample_result, group_prompt_logprobs, - group_sample_logprobs) in zip(sampling_metadata.seq_groups, - maybe_deferred_sample_results, - prompt_logprobs, sample_logprobs): - seq_ids = seq_group.seq_ids - next_token_ids, parent_ids = sample_result - seq_outputs: List[SequenceOutput] = [] - for parent_id, next_token_id, logprobs in zip( - parent_ids, next_token_ids, group_sample_logprobs): - seq_outputs.append( - SequenceOutput(seq_ids[parent_id], next_token_id, - logprobs)) - sampler_output.append( - CompletionSequenceGroupOutput(seq_outputs, - group_prompt_logprobs)) - - # If not specified, store None values in SamplerOutput. - if on_device_tensors is not None: - (sampled_token_probs, logprobs_tensor, - sampled_token_ids) = on_device_tensors - else: - sampled_token_probs, logprobs_tensor, sampled_token_ids = (None, None, - None) - - return SamplerOutput( - outputs=sampler_output, - sampled_token_probs=sampled_token_probs, - sampled_token_ids=sampled_token_ids, - logprobs=logprobs_tensor, - deferred_sample_results_args=deferred_sample_results_args) - - -def _get_next_prompt_tokens(seq_group: SequenceGroupToSample) -> List[int]: - """Get a list of next prompt tokens to compute logprob from a - given sequence group. - - It is used to compute prompt logprob. Imagine you have logprob for each - query token. Query token needs to know the next prompt token id to compute - prompt logprob. This is a helper to obtain next prompt token ids. - - This API has to be used only when the caller knows seq_group is in prefill - stage. - - Returns: - A list of next prompt tokens to compute logprob. - """ - assert seq_group.is_prompt, ( - "Caller should ensure the sequence group is in a prefill stage.") - seq_ids = seq_group.seq_ids - query_len = seq_group.query_len - assert query_len is not None - # prompt has only 1 seq id. - assert len(seq_ids) == 1 - seq_data = seq_group.seq_data[seq_ids[0]] + return prompt_logprobs, top_logprob_idx, selected_logprobs_idx + + +def _get_sampled_logprob_if_needed( + seq_group: SequenceGroupToSample, + sample_result: Tuple[List[int], List[int]], + selected_logprobs: torch.Tensor, + ranks: torch.Tensor, + top_token_ids: torch.Tensor, + top_logprobs: torch.Tensor, + selected_logprobs_idx: int, + top_logprob_idx: int, +): + """Compute the sample logprob if needed.""" + seq_ids = seq_group.seq_ids + num_logprobs = seq_group.sampling_params.logprobs + sampled_logprobs: SampleLogprobs = [] + next_token_ids, parent_seq_ids = sample_result + + if seq_group.do_sample: + assert len(next_token_ids) > 0 + if num_logprobs is None: + for next_token_id in next_token_ids: + # Use a dummy logprob + sampled_logprobs.append({next_token_id: Logprob(inf)}) + else: + # Pre-select items from tensor. tolist() is faster than repetitive + # `.item()` calls. + selected_logprob_items = selected_logprobs[ + selected_logprobs_idx:selected_logprobs_idx + + len(next_token_ids)].tolist() + rank_items = ranks[selected_logprobs_idx:selected_logprobs_idx + + len(next_token_ids)].tolist() + for idx, (next_token_id, parent_id) in enumerate( + zip(next_token_ids, parent_seq_ids)): + # Get the logprob of a sampled token. + sampled_logprobs_dict = { + next_token_id: + (selected_logprob_items[idx], rank_items[idx]) + } + if num_logprobs is not None and num_logprobs > 0: + # Get top K logprobs. + top_ids = top_token_ids[top_logprob_idx + + parent_id, :num_logprobs].tolist() + top_probs = top_logprobs[ + top_logprob_idx + parent_id, :num_logprobs].tolist() + # Top K is already sorted by rank, so we can use 1 ~ + # num_logprobs + 1 for rank. + top_ranks = range(1, num_logprobs + 1) + sampled_logprobs_dict.update({ + top_id: (top_prob, rank) + for top_id, top_prob, rank in zip( + top_ids, top_probs, top_ranks) + }) + + sampled_logprobs.append({ + token_id: Logprob(*logprob_and_rank) + for token_id, logprob_and_rank in + sampled_logprobs_dict.items() + }) + + # NOTE: This part of code is not intuitive. `selected_logprobs` include + # logprobs for the current step, which has len(next_token_ids) tokens + # per sequence group. `logprobs` includes logprobs from the previous + # steps, which has len(seq_ids) tokens per sequence group. + + # Iterate to the next sequence group in a batch. + selected_logprobs_idx += len(next_token_ids) + # Iterate to the next sequence group in a batch. + top_logprob_idx += len(seq_ids) + return sampled_logprobs, top_logprob_idx, selected_logprobs_idx + + +def _modify_greedy_probs_inplace(logprobs: torch.Tensor, probs: torch.Tensor, + sample_indices: torch.Tensor, + greedy_samples: torch.Tensor) -> None: + """Modify the probability distributions of the greedily-sampled tokens such + that each sampled token has a "probability" of 1.0. This is required by + speculative decoding, which depends on the sampling method being encoded + within the probability distribution for correctness. + + # Why do we only need to do this for greedy sampling? + + vLLM's sampler performs the following steps for greedy or multinomial + (random) sampling: + 1. Get logits from model. + 2. Modify logits according to per-sequence sampling parameters. + - Multiply by temperature, top-k and top-p masking, penalize tokens + according to their frequency, etc. + 3. Sample a token. + - Random sampling simply samples from the modified probability + distribution. + - Greedy sampling performs `argmax` to obtain the token with the + highest likelihood. + + Ignoring greedy sampling for a moment, we find that the computed probability + distribution has the following property: we can sample from it independently + and find that the token sampled by the Sampler has a frequency corresponding + to how often we see it in our sampling. In other words, for tokens sampled + with vLLM's random SamplingType, the computed probability distribution + encodes the sampling methodology completely. + + Greedy sampling does not normally have this property. vLLM modifies logits + according to sampling params, then performs `argmax`, then returns the + sampled token and the computed probability distribution. If we sample from + the distribution, we'll find the likelihood of the greedily-sampled token + is not always 1.0. + + Since lossless speculative decoding requires that the sampling methodology + be encoded within the probability distribution, we are motivated to modify + the probability distribution such that the sampled token has probability 1 + when speculative decoding is used. + + NOTE: Alternatively, we could use an extremely low temperature to achieve + greedy sampling using multinomial computation and unite the codepaths. This + has implications on the overall design of the sampler, e.g. how to record + accurate logprobs for the user, so this improvement is deferred to later. + """ + # NOTE: logprobs are not modified so they can be returned to the user. + probs[sample_indices, :] = 0 + probs[sample_indices, greedy_samples] = 1.0 + + +def _build_sampler_output( + maybe_deferred_sample_results: MaybeDeferredSampleResultType, + sampling_metadata: SamplingMetadata, + prompt_logprobs: Optional[List[Optional[PromptLogprobs]]], + sample_logprobs: Optional[List[SampleLogprobs]], + on_device_tensors: Optional[Tuple[torch.Tensor, torch.Tensor, + torch.Tensor]], + skip_sampler_cpu_output: bool = False, +) -> SamplerOutput: + """Construct Python objects with the output of sampling. + + Args: + on_device_tensors: Tuple containing on-device tensors with the + probabilities used in sampling and the sampled token ids. This + allows post-processing without copies to CPU/serialization, e.g. in + speculative decoding rejection sampling. + """ + sampler_output: List[CompletionSequenceGroupOutput] = [] + + if skip_sampler_cpu_output: + assert isinstance(maybe_deferred_sample_results, SampleResultArgsType) + deferred_sample_results_args = maybe_deferred_sample_results + else: + assert prompt_logprobs is not None + assert sample_logprobs is not None + assert not isinstance(maybe_deferred_sample_results, + SampleResultArgsType) + deferred_sample_results_args = None + + for (seq_group, sample_result, group_prompt_logprobs, + group_sample_logprobs) in zip(sampling_metadata.seq_groups, + maybe_deferred_sample_results, + prompt_logprobs, sample_logprobs): + seq_ids = seq_group.seq_ids + next_token_ids, parent_ids = sample_result + seq_outputs: List[SequenceOutput] = [] + for parent_id, next_token_id, logprobs in zip( + parent_ids, next_token_ids, group_sample_logprobs): + seq_outputs.append( + SequenceOutput(seq_ids[parent_id], next_token_id, + logprobs)) + sampler_output.append( + CompletionSequenceGroupOutput(seq_outputs, + group_prompt_logprobs)) + + # If not specified, store None values in SamplerOutput. + if on_device_tensors is not None: + (sampled_token_probs, logprobs_tensor, + sampled_token_ids) = on_device_tensors + else: + sampled_token_probs, logprobs_tensor, sampled_token_ids = (None, None, + None) + + return SamplerOutput( + outputs=sampler_output, + sampled_token_probs=sampled_token_probs, + sampled_token_ids=sampled_token_ids, + logprobs=logprobs_tensor, + deferred_sample_results_args=deferred_sample_results_args) + + +def _get_next_prompt_tokens(seq_group: SequenceGroupToSample) -> List[int]: + """Get a list of next prompt tokens to compute logprob from a + given sequence group. + + It is used to compute prompt logprob. Imagine you have logprob for each + query token. Query token needs to know the next prompt token id to compute + prompt logprob. This is a helper to obtain next prompt token ids. + + This API has to be used only when the caller knows seq_group is in prefill + stage. + + Returns: + A list of next prompt tokens to compute logprob. + """ + assert seq_group.is_prompt, ( + "Caller should ensure the sequence group is in a prefill stage.") + seq_ids = seq_group.seq_ids + query_len = seq_group.query_len + assert query_len is not None + # prompt has only 1 seq id. + assert len(seq_ids) == 1 + seq_data = seq_group.seq_data[seq_ids[0]] computed_len = seq_data.get_num_computed_tokens() prompt_tokens = seq_data.prompt_token_ids next_prompt_tokens = [] diff --git a/vllm_overrides/model_executor/sampling_metadata.py b/vllm_overrides/model_executor/sampling_metadata.py index 56785643..98ffffdc 100644 --- a/vllm_overrides/model_executor/sampling_metadata.py +++ b/vllm_overrides/model_executor/sampling_metadata.py @@ -1,58 +1,58 @@ -from array import array -from dataclasses import dataclass -from typing import Dict, List, Optional, Tuple - -import torch - -from vllm.sampling_params import SamplingParams, SamplingType -from vllm.sequence import (VLLM_TOKEN_ID_ARRAY_TYPE, SequenceData, - SequenceGroupMetadata) -from vllm.utils import (PyObjectCache, async_tensor_h2d, - is_pin_memory_available, make_tensor_with_pad) - -_SAMPLING_EPS = 1e-5 - - -@dataclass -class SequenceGroupToSample: - # |---------- N-1 iteration --------| - # |---------------- N iteration ---------------------| - # |- tokenA -|......................|-- newTokens ---| - # |---------- context_len ----------| - # |-------------------- seq_len ----------------------| - # |-- query_len ---| - - # Sequence ids for the sequence group in a previous step. - seq_ids: List[int] - sampling_params: SamplingParams - # seq_id -> sequence data. - seq_data: Dict[int, SequenceData] - # The length of the sequence (all tokens seen in the past + new token to - # compute attention) of the sequence group. None if it is in a decode - # stage. - seq_len: Optional[int] - # The length of new query tokens to compute in the current step. None if it - # is in a decode stage. The length of query_len <= seq_len if chunked - # prefill is enabled. - query_len: Optional[int] - # A random number generator for sampling. - generator: Optional[torch.Generator] - # True if the sequence group is in prefill stage. False if it is in a - # decode stage. - is_prompt: bool +from array import array +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple + +import torch + +from vllm.sampling_params import SamplingParams, SamplingType +from vllm.sequence import (VLLM_TOKEN_ID_ARRAY_TYPE, SequenceData, + SequenceGroupMetadata) +from vllm.utils import (PyObjectCache, async_tensor_h2d, + is_pin_memory_available, make_tensor_with_pad) + +_SAMPLING_EPS = 1e-5 + + +@dataclass +class SequenceGroupToSample: + # |---------- N-1 iteration --------| + # |---------------- N iteration ---------------------| + # |- tokenA -|......................|-- newTokens ---| + # |---------- context_len ----------| + # |-------------------- seq_len ----------------------| + # |-- query_len ---| + + # Sequence ids for the sequence group in a previous step. + seq_ids: List[int] + sampling_params: SamplingParams + # seq_id -> sequence data. + seq_data: Dict[int, SequenceData] + # The length of the sequence (all tokens seen in the past + new token to + # compute attention) of the sequence group. None if it is in a decode + # stage. + seq_len: Optional[int] + # The length of new query tokens to compute in the current step. None if it + # is in a decode stage. The length of query_len <= seq_len if chunked + # prefill is enabled. + query_len: Optional[int] + # A random number generator for sampling. + generator: Optional[torch.Generator] + # True if the sequence group is in prefill stage. False if it is in a + # decode stage. + is_prompt: bool # Query token indices from logits. to compute prompt logprob. Empty if # prompt logprob is not required. prompt_logprob_indices: List[int] # Output offsets within this prefill chunk. Sparse diagnostic requests use # this to retain the standard full-length prompt-logprob response shape. prompt_logprob_output_indices: List[int] - # Sample token indices from logits. Empty if sampling is not required. - sample_indices: List[int] - - @property - def do_sample(self): - return len(self.sample_indices) > 0 - + # Sample token indices from logits. Empty if sampling is not required. + sample_indices: List[int] + + @property + def do_sample(self): + return len(self.sample_indices) > 0 + def __post_init__(self): if len(self.prompt_logprob_indices) > 0: assert self.sampling_params.prompt_logprobs is not None @@ -66,142 +66,142 @@ class SequenceGroupToSample: assert all( 0 <= index < self.query_len for index in self.prompt_logprob_output_indices) - - -def gen_seq_group_to_sample_builder(num_seqs: int): - return lambda: SequenceGroupToSample( - seq_ids=[0] * num_seqs, - sampling_params=None, - seq_data=None, # type: ignore - seq_len=0, - query_len=0, - generator=None, + + +def gen_seq_group_to_sample_builder(num_seqs: int): + return lambda: SequenceGroupToSample( + seq_ids=[0] * num_seqs, + sampling_params=None, + seq_data=None, # type: ignore + seq_len=0, + query_len=0, + generator=None, is_prompt=True, prompt_logprob_indices=[], prompt_logprob_output_indices=[], sample_indices=[], - ) - - -class SamplingMetadataCache: - """Used to cache SamplingMetadata objects between scheduler iterations""" - - def __init__(self): - self._seq_group_to_sample_cache: Dict[int, PyObjectCache] = {} - - def get_cached_seq_group_to_sample(self, num_seqs): - if num_seqs not in self._seq_group_to_sample_cache: - self._seq_group_to_sample_cache[num_seqs] = PyObjectCache( - gen_seq_group_to_sample_builder(num_seqs)) - - obj = self._seq_group_to_sample_cache[num_seqs].get_object() - return obj - - def reset(self): - for cache in self._seq_group_to_sample_cache.values(): - cache.reset() - - -class SamplingMetadata: - """Metadata for input sequences. Used in sampler. - - The usage is as follow; - ``` - hidden_states = execute_model(...) - logits = hidden_states[sampling_metadata.selected_token_indices] - sample(logits) - - def sample(logits): - # Use categorized_sample_indices for sampling.... - ``` - - Args: - seq_groups: List of batched sequence groups. - selected_token_indices: (num_query_tokens_to_logprob). Indices to find - logits from the initial model output hidden states. - categorized_sample_indices: SamplingType -> token indices to sample. - Each token indices is 2D tensor of (num_indices, num_indices) where - the first item means the sample index within the returned logit - (before pruning padding), and the second item means the sample - index after pruning using selected_token_indices. - For example, if the returned logit is [1, 2, 3], and we select - [1, 2] for sampling, the pruned logit will be [2, 3]. In this case, - The first tuple is [1, 2] (sampled index within original logit), - and the second tuple is [0, 1] (sampled index within pruned logit). - num_prompts: Number of prompt sequence groups in seq_groups. - skip_sampler_cpu_output: Indicates if we want to skip the GPU=>CPU - serialization of token outputs. - reuse_sampling_tensors: Indicates if we want to reuse sampling - tensors that are part of the sampler forward pass. Currently, - it is mainly used for multi-step decode. - - """ - - def __init__( - self, - seq_groups: List[SequenceGroupToSample], - selected_token_indices: torch.Tensor, - categorized_sample_indices: Dict[SamplingType, torch.Tensor], - num_prompts: int, - skip_sampler_cpu_output: bool = False, - reuse_sampling_tensors: bool = False, - ) -> None: - self.seq_groups = seq_groups - self.selected_token_indices = selected_token_indices - self.categorized_sample_indices = categorized_sample_indices - self.num_prompts = num_prompts - self.skip_sampler_cpu_output = skip_sampler_cpu_output - self.reuse_sampling_tensors = reuse_sampling_tensors - - @staticmethod - def prepare( - seq_group_metadata_list: List[SequenceGroupMetadata], - seq_lens: List[int], - query_lens: List[int], - device: str, - pin_memory: bool, - generators: Optional[Dict[str, torch.Generator]] = None, - cache: Optional[SamplingMetadataCache] = None, - ) -> "SamplingMetadata": - ( - seq_groups, - selected_token_indices, - categorized_sample_indices, - num_prompts, - ) = _prepare_seq_groups(seq_group_metadata_list, seq_lens, query_lens, - device, generators, cache) - selected_token_indices = async_tensor_h2d( - selected_token_indices, - dtype=torch.long, - target_device=device, - pin_memory=pin_memory, - ) - categorized_sample_indices = { - t: async_tensor_h2d( - seq_ids, - dtype=torch.int, - target_device=device, - pin_memory=pin_memory, - ) - for t, seq_ids in categorized_sample_indices.items() - } - - sampling_metadata = SamplingMetadata( - seq_groups=seq_groups, - selected_token_indices=selected_token_indices, - categorized_sample_indices=categorized_sample_indices, - num_prompts=num_prompts, - ) - return sampling_metadata - - def __repr__(self) -> str: - return ( - "SamplingMetadata(" - f"seq_groups={self.seq_groups}, " - f"selected_token_indices={self.selected_token_indices}, " - f"categorized_sample_indices={self.categorized_sample_indices}), ") - - + ) + + +class SamplingMetadataCache: + """Used to cache SamplingMetadata objects between scheduler iterations""" + + def __init__(self): + self._seq_group_to_sample_cache: Dict[int, PyObjectCache] = {} + + def get_cached_seq_group_to_sample(self, num_seqs): + if num_seqs not in self._seq_group_to_sample_cache: + self._seq_group_to_sample_cache[num_seqs] = PyObjectCache( + gen_seq_group_to_sample_builder(num_seqs)) + + obj = self._seq_group_to_sample_cache[num_seqs].get_object() + return obj + + def reset(self): + for cache in self._seq_group_to_sample_cache.values(): + cache.reset() + + +class SamplingMetadata: + """Metadata for input sequences. Used in sampler. + + The usage is as follow; + ``` + hidden_states = execute_model(...) + logits = hidden_states[sampling_metadata.selected_token_indices] + sample(logits) + + def sample(logits): + # Use categorized_sample_indices for sampling.... + ``` + + Args: + seq_groups: List of batched sequence groups. + selected_token_indices: (num_query_tokens_to_logprob). Indices to find + logits from the initial model output hidden states. + categorized_sample_indices: SamplingType -> token indices to sample. + Each token indices is 2D tensor of (num_indices, num_indices) where + the first item means the sample index within the returned logit + (before pruning padding), and the second item means the sample + index after pruning using selected_token_indices. + For example, if the returned logit is [1, 2, 3], and we select + [1, 2] for sampling, the pruned logit will be [2, 3]. In this case, + The first tuple is [1, 2] (sampled index within original logit), + and the second tuple is [0, 1] (sampled index within pruned logit). + num_prompts: Number of prompt sequence groups in seq_groups. + skip_sampler_cpu_output: Indicates if we want to skip the GPU=>CPU + serialization of token outputs. + reuse_sampling_tensors: Indicates if we want to reuse sampling + tensors that are part of the sampler forward pass. Currently, + it is mainly used for multi-step decode. + + """ + + def __init__( + self, + seq_groups: List[SequenceGroupToSample], + selected_token_indices: torch.Tensor, + categorized_sample_indices: Dict[SamplingType, torch.Tensor], + num_prompts: int, + skip_sampler_cpu_output: bool = False, + reuse_sampling_tensors: bool = False, + ) -> None: + self.seq_groups = seq_groups + self.selected_token_indices = selected_token_indices + self.categorized_sample_indices = categorized_sample_indices + self.num_prompts = num_prompts + self.skip_sampler_cpu_output = skip_sampler_cpu_output + self.reuse_sampling_tensors = reuse_sampling_tensors + + @staticmethod + def prepare( + seq_group_metadata_list: List[SequenceGroupMetadata], + seq_lens: List[int], + query_lens: List[int], + device: str, + pin_memory: bool, + generators: Optional[Dict[str, torch.Generator]] = None, + cache: Optional[SamplingMetadataCache] = None, + ) -> "SamplingMetadata": + ( + seq_groups, + selected_token_indices, + categorized_sample_indices, + num_prompts, + ) = _prepare_seq_groups(seq_group_metadata_list, seq_lens, query_lens, + device, generators, cache) + selected_token_indices = async_tensor_h2d( + selected_token_indices, + dtype=torch.long, + target_device=device, + pin_memory=pin_memory, + ) + categorized_sample_indices = { + t: async_tensor_h2d( + seq_ids, + dtype=torch.int, + target_device=device, + pin_memory=pin_memory, + ) + for t, seq_ids in categorized_sample_indices.items() + } + + sampling_metadata = SamplingMetadata( + seq_groups=seq_groups, + selected_token_indices=selected_token_indices, + categorized_sample_indices=categorized_sample_indices, + num_prompts=num_prompts, + ) + return sampling_metadata + + def __repr__(self) -> str: + return ( + "SamplingMetadata(" + f"seq_groups={self.seq_groups}, " + f"selected_token_indices={self.selected_token_indices}, " + f"categorized_sample_indices={self.categorized_sample_indices}), ") + + def _get_prompt_logprob_output_indices( sampling_params: SamplingParams, seq_data: SequenceData, @@ -231,106 +231,106 @@ def _get_prompt_logprob_output_indices( def _prepare_seq_groups( - seq_group_metadata_list: List[SequenceGroupMetadata], - seq_lens: List[int], - query_lens: List[int], - device: str, - generators: Optional[Dict[str, torch.Generator]] = None, - cache: Optional[SamplingMetadataCache] = None, -) -> Tuple[List[SequenceGroupToSample], List[int], Dict[SamplingType, - List[int]], int, ]: - """Prepare sequence groups and indices for sampling. - - Args: - seq_group_metadata_list: A list of sequence group to batch. - seq_lens: A list of sequence lens per sequence group. - Index of prompt len should match with seq_group_metadata_list. - query_lens: A list of query lengths. Prompt lens include the length - of entire prompt tokens, and it could be shorter. - device: A device to use for random number generators, - `SequenceGroupToSample.generator`. - generators: A store of per-request random number generators used - for seeded requests. - - Returns: - seq_groups: A list of sequence group to sample. - selected_token_indices: See the definition from `SamplingMetadata`. - categorized_sample_indices: See the definition from `SamplingMetadata`. - num_prompts: Total number of prompts from `seq_group_metadata_list`. - """ - # Batched sequence groups for the current model forward stsep. - seq_groups: List[SequenceGroupToSample] = [] - # A list of token indices to sample/compute logprob. It is used to - # prune the outcome logits from the model for the performance. - selected_token_indices: List[int] = [] - # Used for selected_token_indices. - model_output_idx = 0 - - # Sampling type -> ( - # indices to sample/prompt logprob within pruned output logits, - # indices to sample within pruned logits) - categorized_sample_indices: Dict[SamplingType, List[int]] = { - t: [] - for t in SamplingType - } - # Index of logits to compute logprob. Logits include both prompt logprob - # and sample logprob indices. - logit_idx = 0 - # Total number of prompts from given sequence groups. + seq_group_metadata_list: List[SequenceGroupMetadata], + seq_lens: List[int], + query_lens: List[int], + device: str, + generators: Optional[Dict[str, torch.Generator]] = None, + cache: Optional[SamplingMetadataCache] = None, +) -> Tuple[List[SequenceGroupToSample], List[int], Dict[SamplingType, + List[int]], int, ]: + """Prepare sequence groups and indices for sampling. + + Args: + seq_group_metadata_list: A list of sequence group to batch. + seq_lens: A list of sequence lens per sequence group. + Index of prompt len should match with seq_group_metadata_list. + query_lens: A list of query lengths. Prompt lens include the length + of entire prompt tokens, and it could be shorter. + device: A device to use for random number generators, + `SequenceGroupToSample.generator`. + generators: A store of per-request random number generators used + for seeded requests. + + Returns: + seq_groups: A list of sequence group to sample. + selected_token_indices: See the definition from `SamplingMetadata`. + categorized_sample_indices: See the definition from `SamplingMetadata`. + num_prompts: Total number of prompts from `seq_group_metadata_list`. + """ + # Batched sequence groups for the current model forward stsep. + seq_groups: List[SequenceGroupToSample] = [] + # A list of token indices to sample/compute logprob. It is used to + # prune the outcome logits from the model for the performance. + selected_token_indices: List[int] = [] + # Used for selected_token_indices. + model_output_idx = 0 + + # Sampling type -> ( + # indices to sample/prompt logprob within pruned output logits, + # indices to sample within pruned logits) + categorized_sample_indices: Dict[SamplingType, List[int]] = { + t: [] + for t in SamplingType + } + # Index of logits to compute logprob. Logits include both prompt logprob + # and sample logprob indices. + logit_idx = 0 + # Total number of prompts from given sequence groups. num_prompts = 0 - - for i, seq_group_metadata in enumerate(seq_group_metadata_list): - seq_ids = seq_group_metadata.seq_data.keys() - - if cache is not None: - sample_obj = cache.get_cached_seq_group_to_sample(len(seq_ids)) - - for j, seq_id in enumerate(seq_ids): - sample_obj.seq_ids[j] = seq_id - + + for i, seq_group_metadata in enumerate(seq_group_metadata_list): + seq_ids = seq_group_metadata.seq_data.keys() + + if cache is not None: + sample_obj = cache.get_cached_seq_group_to_sample(len(seq_ids)) + + for j, seq_id in enumerate(seq_ids): + sample_obj.seq_ids[j] = seq_id + sample_obj.prompt_logprob_indices.clear() sample_obj.prompt_logprob_output_indices.clear() sample_obj.sample_indices.clear() - - sampling_params = seq_group_metadata.sampling_params - is_prompt = seq_group_metadata.is_prompt - generator: Optional[torch.Generator] = None - # If the current seq group is in decode stage, it is None. - seq_len: Optional[int] = None - query_len: Optional[int] = None + + sampling_params = seq_group_metadata.sampling_params + is_prompt = seq_group_metadata.is_prompt + generator: Optional[torch.Generator] = None + # If the current seq group is in decode stage, it is None. + seq_len: Optional[int] = None + query_len: Optional[int] = None prompt_logprob_indices: List[int] = (sample_obj.prompt_logprob_indices if cache is not None else []) prompt_logprob_output_indices: List[int] = ( sample_obj.prompt_logprob_output_indices if cache is not None else []) - sample_indices: List[int] = (sample_obj.sample_indices - if cache is not None else []) - do_sample = seq_group_metadata.do_sample - - if seq_group_metadata.is_prompt: - if sampling_params.seed is not None: - generator = torch.Generator(device=device).manual_seed( - sampling_params.seed) - if generators is not None: - generators[seq_group_metadata.request_id] = generator - - num_prompts += 1 - num_prefill_sample = len(seq_ids) - assert num_prefill_sample == 1 - assert query_lens is not None and seq_lens is not None - query_len, seq_len = query_lens[i], seq_lens[i] - # If we need sampling, exclude num_prefill_sample tokens from - # prompt logprob. - prompt_logprob_len = (query_len - num_prefill_sample - if do_sample else query_len) - sample_len = num_prefill_sample if do_sample else 0 + sample_indices: List[int] = (sample_obj.sample_indices + if cache is not None else []) + do_sample = seq_group_metadata.do_sample + + if seq_group_metadata.is_prompt: + if sampling_params.seed is not None: + generator = torch.Generator(device=device).manual_seed( + sampling_params.seed) + if generators is not None: + generators[seq_group_metadata.request_id] = generator + + num_prompts += 1 + num_prefill_sample = len(seq_ids) + assert num_prefill_sample == 1 + assert query_lens is not None and seq_lens is not None + query_len, seq_len = query_lens[i], seq_lens[i] + # If we need sampling, exclude num_prefill_sample tokens from + # prompt logprob. + prompt_logprob_len = (query_len - num_prefill_sample + if do_sample else query_len) + sample_len = num_prefill_sample if do_sample else 0 else: - # Decode - prompt_logprob_len = 0 - query_len = query_lens[i] if query_lens is not None else 1 - sample_len = len(seq_ids) * query_len if do_sample else 0 - - if sampling_params.seed is not None and generators is not None: + # Decode + prompt_logprob_len = 0 + query_len = query_lens[i] if query_lens is not None else 1 + sample_len = len(seq_ids) * query_len if do_sample else 0 + + if sampling_params.seed is not None and generators is not None: generator = generators.get(seq_group_metadata.request_id) seq_data = next(iter(seq_group_metadata.seq_data.values())) @@ -340,65 +340,65 @@ def _prepare_seq_groups( seq_data, prompt_logprob_len, )) - - # Update indices to select from the model output. - """ - This blocks computes selected_token_indices which is used in the - following way. - - hidden_states = model(...) - logits = hidden_states[selected_token_indices] - """ - + + # Update indices to select from the model output. + """ + This blocks computes selected_token_indices which is used in the + following way. + + hidden_states = model(...) + logits = hidden_states[selected_token_indices] + """ + if sampling_params.prompt_logprobs is not None: selected_token_indices.extend( model_output_idx + output_index for output_index in prompt_logprob_output_indices) - model_output_idx += prompt_logprob_len - if do_sample: - selected_token_indices.extend( - range(model_output_idx, model_output_idx + sample_len)) - model_output_idx += sample_len - - # We now find indices for logprob computation and sampling. - """ - This block computes categorized_sample_indices which is used in the - following way. - - hidden_states = model(...) - logits = hidden_states[selected_token_indices] - def sample(logits): - # Use categorized_sample_indices for sampling. - # prompt_logprob_indices to find prompt logprob indices. - # sample_indices to find sample indices. - """ - + model_output_idx += prompt_logprob_len + if do_sample: + selected_token_indices.extend( + range(model_output_idx, model_output_idx + sample_len)) + model_output_idx += sample_len + + # We now find indices for logprob computation and sampling. + """ + This block computes categorized_sample_indices which is used in the + following way. + + hidden_states = model(...) + logits = hidden_states[selected_token_indices] + def sample(logits): + # Use categorized_sample_indices for sampling. + # prompt_logprob_indices to find prompt logprob indices. + # sample_indices to find sample indices. + """ + if sampling_params.prompt_logprobs is not None: prompt_logprob_indices.extend( range(logit_idx, logit_idx + len(prompt_logprob_output_indices))) logit_idx += len(prompt_logprob_output_indices) - if do_sample: - sample_indices.extend(range(logit_idx, logit_idx + sample_len)) - categorized_sample_indices[sampling_params.sampling_type].extend( - list(range(logit_idx, logit_idx + sample_len))) - logit_idx += sample_len - + if do_sample: + sample_indices.extend(range(logit_idx, logit_idx + sample_len)) + categorized_sample_indices[sampling_params.sampling_type].extend( + list(range(logit_idx, logit_idx + sample_len))) + logit_idx += sample_len + if cache is not None: sample_obj.sampling_params = sampling_params sample_obj.seq_data = seq_group_metadata.seq_data - sample_obj.seq_len = seq_len - sample_obj.query_len = query_len + sample_obj.seq_len = seq_len + sample_obj.query_len = query_len sample_obj.generator = generator sample_obj.is_prompt = is_prompt else: - sample_obj = SequenceGroupToSample( - seq_ids=list(seq_ids), - sampling_params=sampling_params, - seq_data=seq_group_metadata.seq_data, - seq_len=seq_len, - query_len=query_len, - generator=generator, + sample_obj = SequenceGroupToSample( + seq_ids=list(seq_ids), + sampling_params=sampling_params, + seq_data=seq_group_metadata.seq_data, + seq_len=seq_len, + query_len=query_len, + generator=generator, is_prompt=is_prompt, prompt_logprob_indices=list(prompt_logprob_indices), prompt_logprob_output_indices=list( @@ -409,236 +409,236 @@ def _prepare_seq_groups( assert (len(sample_obj.prompt_logprob_indices) == len(sample_obj.prompt_logprob_output_indices)) seq_groups.append(sample_obj) - - if cache is not None: - cache.reset() - - return (seq_groups, selected_token_indices, categorized_sample_indices, - num_prompts) - - -@dataclass -class SamplingTensors: - """Tensors for sampling.""" - - temperatures: torch.Tensor - top_ps: torch.Tensor - top_ks: torch.Tensor - min_ps: torch.Tensor - presence_penalties: torch.Tensor - frequency_penalties: torch.Tensor - repetition_penalties: torch.Tensor - prompt_tokens: torch.Tensor - output_tokens: torch.Tensor - - @classmethod - def from_sampling_metadata( - cls, - sampling_metadata: "SamplingMetadata", - vocab_size: int, - device: torch.device, - dtype: torch.dtype, - ) -> Tuple["SamplingTensors", bool, bool, bool]: - prompt_tokens: List[array] = [] - output_tokens: List[array] = [] - top_ks: List[int] = [] - temperatures: List[float] = [] - top_ps: List[float] = [] - min_ps: List[float] = [] - presence_penalties: List[float] = [] - frequency_penalties: List[float] = [] - repetition_penalties: List[float] = [] - do_penalties = False - do_top_p_top_k = False - do_min_p = False - - assert sampling_metadata.seq_groups is not None - for seq_group in sampling_metadata.seq_groups: - seq_ids = seq_group.seq_ids - sampling_params = seq_group.sampling_params - temperature = sampling_params.temperature - p = sampling_params.presence_penalty - f = sampling_params.frequency_penalty - r = sampling_params.repetition_penalty - top_p = sampling_params.top_p - min_p = sampling_params.min_p - - # k should not be greater than the vocab size. - top_k = min(sampling_params.top_k, vocab_size) - top_k = vocab_size if top_k == -1 else top_k - if temperature < _SAMPLING_EPS: - # NOTE: Zero temperature means deterministic sampling - # (i.e., greedy sampling or beam search). - # Set the temperature to 1 to avoid division by zero. - temperature = 1.0 - if not do_top_p_top_k and (top_p < 1.0 - _SAMPLING_EPS - or top_k != vocab_size): - do_top_p_top_k = True - if not do_min_p and min_p > _SAMPLING_EPS: - do_min_p = True - if not do_penalties and (abs(p) >= _SAMPLING_EPS - or abs(f) >= _SAMPLING_EPS - or abs(r - 1.0) >= _SAMPLING_EPS): - do_penalties = True - - is_prompt = seq_group.is_prompt - if is_prompt and sampling_params.prompt_logprobs is not None: - # For tokens in the prompt that we only need to get - # their logprobs - query_len = seq_group.query_len - assert query_len is not None - prefill_len = len(seq_group.prompt_logprob_indices) - temperatures += [temperature] * prefill_len - top_ps += [top_p] * prefill_len - top_ks += [top_k] * prefill_len - min_ps += [min_p] * prefill_len - presence_penalties += [0] * prefill_len - frequency_penalties += [0] * prefill_len - repetition_penalties += [1] * prefill_len - - if seq_group.do_sample: - sample_lens = len(seq_group.sample_indices) - assert sample_lens >= len(seq_ids) - temperatures += [temperature] * sample_lens - top_ps += [top_p] * sample_lens - top_ks += [top_k] * sample_lens - min_ps += [min_p] * sample_lens - presence_penalties += [p] * sample_lens - frequency_penalties += [f] * sample_lens - repetition_penalties += [r] * sample_lens - - if do_penalties: - for seq_group in sampling_metadata.seq_groups: - seq_ids = seq_group.seq_ids - if (seq_group.is_prompt - and sampling_params.prompt_logprobs is not None): - prefill_len = len(seq_group.prompt_logprob_indices) - prompt_tokens.extend( - array(VLLM_TOKEN_ID_ARRAY_TYPE) - for _ in range(prefill_len)) - output_tokens.extend( - array(VLLM_TOKEN_ID_ARRAY_TYPE) - for _ in range(prefill_len)) - if seq_group.do_sample: - for seq_id in seq_ids: - seq_data = seq_group.seq_data[seq_id] - prompt_tokens.append(seq_data.prompt_token_ids_array) - output_tokens.append(seq_data.output_token_ids_array) - - sampling_tensors = SamplingTensors.from_lists( - temperatures, - top_ps, - top_ks, - min_ps, - presence_penalties, - frequency_penalties, - repetition_penalties, - prompt_tokens, - output_tokens, - vocab_size, - device, - dtype, - ) - return (sampling_tensors, do_penalties, do_top_p_top_k, do_min_p) - - @classmethod - def from_lists( - cls, - temperatures: List[float], - top_ps: List[float], - top_ks: List[int], - min_ps: List[float], - presence_penalties: List[float], - frequency_penalties: List[float], - repetition_penalties: List[float], - prompt_tokens: List[array], - output_tokens: List[array], - vocab_size: int, - device: torch.device, - dtype: torch.dtype, - ) -> "SamplingTensors": - # Note that the performance will be very bad without - # pinned memory. - pin_memory = is_pin_memory_available() - - do_penalties = prompt_tokens or output_tokens - - if do_penalties: - prompt_t = make_tensor_with_pad( - prompt_tokens, - vocab_size, - device="cpu", - dtype=torch.int64, - pin_memory=pin_memory, - ) - output_t = make_tensor_with_pad( - output_tokens, - vocab_size, - device="cpu", - dtype=torch.int64, - pin_memory=pin_memory, - ) - else: - empty_tensor = torch.empty(0, device=device, dtype=torch.long) - prompt_t = empty_tensor - output_t = empty_tensor - - temperatures_t = torch.tensor( - temperatures, - device="cpu", - dtype=dtype, - pin_memory=pin_memory, - ) - top_ps_t = torch.tensor( - top_ps, - device="cpu", - dtype=dtype, - pin_memory=pin_memory, - ) - min_ps_t = torch.tensor( - min_ps, - device="cpu", - dtype=dtype, - pin_memory=pin_memory, - ) - presence_penalties_t = torch.tensor( - presence_penalties, - device="cpu", - dtype=dtype, - pin_memory=pin_memory, - ) - frequency_penalties_t = torch.tensor( - frequency_penalties, - device="cpu", - dtype=dtype, - pin_memory=pin_memory, - ) - repetition_penalties_t = torch.tensor( - repetition_penalties, - device="cpu", - dtype=dtype, - pin_memory=pin_memory, - ) - top_ks_t = torch.tensor( - top_ks, - device="cpu", - dtype=torch.int, - pin_memory=pin_memory, - ) - # Because the memory is pinned, we can do non-blocking - # transfer to device. - - return cls( - temperatures=temperatures_t.to(device=device, non_blocking=True), - top_ps=top_ps_t.to(device=device, non_blocking=True), - top_ks=top_ks_t.to(device=device, non_blocking=True), - min_ps=min_ps_t.to(device=device, non_blocking=True), - presence_penalties=presence_penalties_t.to(device=device, - non_blocking=True), - frequency_penalties=frequency_penalties_t.to(device=device, - non_blocking=True), - repetition_penalties=repetition_penalties_t.to(device=device, - non_blocking=True), - prompt_tokens=prompt_t.to(device=device, non_blocking=True), - output_tokens=output_t.to(device=device, non_blocking=True), - ) + + if cache is not None: + cache.reset() + + return (seq_groups, selected_token_indices, categorized_sample_indices, + num_prompts) + + +@dataclass +class SamplingTensors: + """Tensors for sampling.""" + + temperatures: torch.Tensor + top_ps: torch.Tensor + top_ks: torch.Tensor + min_ps: torch.Tensor + presence_penalties: torch.Tensor + frequency_penalties: torch.Tensor + repetition_penalties: torch.Tensor + prompt_tokens: torch.Tensor + output_tokens: torch.Tensor + + @classmethod + def from_sampling_metadata( + cls, + sampling_metadata: "SamplingMetadata", + vocab_size: int, + device: torch.device, + dtype: torch.dtype, + ) -> Tuple["SamplingTensors", bool, bool, bool]: + prompt_tokens: List[array] = [] + output_tokens: List[array] = [] + top_ks: List[int] = [] + temperatures: List[float] = [] + top_ps: List[float] = [] + min_ps: List[float] = [] + presence_penalties: List[float] = [] + frequency_penalties: List[float] = [] + repetition_penalties: List[float] = [] + do_penalties = False + do_top_p_top_k = False + do_min_p = False + + assert sampling_metadata.seq_groups is not None + for seq_group in sampling_metadata.seq_groups: + seq_ids = seq_group.seq_ids + sampling_params = seq_group.sampling_params + temperature = sampling_params.temperature + p = sampling_params.presence_penalty + f = sampling_params.frequency_penalty + r = sampling_params.repetition_penalty + top_p = sampling_params.top_p + min_p = sampling_params.min_p + + # k should not be greater than the vocab size. + top_k = min(sampling_params.top_k, vocab_size) + top_k = vocab_size if top_k == -1 else top_k + if temperature < _SAMPLING_EPS: + # NOTE: Zero temperature means deterministic sampling + # (i.e., greedy sampling or beam search). + # Set the temperature to 1 to avoid division by zero. + temperature = 1.0 + if not do_top_p_top_k and (top_p < 1.0 - _SAMPLING_EPS + or top_k != vocab_size): + do_top_p_top_k = True + if not do_min_p and min_p > _SAMPLING_EPS: + do_min_p = True + if not do_penalties and (abs(p) >= _SAMPLING_EPS + or abs(f) >= _SAMPLING_EPS + or abs(r - 1.0) >= _SAMPLING_EPS): + do_penalties = True + + is_prompt = seq_group.is_prompt + if is_prompt and sampling_params.prompt_logprobs is not None: + # For tokens in the prompt that we only need to get + # their logprobs + query_len = seq_group.query_len + assert query_len is not None + prefill_len = len(seq_group.prompt_logprob_indices) + temperatures += [temperature] * prefill_len + top_ps += [top_p] * prefill_len + top_ks += [top_k] * prefill_len + min_ps += [min_p] * prefill_len + presence_penalties += [0] * prefill_len + frequency_penalties += [0] * prefill_len + repetition_penalties += [1] * prefill_len + + if seq_group.do_sample: + sample_lens = len(seq_group.sample_indices) + assert sample_lens >= len(seq_ids) + temperatures += [temperature] * sample_lens + top_ps += [top_p] * sample_lens + top_ks += [top_k] * sample_lens + min_ps += [min_p] * sample_lens + presence_penalties += [p] * sample_lens + frequency_penalties += [f] * sample_lens + repetition_penalties += [r] * sample_lens + + if do_penalties: + for seq_group in sampling_metadata.seq_groups: + seq_ids = seq_group.seq_ids + if (seq_group.is_prompt + and sampling_params.prompt_logprobs is not None): + prefill_len = len(seq_group.prompt_logprob_indices) + prompt_tokens.extend( + array(VLLM_TOKEN_ID_ARRAY_TYPE) + for _ in range(prefill_len)) + output_tokens.extend( + array(VLLM_TOKEN_ID_ARRAY_TYPE) + for _ in range(prefill_len)) + if seq_group.do_sample: + for seq_id in seq_ids: + seq_data = seq_group.seq_data[seq_id] + prompt_tokens.append(seq_data.prompt_token_ids_array) + output_tokens.append(seq_data.output_token_ids_array) + + sampling_tensors = SamplingTensors.from_lists( + temperatures, + top_ps, + top_ks, + min_ps, + presence_penalties, + frequency_penalties, + repetition_penalties, + prompt_tokens, + output_tokens, + vocab_size, + device, + dtype, + ) + return (sampling_tensors, do_penalties, do_top_p_top_k, do_min_p) + + @classmethod + def from_lists( + cls, + temperatures: List[float], + top_ps: List[float], + top_ks: List[int], + min_ps: List[float], + presence_penalties: List[float], + frequency_penalties: List[float], + repetition_penalties: List[float], + prompt_tokens: List[array], + output_tokens: List[array], + vocab_size: int, + device: torch.device, + dtype: torch.dtype, + ) -> "SamplingTensors": + # Note that the performance will be very bad without + # pinned memory. + pin_memory = is_pin_memory_available() + + do_penalties = prompt_tokens or output_tokens + + if do_penalties: + prompt_t = make_tensor_with_pad( + prompt_tokens, + vocab_size, + device="cpu", + dtype=torch.int64, + pin_memory=pin_memory, + ) + output_t = make_tensor_with_pad( + output_tokens, + vocab_size, + device="cpu", + dtype=torch.int64, + pin_memory=pin_memory, + ) + else: + empty_tensor = torch.empty(0, device=device, dtype=torch.long) + prompt_t = empty_tensor + output_t = empty_tensor + + temperatures_t = torch.tensor( + temperatures, + device="cpu", + dtype=dtype, + pin_memory=pin_memory, + ) + top_ps_t = torch.tensor( + top_ps, + device="cpu", + dtype=dtype, + pin_memory=pin_memory, + ) + min_ps_t = torch.tensor( + min_ps, + device="cpu", + dtype=dtype, + pin_memory=pin_memory, + ) + presence_penalties_t = torch.tensor( + presence_penalties, + device="cpu", + dtype=dtype, + pin_memory=pin_memory, + ) + frequency_penalties_t = torch.tensor( + frequency_penalties, + device="cpu", + dtype=dtype, + pin_memory=pin_memory, + ) + repetition_penalties_t = torch.tensor( + repetition_penalties, + device="cpu", + dtype=dtype, + pin_memory=pin_memory, + ) + top_ks_t = torch.tensor( + top_ks, + device="cpu", + dtype=torch.int, + pin_memory=pin_memory, + ) + # Because the memory is pinned, we can do non-blocking + # transfer to device. + + return cls( + temperatures=temperatures_t.to(device=device, non_blocking=True), + top_ps=top_ps_t.to(device=device, non_blocking=True), + top_ks=top_ks_t.to(device=device, non_blocking=True), + min_ps=min_ps_t.to(device=device, non_blocking=True), + presence_penalties=presence_penalties_t.to(device=device, + non_blocking=True), + frequency_penalties=frequency_penalties_t.to(device=device, + non_blocking=True), + repetition_penalties=repetition_penalties_t.to(device=device, + non_blocking=True), + prompt_tokens=prompt_t.to(device=device, non_blocking=True), + output_tokens=output_t.to(device=device, non_blocking=True), + ) diff --git a/vllm_overrides/sampling_params.py b/vllm_overrides/sampling_params.py index 884d07e0..58790be1 100644 --- a/vllm_overrides/sampling_params.py +++ b/vllm_overrides/sampling_params.py @@ -1,391 +1,391 @@ -"""Sampling parameters for text generation.""" -import copy -from dataclasses import dataclass -from enum import Enum, IntEnum -from functools import cached_property -from typing import Any, Callable, Dict, List, Optional, Set, Union - -import msgspec -import torch -from pydantic import BaseModel -from typing_extensions import Annotated - -from vllm.logger import init_logger - -logger = init_logger(__name__) - -_SAMPLING_EPS = 1e-5 -_MAX_TEMP = 1e-2 - - -class SamplingType(IntEnum): - GREEDY = 0 - RANDOM = 1 - RANDOM_SEED = 2 - - -LogitsProcessor = Union[Callable[[List[int], torch.Tensor], torch.Tensor], - Callable[[List[int], List[int], torch.Tensor], - torch.Tensor]] -"""LogitsProcessor is a function that takes a list -of previously generated tokens, the logits tensor -for the next token and, optionally, prompt tokens as a -first argument, and returns a modified tensor of logits -to sample from.""" - - -# maybe make msgspec? -@dataclass -class GuidedDecodingParams: - """One of these fields will be used to build a logit processor.""" - json: Optional[Union[str, Dict]] = None - regex: Optional[str] = None - choice: Optional[List[str]] = None - grammar: Optional[str] = None - json_object: Optional[bool] = None - """These are other options that can be set""" - backend: Optional[str] = None - whitespace_pattern: Optional[str] = None - - @staticmethod - def from_optional( - json: Optional[Union[Dict, BaseModel, str]], - regex: Optional[str] = None, - choice: Optional[List[str]] = None, - grammar: Optional[str] = None, - json_object: Optional[bool] = None, - backend: Optional[str] = None, - whitespace_pattern: Optional[str] = None, - ) -> "GuidedDecodingParams": - # Extract json schemas from pydantic models - if isinstance(json, (BaseModel, type(BaseModel))): - json = json.model_json_schema() - return GuidedDecodingParams( - json=json, - regex=regex, - choice=choice, - grammar=grammar, - json_object=json_object, - backend=backend, - whitespace_pattern=whitespace_pattern, - ) - - def __post_init__(self): - """Validate that some fields are mutually exclusive.""" - guide_count = sum([ - self.json is not None, self.regex is not None, self.choice - is not None, self.grammar is not None, self.json_object is not None - ]) - if guide_count > 1: - raise ValueError( - "You can only use one kind of guided decoding but multiple are " - f"specified: {self.__dict__}") - - -class RequestOutputKind(Enum): - # Return entire output so far in every RequestOutput - CUMULATIVE = 0 - # Return only deltas in each RequestOutput - DELTA = 1 - # Do not return intermediate RequestOuputs - FINAL_ONLY = 2 - - -class SamplingParams( - msgspec.Struct, - omit_defaults=True, # type: ignore[call-arg] - # required for @cached_property. - dict=True): # type: ignore[call-arg] - """Sampling parameters for text generation. - - Overall, we follow the sampling parameters from the OpenAI text completion - API (https://platform.openai.com/docs/api-reference/completions/create). - In addition, we support beam search, which is not supported by OpenAI. - - Args: - n: Number of output sequences to return for the given prompt. - best_of: Number of output sequences that are generated from the prompt. - From these `best_of` sequences, the top `n` sequences are returned. - `best_of` must be greater than or equal to `n`. By default, - `best_of` is set to `n`. - presence_penalty: Float that penalizes new tokens based on whether they - appear in the generated text so far. Values > 0 encourage the model - to use new tokens, while values < 0 encourage the model to repeat - tokens. - frequency_penalty: Float that penalizes new tokens based on their - frequency in the generated text so far. Values > 0 encourage the - model to use new tokens, while values < 0 encourage the model to - repeat tokens. - repetition_penalty: Float that penalizes new tokens based on whether - they appear in the prompt and the generated text so far. Values > 1 - encourage the model to use new tokens, while values < 1 encourage - the model to repeat tokens. - temperature: Float that controls the randomness of the sampling. Lower - values make the model more deterministic, while higher values make - the model more random. Zero means greedy sampling. - top_p: Float that controls the cumulative probability of the top tokens - to consider. Must be in (0, 1]. Set to 1 to consider all tokens. - top_k: Integer that controls the number of top tokens to consider. Set - to -1 to consider all tokens. - min_p: Float that represents the minimum probability for a token to be - considered, relative to the probability of the most likely token. - Must be in [0, 1]. Set to 0 to disable this. - seed: Random seed to use for the generation. - stop: List of strings that stop the generation when they are generated. - The returned output will not contain the stop strings. - stop_token_ids: List of tokens that stop the generation when they are - generated. The returned output will contain the stop tokens unless - the stop tokens are special tokens. - include_stop_str_in_output: Whether to include the stop strings in - output text. Defaults to False. - ignore_eos: Whether to ignore the EOS token and continue generating - tokens after the EOS token is generated. - max_tokens: Maximum number of tokens to generate per output sequence. - min_tokens: Minimum number of tokens to generate per output sequence - before EOS or stop_token_ids can be generated - logprobs: Number of log probabilities to return per output token. - When set to None, no probability is returned. If set to a non-None - value, the result includes the log probabilities of the specified - number of most likely tokens, as well as the chosen tokens. - Note that the implementation follows the OpenAI API: The API will - always return the log probability of the sampled token, so there - may be up to `logprobs+1` elements in the response. - prompt_logprobs: Number of log probabilities to return per prompt token. - detokenize: Whether to detokenize the output. Defaults to True. - skip_special_tokens: Whether to skip special tokens in the output. - spaces_between_special_tokens: Whether to add spaces between special - tokens in the output. Defaults to True. - logits_processors: List of functions that modify logits based on - previously generated tokens, and optionally prompt tokens as - a first argument. - truncate_prompt_tokens: If set to an integer k, will use only the last k - tokens from the prompt (i.e., left truncation). Defaults to None - (i.e., no truncation). - guided_decoding: If provided, the engine will construct a guided - decoding logits processor from these parameters. Defaults to None. - logit_bias: If provided, the engine will construct a logits processor - that applies these logit biases. Defaults to None. +"""Sampling parameters for text generation.""" +import copy +from dataclasses import dataclass +from enum import Enum, IntEnum +from functools import cached_property +from typing import Any, Callable, Dict, List, Optional, Set, Union + +import msgspec +import torch +from pydantic import BaseModel +from typing_extensions import Annotated + +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_SAMPLING_EPS = 1e-5 +_MAX_TEMP = 1e-2 + + +class SamplingType(IntEnum): + GREEDY = 0 + RANDOM = 1 + RANDOM_SEED = 2 + + +LogitsProcessor = Union[Callable[[List[int], torch.Tensor], torch.Tensor], + Callable[[List[int], List[int], torch.Tensor], + torch.Tensor]] +"""LogitsProcessor is a function that takes a list +of previously generated tokens, the logits tensor +for the next token and, optionally, prompt tokens as a +first argument, and returns a modified tensor of logits +to sample from.""" + + +# maybe make msgspec? +@dataclass +class GuidedDecodingParams: + """One of these fields will be used to build a logit processor.""" + json: Optional[Union[str, Dict]] = None + regex: Optional[str] = None + choice: Optional[List[str]] = None + grammar: Optional[str] = None + json_object: Optional[bool] = None + """These are other options that can be set""" + backend: Optional[str] = None + whitespace_pattern: Optional[str] = None + + @staticmethod + def from_optional( + json: Optional[Union[Dict, BaseModel, str]], + regex: Optional[str] = None, + choice: Optional[List[str]] = None, + grammar: Optional[str] = None, + json_object: Optional[bool] = None, + backend: Optional[str] = None, + whitespace_pattern: Optional[str] = None, + ) -> "GuidedDecodingParams": + # Extract json schemas from pydantic models + if isinstance(json, (BaseModel, type(BaseModel))): + json = json.model_json_schema() + return GuidedDecodingParams( + json=json, + regex=regex, + choice=choice, + grammar=grammar, + json_object=json_object, + backend=backend, + whitespace_pattern=whitespace_pattern, + ) + + def __post_init__(self): + """Validate that some fields are mutually exclusive.""" + guide_count = sum([ + self.json is not None, self.regex is not None, self.choice + is not None, self.grammar is not None, self.json_object is not None + ]) + if guide_count > 1: + raise ValueError( + "You can only use one kind of guided decoding but multiple are " + f"specified: {self.__dict__}") + + +class RequestOutputKind(Enum): + # Return entire output so far in every RequestOutput + CUMULATIVE = 0 + # Return only deltas in each RequestOutput + DELTA = 1 + # Do not return intermediate RequestOuputs + FINAL_ONLY = 2 + + +class SamplingParams( + msgspec.Struct, + omit_defaults=True, # type: ignore[call-arg] + # required for @cached_property. + dict=True): # type: ignore[call-arg] + """Sampling parameters for text generation. + + Overall, we follow the sampling parameters from the OpenAI text completion + API (https://platform.openai.com/docs/api-reference/completions/create). + In addition, we support beam search, which is not supported by OpenAI. + + Args: + n: Number of output sequences to return for the given prompt. + best_of: Number of output sequences that are generated from the prompt. + From these `best_of` sequences, the top `n` sequences are returned. + `best_of` must be greater than or equal to `n`. By default, + `best_of` is set to `n`. + presence_penalty: Float that penalizes new tokens based on whether they + appear in the generated text so far. Values > 0 encourage the model + to use new tokens, while values < 0 encourage the model to repeat + tokens. + frequency_penalty: Float that penalizes new tokens based on their + frequency in the generated text so far. Values > 0 encourage the + model to use new tokens, while values < 0 encourage the model to + repeat tokens. + repetition_penalty: Float that penalizes new tokens based on whether + they appear in the prompt and the generated text so far. Values > 1 + encourage the model to use new tokens, while values < 1 encourage + the model to repeat tokens. + temperature: Float that controls the randomness of the sampling. Lower + values make the model more deterministic, while higher values make + the model more random. Zero means greedy sampling. + top_p: Float that controls the cumulative probability of the top tokens + to consider. Must be in (0, 1]. Set to 1 to consider all tokens. + top_k: Integer that controls the number of top tokens to consider. Set + to -1 to consider all tokens. + min_p: Float that represents the minimum probability for a token to be + considered, relative to the probability of the most likely token. + Must be in [0, 1]. Set to 0 to disable this. + seed: Random seed to use for the generation. + stop: List of strings that stop the generation when they are generated. + The returned output will not contain the stop strings. + stop_token_ids: List of tokens that stop the generation when they are + generated. The returned output will contain the stop tokens unless + the stop tokens are special tokens. + include_stop_str_in_output: Whether to include the stop strings in + output text. Defaults to False. + ignore_eos: Whether to ignore the EOS token and continue generating + tokens after the EOS token is generated. + max_tokens: Maximum number of tokens to generate per output sequence. + min_tokens: Minimum number of tokens to generate per output sequence + before EOS or stop_token_ids can be generated + logprobs: Number of log probabilities to return per output token. + When set to None, no probability is returned. If set to a non-None + value, the result includes the log probabilities of the specified + number of most likely tokens, as well as the chosen tokens. + Note that the implementation follows the OpenAI API: The API will + always return the log probability of the sampled token, so there + may be up to `logprobs+1` elements in the response. + prompt_logprobs: Number of log probabilities to return per prompt token. + detokenize: Whether to detokenize the output. Defaults to True. + skip_special_tokens: Whether to skip special tokens in the output. + spaces_between_special_tokens: Whether to add spaces between special + tokens in the output. Defaults to True. + logits_processors: List of functions that modify logits based on + previously generated tokens, and optionally prompt tokens as + a first argument. + truncate_prompt_tokens: If set to an integer k, will use only the last k + tokens from the prompt (i.e., left truncation). Defaults to None + (i.e., no truncation). + guided_decoding: If provided, the engine will construct a guided + decoding logits processor from these parameters. Defaults to None. + logit_bias: If provided, the engine will construct a logits processor + that applies these logit biases. Defaults to None. allowed_token_ids: If provided, the engine will construct a logits processor which only retains scores for the given token ids. Defaults to None. prompt_logprob_positions: Optional prompt-token positions whose logits should be materialized. None preserves the standard all-position prompt-logprob behavior. - """ - - n: int = 1 - best_of: Optional[int] = None - _real_n: Optional[int] = None - presence_penalty: float = 0.0 - frequency_penalty: float = 0.0 - repetition_penalty: float = 1.0 - temperature: float = 1.0 - top_p: float = 1.0 - top_k: int = -1 - min_p: float = 0.0 - seed: Optional[int] = None - stop: Optional[Union[str, List[str]]] = None - stop_token_ids: Optional[List[int]] = None - ignore_eos: bool = False - max_tokens: Optional[int] = 16 - min_tokens: int = 0 - logprobs: Optional[int] = None - prompt_logprobs: Optional[int] = None - # NOTE: This parameter is only exposed at the engine level for now. - # It is not exposed in the OpenAI API server, as the OpenAI API does - # not support returning only a list of token IDs. - detokenize: bool = True - skip_special_tokens: bool = True - spaces_between_special_tokens: bool = True - # Optional[List[LogitsProcessor]] type. We use Any here because - # Optional[List[LogitsProcessor]] type is not supported by msgspec. - logits_processors: Optional[Any] = None - include_stop_str_in_output: bool = False - truncate_prompt_tokens: Optional[Annotated[int, msgspec.Meta(ge=1)]] = None - output_kind: RequestOutputKind = RequestOutputKind.CUMULATIVE - - # The below fields are not supposed to be used as an input. - # They are set in post_init. - output_text_buffer_length: int = 0 - _all_stop_token_ids: Set[int] = msgspec.field(default_factory=set) - - # Fields used to construct logits processors + """ + + n: int = 1 + best_of: Optional[int] = None + _real_n: Optional[int] = None + presence_penalty: float = 0.0 + frequency_penalty: float = 0.0 + repetition_penalty: float = 1.0 + temperature: float = 1.0 + top_p: float = 1.0 + top_k: int = -1 + min_p: float = 0.0 + seed: Optional[int] = None + stop: Optional[Union[str, List[str]]] = None + stop_token_ids: Optional[List[int]] = None + ignore_eos: bool = False + max_tokens: Optional[int] = 16 + min_tokens: int = 0 + logprobs: Optional[int] = None + prompt_logprobs: Optional[int] = None + # NOTE: This parameter is only exposed at the engine level for now. + # It is not exposed in the OpenAI API server, as the OpenAI API does + # not support returning only a list of token IDs. + detokenize: bool = True + skip_special_tokens: bool = True + spaces_between_special_tokens: bool = True + # Optional[List[LogitsProcessor]] type. We use Any here because + # Optional[List[LogitsProcessor]] type is not supported by msgspec. + logits_processors: Optional[Any] = None + include_stop_str_in_output: bool = False + truncate_prompt_tokens: Optional[Annotated[int, msgspec.Meta(ge=1)]] = None + output_kind: RequestOutputKind = RequestOutputKind.CUMULATIVE + + # The below fields are not supposed to be used as an input. + # They are set in post_init. + output_text_buffer_length: int = 0 + _all_stop_token_ids: Set[int] = msgspec.field(default_factory=set) + + # Fields used to construct logits processors guided_decoding: Optional[GuidedDecodingParams] = None logit_bias: Optional[Dict[int, float]] = None allowed_token_ids: Optional[List[int]] = None prompt_logprob_positions: Optional[List[int]] = None - - @staticmethod - def from_optional( - n: Optional[int] = 1, - best_of: Optional[int] = None, - presence_penalty: Optional[float] = 0.0, - frequency_penalty: Optional[float] = 0.0, - repetition_penalty: Optional[float] = 1.0, - temperature: Optional[float] = 1.0, - top_p: Optional[float] = 1.0, - top_k: int = -1, - min_p: float = 0.0, - seed: Optional[int] = None, - stop: Optional[Union[str, List[str]]] = None, - stop_token_ids: Optional[List[int]] = None, - include_stop_str_in_output: bool = False, - ignore_eos: bool = False, - max_tokens: Optional[int] = 16, - min_tokens: int = 0, - logprobs: Optional[int] = None, - prompt_logprobs: Optional[int] = None, - detokenize: bool = True, - skip_special_tokens: bool = True, - spaces_between_special_tokens: bool = True, - logits_processors: Optional[List[LogitsProcessor]] = None, - truncate_prompt_tokens: Optional[Annotated[int, - msgspec.Meta(ge=1)]] = None, - output_kind: RequestOutputKind = RequestOutputKind.CUMULATIVE, + + @staticmethod + def from_optional( + n: Optional[int] = 1, + best_of: Optional[int] = None, + presence_penalty: Optional[float] = 0.0, + frequency_penalty: Optional[float] = 0.0, + repetition_penalty: Optional[float] = 1.0, + temperature: Optional[float] = 1.0, + top_p: Optional[float] = 1.0, + top_k: int = -1, + min_p: float = 0.0, + seed: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, + stop_token_ids: Optional[List[int]] = None, + include_stop_str_in_output: bool = False, + ignore_eos: bool = False, + max_tokens: Optional[int] = 16, + min_tokens: int = 0, + logprobs: Optional[int] = None, + prompt_logprobs: Optional[int] = None, + detokenize: bool = True, + skip_special_tokens: bool = True, + spaces_between_special_tokens: bool = True, + logits_processors: Optional[List[LogitsProcessor]] = None, + truncate_prompt_tokens: Optional[Annotated[int, + msgspec.Meta(ge=1)]] = None, + output_kind: RequestOutputKind = RequestOutputKind.CUMULATIVE, guided_decoding: Optional[GuidedDecodingParams] = None, logit_bias: Optional[Union[Dict[int, float], Dict[str, float]]] = None, allowed_token_ids: Optional[List[int]] = None, prompt_logprob_positions: Optional[List[int]] = None, - ) -> "SamplingParams": - if logit_bias is not None: - logit_bias = { - int(token): bias - for token, bias in logit_bias.items() - } - - return SamplingParams( - n=1 if n is None else n, - best_of=best_of, - presence_penalty=0.0 - if presence_penalty is None else presence_penalty, - frequency_penalty=0.0 - if frequency_penalty is None else frequency_penalty, - repetition_penalty=1.0 - if repetition_penalty is None else repetition_penalty, - temperature=1.0 if temperature is None else temperature, - top_p=1.0 if top_p is None else top_p, - top_k=top_k, - min_p=min_p, - seed=seed, - stop=stop, - stop_token_ids=stop_token_ids, - include_stop_str_in_output=include_stop_str_in_output, - ignore_eos=ignore_eos, - max_tokens=max_tokens, - min_tokens=min_tokens, - logprobs=logprobs, - prompt_logprobs=prompt_logprobs, - detokenize=detokenize, - skip_special_tokens=skip_special_tokens, - spaces_between_special_tokens=spaces_between_special_tokens, - logits_processors=logits_processors, - truncate_prompt_tokens=truncate_prompt_tokens, - output_kind=output_kind, + ) -> "SamplingParams": + if logit_bias is not None: + logit_bias = { + int(token): bias + for token, bias in logit_bias.items() + } + + return SamplingParams( + n=1 if n is None else n, + best_of=best_of, + presence_penalty=0.0 + if presence_penalty is None else presence_penalty, + frequency_penalty=0.0 + if frequency_penalty is None else frequency_penalty, + repetition_penalty=1.0 + if repetition_penalty is None else repetition_penalty, + temperature=1.0 if temperature is None else temperature, + top_p=1.0 if top_p is None else top_p, + top_k=top_k, + min_p=min_p, + seed=seed, + stop=stop, + stop_token_ids=stop_token_ids, + include_stop_str_in_output=include_stop_str_in_output, + ignore_eos=ignore_eos, + max_tokens=max_tokens, + min_tokens=min_tokens, + logprobs=logprobs, + prompt_logprobs=prompt_logprobs, + detokenize=detokenize, + skip_special_tokens=skip_special_tokens, + spaces_between_special_tokens=spaces_between_special_tokens, + logits_processors=logits_processors, + truncate_prompt_tokens=truncate_prompt_tokens, + output_kind=output_kind, guided_decoding=guided_decoding, logit_bias=logit_bias, allowed_token_ids=allowed_token_ids, prompt_logprob_positions=prompt_logprob_positions, ) - - def __post_init__(self) -> None: - # how we deal with `best_of``: - # if `best_of`` is not set, we default to `n`; - # if `best_of`` is set, we set `n`` to `best_of`, - # and set `_real_n`` to the original `n`. - # when we return the result, we will check - # if we need to return `n` or `_real_n` results - if self.best_of: - if self.best_of < self.n: - raise ValueError( - f"best_of must be greater than or equal to n, " - f"got n={self.n} and best_of={self.best_of}.") - self._real_n = self.n - self.n = self.best_of - if 0 < self.temperature < _MAX_TEMP: - logger.warning( - "temperature %s is less than %s, which may cause numerical " - "errors nan or inf in tensors. We have maxed it out to %s.", - self.temperature, _MAX_TEMP, _MAX_TEMP) - self.temperature = max(self.temperature, _MAX_TEMP) - if self.seed == -1: - self.seed = None - else: - self.seed = self.seed - if self.stop is None: - self.stop = [] - elif isinstance(self.stop, str): - self.stop = [self.stop] - else: - self.stop = list(self.stop) - if self.stop_token_ids is None: - self.stop_token_ids = [] - else: - self.stop_token_ids = list(self.stop_token_ids) - self.logprobs = 1 if self.logprobs is True else self.logprobs + + def __post_init__(self) -> None: + # how we deal with `best_of``: + # if `best_of`` is not set, we default to `n`; + # if `best_of`` is set, we set `n`` to `best_of`, + # and set `_real_n`` to the original `n`. + # when we return the result, we will check + # if we need to return `n` or `_real_n` results + if self.best_of: + if self.best_of < self.n: + raise ValueError( + f"best_of must be greater than or equal to n, " + f"got n={self.n} and best_of={self.best_of}.") + self._real_n = self.n + self.n = self.best_of + if 0 < self.temperature < _MAX_TEMP: + logger.warning( + "temperature %s is less than %s, which may cause numerical " + "errors nan or inf in tensors. We have maxed it out to %s.", + self.temperature, _MAX_TEMP, _MAX_TEMP) + self.temperature = max(self.temperature, _MAX_TEMP) + if self.seed == -1: + self.seed = None + else: + self.seed = self.seed + if self.stop is None: + self.stop = [] + elif isinstance(self.stop, str): + self.stop = [self.stop] + else: + self.stop = list(self.stop) + if self.stop_token_ids is None: + self.stop_token_ids = [] + else: + self.stop_token_ids = list(self.stop_token_ids) + self.logprobs = 1 if self.logprobs is True else self.logprobs self.prompt_logprobs = (1 if self.prompt_logprobs is True else self.prompt_logprobs) if self.prompt_logprob_positions is not None: self.prompt_logprob_positions = list( self.prompt_logprob_positions) - - # Number of characters to hold back for stop string evaluation - # until sequence is finished. - if self.stop and not self.include_stop_str_in_output: - self.output_text_buffer_length = max(len(s) for s in self.stop) - 1 - - self._verify_args() - - if self.temperature < _SAMPLING_EPS: - # Zero temperature means greedy sampling. - self.top_p = 1.0 - self.top_k = -1 - self.min_p = 0.0 - self._verify_greedy_sampling() - # eos_token_id is added to this by the engine - self._all_stop_token_ids = set(self.stop_token_ids) - - def _verify_args(self) -> None: - if not isinstance(self.n, int): - raise ValueError(f"n must be an int, but is of " - f"type {type(self.n)}") - if self.n < 1: - raise ValueError(f"n must be at least 1, got {self.n}.") - if not -2.0 <= self.presence_penalty <= 2.0: - raise ValueError("presence_penalty must be in [-2, 2], got " - f"{self.presence_penalty}.") - if not -2.0 <= self.frequency_penalty <= 2.0: - raise ValueError("frequency_penalty must be in [-2, 2], got " - f"{self.frequency_penalty}.") - if not 0.0 < self.repetition_penalty <= 2.0: - raise ValueError("repetition_penalty must be in (0, 2], got " - f"{self.repetition_penalty}.") - if self.temperature < 0.0: - raise ValueError( - f"temperature must be non-negative, got {self.temperature}.") - if not 0.0 < self.top_p <= 1.0: - raise ValueError(f"top_p must be in (0, 1], got {self.top_p}.") - if self.top_k < -1 or self.top_k == 0: - raise ValueError(f"top_k must be -1 (disable), or at least 1, " - f"got {self.top_k}.") - if not isinstance(self.top_k, int): - raise TypeError( - f"top_k must be an integer, got {type(self.top_k).__name__}") - if not 0.0 <= self.min_p <= 1.0: - raise ValueError("min_p must be in [0, 1], got " - f"{self.min_p}.") - if self.max_tokens is not None and self.max_tokens < 1: - raise ValueError( - f"max_tokens must be at least 1, got {self.max_tokens}.") - if self.min_tokens < 0: - raise ValueError(f"min_tokens must be greater than or equal to 0, " - f"got {self.min_tokens}.") - if self.max_tokens is not None and self.min_tokens > self.max_tokens: - raise ValueError( - f"min_tokens must be less than or equal to " - f"max_tokens={self.max_tokens}, got {self.min_tokens}.") - if self.logprobs is not None and self.logprobs < 0: - raise ValueError( - f"logprobs must be non-negative, got {self.logprobs}.") + + # Number of characters to hold back for stop string evaluation + # until sequence is finished. + if self.stop and not self.include_stop_str_in_output: + self.output_text_buffer_length = max(len(s) for s in self.stop) - 1 + + self._verify_args() + + if self.temperature < _SAMPLING_EPS: + # Zero temperature means greedy sampling. + self.top_p = 1.0 + self.top_k = -1 + self.min_p = 0.0 + self._verify_greedy_sampling() + # eos_token_id is added to this by the engine + self._all_stop_token_ids = set(self.stop_token_ids) + + def _verify_args(self) -> None: + if not isinstance(self.n, int): + raise ValueError(f"n must be an int, but is of " + f"type {type(self.n)}") + if self.n < 1: + raise ValueError(f"n must be at least 1, got {self.n}.") + if not -2.0 <= self.presence_penalty <= 2.0: + raise ValueError("presence_penalty must be in [-2, 2], got " + f"{self.presence_penalty}.") + if not -2.0 <= self.frequency_penalty <= 2.0: + raise ValueError("frequency_penalty must be in [-2, 2], got " + f"{self.frequency_penalty}.") + if not 0.0 < self.repetition_penalty <= 2.0: + raise ValueError("repetition_penalty must be in (0, 2], got " + f"{self.repetition_penalty}.") + if self.temperature < 0.0: + raise ValueError( + f"temperature must be non-negative, got {self.temperature}.") + if not 0.0 < self.top_p <= 1.0: + raise ValueError(f"top_p must be in (0, 1], got {self.top_p}.") + if self.top_k < -1 or self.top_k == 0: + raise ValueError(f"top_k must be -1 (disable), or at least 1, " + f"got {self.top_k}.") + if not isinstance(self.top_k, int): + raise TypeError( + f"top_k must be an integer, got {type(self.top_k).__name__}") + if not 0.0 <= self.min_p <= 1.0: + raise ValueError("min_p must be in [0, 1], got " + f"{self.min_p}.") + if self.max_tokens is not None and self.max_tokens < 1: + raise ValueError( + f"max_tokens must be at least 1, got {self.max_tokens}.") + if self.min_tokens < 0: + raise ValueError(f"min_tokens must be greater than or equal to 0, " + f"got {self.min_tokens}.") + if self.max_tokens is not None and self.min_tokens > self.max_tokens: + raise ValueError( + f"min_tokens must be less than or equal to " + f"max_tokens={self.max_tokens}, got {self.min_tokens}.") + if self.logprobs is not None and self.logprobs < 0: + raise ValueError( + f"logprobs must be non-negative, got {self.logprobs}.") if self.prompt_logprobs is not None and self.prompt_logprobs < 0: raise ValueError(f"prompt_logprobs must be non-negative, got " f"{self.prompt_logprobs}.") @@ -407,114 +407,114 @@ class SamplingParams( raise ValueError( "prompt_logprob_positions must be a sorted unique list " "of positive integers.") - if (self.truncate_prompt_tokens is not None - and self.truncate_prompt_tokens < 1): - raise ValueError(f"truncate_prompt_tokens must be >= 1, " - f"got {self.truncate_prompt_tokens}") - assert isinstance(self.stop, list) - if any(not stop_str for stop_str in self.stop): - raise ValueError("stop cannot contain an empty string.") - if self.stop and not self.detokenize: - raise ValueError( - "stop strings are only supported when detokenize is True. " - "Set detokenize=True to use stop.") - if self.best_of != self._real_n and self.output_kind == ( - RequestOutputKind.DELTA): - raise ValueError("best_of must equal n to use output_kind=DELTA") - - def _verify_greedy_sampling(self) -> None: - if self.n > 1: - raise ValueError("n must be 1 when using greedy sampling, " - f"got {self.n}.") - - def update_from_generation_config( - self, - generation_config: Dict[str, Any], - model_eos_token_id: Optional[int] = None) -> None: - """Update if there are non-default values from generation_config""" - - if model_eos_token_id is not None: - # Add the eos token id into the sampling_params to support - # min_tokens processing. - self._all_stop_token_ids.add(model_eos_token_id) - - # Update eos_token_id for generation - if (eos_ids := generation_config.get("eos_token_id")) is not None: - # it can be either int or list of int - eos_ids = {eos_ids} if isinstance(eos_ids, int) else set(eos_ids) - if model_eos_token_id is not None: - # We don't need to include the primary eos_token_id in - # stop_token_ids since it's handled separately for stopping - # purposes. - eos_ids.discard(model_eos_token_id) - if eos_ids: - self._all_stop_token_ids.update(eos_ids) - if not self.ignore_eos: - eos_ids.update(self.stop_token_ids) - self.stop_token_ids = list(eos_ids) - - @cached_property - def sampling_type(self) -> SamplingType: - if self.temperature < _SAMPLING_EPS: - return SamplingType.GREEDY - if self.seed is not None: - return SamplingType.RANDOM_SEED - return SamplingType.RANDOM - - @property - def all_stop_token_ids(self) -> Set[int]: - return self._all_stop_token_ids - - def clone(self) -> "SamplingParams": - """Deep copy excluding LogitsProcessor objects. - - LogitsProcessor objects are excluded because they may contain an - arbitrary, nontrivial amount of data. - See https://github.com/vllm-project/vllm/issues/3087 - """ - - logit_processor_refs = None if self.logits_processors is None else { - id(lp): lp - for lp in self.logits_processors - } - return copy.deepcopy(self, memo=logit_processor_refs) - - def __repr__(self) -> str: - return ( - f"SamplingParams(n={self.n}, " - f"presence_penalty={self.presence_penalty}, " - f"frequency_penalty={self.frequency_penalty}, " - f"repetition_penalty={self.repetition_penalty}, " - f"temperature={self.temperature}, " - f"top_p={self.top_p}, " - f"top_k={self.top_k}, " - f"min_p={self.min_p}, " - f"seed={self.seed}, " - f"stop={self.stop}, " - f"stop_token_ids={self.stop_token_ids}, " - f"include_stop_str_in_output={self.include_stop_str_in_output}, " - f"ignore_eos={self.ignore_eos}, " - f"max_tokens={self.max_tokens}, " - f"min_tokens={self.min_tokens}, " + if (self.truncate_prompt_tokens is not None + and self.truncate_prompt_tokens < 1): + raise ValueError(f"truncate_prompt_tokens must be >= 1, " + f"got {self.truncate_prompt_tokens}") + assert isinstance(self.stop, list) + if any(not stop_str for stop_str in self.stop): + raise ValueError("stop cannot contain an empty string.") + if self.stop and not self.detokenize: + raise ValueError( + "stop strings are only supported when detokenize is True. " + "Set detokenize=True to use stop.") + if self.best_of != self._real_n and self.output_kind == ( + RequestOutputKind.DELTA): + raise ValueError("best_of must equal n to use output_kind=DELTA") + + def _verify_greedy_sampling(self) -> None: + if self.n > 1: + raise ValueError("n must be 1 when using greedy sampling, " + f"got {self.n}.") + + def update_from_generation_config( + self, + generation_config: Dict[str, Any], + model_eos_token_id: Optional[int] = None) -> None: + """Update if there are non-default values from generation_config""" + + if model_eos_token_id is not None: + # Add the eos token id into the sampling_params to support + # min_tokens processing. + self._all_stop_token_ids.add(model_eos_token_id) + + # Update eos_token_id for generation + if (eos_ids := generation_config.get("eos_token_id")) is not None: + # it can be either int or list of int + eos_ids = {eos_ids} if isinstance(eos_ids, int) else set(eos_ids) + if model_eos_token_id is not None: + # We don't need to include the primary eos_token_id in + # stop_token_ids since it's handled separately for stopping + # purposes. + eos_ids.discard(model_eos_token_id) + if eos_ids: + self._all_stop_token_ids.update(eos_ids) + if not self.ignore_eos: + eos_ids.update(self.stop_token_ids) + self.stop_token_ids = list(eos_ids) + + @cached_property + def sampling_type(self) -> SamplingType: + if self.temperature < _SAMPLING_EPS: + return SamplingType.GREEDY + if self.seed is not None: + return SamplingType.RANDOM_SEED + return SamplingType.RANDOM + + @property + def all_stop_token_ids(self) -> Set[int]: + return self._all_stop_token_ids + + def clone(self) -> "SamplingParams": + """Deep copy excluding LogitsProcessor objects. + + LogitsProcessor objects are excluded because they may contain an + arbitrary, nontrivial amount of data. + See https://github.com/vllm-project/vllm/issues/3087 + """ + + logit_processor_refs = None if self.logits_processors is None else { + id(lp): lp + for lp in self.logits_processors + } + return copy.deepcopy(self, memo=logit_processor_refs) + + def __repr__(self) -> str: + return ( + f"SamplingParams(n={self.n}, " + f"presence_penalty={self.presence_penalty}, " + f"frequency_penalty={self.frequency_penalty}, " + f"repetition_penalty={self.repetition_penalty}, " + f"temperature={self.temperature}, " + f"top_p={self.top_p}, " + f"top_k={self.top_k}, " + f"min_p={self.min_p}, " + f"seed={self.seed}, " + f"stop={self.stop}, " + f"stop_token_ids={self.stop_token_ids}, " + f"include_stop_str_in_output={self.include_stop_str_in_output}, " + f"ignore_eos={self.ignore_eos}, " + f"max_tokens={self.max_tokens}, " + f"min_tokens={self.min_tokens}, " f"logprobs={self.logprobs}, " f"prompt_logprobs={self.prompt_logprobs}, " "prompt_logprob_positions=" f"{self.prompt_logprob_positions}, " - f"skip_special_tokens={self.skip_special_tokens}, " - "spaces_between_special_tokens=" - f"{self.spaces_between_special_tokens}, " - f"truncate_prompt_tokens={self.truncate_prompt_tokens}), " - f"guided_decoding={self.guided_decoding}") - - -class BeamSearchParams( - msgspec.Struct, - omit_defaults=True, # type: ignore[call-arg] - # required for @cached_property. - dict=True): # type: ignore[call-arg] - """Beam search parameters for text generation.""" - beam_width: int - max_tokens: int - ignore_eos: bool = False - temperature: float = 0.0 - length_penalty: float = 1.0 + f"skip_special_tokens={self.skip_special_tokens}, " + "spaces_between_special_tokens=" + f"{self.spaces_between_special_tokens}, " + f"truncate_prompt_tokens={self.truncate_prompt_tokens}), " + f"guided_decoding={self.guided_decoding}") + + +class BeamSearchParams( + msgspec.Struct, + omit_defaults=True, # type: ignore[call-arg] + # required for @cached_property. + dict=True): # type: ignore[call-arg] + """Beam search parameters for text generation.""" + beam_width: int + max_tokens: int + ignore_eos: bool = False + temperature: float = 0.0 + length_penalty: float = 1.0