feat: import CUDA kernels from xllm/CCCL/FLA upstream repos

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)

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,

cccl_upstream/cub/cub/device/dispatch/dispatch_batch_memcpy.cuh

@@ -152,8 +152,8 @@ __launch_bounds__(int(current_policy<PolicySelector>().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<MemcpyOpt == CopyAlg::Memcpy, AliasT, InputBufferT>(input_buffer_it[buffer_id], thread_offset);
           write_item<MemcpyOpt == CopyAlg::Memcpy, AliasT, OutputBufferT>(
             output_buffer_it[buffer_id], thread_offset, value);
         }

cccl_upstream/cub/cub/device/dispatch/dispatch_copy_mdspan.cuh

@@ -21,8 +21,14 @@
 #include <cub/device/dispatch/tuning/tuning_transform.cuh>
 #include <cub/util_debug.cuh>
 
+#include <cuda/__algorithm/copy.h>
 #include <cuda/__functional/always_true_false.h>
+#include <cuda/__stream/get_stream.h>
+#include <cuda/__stream/stream_ref.h>
+#include <cuda/std/__concepts/same_as.h>
 #include <cuda/std/__functional/identity.h>
+#include <cuda/std/__host_stdlib/stdexcept>
+#include <cuda/std/__type_traits/is_callable.h>
 #include <cuda/std/mdspan>
 
 CUB_NAMESPACE_BEGIN
@@ -47,6 +53,46 @@ struct copy_mdspan_t
   }
 };
 
