diff --git a/muh/include/muh/tuning/tuning_batched_topk.cuh b/muh/include/muh/tuning/tuning_batched_topk.cuh index 754c30cb..97d102af 100644 --- a/muh/include/muh/tuning/tuning_batched_topk.cuh +++ b/muh/include/muh/tuning/tuning_batched_topk.cuh @@ -4,7 +4,14 @@ // CCCL: extends topk with batch-level parallelism // // vllm relevance: multi-sequence parallel decode top-k sampling -// SMEM risk: same as topk (handled by bits_per_pass) +// SMEM risk: histogram SMEM = (1 << bits) * sizeof(int) * max_batches +// bits=11: 2048*4*batches — even at batches=1, keys_tile + 8192 can overflow +// bits=8: 256*4*batches — much safer +// Decision: use bits=8 (same as radix_sort) for all key sizes on BI-V100. +// +// SMEM layout: keys_tile (union with values_tile) + histogram per batch +// total = threads * items * max(key_size, value_size) + (1< 0 ? remaining_smem / hist_smem_per_batch : 1; if (max_batches < 1) max_batches = 1; if (max_batches > 32) max_batches = 32; // cap for occupancy - return {base, max_batches}; + // Verify total SMEM + int total_smem = keys_tile + hist_smem_per_batch * max_batches; + while (total_smem > hw.max_shared_memory_per_block && max_batches > 1) { + max_batches--; + total_smem = keys_tile + hist_smem_per_batch * max_batches; + } + + // Override base bits_per_pass to 8 + topk::TopkPolicy adjusted_base = base; + adjusted_base.bits_per_pass = bits; + + return {adjusted_base, max_batches}; } };