diff --git a/fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_sm100a.cuh b/fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_sm100a.cuh index 82be6ba73..3605b8e25 100644 --- a/fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_sm100a.cuh +++ b/fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_sm100a.cuh @@ -174,6 +174,13 @@ __host__ inline cudaError_t make_tma_kv_units(CUtensorMap* map, const __nv_bfloa // subset of the repeatedly reduced dQ lines resident pays (~0.6% at 524K/1M tokens); below it the // policy register costs more than it saves. Threshold scales with accumulator BYTES. constexpr int CACHE_WAVE_MIN_SEQ_LEN = 524288; +// The in-kernel identity order (workitem_remap == nullptr below ORDER_MIN_KV_BLOCKS) has no +// sub-array to offset a chunk into, so the chunked launches that DQ_L2_KEEP enables must only +// ever run with a device-computed order: the L2 transition has to sit at or above the +// order-kernel threshold. keep_dq_l2 <=> S * sizeof(dq_accum_t) >= 2 * CACHE_WAVE_MIN_SEQ_LEN. +static_assert((size_t)CACHE_WAVE_MIN_SEQ_LEN * 2 / sizeof(dq_accum_t) >= + (size_t)ORDER_MIN_KV_BLOCKS * BLOCK, + "DQ_L2_KEEP chunked launches need the device-computed work order"); template __host__ inline cudaError_t launch_main(const BlockSparseVsaBwdArgs& args, const int* work_remap, @@ -205,12 +212,18 @@ __host__ inline cudaError_t launch_main(const BlockSparseVsaBwdArgs& args, const at[0].val.clusterDim.z = 1; cfg.attrs = at; cfg.numAttrs = 1; + // Guarded by the static_assert above; kept as a runtime check for other callers of the + // launch API that pass their own thresholds. + if (work_remap == nullptr && chunk != total) { + return cudaErrorInvalidValue; + } for (int base = 0; base < total; base += chunk) { const int count = std::min(chunk, total - base); cfg.gridDim = dim3((unsigned)count, 1, 1); // A chunk starts at work id `base`: it gets the order's sub-array. e = cudaLaunchKernelEx(&cfg, kernel, tk, tv, tqt, tdot, tdk, tdv, args.dqaccum, args.lse, - args.delta, args.k2q_idx, args.k2q_num, work_remap + base, + args.delta, args.k2q_idx, args.k2q_num, + work_remap ? work_remap + base : nullptr, args.variable_block_sizes, args.max_q_blocks, B, H, S, scale_log2, args.sm_scale); if (e != cudaSuccess) { diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py index 2de50d47c..76d338cd6 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py @@ -33,8 +33,9 @@ BHSD = True def set_extension(module) -> None: """Use an already-loaded extension module exposing ``block_sparse_sm100a_bwd``. - The in-repo tests build the binding standalone (tests/jit_ext.py) and inject it here, so - the backend can be exercised without a fastvideo_kernel install. + A standalone build of ``block_sparse_bwd_sm100a.cu`` (for example through + ``torch.utils.cpp_extension.load`` with a ten-line pybind wrapper) can be injected here, so + the backend can be exercised without rebuilding the fastvideo_kernel wheel. """ global _C, _BWD, _HAS_VSA_BWD_SM100A _C = module