+template <class _MDSpanIn, class _MDSpanOut>
+[[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<void>(__e);
+    return ::cudaErrorInvalidValue;
+  }
+#endif // _CCCL_HOSTED
+  _CCCL_CATCH_ALL
+  {
+    return ::cudaErrorUnknown;
+  }
+
+  return ::cudaSuccess;
+}
+
+template <class _MDSpanIn, class _MDSpanOut, class _Env>
+[[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 <typename T_In,
           typename E_In,
           typename L_In,
@@ -64,13 +110,38 @@ copy(::cuda::std::mdspan<T_In, E_In, L_In, A_In> 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<T_In, T_Out>
+                  && ::cuda::__detail::__can_mdspan_copy_bytes<T_In, E_In, L_In, T_Out, E_Out, L_Out>
+                  && ::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

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
     {

cccl_upstream/cub/cub/device/dispatch/dispatch_reduce_by_key.cuh

@@ -208,14 +208,14 @@ __launch_bounds__(int(current_policy<PolicySelector>().lookback.threads_per_bloc
 {
   static constexpr ReduceByKeyPolicy policy = current_policy<PolicySelector>();
   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.lookback_delay.kind,
-                                    policy.lookback.lookback_delay.delay,
-                                    policy.lookback.lookback_delay.l2_write_latency>>;
+    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.lookback_delay.kind,
+                        policy.lookback.lookback_delay.delay,
+                        policy.lookback.lookback_delay.l2_write_latency>>;
 
   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.lookback_delay.kind,
-                                              policy.lookback.lookback_delay.delay,
-                                              policy.lookback.lookback_delay.l2_write_latency>>;
+    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.lookback_delay.kind,
+                        policy.lookback.lookback_delay.delay,
+                        policy.lookback.lookback_delay.l2_write_latency>>;
   using vsmem_helper_t = vsmem_helper_default_fallback_policy_t<Policy, AgentReduceByKey, Args...>;
   return ::cuda::std::tuple{vsmem_helper_t::agent_policy_t::BLOCK_THREADS,
                             vsmem_helper_t::agent_policy_t::ITEMS_PER_THREAD,

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<ScanOpT,
-                                                                 cub::detail::it_value_t<InputIteratorT>,
-                                                                 ::cuda::std::_If<::cuda::std::is_same_v<InitValueT, NullType>,
-                                                                                  cub::detail::it_value_t<InputIteratorT>,
-                                                                                  typename InitValueT::value_type>>,
+  typename AccumT = ::cuda::std::__accumulator_t<ScanOpT,
+                                                 cub::detail::it_value_t<InputIteratorT>,
+                                                 ::cuda::std::_If<::cuda::std::is_same_v<InitValueT, NullType>,
+                                                                  cub::detail::it_value_t<InputIteratorT>,
+                                                                  typename InitValueT::value_type>>,
   ForceInclusive EnforceInclusive = ForceInclusive::No,
   typename PolicyHub              = detail::scan::
     policy_hub<detail::it_value_t<InputIteratorT>, detail::it_value_t<OutputIteratorT>, 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<int>(
-                     ::cuda::ceil_div(num_items, static_cast<OffsetT>(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<int>(
+          ::cuda::ceil_div(num_items, static_cast<OffsetT>(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<int>(kernel_source.InputSize()),
-                       static_cast<int>(kernel_source.InputAlign()),
-                       static_cast<int>(kernel_source.OutputAlign()),
-                       static_cast<int>(kernel_source.AccumSize()),
-                       static_cast<int>(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<int>(kernel_source.InputSize()),
+            static_cast<int>(kernel_source.InputAlign()),
+            static_cast<int>(kernel_source.OutputAlign()),
+            static_cast<int>(kernel_source.AccumSize()),
+            static_cast<int>(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<int>(
-                   ::cuda::ceil_div(num_items, static_cast<OffsetT>(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<int>(
+        ::cuda::ceil_div(num_items, static_cast<OffsetT>(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<int>(kernel_source.InputSize()),
-                     static_cast<int>(kernel_source.InputAlign()),
-                     static_cast<int>(kernel_source.OutputAlign()),
-                     static_cast<int>(kernel_source.AccumSize()),
-                     static_cast<int>(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<int>(kernel_source.InputSize()),
+          static_cast<int>(kernel_source.InputAlign()),
+          static_cast<int>(kernel_source.OutputAlign()),
+          static_cast<int>(kernel_source.AccumSize()),
+          static_cast<int>(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
   {

cccl_upstream/cub/cub/device/dispatch/dispatch_scan_by_key.cuh

@@ -238,15 +238,15 @@ template <
   typename PolicyHub = policy_hub<KeysInputIteratorT, AccumT, cub::detail::it_value_t<ValuesInputIteratorT>, ScanOpT>,
   typename PolicySelector = policy_selector_from_hub<PolicyHub>,
   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<ValuesInputIteratorT>,
                                                        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<PolicySelector>
@@ -875,15 +875,15 @@ template <
     detail::scan_by_key::policy_hub<KeysInputIteratorT, AccumT, cub::detail::it_value_t<ValuesInputIteratorT>, ScanOpT>,
   typename PolicySelector = detail::scan_by_key::policy_selector_from_hub<PolicyHub>,
   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<

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<KeyT, ValueT, SegmentSizeT>;
   using policy_selector_t         = ::cuda::std::decay_t<
-            ::cuda::std::execution::__query_result_or_t<TuningEnvT, SegmentedRadixSortPolicy, default_policy_selector_t>>;
+    ::cuda::std::execution::__query_result_or_t<TuningEnvT, SegmentedRadixSortPolicy, default_policy_selector_t>>;
 #if _CCCL_HAS_CONCEPTS()
   static_assert(segmented_radix_sort_policy_selector<policy_selector_t>);
 #endif // _CCCL_HAS_CONCEPTS()

cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_reduce.cuh

@@ -503,15 +503,15 @@ template <
     decltype(select_segmented_accum_t<InputIteratorT, InitValueT, ReductionOpT>(static_cast<OverrideAccumT*>(nullptr))),
   typename PolicySelector = policy_selector_from_types<AccumT, OffsetT, ReductionOpT>,
   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<PolicySelector>
@@ -711,13 +711,13 @@ template <typename OverrideAccumT = use_default,
             static_cast<OverrideAccumT*>(nullptr))),
           typename PolicySelector = policy_selector_from_types<AccumT, OffsetT, ReductionOpT>,
           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<OffsetT>, int> = 0>
 #if _CCCL_HAS_CONCEPTS()

cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_scan.cuh

@@ -96,17 +96,17 @@ template <
     common_iterator_value_t<BeginOffsetIteratorInputT, EndOffsetIteratorInputT, BeginOffsetIteratorOutputT>,
   typename PolicySelector = policy_selector_from_types<AccumT>,
   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<PolicySelector>

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.lookback_delay.kind,
-                                               active_policy.lookback.lookback_delay.delay,
-                                               active_policy.lookback.lookback_delay.l2_write_latency>>;
+    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.lookback_delay.kind,
+                        active_policy.lookback.lookback_delay.delay,
+                        active_policy.lookback.lookback_delay.l2_write_latency>,
+    active_policy.lookback._load_prefetch>;
   using type = vsmem_helper_default_fallback_policy_t<
     agent_policy_t,
     bind_selection_opt<SelectionOpt>::template agent_t,

cccl_upstream/cub/cub/device/dispatch/dispatch_three_way_partition.cuh

@@ -444,18 +444,18 @@ template <typename InputIteratorT,
           typename OffsetT,
           typename PolicySelector = policy_selector_from_types<it_value_t<InputIteratorT>, per_partition_offset_t>,
           typename KernelSource   = DeviceThreeWayPartitionKernelSource<
-              PolicySelector,
-              InputIteratorT,
-              FirstOutputIteratorT,
-              SecondOutputIteratorT,
-              UnselectedOutputIteratorT,
-              NumSelectedIteratorT,
-              ScanTileStateT,
-              SelectFirstPartOp,
-              SelectSecondPartOp,
-              per_partition_offset_t,
-              streaming_context_t<OffsetT>,
-              OffsetT>,
+            PolicySelector,
+            InputIteratorT,
+            FirstOutputIteratorT,
+            SecondOutputIteratorT,
+            UnselectedOutputIteratorT,
+            NumSelectedIteratorT,
+            ScanTileStateT,
+            SelectFirstPartOp,
+            SelectSecondPartOp,
+            per_partition_offset_t,
+            streaming_context_t<OffsetT>,
+            OffsetT>,
           typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY>
 #if _CCCL_HAS_CONCEPTS()
   requires three_way_partition_policy_selector<PolicySelector>

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

cccl_upstream/cub/cub/device/dispatch/tuning/common.cuh

@@ -13,13 +13,14 @@
 #  pragma system_header
 #endif // no system header
 
+#include <cub/util_device.cuh>
 #include <cub/util_type.cuh>
 
 #include <thrust/type_traits/is_contiguous_iterator.h>
-#include <thrust/type_traits/is_trivially_relocatable.h>
 
 #include <cuda/__functional/maximum.h>
 #include <cuda/__functional/minimum.h>
+#include <cuda/__type_traits/is_trivially_copyable.h>
 #include <cuda/std/__concepts/same_as.h>
 #include <cuda/std/__functional/operations.h>
 #include <cuda/std/__fwd/format.h>
@@ -130,7 +131,7 @@ template <typename It>
   return iterator_info{
     static_cast<int>(size_of<vt>),
     static_cast<int>(align_of<vt>),
-    THRUST_NS_QUALIFIER::is_trivially_relocatable_v<vt>,
+    ::cuda::is_trivially_copyable_v<vt>,
     THRUST_NS_QUALIFIER::is_contiguous_iterator_v<It>};
 }
 

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);

cccl_upstream/cub/cub/device/dispatch/tuning/tuning_batched_topk.cuh

@@ -16,20 +16,27 @@
 #include <cub/block/block_load.cuh>
 #include <cub/block/block_scan.cuh>
 #include <cub/block/block_store.cuh>
+#include <cub/util_device.cuh>
 
+#include <cuda/__cmath/pow2.h>
 #include <cuda/__device/compute_capability.h>
+#include <cuda/__execution/determinism.h>
+#include <cuda/__execution/tie_break.h>
 #include <cuda/std/__host_stdlib/ostream>
 #include <cuda/std/array>
+#include <cuda/std/cstdint>
 
 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_policy, 6> 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 <typename T>
-concept batched_topk_policy_selector = policy_selector<T, batched_topk_policy>;
+concept topk_policy_selector = policy_selector<T, topk_policy>;
 #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 <typename KeyT, typename ValueT, typename SegmentSizeT, ::cuda::std::int64_t MaxK>
-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<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

cccl_upstream/cub/cub/device/dispatch/tuning/tuning_merge.cuh

@@ -20,9 +20,9 @@
 #include <cub/util_type.cuh>
 
 #include <thrust/type_traits/is_contiguous_iterator.h>
-#include <thrust/type_traits/is_trivially_relocatable.h>
 
 #include <cuda/__device/compute_capability.h>
+#include <cuda/__type_traits/is_trivially_copyable.h>
 #include <cuda/std/__algorithm/clamp.h>
 #include <cuda/std/__host_stdlib/ostream>
 #include <cuda/std/concepts>
@@ -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<key_t>,
+      ::cuda::is_trivially_copyable_v<key_t>,
       THRUST_NS_QUALIFIER::is_contiguous_iterator_v<KeysIt1> && THRUST_NS_QUALIFIER::is_contiguous_iterator_v<KeysIt2>,
       ::cuda::std::is_same_v<key_t, it_value_t<KeysIt2>>,
       ::cuda::std::is_same_v<item_t, NullType> ? 0 : int{sizeof(item_t)},
       int{alignof(item_t)},
-      THRUST_NS_QUALIFIER::is_trivially_relocatable_v<item_t>,
+      ::cuda::is_trivially_copyable_v<item_t>,
       THRUST_NS_QUALIFIER::is_contiguous_iterator_v<ItemsIt1>
         && THRUST_NS_QUALIFIER::is_contiguous_iterator_v<ItemsIt2>,
       ::cuda::std::is_same_v<item_t, it_value_t<ItemsIt2>>,

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);

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);

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);

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);
   }

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);

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<DominantT>(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;

cccl_upstream/cub/cub/device/dispatch/tuning/tuning_select_if.cuh

@@ -18,6 +18,7 @@
 #include <cub/block/block_load.cuh>
 #include <cub/block/block_scan.cuh>
 #include <cub/detail/delay_constructor.cuh>
+#include <cub/detail/prefetch.cuh>
 #include <cub/device/dispatch/tuning/common.cuh>
 #include <cub/util_device.cuh>
 #include <cub/util_math.cuh>
@@ -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);

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);

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",
+]

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

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

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

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)

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

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

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",
+]

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."
+    )

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

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 <c10/cuda/CUDAGuard.h>
+#include <torch/cuda.h>
+
+#include <cstdint>
+
+#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 <typename scalar_t,
+          scalar_t (*ACT_FN)(const scalar_t&),
+          bool act_first>
+__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<uintptr_t>(ptr) & 15) == 0;
+}
+
+// Activation and gating kernel template with 128-bit vectorized access
+// optimization.
+template <typename scalar_t,
+          scalar_t (*ACT_FN)(const scalar_t&),
+          bool act_first>
+__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<const int4*>(x_ptr);
+    const int4* y_vec = reinterpret_cast<const int4*>(y_ptr);
+    int4* out_vec = reinterpret_cast<int4*>(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<scalar_t*>(&x);
+      auto* yp = reinterpret_cast<scalar_t*>(&y);
+      auto* rp = reinterpret_cast<scalar_t*>(&r);
+#pragma unroll
+      for (int j = 0; j < kVecSize; j++) {
+        rp[j] = compute<scalar_t, ACT_FN, act_first>(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<scalar_t, ACT_FN, act_first>(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<scalar_t, ACT_FN, act_first>(x, y);
+    }
+  }
+}
+
+template <typename T>
+__device__ __forceinline__ T silu_kernel(const T& x) {
+  // x * sigmoid(x)
+  const float f = static_cast<float>(x);
+  return static_cast<T>(f / (1.0f + expf(-f)));
+}
+
+template <typename T>
+__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<float>(x);
+  constexpr float kAlpha = M_SQRT1_2;
+  return static_cast<T>(f * 0.5f * (1.0f + ::erf(f * kAlpha)));
+}
+
+template <typename T>
+__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<float>(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<T>(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<scalar_t, KERNEL<scalar_t>, ACT_FIRST>              \
+        <<<grid, block, 0, stream>>>(                                      \
+            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), 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

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 <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+#include <cuda_runtime.h>
+#include <torch/extension.h>
+
+#include <cstdint>
+#include <type_traits>
+
+#include "cuda_ops_api.h"
+#include "utils.h"
+
+namespace xllm::kernel::cuda {
+namespace {
+
+template <typename scalar_t>
+struct VecType;
+
+template <>
+struct VecType<c10::Half> {
+  using type = uint4;
+  static constexpr int32_t vec_width = 8;
+};
+
+template <>
+struct VecType<c10::BFloat16> {
+  using type = uint4;
+  static constexpr int32_t vec_width = 8;
+};
+
+template <>
+struct VecType<float> {
+  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 <typename scalar_t, bool kVectorized>
+__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<int64_t>(blockIdx.x);
+  const int32_t dst_linear_idx = static_cast<int32_t>(blockIdx.y);
+  const int64_t tile_idx = static_cast<int64_t>(blockIdx.z);
+
+  scalar_t* __restrict__ key_cache = reinterpret_cast<scalar_t*>(
+      static_cast<uintptr_t>(key_cache_ptrs[layer_idx]));
+  scalar_t* __restrict__ value_cache = reinterpret_cast<scalar_t*>(
+      static_cast<uintptr_t>(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<int64_t>(src_block) * numel_per_block;
+  const int64_t dst_offset = static_cast<int64_t>(dst_block) * numel_per_block;
+
+  if constexpr (kVectorized) {
+    using VecTypeT = typename VecType<scalar_t>::type;
+    constexpr int32_t kVecWidth = VecType<scalar_t>::vec_width;
+    const int64_t num_vecs_per_block = numel_per_block / kVecWidth;
+    const int64_t vec_idx = tile_idx * static_cast<int64_t>(blockDim.x) +
+                            static_cast<int64_t>(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<const VecTypeT*>(key_cache + src_offset + elem_offset);
+    const auto* value_src_vec = reinterpret_cast<const VecTypeT*>(
+        value_cache + src_offset + elem_offset);
+    auto* key_dst_vec =
+        reinterpret_cast<VecTypeT*>(key_cache + dst_offset + elem_offset);
+    auto* value_dst_vec =
+        reinterpret_cast<VecTypeT*>(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<int64_t>(blockDim.x) +
+                             static_cast<int64_t>(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<int32_t>(key_cache_ptrs.size(0));
+  const int32_t num_groups = static_cast<int32_t>(src_block_indices.size(0));
+  const int32_t num_dst_blocks =
+      static_cast<int32_t>(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<scalar_t, float> ||
+                                 std::is_same_v<scalar_t, c10::Half> ||
+                                 std::is_same_v<scalar_t, c10::BFloat16>;
+
+    if constexpr (kHasVecType) {
+      constexpr int32_t kVecWidth = VecType<scalar_t>::vec_width;
+      if (numel_per_block % kVecWidth == 0) {
+        const int64_t tiles_per_block =
+            ceil_div<int64_t>(numel_per_block / kVecWidth, kThreadsPerBlock);
+        const dim3 grid(num_layers, num_dst_blocks, tiles_per_block);
+        block_copy_kernel<scalar_t, true>
+            <<<grid, kThreadsPerBlock, 0, stream>>>(
+                key_cache_ptrs.data_ptr<int64_t>(),
+                value_cache_ptrs.data_ptr<int64_t>(),
+                src_block_indices.data_ptr<int32_t>(),
+                dst_block_indices.data_ptr<int32_t>(),
+                cum_sum.data_ptr<int32_t>(),
+                num_groups,
+                numel_per_block);
+        C10_CUDA_KERNEL_LAUNCH_CHECK();
+        return;
+      }
+    }
+
+    const int64_t tiles_per_block =
+        ceil_div<int64_t>(numel_per_block, kThreadsPerBlock);
+    const dim3 grid(num_layers, num_dst_blocks, tiles_per_block);
+    block_copy_kernel<scalar_t, false><<<grid, kThreadsPerBlock, 0, stream>>>(
+        key_cache_ptrs.data_ptr<int64_t>(),
+        value_cache_ptrs.data_ptr<int64_t>(),
+        src_block_indices.data_ptr<int32_t>(),
+        dst_block_indices.data_ptr<int32_t>(),
+        cum_sum.data_ptr<int32_t>(),
+        num_groups,
+        numel_per_block);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+  });
+}
+
+}  // namespace xllm::kernel::cuda

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 <ATen/DynamicLibrary.h>
+#include <ATen/core/dispatch/Dispatcher.h>
+#include <glog/logging.h>
+
+#include <optional>
+#include <tuple>
+#include <vector>
+
+#include "utils.h"
+
+namespace xllm::kernel::cuda {
+
+// TODO: add head_size parameter
+void rotary_embedding(torch::Tensor& positions,
+                      torch::Tensor& query,
+                      std::optional<torch::Tensor> 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<int64_t> 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<torch::Tensor>& output_lse,
+                   const std::optional<torch::Tensor>& 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<int64_t> 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<torch::Tensor>& output_lse);
+
+void batch_prefill_non_causal(
+    const std::string& uri,
+    ffi::Array<int64_t> 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<torch::Tensor>& output_lse,
+    const std::optional<torch::Tensor>& mask = std::nullopt);
+
+void batch_chunked_prefill(
+    const std::string& uri,
+    ffi::Array<int64_t> 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<torch::Tensor>& output_lse,
+    std::optional<torch::Tensor> qo_indptr = std::nullopt,
+    bool causal = true);
+
+void batch_decode(const std::string& uri,
+                  ffi::Array<int64_t> 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<torch::Tensor>& output_lse,
+                  bool use_tensor_core,
+                  std::optional<torch::Tensor> 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<torch::Tensor> 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<torch::Tensor> 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<torch::Tensor, torch::Tensor> fp8_scaled_quantize(
+    const torch::Tensor& input,
+    const std::optional<torch::Tensor>& output = std::nullopt,
+    const std::optional<torch::Tensor>& 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<torch::Tensor>& bias = std::nullopt,
+    const std::optional<torch::Tensor>& output = std::nullopt);
+
+std::pair<torch::Tensor, torch::Tensor> 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<torch::Tensor, torch::Tensor> 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<torch::Tensor, torch::Tensor> moe_fused_topk(
+    torch::Tensor& gating_output,
+    int64_t topk,
+    bool renormalize,
+    const std::optional<torch::Tensor>& 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<torch::Tensor>& 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<torch::Tensor>& fc1_expert_biases = std::nullopt,
+    const std::optional<torch::Tensor>& fc2_expert_biases = std::nullopt,
+    const std::optional<torch::Tensor>& input_sf = std::nullopt,
+    const std::optional<torch::Tensor>& swiglu_alpha = std::nullopt,
+    const std::optional<torch::Tensor>& swiglu_beta = std::nullopt,
+    const std::optional<torch::Tensor>& swiglu_limit = std::nullopt,
+    const std::optional<torch::Tensor>& 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<torch::Tensor, torch::Tensor, torch::Tensor> 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

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 <hip/amd_detail/amd_hip_bf16.h>
+
+#include <hipcub/hipcub.hpp>
+
+namespace cub = hipcub;
+#else
+#include <cub/cub.cuh>
+#if CUB_VERSION >= 200800
+#include <cuda/functional>
+#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 <typename T,
+          // Number of elements in the array
+          int N,
+          // Alignment requirement in bytes
+          int Alignment = sizeof(T) * N>
+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 <typename T>
+__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 <typename T>
+__device__ float convert_to_float(T x) {
+  if constexpr (std::is_same_v<T, __half>) {
+    return __half2float(x);
+#if defined(USE_DCU)
+  } else if constexpr (std::is_same_v<T, hip_bfloat16>) {
+    return __bfloat162float(reinterpret_cast<const __hip_bfloat16&>(x));
+#else
+  } else if constexpr (std::is_same_v<T, __nv_bfloat16>) {
+    return __bfloat162float(x);
+#endif
+
+  } else if constexpr (std::is_same_v<T, float>) {
+    return x;
+  } else {
+    return static_cast<float>(x);
+  }
+}
+
+// Constructs some constants needed to partition the work across threads at
+// compile time.
+template <typename T, int EXPERTS, int BYTES_PER_LDG>
+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

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 <c10/util/Float8_e4m3fn.h>
+#include <cmath>
+#include <torch/types.h>
+// clang-format on
+namespace xllm {
+namespace kernel {
+namespace cuda {
+
+// FP8 type max value definitions
+template <typename T,
+          typename = std::enable_if_t<std::is_same_v<T, c10::Float8_e4m3fn> ||
+                                      std::is_same_v<T, int8_t>>>
+struct quant_type_max {
+  static constexpr T val() { return std::numeric_limits<T>::max(); }
+};
+
+template <typename T>
+__host__ __device__ static constexpr T quant_type_max_v =
+    quant_type_max<T>::val();
+
+// Minimum scaling factor for quantization types
+template <typename T,
+          typename = std::enable_if_t<std::is_same_v<T, c10::Float8_e4m3fn> ||
+                                      std::is_same_v<T, int8_t>>>
+struct min_scaling_factor {
+  __device__ __host__ static inline float val() {
+    return 1.0f / (quant_type_max_v<T> * 512.0f);
+  }
+};
+
+template <>
+struct min_scaling_factor<int8_t> {
+  __device__ __host__ static inline float val() {
+    return std::numeric_limits<float>::epsilon();
+  }
+};
+
+// Vectorization containers
+template <typename scalar_t, size_t vec_size>
+struct __align__(vec_size * sizeof(scalar_t)) vec_n_t {
+  scalar_t val[vec_size];
+};
+
+template <typename quant_type_t, size_t vec_size>
+struct __align__(vec_size * sizeof(quant_type_t)) q8_n_t {
+  static_assert(std::is_same_v<quant_type_t, int8_t> ||
+                std::is_same_v<quant_type_t, c10::Float8_e4m3fn>);
+  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 <cuda_fp8.h>
+
+// float -> c10::Float8_e4m3fn conversion
+template <typename Tout, typename Tin>
+__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<c10::Float8_e4m3fn, float>(
+    const float& a,
+    const __nv_fp8_interpretation_t fp8_type) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+  return static_cast<c10::Float8_e4m3fn>(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 <bool is_scale_inverted, typename fp8_type>
+__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<fp8_type>, fminf(x, quant_type_max_v<fp8_type>));
+
+#ifdef ENABLE_FP8
+  // Use hardware cvt instruction for fp8 on nvidia
+  return fp8::vec_conversion<fp8_type, float>(r);
+#else
+  return static_cast<fp8_type>(r);
+#endif
+}
+
+// Vectorization utilities
+template <int VEC_SIZE, typename InT, typename OutT, typename ScaOp>
+struct DefaultVecOp {
+  ScaOp scalar_op;
+
+  __device__ __forceinline__ void operator()(
+      vec_n_t<OutT, VEC_SIZE>& dst,
+      const vec_n_t<InT, VEC_SIZE>& src) const {
+#pragma unroll
+    for (int i = 0; i < VEC_SIZE; ++i) {
+      scalar_op(dst.val[i], src.val[i]);
+    }
+  }
+};
+
+template <int VEC_SIZE,
+          typename InT,
+          typename OutT,
+          typename VecOp,
+          typename ScaOp>
+__device__ inline void vectorize_with_alignment(
+    const InT* in,
+    OutT* out,
+    int len,
+    int tid,
+    int stride,
+    VecOp&& vec_op,       // vec_n_t<InT,16> -> vec_n_t<OutT,16>
+    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<uintptr_t>(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<InT, VEC_SIZE>;
+    using vout_t = vec_n_t<OutT, VEC_SIZE>;
+    auto* v_in = reinterpret_cast<const vin_t*>(in);
+    auto* v_out = reinterpret_cast<vout_t*>(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<InT, VEC_SIZE>;
+  using vout_t = vec_n_t<OutT, VEC_SIZE>;
+  auto* v_in = reinterpret_cast<const vin_t*>(in);
+  auto* v_out = reinterpret_cast<vout_t*>(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 <int VEC_SIZE, typename InT, typename OutT, typename ScaOp>
+__device__ __forceinline__ void vectorize_with_alignment(const InT* in,
+                                                         OutT* out,
+                                                         int len,
+                                                         int tid,
+                                                         int stride,
+                                                         ScaOp&& scalar_op) {
+  using Vec = DefaultVecOp<VEC_SIZE, InT, OutT, std::decay_t<ScaOp>>;
+  vectorize_with_alignment<VEC_SIZE>(in,
+                                     out,
+                                     len,
+                                     tid,
+                                     stride,
+                                     Vec{scalar_op},
+                                     std::forward<ScaOp>(scalar_op));
+}
+
+}  // namespace cuda
+}  // namespace kernel
+}  // namespace xllm

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 <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <torch/all.h>
+
+// 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 <typename torch_type>
+class _typeConvert {
+ public:
+  static constexpr bool exists = false;
+};
+
+template <>
+class _typeConvert<float> {
+ 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<c10::Half> {
+ 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<c10::BFloat16> {
+ 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<c10::BFloat16> {
+ 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 <typename scalar_t, int width>
+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<scalar_t>;
+  using T1 = typename Converter::hip_type;
+  using T2 = typename Converter::packed_hip_type;
+  T1 data[width];
+
+  __device__ _f16Vec& operator+=(const _f16Vec<scalar_t, width>& other) {
+    if constexpr (width % 2 == 0) {
+#pragma unroll
+      for (int i = 0; i < width; i += 2) {
+        if constexpr (std::is_same_v<T2, float2>) {
+          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<scalar_t, width>& other) {
+    if constexpr (width % 2 == 0) {
+#pragma unroll
+      for (int i = 0; i < width; i += 2) {
+        if constexpr (std::is_same_v<T2, float2>) {
+          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

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 <ATen/DynamicLibrary.h>
+#if defined(USE_DCU)
+#include <c10/hip/HIPGuard.h>
+#else
+#include <c10/cuda/CUDAGuard.h>
+#endif
+#include <glog/logging.h>
+#include <torch/torch.h>
+#if !defined(USE_DCU)
+#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/extra/c_env_api.h>
+#include <tvm/ffi/extra/module.h>
+#include <tvm/ffi/optional.h>
+#endif
+
+#include <string>
+#include <tuple>
+#include <type_traits>
+#include <unordered_map>
+
+#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 <typename T>
+HOST_DEVICE_INLINE constexpr std::enable_if_t<std::is_integral_v<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<torch::Tensor, double> 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<ffi::Tensor> to_ffi_optional_tensor(
+    const std::optional<torch::Tensor>& optional);
+
+ffi::Array<ffi::Tensor> to_ffi_array_tensors(
+    const std::vector<torch::Tensor>& torch_tensors);
+
+ffi::Optional<ffi::Array<ffi::Tensor>> to_ffi_optional_array_tensors(
+    const std::optional<std::vector<torch::Tensor>>& 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<void*>(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

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 <c10/cuda/CUDAGuard.h>
+#include <torch/cuda.h>
+
+#include <cstdint>
+#include <cub/cub.cuh>
+
+#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 <cuda/std/functional>
+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 <typename scalar_t>
+__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<float>(input[blockIdx.x * input_stride + idx]);
+    variance += x * x;
+  }
+
+  using BlockReduce = cub::BlockReduce<float, 1024>;
+  __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<float>(input[blockIdx.x * input_stride + idx]);
+    out[blockIdx.x * hidden_size + idx] =
+        (static_cast<scalar_t>(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 <typename scalar_t, int width>
+__global__ std::enable_if_t<(width > 0) && _typeConvert<scalar_t>::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<scalar_t, width>>);
+  static_assert(sizeof(_f16Vec<scalar_t, width>) == 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<scalar_t, width>*>(input);
+  auto* __restrict__ residual_v =
+      reinterpret_cast<_f16Vec<scalar_t, width>*>(residual);
+  auto* __restrict__ weight_v =
+      reinterpret_cast<const _f16Vec<scalar_t, width>*>(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<scalar_t, width> temp = input_v[strided_id];
+    temp += residual_v[id];
+    variance += temp.sum_squares();
+    residual_v[id] = temp;
+  }
+
+  using BlockReduce = cub::BlockReduce<float, 1024>;
+  __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<scalar_t, width> 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 <typename scalar_t, int width>
+__global__ std::enable_if_t<(width == 0) || !_typeConvert<scalar_t>::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<float>(z);
+    variance += x * x;
+    residual[blockIdx.x * hidden_size + idx] = z;
+  }
+
+  using BlockReduce = cub::BlockReduce<float, 1024>;
+  __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<float>(residual[blockIdx.x * hidden_size + idx]);
+    input[blockIdx.x * input_stride + idx] =
+        (static_cast<scalar_t>(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<scalar_t, width>                      \
+            <<<grid, block, 0, stream>>>(input.data_ptr<scalar_t>(),    \
+                                         input_stride,                  \
+                                         residual.data_ptr<scalar_t>(), \
+                                         weight.data_ptr<scalar_t>(),   \
+                                         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 <typename scalar_t, typename fp8_type>
+__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<float>(input_row[idx]);
+    variance += x * x;
+  }
+
+  using BlockReduce = cub::BlockReduce<float, 1024>;
+  __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<float>(input_row[idx]);
+    float out_norm = (static_cast<scalar_t>(x * s_variance)) *
+                     static_cast<float>(weight[idx]);
+    out[blockIdx.x * hidden_size + idx] =
+        xllm::kernel::cuda::scaled_fp8_conversion<true, fp8_type>(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 <typename scalar_t, int width, typename fp8_type>
+__global__ std::enable_if_t<(width > 0) && _typeConvert<scalar_t>::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<scalar_t, width>>);
+  static_assert(sizeof(_f16Vec<scalar_t, width>) == 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<scalar_t, width>*>(input);
+  auto* __restrict__ residual_v =
+      reinterpret_cast<_f16Vec<scalar_t, width>*>(residual);
+  auto* __restrict__ weight_v =
+      reinterpret_cast<const _f16Vec<scalar_t, width>*>(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<scalar_t, width> temp = input_v[strided_id];
+    temp += residual_v[id];
+    variance += temp.sum_squares();
+    residual_v[id] = temp;  // Store updated residual
+  }
+
+  using BlockReduce = cub::BlockReduce<float, 1024>;
+  __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<scalar_t, width> 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<scalar_t>::convert(temp.data[i]);
+      out[id * width + i] =
+          xllm::kernel::cuda::scaled_fp8_conversion<true, fp8_type>(val,
+                                                                    scale_inv);
+    }
+  }
+}
+
+/**
+ * Generic fused add + RMSNorm + FP8 quant kernel (fallback for unaligned data)
+ */
+template <typename scalar_t, int width, typename fp8_type>
+__global__ std::enable_if_t<(width == 0) || !_typeConvert<scalar_t>::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<float>(z);
+    variance += x * x;
+    residual[blockIdx.x * hidden_size + idx] = z;  // Store updated residual
+  }
+
+  using BlockReduce = cub::BlockReduce<float, 1024>;
+  __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<float>(residual[blockIdx.x * hidden_size + idx]);
+    float out_norm = (static_cast<scalar_t>(x * s_variance)) *
+                     static_cast<float>(weight[idx]);
+    out[blockIdx.x * hidden_size + idx] =
+        xllm::kernel::cuda::scaled_fp8_conversion<true, fp8_type>(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<scalar_t,            \
+                                                         width,               \
+                                                         fp8_t>               \
+                  <<<grid, block, 0, stream>>>(out.data_ptr<fp8_t>(),         \
+                                               input.data_ptr<scalar_t>(),    \
+                                               input_stride,                  \
+                                               residual.data_ptr<scalar_t>(), \
+                                               weight.data_ptr<scalar_t>(),   \
+                                               scale.data_ptr<float>(),       \
+                                               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<scalar_t>
+        <<<grid, block, 0, stream>>>(output.data_ptr<scalar_t>(),
+                                     input.data_ptr<scalar_t>(),
+                                     input_stride,
+                                     weight.data_ptr<scalar_t>(),
+                                     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<std::uintptr_t>(input.data_ptr());
+  auto res_ptr = reinterpret_cast<std::uintptr_t>(residual.data_ptr());
+  auto wt_ptr = reinterpret_cast<std::uintptr_t>(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<scalar_t, fp8_t>
+              <<<grid, block, 0, stream>>>(out.data_ptr<fp8_t>(),
+                                           input.data_ptr<scalar_t>(),
+                                           input_stride,
+                                           weight.data_ptr<scalar_t>(),
+                                           scale.data_ptr<float>(),
+                                           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<std::uintptr_t>(input.data_ptr());
+  auto res_ptr = reinterpret_cast<std::uintptr_t>(residual.data_ptr());
+  auto wt_ptr = reinterpret_cast<std::uintptr_t>(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

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 <c10/cuda/CUDAStream.h>
+
+#include "cuda_ops_api.h"
+#include "device_utils.cuh"
+
+namespace xllm::kernel::cuda {
+
+template <typename T>
+__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<int>(n, 1024));
+  DISPATCH_FLOATING_TYPES(
+      keys.scalar_type(), "reshape_paged_cache_kernel", [&] {
+        reshape_paged_cache_kernel<scalar_t>
+            <<<grid, block, 0, c10::cuda::getCurrentCUDAStream()>>>(
+                slot_ids.data_ptr<int>(),
+                keys.data_ptr<scalar_t>(),
+                values.data_ptr<scalar_t>(),
+                key_cache.data_ptr<scalar_t>(),
+                value_cache.data_ptr<scalar_t>(),
+                k_stride,
+                v_stride,
+                n_kv_heads,
+                head_dim,
+                block_size);
+      });
+}
+
+}  // namespace xllm::kernel::cuda

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 <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/all.h>
+
+#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 <typename scalar_t, bool IS_NEOX>
+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 <typename scalar_t, bool IS_NEOX>
+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<scalar_t, IS_NEOX>(
+        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<scalar_t, IS_NEOX>(
+          key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
+    }
+  }
+}
+
+template <typename scalar_t, bool IS_NEOX>
+__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<scalar_t, IS_NEOX>(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<torch::Tensor> 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<int64_t>(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<scalar_t, true><<<grid, block, 0, stream>>>(
+              positions.data_ptr<int64_t>(),
+              query.data_ptr<scalar_t>(),
+              key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
+              cos_sin_cache.data_ptr<scalar_t>(),
+              rot_dim,
+              query_stride,
+              key_stride,
+              head_stride,
+              num_heads,
+              num_kv_heads,
+              head_size);
+        } else {
+          rotary_embedding_kernel<scalar_t, false><<<grid, block, 0, stream>>>(
+              positions.data_ptr<int64_t>(),
+              query.data_ptr<scalar_t>(),
+              key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
+              cos_sin_cache.data_ptr<scalar_t>(),
+              rot_dim,
+              query_stride,
+              key_stride,
+              head_stride,
+              num_heads,
+              num_kv_heads,
+              head_size);
+        }
+      });
+}
+
+}  // namespace xllm::kernel::cuda

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

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<torch::Tensor>& value,
+                         torch::Tensor& key_cache,
+                         std::optional<torch::Tensor>& 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<torch::Tensor>& value,
+                   torch::Tensor& output,
+                   std::optional<torch::Tensor>& output_lse,
+                   const std::optional<torch::Tensor>& q_cu_seq_lens,
+                   const std::optional<torch::Tensor>& kv_cu_seq_lens,
+                   const std::optional<torch::Tensor>& alibi_slope,
+                   const std::optional<torch::Tensor>& attn_bias,
+                   const std::optional<torch::Tensor>& q_quant_scale,
+                   const std::optional<torch::Tensor>& k_quant_scale,
+                   const std::optional<torch::Tensor>& 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<double>(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<torch::Tensor>& v_cache,
+                  std::optional<torch::Tensor>& output_lse,
+                  const std::optional<torch::Tensor>& q_quant_scale,
+                  const std::optional<torch::Tensor>& k_cache_quant_scale,
+                  const std::optional<torch::Tensor>& v_cache_quant_scale,
+                  const std::optional<torch::Tensor>& out_quant_scale,
+                  const std::optional<torch::Tensor>& alibi_slope,
+                  const std::optional<torch::Tensor>& 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

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 <glog/logging.h>
+
+#include "ilu_ops_api.h"
+
+namespace xllm::kernel::ilu {
+
+std::tuple<torch::Tensor, torch::Tensor> moe_active_topk(
+    const torch::Tensor& input,
+    int64_t topk,
+    int64_t num_expert_group,
+    int64_t topk_group,
+    bool normalize,
+    const std::optional<torch::Tensor>& mask,
+    const std::string& normed_by,
+    const std::string& scoring_func,
+    double route_scale,
+    const std::optional<torch::Tensor>& 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<torch::Tensor> 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

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<torch::Tensor>& 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<int64_t>());
+
+  return output;
+}
+
+}  // namespace xllm::kernel::ilu

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 <ATen/DynamicLibrary.h>
+#include <ATen/core/dispatch/Dispatcher.h>
+#include <cuda_runtime.h>
+#include <glog/logging.h>
+#include <torch/all.h>
+
+#include <optional>
+
+#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<torch::Tensor>& value,  // (num_tokens, num_heads, head_size)
+    torch::Tensor& key_cache,  // (num_blocks, num_heads, block_size, head_size)
+    std::optional<torch::Tensor>&
+        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<torch::Tensor>& value,
+                   torch::Tensor& output,
+                   std::optional<torch::Tensor>& output_lse,
+                   const std::optional<torch::Tensor>& q_cu_seq_lens,
+                   const std::optional<torch::Tensor>& kv_cu_seq_lens,
+                   const std::optional<torch::Tensor>& alibi_slope,
+                   const std::optional<torch::Tensor>& attn_bias,
+                   const std::optional<torch::Tensor>& q_quant_scale,
+                   const std::optional<torch::Tensor>& k_quant_scale,
+                   const std::optional<torch::Tensor>& 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<torch::Tensor>& v_cache,
+                  std::optional<torch::Tensor>& output_lse,
+                  const std::optional<torch::Tensor>& q_quant_scale,
+                  const std::optional<torch::Tensor>& k_cache_quant_scale,
+                  const std::optional<torch::Tensor>& v_cache_quant_scale,
+                  const std::optional<torch::Tensor>& out_quant_scale,
+                  const std::optional<torch::Tensor>& alibi_slope,
+                  const std::optional<torch::Tensor>& 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<torch::Tensor>& residual,
+                         torch::Tensor& weight,
+                         std::optional<torch::Tensor>& bias,
+                         std::optional<torch::Tensor>& 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<torch::Tensor> bias);
+
+std::tuple<torch::Tensor, torch::Tensor> moe_active_topk(
+    const torch::Tensor& input,
+    int64_t topk,
+    int64_t num_expert_group,
+    int64_t topk_group,
+    bool normalize,
+    const std::optional<torch::Tensor>& mask,
+    const std::string& normed_by,
+    const std::string& scoring_func,
+    double route_scale,
+    const std::optional<torch::Tensor>& e_score_correction_bias);
+
+std::vector<torch::Tensor> 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<torch::Tensor>& dst_to_src,
+                         torch::Tensor& output);
+
+torch::Tensor moe_combine_result(torch::Tensor& input, torch::Tensor& weight);
+}  // namespace xllm::kernel::ilu

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 <torch/all.h>
+
+#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<torch::Tensor>& alibi_slopes,
+    const std::optional<torch::Tensor>& sinks,
+    std::optional<torch::Tensor>& 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<torch::Tensor>& 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<torch::Tensor>& sinks);
+
+torch::Tensor ixformer_linear(torch::Tensor& input,
+                              torch::Tensor& weight,
+                              int64_t act_type,
+                              const std::optional<torch::Tensor>& bias,
+                              const std::optional<torch::Tensor>& out,
+                              const std::optional<bool> persistent);
+
+torch::Tensor ixformer_linear_ex(torch::Tensor& input,
+                                 torch::Tensor& weight,
+                                 const c10::optional<torch::Tensor>& bias,
+                                 const c10::optional<torch::Tensor>& 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<torch::Tensor>& fused_bias,
+                       double alpha,
+                       double eps,
+                       bool is_post);
+
+void rms_norm(torch::Tensor& input,
+              torch::Tensor& weight,
+              torch::Tensor& output,
+              const std::optional<torch::Tensor>& 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<torch::Tensor>& expert_mask,
+    const c10::optional<torch::Tensor>& expert_sizes_cpu,
+    const c10::optional<torch::Tensor>& 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<torch::Tensor>& 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<torch::Tensor>& dst_to_src,
+                           const c10::optional<torch::Tensor>& 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<torch::Tensor>& mul_weight,
+                           const c10::optional<torch::Tensor>& mask,
+                           const c10::optional<torch::Tensor>& extra_residual,
+                           double scaling_factor);
+}  // namespace ixformer::infer

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<torch::Tensor> bias) {
+  int64_t act_type = -1;
+  bool persistent = false;
+  std::vector<int64_t> 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

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<torch::Tensor>& residual,
+                         torch::Tensor& weight,
+                         std::optional<torch::Tensor>& bias,
+                         std::optional<torch::Tensor>& 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<torch::Tensor> fused_bias = std::nullopt;
+  infer::rms_norm(input, weight, output, fused_bias, eps);
+}
+
+}  // namespace xllm::kernel::ilu

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

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

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']

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

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

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

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()

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