From 090b7e0841a1b81d2b01101dce0404fe1511e55b Mon Sep 17 00:00:00 2001 From: Mingyang Hao Date: Thu, 14 May 2026 00:27:05 -0700 Subject: [PATCH 1/2] [None][perf] DSV4 compressor: enable 4-warp Phase 3 reduction for MTP NEXT_N=3..4 For DSV4 GEN (decode) with MTP-3, each step accepts up to NEXT_N=4 tokens per request, so pagedKvCompressKernel runs with NEXT_N=4. The current dispatch only enables the 4-warp Phase 3 reduction for NEXT_N<=2, leaving the NEXT_N=3..4 path as a single warp doing 128 serial paged loads. With the kernel launched at (batch, HEAD_BLOCKS=4)=tiny grid + 1 warp/block, per-SM resident warps are <14 and the 300-cycle DRAM latency on every load is fully exposed -- slow path balloons to ~55-60 us per call at batch>=256. Relaxing the dispatch condition from `next_n<=2` to `next_n<=4` and instantiating the existing multi-warp template for NEXT_N=3,4 across both head_dim variants drops the slow path 3x at the production batch=16 per-rank shape (20.9us -> 7.2us) and 2x at batch=128 (40.1us -> 13.9us). End-to-end this kernel goes from 2.1% of GEN GPU time to ~0.6% on the user's DSV4-Pro concurrency=512 trace. Phase 3 multi-warp merge logic is per-c (per emitted compressed token) and already handles NEXT_N >= 2 generically -- no algorithmic change. Tested: - tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_kernel.py::test_decode_mtp passes all 4 MTP_CONFIGS (overlap_hd128_next4, overlap_hd512_next3, basic_hd128_multi_batch_next4, basic_hd512_next4) -> bit-identical numerical results. - Microbench (HD=512, KV=bf16, STATE=fp32, R=128) emit-path latency: batch=16 NEXT_N=4: 20.9us -> 7.2us (2.90x) batch=128 NEXT_N=4: 40.1us -> 13.9us (2.88x) batch=256 NEXT_N=4: 51.1us -> 25.6us (2.00x) Signed-off-by: Mingyang Hao --- .../compressorKernels/compressorKernels.cu | 29 ++++++++++++++++--- 1 file changed, 25 insertions(+), 4 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu index cf48dd04b720..869aac6362a8 100644 --- a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu +++ b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu @@ -633,7 +633,7 @@ INST_DECODE_DTYPES(512, 128) // HD=128: ELEM_PER_BLOCK=128, smem = 3*4*128*4 = 6 KB. // HD=512 bf16: ELEM_PER_BLOCK=256, smem = 3*4*256*4 = 12 KB. // HD=512 fp32: ELEM_PER_BLOCK=128, smem = 3*4*128*4 = 6 KB. -// NEXT_N=1 (single-token decode) and NEXT_N=2 (MTP speculative decode). +// NEXT_N=1 (single-token decode), NEXT_N=2..4 (MTP speculative decode; MTP-3 → NEXT_N=4). INST_DECODE(128, 2, 2, 128, 1, 4) INST_DECODE(128, 2, 4, 128, 1, 4) INST_DECODE(128, 4, 2, 128, 1, 4) @@ -642,6 +642,14 @@ INST_DECODE(128, 2, 2, 128, 2, 4) INST_DECODE(128, 2, 4, 128, 2, 4) INST_DECODE(128, 4, 2, 128, 2, 4) INST_DECODE(128, 4, 4, 128, 2, 4) +INST_DECODE(128, 2, 2, 128, 3, 4) +INST_DECODE(128, 2, 4, 128, 3, 4) +INST_DECODE(128, 4, 2, 128, 3, 4) +INST_DECODE(128, 4, 4, 128, 3, 4) +INST_DECODE(128, 2, 2, 128, 4, 4) +INST_DECODE(128, 2, 4, 128, 4, 4) +INST_DECODE(128, 4, 2, 128, 4, 4) +INST_DECODE(128, 4, 4, 128, 4, 4) INST_DECODE(512, 2, 2, 128, 1, 4) INST_DECODE(512, 2, 4, 128, 1, 4) INST_DECODE(512, 4, 2, 128, 1, 4) @@ -650,6 +658,14 @@ INST_DECODE(512, 2, 2, 128, 2, 4) INST_DECODE(512, 2, 4, 128, 2, 4) INST_DECODE(512, 4, 2, 128, 2, 4) INST_DECODE(512, 4, 4, 128, 2, 4) +INST_DECODE(512, 2, 2, 128, 3, 4) +INST_DECODE(512, 2, 4, 128, 3, 4) +INST_DECODE(512, 4, 2, 128, 3, 4) +INST_DECODE(512, 4, 4, 128, 3, 4) +INST_DECODE(512, 2, 2, 128, 4, 4) +INST_DECODE(512, 2, 4, 128, 4, 4) +INST_DECODE(512, 4, 2, 128, 4, 4) +INST_DECODE(512, 4, 4, 128, 4, 4) #undef INST_DECODE_DTYPES #undef INST_DECODE_NN @@ -692,7 +708,10 @@ void pagedKvCompressLaunch(void const* kv_score, float const* ape, void* paged_k // For large compress_ratio, use 4-warp parallel reduction to cut the serial // softmax loop from COMPRESS_RATIO iterations to COMPRESS_RATIO/4 per warp. - // Supported configs: CR=128, (HD=128 or HD=512), NEXT_N=1 or NEXT_N=2. + // Supported configs: CR=128, (HD=128 or HD=512), NEXT_N in 1..4. NEXT_N>2 + // is required for MTP-3 decode (each step accepts up to 4 tokens per request); + // without multi-warp the slow path is a single warp doing 128 serial paged + // loads, which is DRAM-latency-bound (no other warps to hide it). // // smem per block = 3 * MULTI_WARP * ELEM_PER_BLOCK * sizeof(float) // where ELEM_PER_BLOCK = nthreads_inner * vec = HEAD_DIM / HEAD_BLOCKS. @@ -700,7 +719,7 @@ void pagedKvCompressLaunch(void const* kv_score, float const* ape, void* paged_k // HD=512 with max elem size 2 (vec=8, HEAD_BLOCKS=2): ELEM_PER_BLOCK=256 → 12 KB. // HD=512 with max elem size 4 (vec=4, HEAD_BLOCKS=4): ELEM_PER_BLOCK=128 → 6 KB. constexpr int MULTI_WARP = 4; - bool const use_multi_warp = (compress_ratio == 128 && next_n <= 2); + bool const use_multi_warp = (compress_ratio == 128 && next_n <= 4); int const num_red_warps = use_multi_warp ? MULTI_WARP : 1; int const nthreads = nthreads_inner * num_red_warps; int const elem_per_block = nthreads_inner * vec; // = HEAD_DIM / HEAD_BLOCKS @@ -721,8 +740,10 @@ void pagedKvCompressLaunch(void const* kv_score, float const* ape, void* paged_k #define DISPATCH_NN_MW(HD, KV_EB, STATE_EB, CR) \ switch (next_n) \ { \ + case 1: LAUNCH_DECODE_MW(HD, KV_EB, STATE_EB, CR, 1); break; \ case 2: LAUNCH_DECODE_MW(HD, KV_EB, STATE_EB, CR, 2); break; \ - default: LAUNCH_DECODE_MW(HD, KV_EB, STATE_EB, CR, 1); break; \ + case 3: LAUNCH_DECODE_MW(HD, KV_EB, STATE_EB, CR, 3); break; \ + default: LAUNCH_DECODE_MW(HD, KV_EB, STATE_EB, CR, 4); break; \ } #define DISPATCH_NN(HD, KV_EB, STATE_EB, CR) \ From 8c36875f0962452dc442c66be4687ad144c8a4c1 Mon Sep 17 00:00:00 2001 From: Mingyang Hao Date: Thu, 14 May 2026 01:20:39 -0700 Subject: [PATCH 2/2] [None][chore] DSV4 compressor: collapse dispatch macros into single X-macro The decode-kernel side of compressorKernels.cu had grown 5 layered macros (INST_DECODE / INST_DECODE_NN / INST_DECODE_DTYPES / LAUNCH_DECODE / LAUNCH_DECODE_MW / DISPATCH_NN / DISPATCH_NN_MW / DISPATCH_DTYPE) plus a 7-arm if-else cascade in the launcher. Adding a new (HD, KV_EB, STATE_EB, CR, NN, NRW) config required touching 5 places, and the multi-warp NEXT_N=3..4 patch from the previous commit had to hand-write 16 INST_DECODE lines because the layered macros could not express "this dtype combo exists at NRW=4 only". Replace the layered dispatch with a single X-macro FOREACH_DECODE_CONFIG(F) listing every valid tuple once. Both the explicit template instantiations and the runtime dispatcher walk the same list -- adding/removing a config is a one-line edit, and instantiation/dispatch can never drift out of sync. Two small fan-out helpers (FOREACH_DECODE_NN, FOREACH_DECODE_DTYPE) keep the master list one row per (HD, CR, NRW) bucket. Net: -141 / +63 lines (78 fewer); the launcher dispatch shrinks from ~90 lines of nested switch/if-else to a single TRY_LAUNCH walk over the config list. No semantic change. Verified: - libtensorrt_llm.so links cleanly (compile + relink with sm_100f). - test_compressor_kernel.py full suite: 63 passed, 22 skipped (no regressions). - test_decode_mtp 4/4 passes for all MTP_CONFIGS -- bit-identical to pre-refactor. Signed-off-by: Mingyang Hao --- .../compressorKernels/compressorKernels.cu | 206 ++++++------------ 1 file changed, 65 insertions(+), 141 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu index 869aac6362a8..24e88c604c17 100644 --- a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu +++ b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu @@ -64,6 +64,7 @@ #include "tensorrt_llm/kernels/compressorKernels/compressorKernels.h" #include "tensorrt_llm/common/assert.h" +#include #include #include #include @@ -610,65 +611,52 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo } } -// Explicit instantiations for decode kernel (NUM_RED_WARPS defaults to 1) +// ============================================================================ +// Decode kernel configuration matrix — single source of truth. +// +// X-macro listing every supported (HD, KV_EB, STATE_EB, CR, NN, NRW) tuple. +// Both the explicit template instantiations below AND the runtime dispatcher +// in pagedKvCompressLaunch() walk this list, so adding/removing a config is +// a one-line edit. +// +// HD — HEAD_DIM in {128, 512} +// KV_EB — kv_score element bytes in {2 (bf16), 4 (fp32)} +// STATE_EB — paged state element bytes in {2 (bf16), 4 (fp32)} +// CR — COMPRESS_RATIO in {4, 128} +// NN — NEXT_N (new tokens / decode step) in {1..4} +// NRW — NUM_RED_WARPS — 4 only when CR=128 (multi-warp Phase 3 reduce +// hides DRAM latency for the heavier R=128 chunk); 1 otherwise. +// +// Multi-warp SMEM budget (per block): 3 * NRW * ELEM_PER_BLOCK * sizeof(float). +// HD=128: ELEM_PER_BLOCK=128 → 6 KB +// HD=512 bf16: ELEM_PER_BLOCK=256 → 12 KB +// HD=512 fp32: ELEM_PER_BLOCK=128 → 6 KB +// ============================================================================ + +// Per-axis fan-outs (used to keep the master list compact). +#define FOREACH_DECODE_NN(F, HD, KV, ST, CR, NRW) \ + F(HD, KV, ST, CR, 1, NRW) F(HD, KV, ST, CR, 2, NRW) F(HD, KV, ST, CR, 3, NRW) F(HD, KV, ST, CR, 4, NRW) +#define FOREACH_DECODE_DTYPE(F, HD, CR, NRW) \ + FOREACH_DECODE_NN(F, HD, 2, 2, CR, NRW) \ + FOREACH_DECODE_NN(F, HD, 2, 4, CR, NRW) \ + FOREACH_DECODE_NN(F, HD, 4, 2, CR, NRW) FOREACH_DECODE_NN(F, HD, 4, 4, CR, NRW) + +// Master list. Order does not matter; the dispatcher walks linearly. +// clang-format off +#define FOREACH_DECODE_CONFIG(F) \ + /* CR=4: single-warp only (small reduction; multi-warp would over-subscribe). */ \ + FOREACH_DECODE_DTYPE(F, 128, 4, 1) FOREACH_DECODE_DTYPE(F, 512, 4, 1) \ + /* CR=128: single-warp fallback (covers next_n>4 path which currently isn't reached). */ \ + FOREACH_DECODE_DTYPE(F, 128, 128, 1) FOREACH_DECODE_DTYPE(F, 512, 128, 1) \ + /* CR=128: multi-warp fast path. Used whenever next_n <= 4 (i.e. MTP-3 and below). */ \ + FOREACH_DECODE_DTYPE(F, 128, 128, 4) FOREACH_DECODE_DTYPE(F, 512, 128, 4) +// clang-format on + +// Generate explicit template instantiations. #define INST_DECODE(HD, KV_EB, STATE_EB, CR, NN, NRW) \ template __global__ void pagedKvCompressKernel(void const*, float const*, void*, \ void*, int32_t const*, int32_t const*, void*, int32_t const*, int32_t const*, int32_t const*, int, int, int); - -#define INST_DECODE_NN(HD, KV_EB, STATE_EB, CR) \ - INST_DECODE(HD, KV_EB, STATE_EB, CR, 1, 1) \ - INST_DECODE(HD, KV_EB, STATE_EB, CR, 2, 1) \ - INST_DECODE(HD, KV_EB, STATE_EB, CR, 3, 1) INST_DECODE(HD, KV_EB, STATE_EB, CR, 4, 1) - -#define INST_DECODE_DTYPES(HD, CR) \ - INST_DECODE_NN(HD, 2, 2, CR) \ - INST_DECODE_NN(HD, 2, 4, CR) INST_DECODE_NN(HD, 4, 2, CR) INST_DECODE_NN(HD, 4, 4, CR) - -INST_DECODE_DTYPES(128, 4) -INST_DECODE_DTYPES(128, 128) -INST_DECODE_DTYPES(512, 4) -INST_DECODE_DTYPES(512, 128) - -// 4-warp parallel reduction variants for large compress_ratio. -// HD=128: ELEM_PER_BLOCK=128, smem = 3*4*128*4 = 6 KB. -// HD=512 bf16: ELEM_PER_BLOCK=256, smem = 3*4*256*4 = 12 KB. -// HD=512 fp32: ELEM_PER_BLOCK=128, smem = 3*4*128*4 = 6 KB. -// NEXT_N=1 (single-token decode), NEXT_N=2..4 (MTP speculative decode; MTP-3 → NEXT_N=4). -INST_DECODE(128, 2, 2, 128, 1, 4) -INST_DECODE(128, 2, 4, 128, 1, 4) -INST_DECODE(128, 4, 2, 128, 1, 4) -INST_DECODE(128, 4, 4, 128, 1, 4) -INST_DECODE(128, 2, 2, 128, 2, 4) -INST_DECODE(128, 2, 4, 128, 2, 4) -INST_DECODE(128, 4, 2, 128, 2, 4) -INST_DECODE(128, 4, 4, 128, 2, 4) -INST_DECODE(128, 2, 2, 128, 3, 4) -INST_DECODE(128, 2, 4, 128, 3, 4) -INST_DECODE(128, 4, 2, 128, 3, 4) -INST_DECODE(128, 4, 4, 128, 3, 4) -INST_DECODE(128, 2, 2, 128, 4, 4) -INST_DECODE(128, 2, 4, 128, 4, 4) -INST_DECODE(128, 4, 2, 128, 4, 4) -INST_DECODE(128, 4, 4, 128, 4, 4) -INST_DECODE(512, 2, 2, 128, 1, 4) -INST_DECODE(512, 2, 4, 128, 1, 4) -INST_DECODE(512, 4, 2, 128, 1, 4) -INST_DECODE(512, 4, 4, 128, 1, 4) -INST_DECODE(512, 2, 2, 128, 2, 4) -INST_DECODE(512, 2, 4, 128, 2, 4) -INST_DECODE(512, 4, 2, 128, 2, 4) -INST_DECODE(512, 4, 4, 128, 2, 4) -INST_DECODE(512, 2, 2, 128, 3, 4) -INST_DECODE(512, 2, 4, 128, 3, 4) -INST_DECODE(512, 4, 2, 128, 3, 4) -INST_DECODE(512, 4, 4, 128, 3, 4) -INST_DECODE(512, 2, 2, 128, 4, 4) -INST_DECODE(512, 2, 4, 128, 4, 4) -INST_DECODE(512, 4, 2, 128, 4, 4) -INST_DECODE(512, 4, 4, 128, 4, 4) - -#undef INST_DECODE_DTYPES -#undef INST_DECODE_NN +FOREACH_DECODE_CONFIG(INST_DECODE) #undef INST_DECODE // ============================================================================ @@ -727,97 +715,33 @@ void pagedKvCompressLaunch(void const* kv_score, float const* ape, void* paged_k dim3 grid(batch_size, head_blocks); -#define LAUNCH_DECODE(HD, KV_EB, STATE_EB, CR, NN) \ - pagedKvCompressKernel<<>>(kv_score, ape, \ - paged_kv, paged_score, block_table_kv, block_table_score, output, kv_lens, cu_seq_lens, cu_kv_comp, page_size, \ - max_blocks, out_elem_bytes) + // Clamp the runtime next_n into the supported range; configs above 4 fall + // back to the NN=4 instantiation (matches the prior `default:` arm). + int const next_n_dispatch = std::min(next_n, 4); -#define LAUNCH_DECODE_MW(HD, KV_EB, STATE_EB, CR, NN) \ - pagedKvCompressKernel<<>>(kv_score, \ - ape, paged_kv, paged_score, block_table_kv, block_table_score, output, kv_lens, cu_seq_lens, cu_kv_comp, \ - page_size, max_blocks, out_elem_bytes) - -#define DISPATCH_NN_MW(HD, KV_EB, STATE_EB, CR) \ - switch (next_n) \ - { \ - case 1: LAUNCH_DECODE_MW(HD, KV_EB, STATE_EB, CR, 1); break; \ - case 2: LAUNCH_DECODE_MW(HD, KV_EB, STATE_EB, CR, 2); break; \ - case 3: LAUNCH_DECODE_MW(HD, KV_EB, STATE_EB, CR, 3); break; \ - default: LAUNCH_DECODE_MW(HD, KV_EB, STATE_EB, CR, 4); break; \ - } - -#define DISPATCH_NN(HD, KV_EB, STATE_EB, CR) \ - switch (next_n) \ + // Walk FOREACH_DECODE_CONFIG until we find a matching (HD, KV, ST, CR, NN, NRW) + // tuple, then launch that instantiation. Any unsupported tuple bails via TLLM_THROW. +#define TRY_LAUNCH(HD, KV_EB, STATE_EB, CR, NN, NRW) \ + if (head_dim == HD && kv_score_elem_bytes == KV_EB && state_elem_bytes == STATE_EB && compress_ratio == CR \ + && next_n_dispatch == NN && num_red_warps == NRW) \ { \ - case 1: LAUNCH_DECODE(HD, KV_EB, STATE_EB, CR, 1); break; \ - case 2: LAUNCH_DECODE(HD, KV_EB, STATE_EB, CR, 2); break; \ - case 3: LAUNCH_DECODE(HD, KV_EB, STATE_EB, CR, 3); break; \ - default: LAUNCH_DECODE(HD, KV_EB, STATE_EB, CR, 4); break; \ + pagedKvCompressKernel<<>>(kv_score, ape, \ + paged_kv, paged_score, block_table_kv, block_table_score, output, kv_lens, cu_seq_lens, cu_kv_comp, \ + page_size, max_blocks, out_elem_bytes); \ + return; \ } + FOREACH_DECODE_CONFIG(TRY_LAUNCH) +#undef TRY_LAUNCH -#define DISPATCH_DTYPE(HD, CR, DISPATCH_MACRO) \ - do \ - { \ - if (kv_score_elem_bytes == 4 && state_elem_bytes == 4) \ - { \ - DISPATCH_MACRO(HD, 4, 4, CR); \ - } \ - else if (kv_score_elem_bytes == 2 && state_elem_bytes == 4) \ - { \ - DISPATCH_MACRO(HD, 2, 4, CR); \ - } \ - else if (kv_score_elem_bytes == 4 && state_elem_bytes == 2) \ - { \ - DISPATCH_MACRO(HD, 4, 2, CR); \ - } \ - else \ - { \ - DISPATCH_MACRO(HD, 2, 2, CR); \ - } \ - } while (false) - - if (use_multi_warp) - { - // Multi-warp path: HD=128 or HD=512, NEXT_N=1 or NEXT_N=2, CR=128. - if (head_dim == 512) - { - DISPATCH_DTYPE(512, 128, DISPATCH_NN_MW); - } - else - { - DISPATCH_DTYPE(128, 128, DISPATCH_NN_MW); - } - } - else if (compress_ratio == 4) - { - if (head_dim == 512) - { - DISPATCH_DTYPE(512, 4, DISPATCH_NN); - } - else - { - DISPATCH_DTYPE(128, 4, DISPATCH_NN); - } - } - else - { - if (head_dim == 512) - { - DISPATCH_DTYPE(512, 128, DISPATCH_NN); - } - else - { - DISPATCH_DTYPE(128, 128, DISPATCH_NN); - } - } - -#undef DISPATCH_DTYPE -#undef DISPATCH_NN -#undef DISPATCH_NN_MW -#undef LAUNCH_DECODE_MW -#undef LAUNCH_DECODE + TLLM_THROW( + "pagedKvCompressLaunch: no matching instantiation for HD=%d, kv_eb=%d, state_eb=%d, CR=%d, NN=%d, NRW=%d", + head_dim, kv_score_elem_bytes, state_elem_bytes, compress_ratio, next_n_dispatch, num_red_warps); } +#undef FOREACH_DECODE_CONFIG +#undef FOREACH_DECODE_DTYPE +#undef FOREACH_DECODE_NN + // ============================================================================ // Prefill Kernel: prefillReductionKernel //