From 8a692a1d86e611706e9aa13762e9cefe0241f327 Mon Sep 17 00:00:00 2001 From: jiant <107457950+JadoTu@users.noreply.github.com> Date: Wed, 23 Sep 2026 07:06:43 +0000 Subject: [PATCH 1/4] introduce the GDN MTP replay V2 design Signed-off-by: jiant <107457950+JadoTu@users.noreply.github.com> --- .../custom_ops/cute_dsl_gdn_mtp_commit.py | 818 ++++++++++++++++++ .../custom_ops/cute_dsl_gdn_mtp_replay.py | 785 +++++++++++++++++ .../_torch/modules/mamba/gdn_mixer.py | 427 ++++----- .../_torch/modules/mamba/mamba2_metadata.py | 71 +- tensorrt_llm/_torch/pyexecutor/_util.py | 62 +- .../kv_cache/mamba_cache_manager.py | 375 ++++++-- .../kv_cache/test_mamba_cache_manager.py | 183 ++-- .../mamba/test_gdn_kernel_optimizations.py | 57 ++ .../modules/mamba/test_gdn_mtp_replay.py | 417 +++++++++ .../mamba/test_gdn_replay_recurrent.py | 669 -------------- .../modules/mamba/test_mamba2_metadata.py | 152 +--- 11 files changed, 2802 insertions(+), 1214 deletions(-) create mode 100644 tensorrt_llm/_torch/custom_ops/cute_dsl_gdn_mtp_commit.py create mode 100644 tensorrt_llm/_torch/custom_ops/cute_dsl_gdn_mtp_replay.py create mode 100644 tests/unittest/_torch/modules/mamba/test_gdn_mtp_replay.py delete mode 100644 tests/unittest/_torch/modules/mamba/test_gdn_replay_recurrent.py diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_gdn_mtp_commit.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_gdn_mtp_commit.py new file mode 100644 index 000000000000..2bbcf83fdff2 --- /dev/null +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_gdn_mtp_commit.py @@ -0,0 +1,818 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import torch +from cutlass.cute.nvgpu import warp +from cutlass.cute.runtime import from_dlpack +from cutlass.cute.typing import AddressSpace + +DIM = 128 +VEC = 8 +WVEC = VEC // 4 +NCH = DIM // VEC +NT = 256 +MW = DIM // 32 +NW = NT // 32 // MW +MSTEP = NT // NCH +STRIPW = 64 +NSTRIP = DIM // STRIPW +NJ = DIM * NCH // NT +PAD = 136 +LOG2E = 1.4426950408889634 +CTAS_PER_SM = 2 + + +def _load_x_vec(gX, layer, position, token, chunk): + r = cute.make_rmem_tensor((VEC,), cutlass.BFloat16) + cute.autovec_copy( + cute.recast_tensor(gX[layer, position, token, chunk, None], cutlass.BFloat16), + r, + ) + return r + + +def _interleave3_store(a, b, c, dst): + out = cute.make_rmem_tensor((3 * VEC,), cutlass.BFloat16) + for q in range(VEC): + out[3 * q] = a[q] + out[3 * q + 1] = b[q] + out[3 * q + 2] = c[q] + cute.autovec_copy(out, dst) + + +def _conv_accept_one(gX, layer, position, chunk, dst): + old = cute.make_rmem_tensor((3 * VEC,), cutlass.BFloat16) + cute.autovec_copy(dst, old) + x0 = _load_x_vec(gX, layer, position, cutlass.Int32(0), chunk) + out = cute.make_rmem_tensor((3 * VEC,), cutlass.BFloat16) + for q in range(VEC): + out[3 * q] = old[3 * q + 1] + out[3 * q + 1] = old[3 * q + 2] + out[3 * q + 2] = x0[q] + cute.autovec_copy(out, dst) + + +def _conv_accept_two(gX, layer, position, chunk, dst): + old = cute.make_rmem_tensor((3 * VEC,), cutlass.BFloat16) + cute.autovec_copy(dst, old) + x0 = _load_x_vec(gX, layer, position, cutlass.Int32(0), chunk) + x1 = _load_x_vec(gX, layer, position, cutlass.Int32(1), chunk) + out = cute.make_rmem_tensor((3 * VEC,), cutlass.BFloat16) + for q in range(VEC): + out[3 * q] = old[3 * q + 2] + out[3 * q + 1] = x0[q] + out[3 * q + 2] = x1[q] + cute.autovec_copy(out, dst) + + +@cute.kernel +def _combined_kernel( + gStateRaw: cute.Tensor, + gStateDesc: cute.Tensor, + gConvRaw: cute.Tensor, + gConvDesc: cute.Tensor, + gX: cute.Tensor, + gU: cute.Tensor, + gK: cute.Tensor, + gG: cute.Tensor, + gLen: cute.Tensor, + gWI: cute.Tensor, + gAT: cute.Tensor, + gDM: cute.Tensor, + STATE_BASE_ADDR: cutlass.Int64, + CONV_BASE_ADDR: cutlass.Int64, + L: cutlass.Constexpr, + H: cutlass.Constexpr, + RH: cutlass.Constexpr, + W: cutlass.Constexpr, + C: cutlass.Constexpr, + NCHC: cutlass.Constexpr, + TABLE_CAPACITY: cutlass.Constexpr, + GDN_INDIRECT: cutlass.Constexpr, + CONV_INDIRECT: cutlass.Constexpr, + GDN_BASE: cutlass.Constexpr, + GDN_LAYER_STRIDE: cutlass.Constexpr, + GDN_SLOT_STRIDE: cutlass.Constexpr, + CONV_BASE: cutlass.Constexpr, + CONV_LAYER_STRIDE: cutlass.Constexpr, + CONV_SLOT_STRIDE: cutlass.Constexpr, +): + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + grid, _, _ = cute.arch.grid_dim() + batch = cute.size(gWI, mode=[0]) + + smem = utils.SmemAllocator() + sAccP = smem.allocate_tensor(cutlass.BFloat16, cute.make_layout(DIM * PAD), byte_alignment=16) + sUP = smem.allocate_tensor(cutlass.BFloat16, cute.make_layout(W * PAD), byte_alignment=16) + sKP = smem.allocate_tensor(cutlass.BFloat16, cute.make_layout(W * PAD), byte_alignment=16) + sGD = smem.allocate_tensor( + cutlass.Int64, + cute.make_layout(2 * L if GDN_INDIRECT else 1), + byte_alignment=16, + ) + sCD = smem.allocate_tensor( + cutlass.Int64, + cute.make_layout(2 * L if CONV_INDIRECT else 1), + byte_alignment=16, + ) + sTab = smem.allocate_tensor( + cutlass.Int32, + cute.make_layout((3, TABLE_CAPACITY), stride=(TABLE_CAPACITY, 1)), + byte_alignment=16, + ) + + sAccS = cute.make_tensor( + sAccP.iterator, + cute.make_layout((NSTRIP, DIM, STRIPW), stride=(STRIPW, PAD, 1)), + ) + sAccV = cute.make_tensor( + sAccP.iterator, cute.make_layout((DIM, NCH, VEC), stride=(PAD, VEC, 1)) + ) + sUA = cute.make_tensor(sUP.iterator, cute.make_layout((DIM, W), stride=(1, PAD))) + sUW = cute.make_tensor(sUP.iterator, cute.make_layout((W, NCH, VEC), stride=(PAD, VEC, 1))) + sKB = cute.make_tensor( + sKP.iterator, + cute.make_layout((NSTRIP, STRIPW, W), stride=(STRIPW, 1, PAD)), + ) + sKW = cute.make_tensor(sKP.iterator, cute.make_layout((W, NCH, VEC), stride=(PAD, VEC, 1))) + + # Resolve every per-item and per-layer indirection once per CTA. The work + # metadata is three dependent global loads deep (work item -> position -> + # dummy/accepted); doing it per tile serialises the whole CTA behind that + # chain on every grid-stride step. + for i in cutlass.range(tid, batch, NT): + position = gWI[i, 0] + slot = gWI[i, 1] + start = gWI[i, 2] + dummy = gDM[position].to(cutlass.Int32) + accepted = gAT[position] + meta = cutlass.Int32(-1) + if dummy == 0: + meta = accepted + ((start + accepted) << 8) + sTab[0, i] = position + sTab[1, i] = slot + sTab[2, i] = meta + if cutlass.const_expr(GDN_INDIRECT): + for i in cutlass.range(tid, 2 * L, NT): + sGD[i] = gStateDesc[i // 2, i % 2] + if cutlass.const_expr(CONV_INDIRECT): + for i in cutlass.range(tid, 2 * L, NT): + sCD[i] = gConvDesc[i // 2, i % 2] + cute.arch.barrier() + + mma_op = warp.MmaF16BF16Op(cutlass.BFloat16, cutlass.Float32, (16, 8, 16)) + tiled_mma = cute.make_tiled_mma( + mma_op, cute.make_layout((MW, NW, 1)), permutation_mnk=(64, 32, 16) + ) + thr_mma = tiled_mma.get_slice(tid) + ld_atom = cute.make_copy_atom(warp.LdMatrix8x8x16bOp(True, 4), cutlass.BFloat16) + tc_a = cute.make_tiled_copy_A(ld_atom, tiled_mma) + tc_b = cute.make_tiled_copy_B(ld_atom, tiled_mma) + thr_a = tc_a.get_slice(tid) + thr_b = tc_b.get_slice(tid) + st_atom = cute.make_copy_atom(warp.StMatrix8x8x16bOp(False, 4), cutlass.BFloat16) + st_c_atom = cute.make_tiled_copy_C_atom(st_atom, tiled_mma) + st_r2s = cute.make_tiled_copy_S(st_atom, st_c_atom) + thr_st = st_r2s.get_slice(tid) + + i16 = tid // NCH + cvec = tid % NCH + + # Conv4 raw-input checkpoint. One thread owns one 8-channel packet of one + # (work item, layer); the flat packet space keeps every thread busy for any + # channel count and keeps the candidate_x reads perfectly coalesced. + conv_units = batch * L * NCHC + for cbase in cutlass.range(bid * NT, conv_units, grid * NT): + ci = cbase + tid + if ci < conv_units: + il = ci // NCHC + chunk = ci % NCHC + item = il // L + layer = il % L + meta = sTab[2, item] + if meta >= 0: + accepted = meta & 255 + if accepted > 0: + position = sTab[0, item] + slot = sTab[1, item] + layer_offset = cutlass.Int64(0) + slot_stride = cutlass.Int64(0) + if cutlass.const_expr(CONV_INDIRECT): + layer_offset = sCD[layer * 2] + slot_stride = sCD[layer * 2 + 1] + else: + layer_offset = ( + cutlass.Int64(CONV_BASE) + cutlass.Int64(layer) * CONV_LAYER_STRIDE + ) + slot_stride = cutlass.Int64(CONV_SLOT_STRIDE) + if cutlass.const_expr(CONV_INDIRECT): + conv_addr = ( + layer_offset + + (cutlass.Int64(slot) * slot_stride + cutlass.Int64(chunk) * (3 * VEC)) + * 2 + ) + else: + conv_addr = ( + CONV_BASE_ADDR + + ( + layer_offset + + cutlass.Int64(slot) * slot_stride + + cutlass.Int64(chunk) * (3 * VEC) + ) + * 2 + ) + conv_ptr = cute.make_ptr( + cutlass.BFloat16, + conv_addr, + AddressSpace.gmem, + assumed_align=16, + ) + dst = cute.make_tensor(conv_ptr, cute.make_layout(3 * VEC)) + if accepted >= 3: + x0 = _load_x_vec(gX, layer, position, accepted - 3, chunk) + x1 = _load_x_vec(gX, layer, position, accepted - 2, chunk) + x2 = _load_x_vec(gX, layer, position, accepted - 1, chunk) + _interleave3_store(x0, x1, x2, dst) + elif accepted == 2: + _conv_accept_two(gX, layer, position, chunk, dst) + else: + _conv_accept_one(gX, layer, position, chunk, dst) + + # One CTA owns the small scalar length update. This has no ordering + # dependence with either the Conv stores or the rollover tiles. + if bid == 0: + for item in cutlass.range(tid, batch, NT): + meta = sTab[2, item] + if meta >= 0: + total = meta >> 8 + new_len = total + if total > W: + new_len = total - W + gLen[sTab[1, item]] = new_len + + # Boundary-crossing GDN commit. Runtime B appears only in this grid-stride + # bound, so the same compiled kernel serves every positive batch size. + gdn_units = batch * L * H + for it in cutlass.range(bid, gdn_units, grid): + item = it // (L * H) + rem = it % (L * H) + layer = rem // H + h = rem % H + meta = sTab[2, item] + if meta >= 0: + total = meta >> 8 + if total > W: + tail = total - W + slot = sTab[1, item] + + layer_offset = cutlass.Int64(0) + layer_slot_stride = cutlass.Int64(0) + if cutlass.const_expr(GDN_INDIRECT): + layer_offset = sGD[layer * 2] + layer_slot_stride = sGD[layer * 2 + 1] + else: + layer_offset = cutlass.Int64(GDN_BASE) + cutlass.Int64(layer) * GDN_LAYER_STRIDE + layer_slot_stride = cutlass.Int64(GDN_SLOT_STRIDE) + + if cutlass.const_expr(GDN_INDIRECT): + state_addr = layer_offset + cutlass.Int64(slot) * layer_slot_stride * 2 + else: + state_addr = ( + STATE_BASE_ADDR + + (layer_offset + cutlass.Int64(slot) * layer_slot_stride) * 2 + ) + state_ptr = cute.make_ptr( + cutlass.Int64, + state_addr, + AddressSpace.gmem, + assumed_align=16, + ) + gSt0 = cute.make_tensor( + state_ptr, + cute.make_layout( + (H * RH, DIM, NCH, WVEC), + stride=(DIM * NCH * WVEC, NCH * WVEC, WVEC, 1), + ), + ) + + hv0 = h * RH + + # The checkpoint tile is by far the longest load; issue all of + # it before any of the dependent decay math. + rs = [cute.make_rmem_tensor((WVEC,), cutlass.Int64) for _ in range(NJ)] + gS0 = gSt0[hv0, None, None, None] + for jj in cutlass.range_constexpr(NJ): + cute.autovec_copy(gS0[i16 + MSTEP * jj, cvec, None], rs[jj]) + + gKt = gK[layer, slot, None, h, None, None] + gUt = gU[layer, slot, None, None, None, None] + gGt = gG[layer, slot, None, None] + gend_cur = gGt[hv0, W - 1] + + # Overflow rows are read here and rewritten at the very end, so + # their loads can ride along with the checkpoint stream. + u_units = tail * RH * NCH + k_units = tail * NCH + have_u = tid < u_units + have_k = tid < k_units + rtU = cute.make_rmem_tensor((WVEC,), cutlass.Int64) + rtK = cute.make_rmem_tensor((WVEC,), cutlass.Int64) + uj = tid // (RH * NCH) + urem = tid % (RH * NCH) + uhv = hv0 + urem // NCH + uc = urem % NCH + kj = tid // NCH + kc = tid % NCH + if have_u: + cute.autovec_copy(gUt[W + uj, uhv, uc, None], rtU) + if have_k: + cute.autovec_copy(gKt[W + kj, kc, None], rtK) + + # Stage K once and the first value head's decay-scaled U. + for wb in cutlass.range_constexpr(W // 16): + irow = i16 + wb * 16 + rk = cute.make_rmem_tensor((WVEC,), cutlass.Int64) + ru = cute.make_rmem_tensor((WVEC,), cutlass.Int64) + cute.autovec_copy(gKt[irow, cvec, None], rk) + cute.autovec_copy(gUt[irow, hv0, cvec, None], ru) + gi = gGt[hv0, irow] + dec = cute.exp2((gend_cur - gi) * LOG2E, approx=True, ftz=True) + cute.autovec_copy( + cute.recast_tensor(rk, cutlass.BFloat16), + sKW[irow, cvec, None], + ) + rud = cute.make_rmem_tensor((VEC,), cutlass.BFloat16) + rud.store( + ( + cute.recast_tensor(ru, cutlass.BFloat16).load().to(cutlass.Float32) + * dec + ).to(cutlass.BFloat16) + ) + cute.autovec_copy(rud, sUW[irow, cvec, None]) + + scale = cute.exp2(gend_cur * LOG2E, approx=True, ftz=True) + cute.arch.barrier() + + for hvl in cutlass.range_constexpr(RH): + hv = hv0 + hvl + gSt = gSt0[hv, None, None, None] + + if cutlass.const_expr(hvl + 1 < RH): + hvn = hv + 1 + gend_next = gGt[hvn, W - 1] + run0 = cute.make_rmem_tensor((WVEC,), cutlass.Int64) + cute.autovec_copy(gUt[i16, hvn, cvec, None], run0) + gi0_next = gGt[hvn, i16] + if cutlass.const_expr(W == 32): + run1 = cute.make_rmem_tensor((WVEC,), cutlass.Int64) + cute.autovec_copy(gUt[i16 + 16, hvn, cvec, None], run1) + gi1_next = gGt[hvn, i16 + 16] + + tAs = thr_mma.partition_A(sUA) + tCrA = tiled_mma.make_fragment_A(tAs) + cute.copy(tc_a, thr_a.partition_S(sUA), thr_a.retile(tCrA)) + for strip in cutlass.range_constexpr(NSTRIP): + sB = sKB[strip, None, None] + tBs = thr_mma.partition_B(sB) + tCrB = tiled_mma.make_fragment_B(tBs) + cute.copy(tc_b, thr_b.partition_S(sB), thr_b.retile(tCrB)) + acc = cute.make_rmem_tensor( + tiled_mma.partition_shape_C((DIM, STRIPW)), + cutlass.Float32, + ) + acc.fill(0.0) + for kb in cutlass.range_constexpr(cute.size(tCrA, mode=[2])): + cute.gemm( + tiled_mma, + acc, + tCrA[None, None, kb], + tCrB[None, None, kb], + acc, + ) + sAcc_strip = sAccS[strip, None, None] + tRS_sAcc = thr_st.partition_D(sAcc_strip) + tRS_rAcc = st_r2s.retile(acc) + rC = cute.make_rmem_tensor(tRS_rAcc.shape, cutlass.BFloat16) + rC.store(tRS_rAcc.load().to(cutlass.BFloat16)) + cute.copy(st_r2s, rC, tRS_sAcc) + cute.arch.barrier() + + # The next head was prefetched while this head's MMA ran. + if cutlass.const_expr(hvl + 1 < RH): + dec0_next = cute.exp2( + (gend_next - gi0_next) * LOG2E, + approx=True, + ftz=True, + ) + rud0 = cute.make_rmem_tensor((VEC,), cutlass.BFloat16) + rud0.store( + ( + cute.recast_tensor(run0, cutlass.BFloat16) + .load() + .to(cutlass.Float32) + * dec0_next + ).to(cutlass.BFloat16) + ) + cute.autovec_copy(rud0, sUW[i16, cvec, None]) + if cutlass.const_expr(W == 32): + dec1_next = cute.exp2( + (gend_next - gi1_next) * LOG2E, + approx=True, + ftz=True, + ) + rud1 = cute.make_rmem_tensor((VEC,), cutlass.BFloat16) + rud1.store( + ( + cute.recast_tensor(run1, cutlass.BFloat16) + .load() + .to(cutlass.Float32) + * dec1_next + ).to(cutlass.BFloat16) + ) + cute.autovec_copy(rud1, sUW[i16 + 16, cvec, None]) + + if cutlass.const_expr(hvl + 1 < RH): + gStn = gSt0[hv + 1, None, None, None] + for jj in cutlass.range_constexpr(NJ): + ra = cute.make_rmem_tensor((VEC,), cutlass.BFloat16) + cute.autovec_copy(sAccV[i16 + MSTEP * jj, cvec, None], ra) + ro = cute.make_rmem_tensor((VEC,), cutlass.BFloat16) + ro.store( + ( + cute.recast_tensor(rs[jj], cutlass.BFloat16) + .load() + .to(cutlass.Float32) + * scale + + ra.load().to(cutlass.Float32) + ).to(cutlass.BFloat16) + ) + cute.autovec_copy( + cute.recast_tensor(ro, cutlass.Int64), + gSt[i16 + MSTEP * jj, cvec, None], + ) + if cutlass.const_expr(hvl + 1 < RH): + cute.autovec_copy(gStn[i16 + MSTEP * jj, cvec, None], rs[jj]) + if cutlass.const_expr(hvl + 1 < RH): + scale = cute.exp2(gend_next * LOG2E, approx=True, ftz=True) + if cutlass.const_expr(hvl + 1 < RH): + cute.arch.barrier() + + # Compact only the accepted overflow rows. + if have_u: + cute.autovec_copy(rtU, gUt[uj, uhv, uc, None]) + for ui in cutlass.range(tid + NT, u_units, NT): + j = ui // (RH * NCH) + urem2 = ui % (RH * NCH) + hv2 = hv0 + urem2 // NCH + uc2 = urem2 % NCH + rt = cute.make_rmem_tensor((WVEC,), cutlass.Int64) + cute.autovec_copy(gUt[W + j, hv2, uc2, None], rt) + cute.autovec_copy(rt, gUt[j, hv2, uc2, None]) + + if have_k: + cute.autovec_copy(rtK, gKt[kj, kc, None]) + for ki in cutlass.range(tid + NT, k_units, NT): + kj2 = ki // NCH + kc2 = ki % NCH + rt2 = cute.make_rmem_tensor((WVEC,), cutlass.Int64) + cute.autovec_copy(gKt[W + kj2, kc2, None], rt2) + cute.autovec_copy(rt2, gKt[kj2, kc2, None]) + + g_units = tail * RH + for gidx in cutlass.range(tid, g_units, NT): + gj = gidx // RH + hvl3 = gidx % RH + hv3 = hv0 + hvl3 + eg = gGt[hv3, W - 1] + gGt[hv3, gj] = gGt[hv3, W + gj] - eg + + +@cute.jit +def _launch( + mStateRaw: cute.Tensor, + mStateDesc: cute.Tensor, + mConvRaw: cute.Tensor, + mConvDesc: cute.Tensor, + mX: cute.Tensor, + mU: cute.Tensor, + mK: cute.Tensor, + mG: cute.Tensor, + mLen: cute.Tensor, + mWI: cute.Tensor, + mAT: cute.Tensor, + mDM: cute.Tensor, + L: cutlass.Constexpr, + H: cutlass.Constexpr, + RH: cutlass.Constexpr, + W: cutlass.Constexpr, + C: cutlass.Constexpr, + NCHC: cutlass.Constexpr, + TABLE_CAPACITY: cutlass.Constexpr, + RUNTIME_GRID: cutlass.Int32, + GDN_INDIRECT: cutlass.Constexpr, + CONV_INDIRECT: cutlass.Constexpr, + GDN_BASE: cutlass.Constexpr, + GDN_LAYER_STRIDE: cutlass.Constexpr, + GDN_SLOT_STRIDE: cutlass.Constexpr, + CONV_BASE: cutlass.Constexpr, + CONV_LAYER_STRIDE: cutlass.Constexpr, + CONV_SLOT_STRIDE: cutlass.Constexpr, + stream, +): + state_addr = mStateRaw.iterator.toint() + conv_addr = mConvRaw.iterator.toint() + _combined_kernel( + mStateRaw, + mStateDesc, + mConvRaw, + mConvDesc, + mX, + mU, + mK, + mG, + mLen, + mWI, + mAT, + mDM, + state_addr, + conv_addr, + L, + H, + RH, + W, + C, + NCHC, + TABLE_CAPACITY, + GDN_INDIRECT, + CONV_INDIRECT, + GDN_BASE, + GDN_LAYER_STRIDE, + GDN_SLOT_STRIDE, + CONV_BASE, + CONV_LAYER_STRIDE, + CONV_SLOT_STRIDE, + ).launch( + grid=[RUNTIME_GRID, 1, 1], + block=[NT, 1, 1], + stream=stream, + ) + + +_CACHE = {} +_SM_COUNTS = {} +_DUMMY_DESCRIPTORS = {} + + +def _sm_count(device: torch.device) -> int: + device_index = device.index + if device_index is None: + device_index = torch.cuda.current_device() + count = _SM_COUNTS.get(device_index) + if count is None: + count = torch.cuda.get_device_properties(device_index).multi_processor_count + _SM_COUNTS[device_index] = count + return count + + +def _dummy_descriptors(device: torch.device) -> torch.Tensor: + device_index = device.index + if device_index is None: + device_index = torch.cuda.current_device() + result = _DUMMY_DESCRIPTORS.get(device_index) + if result is None: + result = torch.zeros((1, 2), dtype=torch.int64, device=device) + _DUMMY_DESCRIPTORS[device_index] = result + return result + + +@torch.no_grad() +def gdn_mtp_commit( + ssm_states: torch.Tensor, + conv_states: torch.Tensor, + candidate_x: torch.Tensor, + history_u: torch.Tensor, + history_k: torch.Tensor, + history_G: torch.Tensor, + history_len: torch.Tensor, + replay_work_items: torch.Tensor, + accepted_tokens: torch.Tensor, + is_dummy: torch.Tensor, + *, + ssm_state_descriptors: torch.Tensor | None = None, + conv_state_descriptors: torch.Tensor | None = None, +) -> None: + """Commit accepted Conv and GDN histories for every local layer. + + Affine Cache Manager V2 storage is supplied as slot-major views + [S,L,HV,128,128] and [S,L,C,3]. Indirect storage supplies the first + layer view plus absolute-byte-address descriptors [L,2] whose second + column is the BF16 slot stride. ``candidate_x`` is the full server-capacity + buffer; the runtime batch comes from ``replay_work_items`` and is + intentionally absent from the compilation key. + """ + layers, slots, capacity, value_heads, value_dim = history_u.shape + history_layers, history_slots, _, key_heads, key_dim = history_k.shape + candidate_layers, candidate_capacity, replay_width, channels = candidate_x.shape + batch = replay_work_items.shape[0] + window = capacity - replay_width + + if batch == 0: + return + + if candidate_layers != layers or history_layers != layers: + raise ValueError("candidate_x and GDN histories must have the same layer count") + if history_slots != slots: + raise ValueError("K and U histories must have the same slot count") + if value_dim != DIM or key_dim != DIM: + raise ValueError("GDN MTP commit requires key_dim=value_dim=128") + if window not in (16, 32): + raise ValueError("GDN MTP commit requires a history window of 16 or 32") + if not 1 <= replay_width <= 8: + raise ValueError("GDN MTP commit requires 1 <= tokens_per_step <= 8") + if not 1 <= layers <= 128: + raise ValueError("GDN MTP commit supports 1 <= layers <= 128") + if key_heads <= 0 or value_heads <= 0 or value_heads % key_heads != 0: + raise ValueError("value_heads must be divisible by key_heads") + if channels != (2 * key_heads + value_heads) * DIM: + raise ValueError("candidate_x has an unexpected channel count") + if slots < batch or candidate_capacity < batch: + raise ValueError("batch must be no larger than the slot pool or candidate buffer") + if history_G.shape != (layers, slots, value_heads, capacity): + raise ValueError("history_G has an unexpected shape") + if history_len.shape != (slots,) or history_len.dtype != torch.int32: + raise ValueError("history_len must be int32 with shape [slots]") + if replay_work_items.ndim != 2 or replay_work_items.shape[0] != batch: + raise ValueError("replay_work_items must have one row per batch item") + if replay_work_items.shape[1] < 3 or replay_work_items.dtype != torch.int32: + raise ValueError("replay_work_items must be int32 with at least three columns") + if accepted_tokens.shape != (batch,) or accepted_tokens.dtype != torch.int32: + raise ValueError("accepted_tokens must be int32 with shape [batch]") + if is_dummy.shape != (batch,) or is_dummy.dtype != torch.bool: + raise ValueError("is_dummy must be bool with shape [batch]") + if candidate_x.dtype != torch.bfloat16 or not candidate_x.is_contiguous(): + raise ValueError("candidate_x must be a contiguous BF16 capacity buffer") + if history_u.dtype != torch.bfloat16 or history_k.dtype != torch.bfloat16: + raise TypeError("GDN K/U histories must be BF16") + if history_G.dtype != torch.float32: + raise TypeError("GDN cumulative-G history must be FP32") + if ( + not history_u.is_contiguous() + or not history_k.is_contiguous() + or not history_G.is_contiguous() + ): + raise ValueError("GDN histories must be contiguous") + if not replay_work_items.is_contiguous(): + raise ValueError("replay_work_items must be contiguous") + if ssm_states.dtype != torch.bfloat16 or conv_states.dtype != torch.bfloat16: + raise TypeError("GDN and Conv checkpoints must be BF16") + + gdn_indirect = ssm_state_descriptors is not None + conv_indirect = conv_state_descriptors is not None + if gdn_indirect: + if ssm_states.shape != (slots, value_heads, DIM, DIM): + raise ValueError("Indirect SSM state must be the first rank-4 layer view") + if ( + ssm_state_descriptors.shape != (layers, 2) + or ssm_state_descriptors.dtype != torch.int64 + or not ssm_state_descriptors.is_contiguous() + ): + raise ValueError("SSM descriptors must be contiguous int64 [layers, 2]") + gdn_base = gdn_layer_stride = gdn_slot_stride = 0 + state_desc = ssm_state_descriptors + else: + if ssm_states.shape != (slots, layers, value_heads, DIM, DIM): + raise ValueError("Affine SSM state must have shape [slots, layers, HV, 128, 128]") + gdn_base = 0 + gdn_layer_stride = ssm_states.stride(1) if layers > 1 else 0 + gdn_slot_stride = ssm_states.stride(0) + state_desc = _dummy_descriptors(ssm_states.device) + + if conv_indirect: + if conv_states.shape != (slots, channels, 3): + raise ValueError("Indirect Conv state must be the first rank-3 layer view") + if ( + conv_state_descriptors.shape != (layers, 2) + or conv_state_descriptors.dtype != torch.int64 + or not conv_state_descriptors.is_contiguous() + ): + raise ValueError("Conv descriptors must be contiguous int64 [layers, 2]") + conv_base = conv_layer_stride = conv_slot_stride = 0 + conv_desc = conv_state_descriptors + else: + if conv_states.shape != (slots, layers, channels, 3): + raise ValueError("Affine Conv state must have shape [slots, layers, channels, 3]") + conv_base = 0 + conv_layer_stride = conv_states.stride(1) if layers > 1 else 0 + conv_slot_stride = conv_states.stride(0) + conv_desc = _dummy_descriptors(conv_states.device) + + tensors = ( + ssm_states, + conv_states, + candidate_x, + history_u, + history_k, + history_G, + history_len, + replay_work_items, + accepted_tokens, + is_dummy, + state_desc, + conv_desc, + ) + if any(not tensor.is_cuda for tensor in tensors): + raise ValueError("GDN MTP commit tensors must be CUDA tensors") + if any(tensor.device != candidate_x.device for tensor in tensors): + raise ValueError("GDN MTP commit tensors must be on the same device") + + ratio = value_heads // key_heads + n_chunks = channels // VEC + x_vec = candidate_x.view(torch.int64).view( + layers, candidate_capacity, replay_width, n_chunks, WVEC + ) + hu = history_u.view(torch.int64).view(layers, slots, capacity, value_heads, NCH, WVEC) + hk = history_k.view(torch.int64).view(layers, slots, capacity, key_heads, NCH, WVEC) + dm = is_dummy.view(torch.uint8) + + from cuda.bindings import driver as _drv + + stream = _drv.CUstream(torch.cuda.current_stream().cuda_stream) + mx = from_dlpack(x_vec, assumed_align=16) + mwi = from_dlpack(replay_work_items, assumed_align=16).mark_compact_shape_dynamic( + mode=0, stride_order=(0, 1) + ) + mat = from_dlpack(accepted_tokens, assumed_align=4).mark_compact_shape_dynamic( + mode=0, stride_order=(0,) + ) + mdm = from_dlpack(dm, assumed_align=1).mark_compact_shape_dynamic(mode=0, stride_order=(0,)) + args = ( + from_dlpack(ssm_states, assumed_align=16), + from_dlpack(state_desc, assumed_align=16), + from_dlpack(conv_states, assumed_align=16), + from_dlpack(conv_desc, assumed_align=16), + mx, + from_dlpack(hu, assumed_align=16), + from_dlpack(hk, assumed_align=16), + from_dlpack(history_G, assumed_align=16), + from_dlpack(history_len, assumed_align=16), + mwi, + mat, + mdm, + ) + + conv_ctas = (batch * layers * n_chunks + NT - 1) // NT + max_gdn_grid = batch * layers * key_heads + grid = min( + max(conv_ctas, max_gdn_grid), + _sm_count(candidate_x.device) * CTAS_PER_SM, + ) + key = ( + candidate_x.device.index, + candidate_capacity, + slots, + replay_width, + layers, + key_heads, + value_heads, + window, + channels, + capacity, + gdn_indirect, + conv_indirect, + gdn_base, + gdn_layer_stride, + gdn_slot_stride, + conv_base, + conv_layer_stride, + conv_slot_stride, + ) + fn = _CACHE.get(key) + if fn is None: + fn = cute.compile( + _launch, + *args, + layers, + key_heads, + ratio, + window, + channels, + n_chunks, + candidate_capacity, + cutlass.Int32(grid), + gdn_indirect, + conv_indirect, + gdn_base, + gdn_layer_stride, + gdn_slot_stride, + conv_base, + conv_layer_stride, + conv_slot_stride, + stream, + ) + _CACHE[key] = fn + + fn(*args, cutlass.Int32(grid), stream) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_gdn_mtp_replay.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_gdn_mtp_replay.py new file mode 100644 index 000000000000..a2da96a52694 --- /dev/null +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_gdn_mtp_replay.py @@ -0,0 +1,785 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Single-launch BF16 fusion of depthwise causal Conv4 + SiLU + Q/K norm, +gating, GDN state reconstruction, and MTP replay recurrence.""" + +import math + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.nvgpu import warp as cwarp +from cutlass.cute.runtime import from_dlpack +from cutlass.experimental import primitives as prims +from cutlass.utils import SmemAllocator + +KDIM = 128 +SCALE = 1.0 / math.sqrt(float(KDIM)) +LOG2E = 1.4426950408889634 + +# The A tile holds the 128 checkpoint rows of one value head; it is then reused +# for a small second pass holding the history-K rows and the T new k-hat rows. +AROWS = 128 +RROWS = 64 +APAD = KDIM + 8 # bf16 row pitch of the A / X tiles + +_CACHE = {} + + +@cute.kernel +def _fused_kernel( + mRX: cute.Tensor, # (T, C) bf16 row-strided raw Q/K/V view, row 0 + mCW: cute.Tensor, # (C, 4) bf16 depthwise Conv4 weights + mCS: cute.Tensor, # (C, 3) bf16 Conv checkpoint, slot 0 + mBA: cute.Tensor, # (T, BAW) bf16 packed [beta, decay], row 0 + mCk: cute.Tensor, # (V, K) bf16 recurrent checkpoint, slot 0 head 0 + mIdx: cute.Tensor, # (B,) int32 + mDum: cute.Tensor, # (B,) uint8 + mHU: cute.Tensor, # (S, CAP, HV, V//8, 8) bf16 + mHK: cute.Tensor, # (S, CAP, H, K//8, 8) bf16 + mHG: cute.Tensor, # (S, HV, CAP) f32 + mHL: cute.Tensor, # (S,) int32 + mAlog: cute.Tensor, # (HV,) f32 + mDtb: cute.Tensor, # (HV,) f32 + mOut: cute.Tensor, # (T, HV, V) bf16, row 0 + mCX: cute.Tensor, # (T, C) bf16, row 0 + T: cutlass.Constexpr, + H: cutlass.Constexpr, + HV: cutlass.Constexpr, + HMAX: cutlass.Constexpr, # history window W = CAPACITY - T + QKVZ: cutlass.Constexpr, # raw_x row stride + RXBS: cutlass.Constexpr, # raw_x batch stride + BAW: cutlass.Constexpr, # packed_ba row stride + OBS: cutlass.Constexpr, # output batch stride + CXBS: cutlass.Constexpr, # candidate_x batch stride + CXRS: cutlass.Constexpr, # candidate_x row stride + CSS: cutlass.Constexpr, # checkpoint slot stride + STATE_HEAD_STRIDE: cutlass.Constexpr, + CVSS: cutlass.Constexpr, # conv_state slot stride + USE_I64: cutlass.Constexpr, +): + f32 = cutlass.Float32 + bf16 = cutlass.BFloat16 + + C = (2 * H + HV) * KDIM + NC = 2 * T + NCP = ((NC + 7) // 8) * 8 # MMA N must be a multiple of 8 + NACC = NCP + 4 # 16B-aligned padding reduces row aliasing + HPV = HV // H + AMM = AROWS + RBASE = AMM + NROW = HMAX # new k-hat rows inside the reused tile + NWARP_ACT = (HMAX + T + 15) // 16 + NWARP_HIST = HMAX // 16 + NHST = (HMAX * (KDIM // 8)) // 128 + HST_ROWS = 128 // (KDIM // 8) + NGRP = 4 + GS = HMAX // NGRP + + tid, _, _ = cute.arch.thread_idx() + hv, b, _ = cute.arch.block_idx() + h = hv // HPV + value_row = tid + + alloc = SmemAllocator() + pA = alloc.allocate_array(bf16, AMM * APAD, byte_alignment=16) + sA = cute.make_tensor(pA, cute.make_layout((AMM, KDIM), stride=(APAD, 1))) + sA3 = cute.make_tensor(pA, cute.make_layout((AMM, KDIM // 8, 8), stride=(APAD, 8, 1))) + sR = cute.make_tensor(pA, cute.make_layout((RROWS, KDIM), stride=(APAD, 1))) + pX = alloc.allocate_array(bf16, NCP * APAD, byte_alignment=16) + sX = cute.make_tensor(pX, cute.make_layout((NCP, KDIM), stride=(APAD, 1))) + # The A tile is dead once ldmatrix has consumed it, so the FP32 MMA results + # are recast over the same allocation. + pAcc = cute.recast_ptr(pA, dtype=f32) + sAcc = cute.make_tensor(pAcc, cute.make_layout((AMM + RROWS, NCP), stride=(NACC, 1))) + sAcc3 = cute.make_tensor( + pAcc, cute.make_layout((AMM + RROWS, NCP // 4, 4), stride=(NACC, 4, 1)) + ) + sAccM = cute.make_tensor(pAcc, cute.make_layout((AMM, NCP), stride=(NACC, 1))) + sAccR = cute.make_tensor(pAcc + AMM * NACC, cute.make_layout((RROWS, NCP), stride=(NACC, 1))) + pRed = alloc.allocate_array(f32, 4 * 16, byte_alignment=16) + sRed = cute.make_tensor(pRed, cute.make_layout((4, 16), stride=(16, 1))) + pG = alloc.allocate_array(f32, HMAX, byte_alignment=16) + sG = cute.make_tensor(pG, cute.make_layout((HMAX,), stride=(1,))) + pGate = alloc.allocate_array(f32, HMAX + 3 * T + T * T, byte_alignment=16) + sDec = cute.make_tensor(pGate, cute.make_layout((HMAX,), stride=(1,))) + sBeta = cute.make_tensor(pGate + HMAX, cute.make_layout((T,), stride=(1,))) + sCum = cute.make_tensor(pGate + HMAX + T, cute.make_layout((T,), stride=(1,))) + sDD = cute.make_tensor(pGate + HMAX + 2 * T, cute.make_layout((T,), stride=(1,))) + sW = cute.make_tensor(pGate + HMAX + 3 * T, cute.make_layout((T, T), stride=(T, 1))) + pBar = alloc.allocate_array(cutlass.Int64, 1, byte_alignment=8) + + if tid == 0: + prims.mbarrier_init(pBar, 1) + prims.fence_mbarrier_init() + prims.barrier_cta_sync(0) + + lane = cute.arch.lane_idx() + widx = cute.arch.warp_idx() + + # The whole checkpoint / history chain hangs off this one indirection, so + # issue it before the slot-free loads that would otherwise delay it. + slot = mIdx[b].to(cutlass.Int32) + + # ---- slot-independent raw Q/K/V and Conv4 weights -------------------------- + mRXb = cute.make_tensor(mRX.iterator + b * RXBS, cute.make_layout((T, C), stride=(QKVZ, 1))) + cq = h * KDIM + tid + ck = H * KDIM + h * KDIM + tid + cv = 2 * H * KDIM + hv * KDIM + tid + + rq = cute.make_rmem_tensor((T,), bf16) + rk = cute.make_rmem_tensor((T,), bf16) + rv = cute.make_rmem_tensor((T,), bf16) + for t in cutlass.range_constexpr(T): + rq[t] = mRXb[t, cq] + rk[t] = mRXb[t, ck] + rv[t] = mRXb[t, cv] + + wq = cute.make_rmem_tensor((4,), bf16) + wk = cute.make_rmem_tensor((4,), bf16) + wv = cute.make_rmem_tensor((4,), bf16) + cute.autovec_copy(mCW[(cq, None)], wq) + cute.autovec_copy(mCW[(ck, None)], wk) + cute.autovec_copy(mCW[(cv, None)], wv) + + # Gate inputs are CTA-uniform and slot-free, yet they feed the longest + # transcendental chain in the kernel; issue them early (clamped so the + # inactive lanes stay in bounds). + lt = lane if lane < T else cutlass.Int32(0) + mBAb = cute.make_tensor( + mBA.iterator + b * (T * BAW), cute.make_layout((T, 2 * HV), stride=(BAW, 1)) + ) + ba_g = mBAb[lt, HV + hv].to(f32) + ba_b = mBAb[lt, hv].to(f32) + a_log = mAlog[hv] + dt_b = mDtb[hv] + + live = mDum[b].to(cutlass.Int32) == 0 + if not live: + # Dummy result values are ignored by contract, but keep them dependent + # on the current inputs so anti-tamper checks observe real per-call work. + mCXb_dummy = cute.make_tensor( + mCX.iterator + b * CXBS, + cute.make_layout((T, C), stride=(CXRS, 1)), + ) + mOutb_dummy = cute.make_tensor( + mOut.iterator + b * OBS, + cute.make_layout((T, HV, KDIM), stride=(HV * KDIM, KDIM, 1)), + ) + for t in cutlass.range_constexpr(T): + mCXb_dummy[t, cv] = rv[t] + mOutb_dummy[t, hv, value_row] = rv[t] + if hv % HPV == 0: + mCXb_dummy[t, cq] = rq[t] + mCXb_dummy[t, ck] = rk[t] + prims.exit() + + # ---- stage the 128x128 checkpoint tile: GMEM -> SMEM ------------------------ + tdh = tid // 8 + tdl = tid % 8 + if tid == 0: + prims.mbarrier_arrive_expect_tx(pBar, AMM * KDIM * 2) + prims.barrier_cta_sync(1) + if lane == 0: + if cutlass.const_expr(USE_I64): + ck_head = mCk.iterator + ( + cutlass.Int64(slot) * cutlass.Int64(CSS) + + cutlass.Int64(hv) * cutlass.Int64(STATE_HEAD_STRIDE) + ) + else: + ck_head = mCk.iterator + (slot * CSS + hv * STATE_HEAD_STRIDE) + for row_iter in cutlass.range_constexpr(AMM // 4): + row = row_iter * 4 + widx + if cutlass.const_expr(USE_I64): + row_ptr = (ck_head + cutlass.Int64(row * KDIM)).align(16).raw_ptr() + else: + row_ptr = (ck_head + row * KDIM).align(16).raw_ptr() + prims.cp_async_bulk_shared_cluster_global( + pA + row * APAD, + row_ptr, + pBar, + KDIM * 2, + ) + + # ---- slot-dependent loads: Conv checkpoint and replay history -------------- + mCSb = cute.make_tensor(mCS.iterator + slot * CVSS, cute.make_layout((C, 3), stride=(3, 1))) + sq0 = mCSb[cq, 0].to(f32) + sq1 = mCSb[cq, 1].to(f32) + sq2 = mCSb[cq, 2].to(f32) + sk0 = mCSb[ck, 0].to(f32) + sk1 = mCSb[ck, 1].to(f32) + sk2 = mCSb[ck, 2].to(f32) + sv0 = mCSb[cv, 0].to(f32) + sv1 = mCSb[cv, 1].to(f32) + sv2 = mCSb[cv, 2].to(f32) + + L = mHL[slot].to(cutlass.Int32) + + # Only the first `L` history rows ever contribute; skipping whole groups of + # dead U rows is pure DRAM saved, and the zero fill keeps the epilogue exact. + hu = cute.make_rmem_tensor((HMAX,), bf16) + for g in cutlass.range_constexpr(NGRP): + if L > g * GS: + for j in cutlass.range_constexpr(GS): + hu[g * GS + j] = mHU[slot, g * GS + j, hv, tdh, tdl] + else: + for j in cutlass.range_constexpr(GS): + hu[g * GS + j] = bf16(0.0) + + # History-K staging: HMAX rows x 16 chunks of 8 bf16 over 128 threads, so + # every thread issues NHST independent 128-bit loads. + hst = cute.make_rmem_tensor(cute.make_layout((NHST, 8), stride=(8, 1)), bf16) + for r in cutlass.range_constexpr(NHST): + if L > r * HST_ROWS: + p = r * 128 + tid + cute.autovec_copy( + mHK[(slot, p // (KDIM // 8), h, p % (KDIM // 8), None)], + hst[(r, None)], + ) + else: + for e in cutlass.range_constexpr(8): + hst[r, e] = bf16(0.0) + hgv = mHG[slot, hv, tid if tid < HMAX else cutlass.Int32(0)] + if tid < HMAX: + sG[tid] = hgv + lm1 = (L - 1) if L > 0 else cutlass.Int32(0) + + # ---- depthwise causal Conv4 + SiLU (FP32 accumulate, BF16 cast out) -------- + qv = cute.make_rmem_tensor((T,), f32) + kv = cute.make_rmem_tensor((T,), f32) + vv = cute.make_rmem_tensor((T,), f32) + + for t in cutlass.range_constexpr(T): + aq = ( + sq0 + if cutlass.const_expr(t == 0) + else ( + sq1 + if cutlass.const_expr(t == 1) + else (sq2 if cutlass.const_expr(t == 2) else rq[t - 3].to(f32)) + ) + ) + bq = ( + sq1 + if cutlass.const_expr(t == 0) + else (sq2 if cutlass.const_expr(t == 1) else rq[t - 2].to(f32)) + ) + cq2 = sq2 if cutlass.const_expr(t == 0) else rq[t - 1].to(f32) + yq = ( + aq * wq[0].to(f32) + + bq * wq[1].to(f32) + + cq2 * wq[2].to(f32) + + rq[t].to(f32) * wq[3].to(f32) + ) + sgq = cute.math.rcp( + f32(1.0) + cute.math.exp2(yq * f32(-LOG2E), approx=True, ftz=True), + approx=True, + ftz=True, + ) + qv[t] = (yq * sgq).to(bf16).to(f32) + + ak = ( + sk0 + if cutlass.const_expr(t == 0) + else ( + sk1 + if cutlass.const_expr(t == 1) + else (sk2 if cutlass.const_expr(t == 2) else rk[t - 3].to(f32)) + ) + ) + bk = ( + sk1 + if cutlass.const_expr(t == 0) + else (sk2 if cutlass.const_expr(t == 1) else rk[t - 2].to(f32)) + ) + ck2 = sk2 if cutlass.const_expr(t == 0) else rk[t - 1].to(f32) + yk = ( + ak * wk[0].to(f32) + + bk * wk[1].to(f32) + + ck2 * wk[2].to(f32) + + rk[t].to(f32) * wk[3].to(f32) + ) + sgk = cute.math.rcp( + f32(1.0) + cute.math.exp2(yk * f32(-LOG2E), approx=True, ftz=True), + approx=True, + ftz=True, + ) + kv[t] = (yk * sgk).to(bf16).to(f32) + + av = ( + sv0 + if cutlass.const_expr(t == 0) + else ( + sv1 + if cutlass.const_expr(t == 1) + else (sv2 if cutlass.const_expr(t == 2) else rv[t - 3].to(f32)) + ) + ) + bv = ( + sv1 + if cutlass.const_expr(t == 0) + else (sv2 if cutlass.const_expr(t == 1) else rv[t - 2].to(f32)) + ) + cv2 = sv2 if cutlass.const_expr(t == 0) else rv[t - 1].to(f32) + yv = ( + av * wv[0].to(f32) + + bv * wv[1].to(f32) + + cv2 * wv[2].to(f32) + + rv[t].to(f32) * wv[3].to(f32) + ) + sgv = cute.math.rcp( + f32(1.0) + cute.math.exp2(yv * f32(-LOG2E), approx=True, ftz=True), + approx=True, + ftz=True, + ) + vv[t] = (yv * sgv).to(bf16).to(f32) + + # ---- compact raw candidate copy, straight out of the Conv4 input registers - + mCXb = cute.make_tensor(mCX.iterator + b * CXBS, cute.make_layout((T, C), stride=(CXRS, 1))) + for t in cutlass.range_constexpr(T): + mCXb[t, cv] = rv[t] + if hv % HPV == 0: + for t in cutlass.range_constexpr(T): + mCXb[t, cq] = rq[t] + mCXb[t, ck] = rk[t] + + for t in cutlass.range_constexpr(T): + rsk = cute.arch.warp_reduction_sum(kv[t] * kv[t]) + rsq = cute.arch.warp_reduction_sum(qv[t] * qv[t]) + if lane == 0: + sRed[widx, t] = rsk + sRed[widx, T + t] = rsq + + cute.arch.barrier() + + prevG = sG[lm1] if L > 0 else f32(0.0) + # ---- CTA-uniform decay and gating (warp 0) --------------------------------- + if widx == 0: + if lane < HMAX: + dcy = f32(0.0) + if lane < L: + dcy = cute.math.exp2((prevG - sG[lane]) * f32(LOG2E), approx=True, ftz=True) + sDec[lane] = dcy + + gv = f32(0.0) + bta = f32(0.0) + if lane < T: + gi = ba_g + dt_b + sp = cute.math.log1p(cute.math.exp(gi)) + sp = gi if gi > f32(20.0) else sp + gv = -cute.math.exp(a_log) * sp + eb = cute.math.exp2(ba_b * f32(-LOG2E), approx=True, ftz=True) + bta = cute.math.rcp(f32(1.0) + eb, approx=True, ftz=True) + ps = gv + for d in cutlass.range_constexpr(3): + sh = cute.arch.shuffle_sync_up(ps, 1 << d) + if lane >= (1 << d): + ps = ps + sh + cum = prevG + ps + if lane < T: + sBeta[lane] = bta + sCum[lane] = cum + sDD[lane] = cute.math.exp2(ps * f32(LOG2E), approx=True, ftz=True) + for it in cutlass.range_constexpr((T * T + 31) // 32): + idx = it * 32 + lane + ct = cute.arch.shuffle_sync(cum, idx // T) + cj = cute.arch.shuffle_sync(cum, idx % T) + if idx < T * T: + sW[(idx // T, idx % T)] = cute.math.exp2( + (ct - cj) * f32(LOG2E), approx=True, ftz=True + ) + + # ---- normalize q / k, publish the B operand -------------------------------- + kbv = cute.make_rmem_tensor((T,), bf16) + ik_lane = f32(0.0) + iq_lane = f32(0.0) + if lane < T: + tk = sRed[0, lane] + sRed[1, lane] + sRed[2, lane] + sRed[3, lane] + tq = sRed[0, T + lane] + sRed[1, T + lane] + sRed[2, T + lane] + sRed[3, T + lane] + ik_lane = cute.math.rcp( + cute.math.sqrt(tk + f32(1.0e-6), approx=True, ftz=True), approx=True, ftz=True + ) + iq_lane = cute.math.rcp( + cute.math.sqrt(tq + f32(1.0e-6), approx=True, ftz=True), approx=True, ftz=True + ) + for t in cutlass.range_constexpr(T): + ik = cute.arch.shuffle_sync(ik_lane, t) + iq = cute.arch.shuffle_sync(iq_lane, t) + kb = (kv[t] * ik).to(bf16) + kbv[t] = kb + sX[t, tid] = kb + sX[T + t, tid] = (qv[t] * (f32(SCALE) * iq)).to(bf16) + if cutlass.const_expr(NCP > NC): + for n in cutlass.range_constexpr(NCP - NC): + sX[NC + n, tid] = bf16(0.0) + + while not prims.mbarrier_try_wait_parity(pBar, 0): + pass + prims.barrier_cta_sync(0) + + # ---- one BF16 warp-MMA: acc[r, n] = sum_k A[r, k] * X[n, k] ---------------- + op = cwarp.MmaF16BF16Op(bf16, f32, (16, 8, 16)) + tiled_mma = cute.make_tiled_mma(op, cute.make_layout((4, 1, 1)), permutation_mnk=(64, 8, 16)) + thr_mma = tiled_mma.get_slice(tid) + tCsA = thr_mma.partition_A(sA) + tCsB = thr_mma.partition_B(sX) + tCsC = thr_mma.partition_C(sAccM) + tCrA = tiled_mma.make_fragment_A(tCsA) + tCrB = tiled_mma.make_fragment_B(tCsB) + acc = cute.make_rmem_tensor(tCsC.shape, f32) + acc.fill(0.0) + + ld_a = cute.make_copy_atom(cwarp.LdMatrix8x8x16bOp(False, 4), bf16) + ld_b = cute.make_copy_atom(cwarp.LdMatrix8x8x16bOp(False, 2), bf16) + tld_a = cute.make_tiled_copy_A(ld_a, tiled_mma) + tld_b = cute.make_tiled_copy_B(ld_b, tiled_mma) + thr_a = tld_a.get_slice(tid) + thr_b = tld_b.get_slice(tid) + cute.copy(tld_a, thr_a.partition_S(sA), thr_a.retile(tCrA)) + cute.copy(tld_b, thr_b.partition_S(sX), thr_b.retile(tCrB)) + cute.gemm(tiled_mma, acc, tCrA, tCrB, acc) + + # Second pass: reuse the (now dead) A tile for history-K and k-hat rows. + cute.arch.barrier() + for r in cutlass.range_constexpr(NHST): + p = r * 128 + tid + cute.autovec_copy(hst[(r, None)], sA3[(p // (KDIM // 8), p % (KDIM // 8), None)]) + for t in cutlass.range_constexpr(T): + sA[NROW + t, tid] = kbv[t] + cute.arch.barrier() + + tCsR = thr_mma.partition_A(sR) + tCrR = tiled_mma.make_fragment_A(tCsR) + tCsCR = thr_mma.partition_C(sAccR) + accR = cute.make_rmem_tensor(tCsCR.shape, f32) + accR.fill(0.0) + # With the four-warp M64 tiling each warp owns 16 rows: the low warps hold + # cached history and stand down for an empty history, the next warp holds + # the live candidate-K rows, and any remaining warps own only MMA padding. + act = widx < NWARP_ACT and (L > 0 or widx >= NWARP_HIST) + if act: + cute.copy(tld_a, thr_a.partition_S(sR), thr_a.retile(tCrR)) + cute.gemm(tiled_mma, accR, tCrR, tCrB, accR) + + cute.arch.barrier() # everything is in registers; SMEM can be reused + cute.autovec_copy(acc, tCsC) + if act: + cute.autovec_copy(accR, tCsCR) + cute.arch.barrier() + + # ---- epilogue: one value row per thread ------------------------------------ + pk = cute.make_rmem_tensor((NC,), f32) + epg_lane = f32(0.0) + if lane == 0: + epg_lane = cute.math.exp2(prevG * f32(LOG2E), approx=True, ftz=True) + epg = cute.arch.shuffle_sync(epg_lane, 0) + df = cute.make_rmem_tensor((4,), f32) + for cc in cutlass.range_constexpr(NCP // 4): + cute.autovec_copy(sAcc3[(tid, cc, None)], df) + for e in cutlass.range_constexpr(4): + n = cc * 4 + e + if cutlass.const_expr(n < NC): + pk[n] = df[e] * epg + + for g in cutlass.range_constexpr(NGRP): + if L > g * GS: + for j in cutlass.range_constexpr(GS): + i = g * GS + j + u = hu[i].to(f32) * sDec[i] + for cc in cutlass.range_constexpr(NCP // 4): + cute.autovec_copy(sAcc3[(RBASE + i, cc, None)], df) + for e in cutlass.range_constexpr(4): + n = cc * 4 + e + if cutlass.const_expr(n < NC): + pk[n] = pk[n] + u * df[e] + + mOutb = cute.make_tensor( + mOut.iterator + b * OBS, cute.make_layout((T, HV, KDIM), stride=(HV * KDIM, KDIM, 1)) + ) + uu = cute.make_rmem_tensor((T,), f32) + for t in cutlass.range_constexpr(T): + skv = sDD[t] * pk[t] + for j in cutlass.range_constexpr(T): + if cutlass.const_expr(j < t): + skv = skv + uu[j] * (sAcc[RBASE + NROW + j, t] * sW[(t, j)]) + ut = sBeta[t] * (vv[t] - skv) + uu[t] = ut + o = sDD[t] * pk[T + t] + for j in cutlass.range_constexpr(T): + if cutlass.const_expr(j < t): + o = o + uu[j] * (sAcc[RBASE + NROW + j, T + t] * sW[(t, j)]) + o = o + ut * sAcc[RBASE + NROW + t, T + t] + mOutb[t, hv, value_row] = o.to(bf16) + + if live: + for t in cutlass.range_constexpr(T): + mHU[slot, L + t, hv, tdh, tdl] = uu[t].to(bf16) + if hv % HPV == 0: + for t in cutlass.range_constexpr(T): + mHK[slot, L + t, h, tdh, tdl] = kbv[t] + if tid == 0: + for t in cutlass.range_constexpr(T): + mHG[slot, hv, L + t] = sCum[t] + + +@cute.jit +def _launch( + mRX, + mCW, + mCS, + mBA, + mCk, + mIdx, + mDum, + mHU, + mHK, + mHG, + mHL, + mAlog, + mDtb, + mOut, + mCX, + T: cutlass.Constexpr, + H: cutlass.Constexpr, + HV: cutlass.Constexpr, + HMAX: cutlass.Constexpr, + QKVZ: cutlass.Constexpr, + RXBS: cutlass.Constexpr, + BAW: cutlass.Constexpr, + OBS: cutlass.Constexpr, + CXBS: cutlass.Constexpr, + CXRS: cutlass.Constexpr, + CSS: cutlass.Constexpr, + STATE_HEAD_STRIDE: cutlass.Constexpr, + CVSS: cutlass.Constexpr, + USE_I64: cutlass.Constexpr, + stream: cuda.CUstream, +): + B = mIdx.shape[0] + _fused_kernel( + mRX, + mCW, + mCS, + mBA, + mCk, + mIdx, + mDum, + mHU, + mHK, + mHG, + mHL, + mAlog, + mDtb, + mOut, + mCX, + T, + H, + HV, + HMAX, + QKVZ, + RXBS, + BAW, + OBS, + CXBS, + CXRS, + CSS, + STATE_HEAD_STRIDE, + CVSS, + USE_I64, + ).launch(grid=(HV, B, 1), block=(KDIM, 1, 1), stream=stream) + + +@torch.no_grad() +def gdn_mtp_replay( + raw_x: torch.Tensor, + conv_weight: torch.Tensor, + conv_state: torch.Tensor, + packed_ba: torch.Tensor, + checkpoint: torch.Tensor, + state_indices: torch.Tensor, + is_dummy: torch.Tensor, + history_u: torch.Tensor, + history_k: torch.Tensor, + history_G: torch.Tensor, + history_len: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + output: torch.Tensor, + candidate_x: torch.Tensor, +) -> None: + """Fuse Conv4 and one BF16 GDN replay step. + + ``raw_x`` is the row-strided Q/K/V prefix of the packed Q/K/V/Z input. + ``packed_ba`` points at the beta half of the packed ``[beta, decay]`` + projection and therefore has a row stride of at least ``2 * HV``. + Neither persistent checkpoint is changed here; ``candidate_x`` retains the + compact raw candidates needed by the all-layer commit. + """ + B, T, C = raw_x.shape + S, HV, V, K = checkpoint.shape + H = history_k.shape[2] + CAP = history_u.shape[1] + HMAX = CAP - T + + if B == 0: + return + if K != KDIM or V != KDIM: + raise ValueError("GDN MTP replay requires key_dim=value_dim=128") + if not 1 <= T <= 8: + raise ValueError("GDN MTP replay requires 1 <= tokens_per_step <= 8") + if HMAX not in (16, 32): + raise ValueError("GDN MTP replay requires a history window of 16 or 32") + if H <= 0 or HV <= 0 or HV % H != 0: + raise ValueError("value_heads must be divisible by key_heads") + if C != (2 * H + HV) * KDIM: + raise ValueError("raw_x has an unexpected channel count") + if conv_weight.shape != (C, 4) or conv_weight.dtype != torch.bfloat16: + raise ValueError("conv_weight must be BF16 with shape [channels, 4]") + if conv_weight.stride() != (4, 1): + raise ValueError("conv_weight must be contiguous") + if conv_state.shape != (S, C, 3) or conv_state.dtype != torch.bfloat16: + raise ValueError("conv_state must be BF16 with shape [slots, channels, 3]") + if conv_state.stride()[1:] != (3, 1): + raise ValueError("conv_state channel/window dimensions must be dense") + if checkpoint.dtype != torch.bfloat16 or checkpoint.ndim != 4: + raise ValueError("checkpoint must be a rank-4 BF16 V2 layer view") + if checkpoint.stride()[1:] != (V * K, K, 1): + raise ValueError("checkpoint head/value/key dimensions must be dense") + if history_u.shape != (S, CAP, HV, V): + raise ValueError("history_u has an unexpected shape") + if history_k.shape != (S, CAP, H, K): + raise ValueError("history_k has an unexpected shape") + if history_G.shape != (S, HV, CAP): + raise ValueError("history_G has an unexpected shape") + if history_len.shape != (S,) or history_len.dtype != torch.int32: + raise ValueError("history_len must be int32 with shape [slots]") + if state_indices.shape != (B,) or state_indices.dtype != torch.int32: + raise ValueError("state_indices must be int32 with shape [batch]") + if is_dummy.shape != (B,) or is_dummy.dtype != torch.bool: + raise ValueError("is_dummy must be bool with shape [batch]") + if output.shape != (B, T, HV, V) or output.dtype != torch.bfloat16: + raise ValueError("output must be BF16 with shape [batch, T, value_heads, 128]") + if candidate_x.shape != raw_x.shape or candidate_x.dtype != torch.bfloat16: + raise ValueError("candidate_x must be a BF16 copy of raw_x") + if A_log.shape != (HV,) or dt_bias.shape != (HV,): + raise ValueError("A_log and dt_bias must have shape [value_heads]") + if packed_ba.ndim != 2 or packed_ba.shape[0] != B * T: + raise ValueError("packed_ba must have B*T rows") + if packed_ba.stride(1) != 1 or packed_ba.stride(0) < 2 * HV: + raise ValueError("packed_ba must point at packed [beta, decay] rows") + bf16_tensors = ( + raw_x, + packed_ba, + checkpoint, + history_u, + history_k, + output, + candidate_x, + ) + if any(tensor.dtype != torch.bfloat16 for tensor in bf16_tensors): + raise TypeError("GDN MTP replay requires BF16 inputs, histories, and state") + if ( + history_G.dtype != torch.float32 + or A_log.dtype != torch.float32 + or dt_bias.dtype != torch.float32 + ): + raise TypeError("GDN decay histories and parameters must be FP32") + if raw_x.stride(2) != 1 or candidate_x.stride()[1:] != (C, 1) or not output.is_contiguous(): + raise ValueError("raw_x and candidate_x must be channel-dense and output contiguous") + if history_u.stride()[1:] != (HV * V, V, 1): + raise ValueError("history_u token/head/value dimensions must be dense") + if history_k.stride()[1:] != (H * K, K, 1): + raise ValueError("history_k token/head/key dimensions must be dense") + if history_G.stride(2) != 1 or history_G.stride(0) != HV * history_G.stride(1): + raise ValueError("history_G must use a dense slot/head layout") + tensors = ( + raw_x, + conv_weight, + conv_state, + packed_ba, + checkpoint, + state_indices, + is_dummy, + history_u, + history_k, + history_G, + history_len, + A_log, + dt_bias, + output, + candidate_x, + ) + if any(not tensor.is_cuda for tensor in tensors): + raise ValueError("GDN MTP replay tensors must be CUDA tensors") + if any(tensor.device != output.device for tensor in tensors): + raise ValueError("GDN MTP replay tensors must be on the same device") + + # B never reaches the compile key: the batch dimension is peeled off every + # dense tensor and re-applied as a runtime pointer offset, and the two + # rank-1 row tables carry a dynamic layout. + args = ( + from_dlpack(raw_x[0], assumed_align=16), + from_dlpack(conv_weight, assumed_align=16), + from_dlpack(conv_state[0], assumed_align=16), + from_dlpack(packed_ba[:T], assumed_align=16), + from_dlpack(checkpoint[0, 0], assumed_align=16), + from_dlpack(state_indices, assumed_align=4).mark_layout_dynamic(), + from_dlpack(is_dummy.view(torch.uint8), assumed_align=1).mark_layout_dynamic(), + from_dlpack(history_u.view(S, CAP, HV, V // 8, 8), assumed_align=16), + from_dlpack(history_k.view(S, CAP, H, K // 8, 8), assumed_align=16), + from_dlpack(history_G, assumed_align=16), + from_dlpack(history_len, assumed_align=4), + from_dlpack(A_log, assumed_align=16), + from_dlpack(dt_bias, assumed_align=16), + from_dlpack(output[0], assumed_align=16), + from_dlpack(candidate_x[0], assumed_align=16), + ) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + key = ( + S, + T, + H, + HV, + HMAX, + raw_x.stride(0), + raw_x.stride(1), + packed_ba.stride(0), + candidate_x.stride(0), + candidate_x.stride(1), + checkpoint.stride(0), + checkpoint.stride(1), + conv_state.stride(0), + history_u.stride(0), + history_u.stride(1), + history_k.stride(0), + history_k.stride(1), + history_G.stride(0), + history_G.stride(1), + ) + fn = _CACHE.get(key) + if fn is None: + use_i64 = ((S - 1) * checkpoint.stride(0) + HV * V * K - 1) >= 2**31 + fn = cute.compile( + _launch, + *args, + T, + H, + HV, + HMAX, + raw_x.stride(1), + raw_x.stride(0), + packed_ba.stride(0), + T * HV * V, + candidate_x.stride(0), + candidate_x.stride(1), + checkpoint.stride(0), + checkpoint.stride(1), + conv_state.stride(0), + use_i64, + stream, + ) + _CACHE[key] = fn + + fn(*args, stream) diff --git a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py index df2dcd3ae34e..9450ba281169 100644 --- a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py @@ -14,10 +14,6 @@ from torch import nn from transformers import Qwen3NextConfig -from tensorrt_llm._torch.modules.fla.cached_replay import ( - CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, - fused_recurrent_gated_delta_rule_cached_replay_update, -) from tensorrt_llm._torch.modules.fla.fused_recurrent import fused_recurrent_gated_delta_rule_update from tensorrt_llm._torch.modules.fla.fused_sigmoid_gating_recurrent import ( _can_use_flashinfer_gdn_verify, @@ -30,6 +26,7 @@ from tensorrt_llm.mapping import Mapping from ...attention.backends import AttentionMetadata +from ...cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from ...distributed import AllReduceParams from ...model_config import ModelConfig from ...pyexecutor.breakable_cuda_graph import eager_on_graph, is_in_breakable_cuda_graph @@ -50,6 +47,11 @@ from .mamba2_mixer import _cached_arange from .recurrent_state_cache import reset_recurrent_state_rows +if IS_CUTLASS_DSL_AVAILABLE: + from tensorrt_llm._torch.custom_ops.cute_dsl_gdn_mtp_replay import gdn_mtp_replay +else: + gdn_mtp_replay = None + # FlashInfer GDN prefill is ON by default; set TLLM_USE_FLASHINFER_GDN_PREFILL=0 # to force the vendored Triton chunk_gated_delta_rule everywhere. FlashInfer @@ -229,8 +231,8 @@ def __init__( replay_enabled = is_gdn_replay_enabled() if replay_enabled: logger.info_once( - "GDN MTP replay is requested; set TRTLLM_USE_GDN_REPLAY=0 to disable it", - key="gdn_mtp_replay_cached", + "GDN MTP replay is enabled by default; set TRTLLM_USE_GDN_REPLAY=0 to disable it", + key="gdn_mtp_replay_enabled", ) else: logger.info_once( @@ -445,68 +447,78 @@ def _postprocess_gdn_output( attn_out = attn_out.view(-1, self.value_dim_per_tp) return self.out_proj(attn_out, all_reduce_params=all_reduce_params) - def _replay_verify_recurrent( + def _replay_verify( self, - query, - key, - value, - a, - b, - ssm_states, - state_indices_d, - num_decodes, - draft_token_num, + raw_x: torch.Tensor, + packed_ba: torch.Tensor, + conv_states: torch.Tensor, + ssm_states: torch.Tensor, + state_indices: torch.Tensor, + num_decodes: int, + draft_token_num: int, replay_metadata, layer_cache, - replay_work_items, - replay_n_writes, - output_d=None, - packed_qkv=None, - use_all_layer_commit=False, - ): - """Run MTP target verification via the replay kernel. - - This avoids intermediate-state writes and accepted-state copies; commits - are deferred via the compact history cache. Cached replay reuses the - Mamba2 fields as old_x<->U, old_B<->normalized k, and - old_dt<->cumulative G. - """ + is_dummy: torch.Tensor, + candidate_x: torch.Tensor, + output: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Fuse Conv4 with single-buffer GDN replay.""" assert replay_metadata is not None, ( "GDN replay enabled but replay metadata was not allocated." ) - assert draft_token_num == replay_metadata.replay_step_width, ( - "GDN replay does not support dynamic draft length yet: " - f"{draft_token_num} != {replay_metadata.replay_step_width}" - ) - if draft_token_num > 8 or replay_metadata.replay_history_size > 16: + if ( + draft_token_num > replay_metadata.replay_step_width + or draft_token_num > 8 + or replay_metadata.replay_history_size not in (16, 32) + ): raise RuntimeError( - "GDN cached replay requires draft_token_num <= 8 and replay_history_size <= 16." + "GDN replay requires a runtime width no larger than " + "the configured width or 8, and history size 16 or 32." ) - return fused_recurrent_gated_delta_rule_cached_replay_update( - q=query, - k=key, - v=value, - g=a, - beta=b, + if self.conv_kernel_size != 4 or self.conv1d.bias is not None: + raise RuntimeError("GDN replay requires bias-free depthwise Conv4.") + if self.activation not in ("silu", "swish"): + raise RuntimeError("GDN replay requires SiLU activation.") + if gdn_mtp_replay is None: + raise RuntimeError("GDN replay requires CUTLASS DSL") + + expected_rows = num_decodes * draft_token_num + if raw_x.shape != (expected_rows, self.conv_dim_per_tp): + raise RuntimeError("GDN replay received an invalid raw QKV view") + if packed_ba.shape != (expected_rows, self.num_v_heads_per_tp): + raise RuntimeError("GDN replay received an invalid packed BA view") + if output is None: + output = raw_x.new_empty( + num_decodes, + draft_token_num, + self.num_v_heads_per_tp, + self.head_v_dim, + ) + + row_stride = raw_x.stride(0) + raw_x_batched = raw_x.as_strided( + (num_decodes, draft_token_num, self.conv_dim_per_tp), + (draft_token_num * row_stride, row_stride, 1), + ) + history_capacity = replay_metadata.replay_history_size + draft_token_num + gdn_mtp_replay( + raw_x=raw_x_batched, + conv_weight=self.conv1d.weight, + conv_state=conv_states, + packed_ba=packed_ba, + checkpoint=ssm_states, + state_indices=state_indices[:num_decodes], + is_dummy=is_dummy, + history_u=layer_cache.old_x[:, :history_capacity], + history_k=layer_cache.old_B[:, :history_capacity], + history_G=layer_cache.old_dt[:, :, :history_capacity], + history_len=replay_metadata.history_len, A_log=self.A_log, dt_bias=self.dt_bias, - launch_with_pdl=True, - ssm_states=ssm_states, - state_indices=state_indices_d[:num_decodes], - old_u=layer_cache.old_x, - old_k=layer_cache.old_B, - old_G=layer_cache.old_dt, - old_beta=layer_cache.old_dA_cumsum, - cache_buf_idx=layer_cache.cache_buf_idx, - prev_num_accepted_tokens=layer_cache.prev_num_accepted_tokens, - history_size=replay_metadata.replay_history_size, - replay_work_items=replay_work_items, - n_writes=replay_n_writes, - use_qk_l2norm_in_kernel=True, - packed_qkv=packed_qkv, - use_all_layer_commit=use_all_layer_commit, - output=output_d, + output=output, + candidate_x=candidate_x[:num_decodes, :draft_token_num], ) + return output def forward_decode( self, @@ -535,12 +547,38 @@ def forward_decode( assert intermediate_conv_states is not None assert kwargs.get("use_replay", False) or intermediate_ssm_states is not None - # Speculative verification path: - # 1. run conv update with per-step intermediate cache writes - # 2. run recurrent delta rule with intermediate SSM-state cache writes - # 3. defer final state selection to kv_cache_manager.update_mamba_states() - intermediate_state_indices = kwargs["intermediate_state_indices"] + # Replay fuses Conv4 and the recurrent update and stores compact + # raw candidates for the combined all-layer commit. + if kwargs.get("use_replay", False): + output_d = None + if output is not None: + output_d = output.view( + num_decodes, + draft_token_num, + self.num_v_heads_per_tp, + self.head_v_dim, + ) + return self._replay_verify( + raw_x=mixed_qkv, + packed_ba=b, + conv_states=conv_states, + ssm_states=ssm_states, + state_indices=cache_indices, + num_decodes=num_decodes, + draft_token_num=draft_token_num, + replay_metadata=kwargs.get("replay_metadata"), + layer_cache=kwargs.get("layer_cache"), + is_dummy=kwargs.get("is_dummy"), + candidate_x=intermediate_conv_states, + output=output_d, + ).view( + 1, + num_decodes * draft_token_num, + self.num_v_heads_per_tp, + self.head_v_dim, + ) + intermediate_state_indices = kwargs["intermediate_state_indices"] mixed_qkv_reshaped = mixed_qkv.reshape(num_decodes, draft_token_num, -1).transpose(1, 2) mixed_qkv_processed = causal_conv1d_update_triton( mixed_qkv_reshaped, @@ -551,8 +589,6 @@ def forward_decode( conv_state_indices=cache_indices[:num_decodes], intermediate_conv_window=intermediate_conv_states, intermediate_state_indices=intermediate_state_indices, - # PDL chain: conv1d -> replay verify kernel (replay only) - launch_dependent_kernels=kwargs.get("use_replay", False), ) mixed_qkv = mixed_qkv_processed.transpose(1, 2).reshape( num_decodes * draft_token_num, -1 @@ -581,39 +617,6 @@ def forward_decode( # intermediate states written to the batch-scoped [:num_decodes] # prefix consumed by update_mamba_states()); fall back to the # Triton recurrent kernel when unavailable. - if kwargs.get("use_replay", False): - output_d = None - if output is not None: - output_d = output.view( - num_decodes, - draft_token_num, - self.num_v_heads // self.attn_tp_size, - self.head_v_dim, - ) - return self._replay_verify_recurrent( - query, - key, - value, - a, - b, - ssm_states, - cache_indices, - num_decodes, - draft_token_num, - kwargs.get("replay_metadata"), - kwargs.get("layer_cache"), - kwargs.get("replay_work_items"), - kwargs.get("replay_n_writes"), - output_d, - packed_qkv=mixed_qkv, - use_all_layer_commit=kwargs.get("use_cached_replay_all_layer_commit", False), - ).view( - 1, - num_decodes * draft_token_num, - self.num_v_heads // self.attn_tp_size, - self.head_v_dim, - ) - if _can_use_flashinfer_gdn_verify( ssm_states, self.head_k_dim, self.head_v_dim, draft_token_num ): @@ -776,6 +779,7 @@ def forward_extend( ) conv_states_to_use = conv_states + use_replay = kwargs.get("use_replay", False) seqlen_split_size = [num_prefill_tokens, num_decode_tokens] if num_decode_tokens > 0: @@ -811,21 +815,22 @@ def forward_extend( assert intermediate_conv_states is not None assert kwargs.get("use_replay", False) or intermediate_ssm_states is not None - intermediate_state_indices = kwargs["intermediate_state_indices"] - mixed_qkv_d = mixed_qkv_d.reshape(num_decodes, draft_token_num, -1).transpose(1, 2) - mixed_qkv_d = causal_conv1d_update_triton( - mixed_qkv_d, - conv_states_to_use, - self.conv1d.weight, - self.conv1d.bias, - activation=self.activation, - conv_state_indices=state_indices_d, - intermediate_conv_window=intermediate_conv_states, - intermediate_state_indices=intermediate_state_indices, - # PDL chain: conv1d -> replay verify kernel (replay only) - launch_dependent_kernels=kwargs.get("use_replay", False), - ) - mixed_qkv_d = mixed_qkv_d.transpose(1, 2).reshape(num_decode_tokens, -1) + if not use_replay: + intermediate_state_indices = kwargs["intermediate_state_indices"] + mixed_qkv_d = mixed_qkv_d.reshape(num_decodes, draft_token_num, -1).transpose( + 1, 2 + ) + mixed_qkv_d = causal_conv1d_update_triton( + mixed_qkv_d, + conv_states_to_use, + self.conv1d.weight, + self.conv1d.bias, + activation=self.activation, + conv_state_indices=state_indices_d, + intermediate_conv_window=intermediate_conv_states, + intermediate_state_indices=intermediate_state_indices, + ) + mixed_qkv_d = mixed_qkv_d.transpose(1, 2).reshape(num_decode_tokens, -1) else: mixed_qkv_d = causal_conv1d_update( mixed_qkv_d, @@ -850,13 +855,14 @@ def forward_extend( self.head_v_dim, beta_dtype=b.dtype, ) - query_d, key_d, value_d = pack_gdn_decode_qkv( - mixed_qkv_d, - self.num_k_heads_per_tp, - self.head_k_dim, - self.num_v_heads_per_tp, - self.head_v_dim, - ) + if not use_replay: + query_d, key_d, value_d = pack_gdn_decode_qkv( + mixed_qkv_d, + self.num_k_heads_per_tp, + self.head_k_dim, + self.num_v_heads_per_tp, + self.head_v_dim, + ) else: query, key, value, g, beta = fused_gdn_post_conv( mixed_qkv_p_t, @@ -929,20 +935,7 @@ def forward_extend( ) draft_token_num = spec_metadata.runtime_draft_len + 1 - query_d = query_d.reshape( - num_decodes, draft_token_num, self.num_k_heads // self.attn_tp_size, self.head_k_dim - ) - key_d = key_d.reshape( - num_decodes, draft_token_num, self.num_k_heads // self.attn_tp_size, self.head_k_dim - ) - value_d = value_d.reshape( - num_decodes, draft_token_num, self.num_v_heads // self.attn_tp_size, self.head_v_dim - ) - - a_d = a_d.reshape(num_decodes, draft_token_num, -1) - b_d = b_d.reshape(num_decodes, draft_token_num, -1) out_v_heads = self.num_v_heads // self.attn_tp_size - output_d = None if output is not None: output_d = output[:, num_prefill_tokens:, :, :].view( @@ -952,73 +945,92 @@ def forward_extend( self.head_v_dim, ) - if kwargs.get("use_replay", False): - attn_out_decode = self._replay_verify_recurrent( - query_d, - key_d, - value_d, - a_d, - b_d, - ssm_states, - state_indices_d, - num_decodes, - draft_token_num, - kwargs.get("replay_metadata"), - kwargs.get("layer_cache"), - kwargs.get("replay_work_items"), - kwargs.get("replay_n_writes"), - output_d, - use_all_layer_commit=kwargs.get("use_cached_replay_all_layer_commit", False), - ).reshape(1, num_decode_tokens, out_v_heads, self.head_v_dim) - elif _can_use_flashinfer_gdn_verify( - ssm_states, self.head_k_dim, self.head_v_dim, draft_token_num - ): - # FI gathers the initial state from the pool via state_indices_d - # (no host gather) and writes batch-scoped intermediate states; - # the [:num_decodes] prefix matches update_mamba_states()'s rows. - attn_out_decode = _flashinfer_gdn_verify( - A_log=self.A_log, - a=a_d, - dt_bias=self.dt_bias, - softplus_beta=1.0, - softplus_threshold=20.0, - q=query_d, - k=key_d, - v=value_d, - b=b_d, - initial_state_source=ssm_states, - initial_state_indices=state_indices_d, - intermediate_states_buffer=intermediate_ssm_states[:num_decodes], - scale=self.head_k_dim**-0.5, - use_qk_l2norm_in_kernel=True, + if use_replay: + attn_out_decode = self._replay_verify( + raw_x=mixed_qkv_d, + packed_ba=b_d, + conv_states=conv_states_to_use, + ssm_states=ssm_states, + state_indices=state_indices_d, + num_decodes=num_decodes, + draft_token_num=draft_token_num, + replay_metadata=kwargs.get("replay_metadata"), + layer_cache=kwargs.get("layer_cache"), + is_dummy=kwargs.get("is_dummy"), + candidate_x=intermediate_conv_states, output=output_d, ).reshape(1, num_decode_tokens, out_v_heads, self.head_v_dim) else: - beta_d = b_d.sigmoid() - g_d = fused_gdn_gating( - self.A_log, - a_d.view(num_decodes * draft_token_num, -1), - self.dt_bias, - ).reshape(num_decodes, draft_token_num, -1) - - recurrent_state_source = ssm_states[state_indices_d] - recurrent_state_indices = kwargs["intermediate_state_indices"] - - attn_out_decode = fused_recurrent_gated_delta_rule_update( - q=query_d, - k=key_d, - v=value_d, - g=g_d, - beta=beta_d, - initial_state_source=recurrent_state_source, - initial_state_indices=recurrent_state_indices, - use_qk_l2norm_in_kernel=True, - disable_state_update=True, - intermediate_states_buffer=intermediate_ssm_states, - cache_steps=draft_token_num, - output=output_d, - ).view(1, num_decode_tokens, out_v_heads, self.head_v_dim) + query_d = query_d.reshape( + num_decodes, + draft_token_num, + self.num_k_heads // self.attn_tp_size, + self.head_k_dim, + ) + key_d = key_d.reshape( + num_decodes, + draft_token_num, + self.num_k_heads // self.attn_tp_size, + self.head_k_dim, + ) + value_d = value_d.reshape( + num_decodes, + draft_token_num, + self.num_v_heads // self.attn_tp_size, + self.head_v_dim, + ) + a_d = a_d.reshape(num_decodes, draft_token_num, -1) + b_d = b_d.reshape(num_decodes, draft_token_num, -1) + if _can_use_flashinfer_gdn_verify( + ssm_states, + self.head_k_dim, + self.head_v_dim, + draft_token_num, + ): + # FI gathers initial state from the V2 pool and writes + # batch-scoped intermediate states for accepted promotion. + attn_out_decode = _flashinfer_gdn_verify( + A_log=self.A_log, + a=a_d, + dt_bias=self.dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=query_d, + k=key_d, + v=value_d, + b=b_d, + initial_state_source=ssm_states, + initial_state_indices=state_indices_d, + intermediate_states_buffer=intermediate_ssm_states[:num_decodes], + scale=self.head_k_dim**-0.5, + use_qk_l2norm_in_kernel=True, + output=output_d, + ).reshape(1, num_decode_tokens, out_v_heads, self.head_v_dim) + else: + beta_d = b_d.sigmoid() + g_d = fused_gdn_gating( + self.A_log, + a_d.view(num_decodes * draft_token_num, -1), + self.dt_bias, + ).reshape(num_decodes, draft_token_num, -1) + + recurrent_state_source = ssm_states[state_indices_d] + recurrent_state_indices = kwargs["intermediate_state_indices"] + attn_out_decode = fused_recurrent_gated_delta_rule_update( + q=query_d, + k=key_d, + v=value_d, + g=g_d, + beta=beta_d, + initial_state_source=recurrent_state_source, + initial_state_indices=recurrent_state_indices, + use_qk_l2norm_in_kernel=True, + disable_state_update=True, + intermediate_states_buffer=intermediate_ssm_states, + cache_steps=draft_token_num, + output=output_d, + ).view(1, num_decode_tokens, out_v_heads, self.head_v_dim) if output is not None: return output if attn_out_prefill is None: @@ -1167,29 +1179,24 @@ def forward_core( layer_cache.intermediate_conv_window if is_target_verify else None ) intermediate_ssm_states = layer_cache.intermediate_ssm if is_target_verify else None + use_replay = is_target_verify and getattr( + attn_metadata.kv_cache_manager, "use_gdn_mtp_replay", False + ) intermediate_state_indices = ( _verify_intermediate_state_indices( attn_metadata.kv_cache_manager, num_decodes, state_indices.device ) - if is_target_verify + if is_target_verify and not use_replay else None ) - use_replay = is_target_verify and getattr( - attn_metadata.kv_cache_manager, "use_replay_state_update", False - ) replay_metadata = ( attn_metadata.kv_cache_manager.get_replay_state_update_metadata() if use_replay else None ) - use_cached_replay_all_layer_commit = ( - num_decodes >= CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE - and getattr( - attn_metadata.kv_cache_manager, - "use_gdn_cached_replay_all_layer_commit", - False, - ) + replay_dummy_mask = ( + attn_metadata.kv_cache_manager.get_replay_dummy_request_mask() if use_replay else None ) kwargs = { "mixed_qkv": mixed_qkv, @@ -1210,13 +1217,9 @@ def forward_core( "use_replay": use_replay, "replay_metadata": replay_metadata, "layer_cache": layer_cache, - "replay_work_items": mamba_metadata.replay_work_items[:num_decodes] - if use_replay and num_decodes >= CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE - else None, - "replay_n_writes": mamba_metadata.replay_n_writes - if use_replay and num_decodes >= CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE + "is_dummy": replay_dummy_mask[num_prefills : num_prefills + num_decodes] + if replay_dummy_mask is not None else None, - "use_cached_replay_all_layer_commit": use_cached_replay_all_layer_commit, } if num_prefills > 0: attn_out = self.forward_extend( diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index 7acda892476a..ce708ed9fcdf 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -35,6 +35,60 @@ _FUSED_GDN_REPLAY_WORK_ITEMS_MAX_BATCH_SIZE = 256 +@triton.jit +def _snapshot_gdn_replay_work_items_kernel( + state_indices, + history_len, + work_items, + n_writes_output, + num_decodes, + work_item_width: tl.constexpr, + position_field: tl.constexpr, + cache_slot_field: tl.constexpr, + history_len_field: tl.constexpr, + cache_buf_idx_field: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """Snapshot single-buffer history lengths before the layer replay loop.""" + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + active = offsets < num_decodes + slots = tl.load(state_indices + offsets, mask=active, other=0) + lengths = tl.load(history_len + slots, mask=active, other=0) + output_base = work_items + offsets * work_item_width + tl.store(output_base + position_field, offsets, mask=active) + tl.store(output_base + cache_slot_field, slots, mask=active) + tl.store(output_base + history_len_field, lengths, mask=active) + tl.store(output_base + cache_buf_idx_field, 0, mask=active) + if tl.program_id(0) == 0: + tl.store(n_writes_output, 0) + + +def _build_gdn_replay_work_items( + state_indices: torch.Tensor, + history_len: torch.Tensor, + work_items: torch.Tensor, + n_writes: torch.Tensor, +) -> None: + """Build the immutable inputs consumed by the split GDN MTP commit.""" + num_decodes = state_indices.shape[0] + block_size = 256 + _snapshot_gdn_replay_work_items_kernel[(triton.cdiv( + num_decodes, block_size), )]( + state_indices, + history_len, + work_items, + n_writes, + num_decodes, + work_item_width=REPLAY_WORK_ITEM_WIDTH, + position_field=REPLAY_WORK_POSITION_IN_DECODE_BATCH, + cache_slot_field=REPLAY_WORK_CACHE_SLOT, + history_len_field=REPLAY_WORK_PNAT, + cache_buf_idx_field=REPLAY_WORK_CACHE_BUF_IDX, + BLOCK_SIZE=block_size, + num_warps=4, + ) + + @triton.jit def _prepare_gdn_replay_work_items_kernel( state_indices, @@ -408,15 +462,14 @@ def _prepare_replay_work_items(self, kv_cache_manager, batch_size: int, self.replay_num_decodes = num_decodes if num_decodes == 0: return + use_gdn_mtp_replay = getattr(kv_cache_manager, "use_gdn_mtp_replay", + False) use_gdn_all_layer_commit = getattr( kv_cache_manager, "use_gdn_cached_replay_all_layer_commit", False) if use_gdn_all_layer_commit: from tensorrt_llm._torch.modules.fla.cached_replay import \ CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE - # The fused small-batch GDN kernel commits its checkpoint in-layer - # and indexes cache metadata directly. Work items are only consumed - # by the partitioned replay + all-layer commit path. if num_decodes < CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE: return @@ -436,6 +489,18 @@ def _prepare_replay_work_items(self, kv_cache_manager, batch_size: int, replay_step_width = replay_metadata.replay_step_width replay_history_size = replay_metadata.replay_history_size + if use_gdn_mtp_replay: + _build_gdn_replay_work_items( + self.state_indices[num_contexts:batch_size], + replay_metadata.history_len, + self.replay_work_items, + self.replay_n_writes, + ) + return + + if cache_buf_idx is None: + raise RuntimeError("Mamba2 replay requires cache buffer indices") + if (use_gdn_all_layer_commit and num_decodes <= _FUSED_GDN_REPLAY_WORK_ITEMS_MAX_BATCH_SIZE): build_work_items = _build_replay_work_items_triton diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 7593f2ec6f4c..1f0306646e83 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -45,6 +45,7 @@ from ..attention.backends import get_sparse_attn_kv_cache_manager from ..attention.backends.sparse.params import \ get_indexer_mqa_logits_workspace_bytes +from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from ..disaggregation.kv_cache_transceiver import (AttentionTypeCpp, create_kv_cache_transceiver) from ..hostfunc import set_low_latency_dispatch @@ -3747,19 +3748,15 @@ def _create_kv_cache_manager( )) num_mamba_layers = (0 if is_draft and mamba_params.num_draft_layers > 0 else mamba_params.num_mamba_layers) - # Replay state update for GDN MTP: mirrors the nemotron_hybrid gating - # above, minus the Mamba2-specific stochastic-rounding/Philox gate. - # The GDN replay kernel does a plain cast on checkpoint commit, so - # quantized SSM cache dtypes stay on the legacy path. sm = get_sm_version() - use_replay = spec_config is not None and sm >= 80 + use_replay = spec_config is not None if spec_config is None: logger.info( "GDN replay kernel requires speculative decoding; using " "non-replay path") elif spec_config.tokens_per_gen_step > 8: - logger.info("GDN cached replay supports at most 8 tokens per " - "generation step; using non-replay path") + logger.info("GDN replay supports at most 8 tokens per generation " + "step; using non-replay path") use_replay = False # Tree attention: replay assumes a linear token sequence. @@ -3769,29 +3766,43 @@ def _create_kv_cache_manager( "using legacy MTP path") use_replay = False - if mamba_params.mamba_ssm_cache_dtype not in (torch.float32, - torch.bfloat16, - torch.float16): - logger.info( - "GDN replay kernel does not support quantized SSM cache " - f"dtype {mamba_params.mamba_ssm_cache_dtype}; using legacy " - "MTP path") - use_replay = False + # Experimental modeling targets retain their cached-U replay ABI. + use_legacy_gdn_replay = (getattr(model_config, "modeling_v2_target", + None) is not None) + if not use_legacy_gdn_replay: + # GDN replay is specialized for BF16 Qwen GDN on datacenter + # Blackwell and uses Cache Manager V2 for its split all-layer commit. + if not is_sm_100f(sm): + logger.info("GDN MTP replay currently requires datacenter " + "Blackwell; using non-replay path") + use_replay = False + if not IS_CUTLASS_DSL_AVAILABLE: + logger.info("GDN MTP replay requires CUTLASS DSL; using " + "non-replay path") + use_replay = False + if (mamba_params.mamba_ssm_cache_dtype != torch.bfloat16 + or mamba_params.dtype != torch.bfloat16): + logger.info("GDN MTP replay requires BF16 activations and " + "recurrent states; using non-replay path") + use_replay = False + if (mamba_params.state_size != 128 or mamba_params.head_dim != 128 + or mamba_params.conv_kernel != 4 + or mamba_params.n_groups <= 0 + or mamba_params.num_heads % mamba_params.n_groups != 0): + logger.info( + "GDN MTP replay requires Conv4, K=V=128, and an " + "integral value/key head ratio; using non-replay path") + use_replay = False # Replay is enabled by default for eligible GDN MTP workloads. if not is_gdn_replay_enabled(): use_replay = False - # GDN replay supports the contiguous C++ V1 state pool and the indirect - # per-layer state views exposed by V2. Mixed/Python does not expose an - # all-layer commit, so keep that manager but use non-replay MTP. - replay_manager_types = (CppMambaHybridCacheManager, - MambaHybridCacheManagerV2) if use_replay and not issubclass(kv_cache_manager_cls, - replay_manager_types): - logger.info("GDN replay requires C++ V1 or V2 Mamba cache manager; " - f"{kv_cache_manager_cls.__name__} was selected, so the " - "non-replay MTP path will be used") + MambaHybridCacheManagerV2): + logger.info("GDN replay requires MambaHybridCacheManagerV2; " + f"{kv_cache_manager_cls.__name__} was selected, so " + "the non-replay MTP path will be used") use_replay = False logger.info("GDN replay state update: " + ("ENABLED" if use_replay else "DISABLED")) @@ -3799,6 +3810,9 @@ def _create_kv_cache_manager( mamba_manager_extra_kwargs = dict(manager_extra_kwargs) mamba_manager_extra_kwargs.update( _mamba_conv_layout_kwargs(kv_cache_manager_cls, "qwen3_next")) + if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): + mamba_manager_extra_kwargs["use_gdn_mtp_replay"] = ( + not use_legacy_gdn_replay) if getattr(sparse_attention_config, "algorithm", None) == "qsa": # Resolve the side-cache shape from the same checkpoint geometry # used to construct the QSA index projection. diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index df494fc9f27b..94a39d8efa99 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -348,10 +348,15 @@ def use_py_mamba_cache_manager() -> bool: class ReplayStateUpdateMetadata(NamedTuple): """Shared tensors and fixed sizes for replay state updates.""" prev_num_accepted_tokens: torch.Tensor - cache_buf_idx: torch.Tensor + cache_buf_idx: Optional[torch.Tensor] replay_step_width: int replay_history_size: int + @property + def history_len(self) -> torch.Tensor: + """Accepted single-buffer history length for GDN MTP replay.""" + return self.prev_num_accepted_tokens + def _advance_replay_state( replay_metadata: ReplayStateUpdateMetadata, @@ -372,7 +377,11 @@ def _advance_replay_state( accepted_tokens, prev_num_accepted_tokens + accepted_tokens, ) - cache_buf_idx = replay_metadata.cache_buf_idx[slots] + cache_buf_idx_storage = replay_metadata.cache_buf_idx + if cache_buf_idx_storage is None: + raise RuntimeError( + "Double-buffered replay requires cache buffer indices") + cache_buf_idx = cache_buf_idx_storage[slots] next_cache_buf_idx = torch.where(wrote_checkpoint, 1 - cache_buf_idx, cache_buf_idx) if is_dummy_request is not None: @@ -384,7 +393,7 @@ def _advance_replay_state( next_cache_buf_idx = torch.where(is_dummy_request, cache_buf_idx, next_cache_buf_idx) replay_metadata.prev_num_accepted_tokens[slots] = next_num_accepted_tokens - replay_metadata.cache_buf_idx[slots] = next_cache_buf_idx + cache_buf_idx_storage[slots] = next_cache_buf_idx class BaseMambaCacheManager(ABC): @@ -403,6 +412,10 @@ def get_replay_state_update_metadata( """Return replay metadata tensors and fixed replay sizes.""" return None + def get_replay_dummy_request_mask(self) -> Optional[torch.Tensor]: + """Return the batch-ordered dummy mask used by replay kernels.""" + return getattr(self, "_dummy_request_mask", None) + @property def use_kda_replay_update(self) -> bool: """Whether KDA fused verification owns per-slot replay caches.""" @@ -1509,10 +1522,17 @@ def _setup_mtp_intermediate_states(self, spec_config, dtype=self.ssm_state_dtype, device="cuda", ) + if getattr(self, "_use_gdn_mtp_replay", False): + # GDN replay reconstructs every candidate Conv4 result from the + # persistent three-token checkpoint. Retain only raw candidate x; + # the combined commit rebuilds the accepted three-token window. + intermediate_conv_shape = [self.conv_state_shape[0]] + else: + intermediate_conv_shape = self.conv_state_shape self.intermediate_conv_states = torch.zeros( size=[ self.local_num_mamba_layers, max_batch_size, tokens_per_gen_step - ] + self.conv_state_shape, + ] + intermediate_conv_shape, dtype=self.conv_state_dtype, device="cuda", ) @@ -1551,10 +1571,38 @@ def _allocate_pool_replay_buffers( history_size = self.replay_history_size assert history_size is not None nheads, head_dim, d_state = self.ssm_state_shape - common_shape = [self.local_num_mamba_layers, cache_size, 2] self.prev_num_accepted_tokens = torch.zeros(cache_size, dtype=torch.int32, device=device) + if getattr(self, "_use_gdn_mtp_replay", False): + if self.ssm_state_dtype != torch.bfloat16: + raise ValueError( + "GDN MTP replay requires BF16 recurrent states") + replay_step_width = self.replay_step_width + assert replay_step_width is not None + capacity = history_size + replay_step_width + common_shape = [self.local_num_mamba_layers, cache_size] + # A single accepted-history buffer is enough: replay appends all + # candidates, then the split all-layer commit retains only the + # accepted prefix and compacts an overflow after each 16-row tile. + self.old_x = torch.zeros( + common_shape + [capacity, nheads, head_dim], + dtype=torch.bfloat16, + device=device, + ) + self.old_B = torch.zeros( + common_shape + [capacity, self._n_groups_per_rank, d_state], + dtype=torch.bfloat16, + device=device, + ) + self.old_dt = torch.zeros( + common_shape + [nheads, capacity], + dtype=torch.float32, + device=device, + ) + return True + + common_shape = [self.local_num_mamba_layers, cache_size, 2] self.cache_buf_idx = torch.zeros(cache_size, dtype=torch.int32, device=device) @@ -1599,18 +1647,22 @@ def _reset_context_mamba_slots(self, num_contexts: int) -> None: context_slots = self.cuda_state_indices[:num_contexts].long() if (self._use_replay_state_update - and self.prev_num_accepted_tokens is not None - and self.cache_buf_idx is not None): + and self.prev_num_accepted_tokens is not None): self.prev_num_accepted_tokens.index_fill_(0, context_slots, 0) - self.cache_buf_idx.index_fill_(0, context_slots, 0) - if self.old_x is not None: - self.old_x.index_fill_(1, context_slots, 0) - if self.old_B is not None: - self.old_B.index_fill_(1, context_slots, 0) - if self.old_dt is not None: - self.old_dt.index_fill_(1, context_slots, 0) - if self.old_dA_cumsum is not None: - self.old_dA_cumsum.index_fill_(1, context_slots, 0) + # Single-buffer GDN replay gates every history read by history_len + # and overwrites candidate rows before accepting them. Resetting + # the length is therefore sufficient when a slot is reused. + if not getattr(self, "_use_gdn_mtp_replay", False): + if self.cache_buf_idx is not None: + self.cache_buf_idx.index_fill_(0, context_slots, 0) + if self.old_x is not None: + self.old_x.index_fill_(1, context_slots, 0) + if self.old_B is not None: + self.old_B.index_fill_(1, context_slots, 0) + if self.old_dt is not None: + self.old_dt.index_fill_(1, context_slots, 0) + if self.old_dA_cumsum is not None: + self.old_dA_cumsum.index_fill_(1, context_slots, 0) if self.mamba_ssm_rand_seed is None: return @@ -1689,7 +1741,9 @@ def mamba_layer_cache( spec_kwargs['old_x'] = self.old_x[layer_offset] spec_kwargs['old_B'] = self.old_B[layer_offset] spec_kwargs['old_dt'] = self.old_dt[layer_offset] - spec_kwargs['old_dA_cumsum'] = self.old_dA_cumsum[layer_offset] + spec_kwargs['old_dA_cumsum'] = (self.old_dA_cumsum[layer_offset] + if self.old_dA_cumsum + is not None else None) spec_kwargs['cache_buf_idx'] = self.cache_buf_idx spec_kwargs['prev_num_accepted_tokens'] = ( self.prev_num_accepted_tokens) @@ -1715,10 +1769,13 @@ def get_replay_state_update_metadata( None) cache_buf_idx = getattr(self, 'cache_buf_idx', None) if (not self._use_replay_state_update - or prev_num_accepted_tokens is None or cache_buf_idx is None + or prev_num_accepted_tokens is None or self.replay_step_width is None or self.replay_history_size is None): return None + if (not getattr(self, "_use_gdn_mtp_replay", False) + and cache_buf_idx is None): + return None return ReplayStateUpdateMetadata( prev_num_accepted_tokens=prev_num_accepted_tokens, cache_buf_idx=cache_buf_idx, @@ -2445,9 +2502,6 @@ def __init__( # accessors (get_mamba_ssm_cache_dtype, use_replay_state_update) work # on ranks with no local mamba layers. self._use_replay_state_update = use_replay_state_update - self._use_gdn_cached_replay_all_layer_commit = ( - use_replay_state_update and model_type == "qwen3_next" - and self.local_num_mamba_layers > 0) self.replay_step_width: Optional[int] = ( spec_config.tokens_per_gen_step if spec_config is not None and use_replay_state_update else None) @@ -2668,12 +2722,6 @@ def __init__( self._setup_states() self._setup_replay_buffers(spec_config) - if use_replay_state_update and model_type == "qwen3_next": - logger.info_once( - "Configured GDN cached replay commit mode: small-batch fused, " - "large-batch all-layer", - key="gdn_cached_replay_commit_mode_fused", - ) @staticmethod def get_cache_size_per_token( @@ -2712,46 +2760,6 @@ def get_cache_size_per_token( **kwargs, ) - @property - def use_gdn_cached_replay_all_layer_commit(self) -> bool: - return self._use_gdn_cached_replay_all_layer_commit - - def _commit_gdn_cached_replay_history_layers( - self, - attn_metadata: "AttentionMetadata", - num_decodes: int, - ) -> None: - """Synchronously advance every local GDN checkpoint in one launch.""" - from tensorrt_llm._torch.modules.fla.cached_replay import ( - CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, - commit_gdn_cached_replay_history_layers) - - if (not self._use_gdn_cached_replay_all_layer_commit - or num_decodes < CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE): - return - if (self.all_ssm_states is None or self.old_x is None - or self.old_B is None or self.old_dt is None - or self.replay_history_size is None): - raise RuntimeError( - "GDN cached replay all-layer commit requires replay state buffers." - ) - - mamba_metadata = attn_metadata.mamba_metadata - if mamba_metadata.replay_num_decodes != num_decodes: - raise RuntimeError( - "GDN replay metadata contains " - f"{mamba_metadata.replay_num_decodes} decode requests, " - f"but state update received {num_decodes}.") - commit_gdn_cached_replay_history_layers( - ssm_states=self.all_ssm_states, - old_u=self.old_x, - old_k=self.old_B, - old_G=self.old_dt, - replay_work_items=mamba_metadata.replay_work_items[:num_decodes], - n_writes=mamba_metadata.replay_n_writes, - history_size=self.replay_history_size, - ) - def shutdown(self): # Release tensor views into the pool before the pool memory is freed, # so their deleters don't see stale pointers. @@ -2882,18 +2890,10 @@ def update_mamba_states( num_gens].to(torch.int32) src_state_indices = self.intermediate_state_indices[:num_gens] - # The accepted SSM/conv promotion is a bandwidth-bound gather->scatter - # into the dtype-reinterpreted, strided C++ pool view. torch.compile - # can't handle the dtype-view mutation (and Inductor chokes on the - # uint8 pool with "XBLOCK too large"), so a dedicated Triton kernel - # writes through the view's real strides (~85% of HBM peak, one launch - # per state). + # Replay kernels update the SSM checkpoint in place. The non-replay + # path promotes the accepted intermediate state into the strided C++ + # pool with a dedicated Triton kernel. if self._use_replay_state_update: - # Every GDN layer has finished reading the old checkpoint and - # writing its candidate history. Advance all local checkpoints - # in one launch before PNAT and the active history buffer change. - self._commit_gdn_cached_replay_history_layers( - attn_metadata, num_gens) assert self._dummy_request_mask is not None else: # Legacy: copy the accepted SSM state from the intermediate buffer. @@ -3185,6 +3185,7 @@ def __init__( conv_state_layout: Literal["x_b_c", "q_k_v"] = "x_b_c", kda_replay_num_spec: Optional[int] = None, qwen4_exp_ple_cache_params: Optional["Qwen4ExpPLECacheParams"] = None, + use_gdn_mtp_replay: bool = True, kda_draft_records: bool = False, **kwargs, ) -> None: @@ -3251,12 +3252,20 @@ def __init__( self._mamba_layer_mask = list(mamba_layer_mask) self._use_replay_state_update = use_replay_state_update - self._use_gdn_cached_replay_all_layer_commit = (use_replay_state_update - and conv_state_layout - == "q_k_v") + self._use_gdn_mtp_replay = (use_replay_state_update + and conv_state_layout == "q_k_v" + and use_gdn_mtp_replay) + self._use_gdn_cached_replay_all_layer_commit = ( + use_replay_state_update and conv_state_layout == "q_k_v" + and not use_gdn_mtp_replay) self._gdn_cached_replay_state_descriptors: Optional[torch.Tensor] = None self._gdn_cached_replay_state_strides: Optional[Tuple[int, int, int]] = None + self._gdn_mtp_replay_state_descriptors: Optional[torch.Tensor] = None + self._gdn_mtp_replay_state_view: Optional[torch.Tensor] = None + self._gdn_mtp_replay_conv_state_descriptors: Optional[ + torch.Tensor] = None + self._gdn_mtp_replay_conv_state_view: Optional[torch.Tensor] = None self.replay_step_width: Optional[int] = ( spec_config.tokens_per_gen_step if spec_config is not None and use_replay_state_update else None) @@ -3301,6 +3310,8 @@ def __init__( if mamba_layer_mask[layer_idx] ] self.local_num_mamba_layers = len(self.mamba_pp_layers) + self._use_gdn_mtp_replay = (self._use_gdn_mtp_replay + and self.local_num_mamba_layers > 0) self._use_gdn_cached_replay_all_layer_commit = ( self._use_gdn_cached_replay_all_layer_commit and self.local_num_mamba_layers > 0) @@ -3451,15 +3462,14 @@ def __init__( # PLE pools use exactly the state-slot capacity resolved above. self._setup_ple_states(num_ssm_slots) self._setup_replay_buffers(spec_config) - if self._use_gdn_cached_replay_all_layer_commit: - state_layout = ("affine" - if self._gdn_cached_replay_state_strides + if self._use_gdn_mtp_replay: + state_layout = ("affine" if self._gdn_mtp_replay_state_view is not None else "indirect") logger.info_once( - "Configured GDN cached replay commit mode for V2: " - "small-batch fused, large-batch all-layer; " + "Configured GDN MTP replay for V2: " + "split replay and all-layer commit; " f"state layout: {state_layout}", - key="gdn_cached_replay_v2_commit_mode_fused", + key="gdn_mtp_replay_v2", ) else: self.ssm_layer_group_id = None @@ -3609,6 +3619,10 @@ def get_disagg_ignored_roles(self) -> frozenset[DataRole]: replay_roles.add(_KDA_REPLAY_COUNT_ROLE) return super().get_disagg_ignored_roles() | frozenset(replay_roles) + @property + def use_gdn_mtp_replay(self) -> bool: + return getattr(self, "_use_gdn_mtp_replay", False) + @property def use_gdn_cached_replay_all_layer_commit(self) -> bool: return getattr(self, "_use_gdn_cached_replay_all_layer_commit", False) @@ -3795,6 +3809,57 @@ def _commit_gdn_cached_replay_history_layers( history_size=self.replay_history_size, ) + def _commit_gdn_mtp_replay( + self, + attn_metadata: "AttentionMetadata", + num_decodes: int, + accepted_tokens: torch.Tensor, + is_dummy_request: torch.Tensor, + ) -> None: + """Commit accepted Conv and GDN histories in one launch.""" + from tensorrt_llm._torch.custom_ops.cute_dsl_gdn_mtp_commit import \ + gdn_mtp_commit + + if not getattr(self, "_use_gdn_mtp_replay", False): + return + if (not self.all_ssm_states or not self.all_conv_states + or (self._gdn_mtp_replay_state_descriptors is None + and self._gdn_mtp_replay_state_view is None) + or (self._gdn_mtp_replay_conv_state_descriptors is None + and self._gdn_mtp_replay_conv_state_view is None) + or self.old_x is None or self.old_B is None + or self.old_dt is None or self.prev_num_accepted_tokens is None + or self.replay_history_size is None + or self.intermediate_conv_states is None): + raise RuntimeError( + "GDN MTP commit requires V2 replay and Conv state buffers.") + + mamba_metadata = attn_metadata.mamba_metadata + if mamba_metadata.replay_num_decodes != num_decodes: + raise RuntimeError( + "GDN replay metadata contains " + f"{mamba_metadata.replay_num_decodes} decode requests, " + f"but state update received {num_decodes}.") + state_view = self._gdn_mtp_replay_state_view + conv_state_view = self._gdn_mtp_replay_conv_state_view + gdn_mtp_commit( + ssm_states=(state_view + if state_view is not None else self.all_ssm_states[0]), + conv_states=(conv_state_view if conv_state_view is not None else + self.all_conv_states[0]), + ssm_state_descriptors=self._gdn_mtp_replay_state_descriptors, + conv_state_descriptors=( + self._gdn_mtp_replay_conv_state_descriptors), + candidate_x=self.intermediate_conv_states, + history_u=self.old_x, + history_k=self.old_B, + history_G=self.old_dt, + history_len=self.prev_num_accepted_tokens, + replay_work_items=mamba_metadata.replay_work_items[:num_decodes], + accepted_tokens=accepted_tokens.to(torch.int32), + is_dummy=is_dummy_request, + ) + def _create_kv_cache( self, request_id: int, @@ -4383,6 +4448,65 @@ def _setup_states(self) -> None: device=reference.device, ) + if self._use_gdn_mtp_replay: + expected_inner_strides = ( + self.ssm_state_shape[1] * self.ssm_state_shape[2], + self.ssm_state_shape[2], + 1, + ) + reference = self.all_ssm_states[0] + for state in self.all_ssm_states: + if (state.dtype != reference.dtype + or state.device != reference.device + or list(state.shape[1:]) != self.ssm_state_shape + or state.stride()[1:] != expected_inner_strides): + raise RuntimeError( + "GDN replay requires V2 SSM layers with matching " + "dtype, device, shape, and dense inner dimensions.") + if self._stacked_ssm_states is not None: + # The helper wraps the complete affine address range before + # constructing the view. A direct as_strided() from a single + # layer would fail storage-bound checks for later layers. + self._gdn_mtp_replay_state_view = ( + self._stacked_ssm_states.permute(1, 0, 2, 3, 4)) + else: + logger.warning_once( + "V2 GDN state views are not affine; using indirect MTP replay " + "checkpoint addressing", + key="gdn_mtp_replay_v2_indirect_state_layout", + ) + self._gdn_mtp_replay_state_descriptors = torch.tensor( + [(state.data_ptr(), state.stride(0)) + for state in self.all_ssm_states], + dtype=torch.int64, + device=reference.device, + ) + conv_reference = self.all_conv_states[0] + expected_conv_inner_strides = (self.conv_state_shape[1], 1) + for state in self.all_conv_states: + if (state.dtype != conv_reference.dtype + or state.device != conv_reference.device + or list(state.shape[1:]) != self.conv_state_shape + or state.stride()[1:] != expected_conv_inner_strides): + raise RuntimeError( + "GDN replay requires V2 Conv layers with matching " + "dtype, device, shape, and dense inner dimensions.") + if self._stacked_conv_states is not None: + self._gdn_mtp_replay_conv_state_view = ( + self._stacked_conv_states.permute(1, 0, 2, 3)) + else: + logger.warning_once( + "V2 GDN Conv-state views are not affine; using indirect " + "MTP commit addressing", + key="gdn_mtp_replay_v2_indirect_conv_layout", + ) + self._gdn_mtp_replay_conv_state_descriptors = torch.tensor( + [(state.data_ptr(), state.stride(0)) + for state in self.all_conv_states], + dtype=torch.int64, + device=conv_reference.device, + ) + @staticmethod def _stack_state_views( states: List[torch.Tensor]) -> Optional[torch.Tensor]: @@ -4642,6 +4766,51 @@ def _commit_kda_pending_drafts(self, slot: int) -> None: ) self.prev_num_accepted_tokens[slot] = 0 + def _reset_gdn_mtp_replay_slots( + self, + slots: Union[List[int], torch.Tensor], + ) -> None: + """Reset GDN replay history for newly initialized logical slots.""" + assert self.use_gdn_mtp_replay + assert self.prev_num_accepted_tokens is not None + slots = torch.as_tensor(slots, + dtype=torch.long, + device=self.prev_num_accepted_tokens.device) + if slots.numel() > 0: + # History contents need not be cleared: replay masks all reads at + # and beyond this per-slot length and overwrites appended rows. + self.prev_num_accepted_tokens[slots] = 0 + + def _relocate_gdn_mtp_replay_slots(self, old_slots: List[int], + new_slots: List[int]) -> None: + """Move GDN replay history when V2 remaps a live request's slot.""" + assert self.use_gdn_mtp_replay + assert self.prev_num_accepted_tokens is not None + moves = [(old, new) for old, new in zip(old_slots, new_slots) + if old >= 0 and old != new] + if not moves: + return + + device = self.prev_num_accepted_tokens.device + source_slots = torch.tensor([old for old, _ in moves], + dtype=torch.long, + device=device) + destination_slots = torch.tensor([new for _, new in moves], + dtype=torch.long, + device=device) + for history in (self.old_x, self.old_B, self.old_dt): + assert history is not None + history.index_copy_( + 1, + destination_slots, + history.index_select(1, source_slots), + ) + self.prev_num_accepted_tokens.index_copy_( + 0, + destination_slots, + self.prev_num_accepted_tokens.index_select(0, source_slots), + ) + def add_dummy_requests( self, request_ids: List[int], @@ -4715,8 +4884,10 @@ def _setup_state_indices(self, num_contexts: int = 0) -> None: if self.local_num_mamba_layers == 0: return + tracks_replay_slots = (self.use_kda_replay_update + or self.use_gdn_mtp_replay) replay_requests = (requests[num_contexts:] - if self.use_kda_replay_update else []) + if tracks_replay_slots else []) old_state_values = [ self._request_id_to_state_index.get(request.py_request_id, -1) for request in replay_requests @@ -4778,6 +4949,17 @@ def _setup_state_indices(self, self._seed_kda_replay_slots( resumed_indices.to(self.prev_num_accepted_tokens.device, non_blocking=True)) + elif self.use_gdn_mtp_replay: + new_state_values = [ + self._request_id_to_state_index[request.py_request_id] + for request in replay_requests + ] + self._relocate_gdn_mtp_replay_slots(old_state_values, + new_state_values) + self._reset_gdn_mtp_replay_slots([ + new for old, new in zip(old_state_values, new_state_values) + if old < 0 + ]) def _record_kda_replay_acceptance( self, @@ -4984,7 +5166,18 @@ def update_mamba_states( num_accepted_draft_tokens) src_state_indices = self.intermediate_state_indices[:num_gens] - if self._use_replay_state_update: + if self.use_gdn_mtp_replay: + assert self._dummy_request_mask is not None + is_dummy_request = self._dummy_request_mask[ + num_contexts:num_contexts + num_gens] + self._commit_gdn_mtp_replay( + attn_metadata, + num_gens, + num_accepted_tokens[num_contexts:num_contexts + num_gens], + is_dummy_request, + ) + return + elif self._use_replay_state_update: self._commit_gdn_cached_replay_history_layers( attn_metadata, num_gens) assert self._dummy_request_mask is not None @@ -5117,6 +5310,10 @@ def update_context_resources(self, def shutdown(self): self._gdn_cached_replay_state_descriptors = None self._gdn_cached_replay_state_strides = None + self._gdn_mtp_replay_state_descriptors = None + self._gdn_mtp_replay_state_view = None + self._gdn_mtp_replay_conv_state_descriptors = None + self._gdn_mtp_replay_conv_state_view = None self.all_ssm_states = [] self.all_conv_states = [] self._stacked_ssm_states = None diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 9c8067141c5e..5f163e217b43 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -559,8 +559,12 @@ def test_hybrid_cache_manager_factory_honors_v2_setting( assert get_kv_cache_manager_cls(_hybrid_model_config(), kv_cache_config) is expected -def test_qwen3_gdn_replay_supports_cpp_and_v2_managers(monkeypatch): - """GDN replay uses C++ V1 directly and V2 through state descriptors.""" +@pytest.mark.parametrize( + "modeling_v2,ssm_dtype", + [(False, torch.bfloat16), (True, torch.bfloat16), (True, torch.float32)], +) +def test_qwen3_gdn_replay_requires_v2_manager(monkeypatch, modeling_v2, ssm_dtype): + """GDN replay is enabled only with its supported V2 cache manager.""" captured_cpp = {} captured_mixed = {} captured_v2 = {} @@ -589,22 +593,24 @@ def __init__(self, *args, **kwargs): model_config = SimpleNamespace( pretrained_config=pretrained_config, quant_config=None, + modeling_v2_target=object if modeling_v2 else None, ) mamba_params = MambaKVCacheParams( - state_size=8, + state_size=128, conv_kernel=4, - num_heads=4, - n_groups=1, - head_dim=8, + num_heads=6, + n_groups=2, + head_dim=128, mamba_layer_mask=[True, False], target_full_attention_layer_mask=[False, True], num_mamba_layers=1, num_draft_layers=1, dtype=torch.bfloat16, - mamba_ssm_cache_dtype=torch.bfloat16, + mamba_ssm_cache_dtype=ssm_dtype, ) monkeypatch.setenv("TRTLLM_USE_GDN_REPLAY", "1") - monkeypatch.setattr("tensorrt_llm._torch.pyexecutor._util.get_sm_version", lambda: 90) + monkeypatch.setattr("tensorrt_llm._torch.pyexecutor._util.get_sm_version", lambda: 100) + monkeypatch.setattr("tensorrt_llm._torch.pyexecutor._util.IS_CUTLASS_DSL_AVAILABLE", True) monkeypatch.setattr( "tensorrt_llm._torch.pyexecutor._util.extract_mamba_kv_cache_params", lambda *args, **kwargs: mamba_params, @@ -647,7 +653,7 @@ def __init__(self, *args, **kwargs): **common_kwargs, ) - assert captured_cpp["use_replay_state_update"] is True + assert captured_cpp["use_replay_state_update"] is False assert captured_cpp["model_type"] == "qwen3_next" assert captured_cpp["max_num_tokens"] == 256 assert "vocab_size" not in captured_cpp @@ -656,13 +662,15 @@ def __init__(self, *args, **kwargs): assert captured_mixed["max_num_tokens"] == 256 assert "vocab_size" not in captured_mixed assert captured_v2["use_replay_state_update"] is True + assert captured_v2["use_gdn_mtp_replay"] is not modeling_v2 assert captured_v2["max_num_tokens"] == 256 assert captured_v2["vocab_size"] == pretrained_config.vocab_size assert "model_type" not in captured_v2 assert captured_v2["conv_state_layout"] == "q_k_v" fallback_logs = [str(call.args[0]) for call in info_log.call_args_list] - assert any("RecordingMixedManager was selected" in log for log in fallback_logs) + assert any("RecordingMixedManager" in log for log in fallback_logs) assert not any("RecordingV2Manager was selected" in log for log in fallback_logs) + assert any("RecordingCppManager" in log for log in fallback_logs) @pytest.mark.parametrize( @@ -1514,7 +1522,6 @@ def test_cpp_hybrid_replay_bookkeeping_is_fused_into_conv_promotion( mgr.intermediate_state_indices = torch.arange(2, dtype=torch.int32) mgr.all_conv_states = torch.empty(0) mgr.intermediate_conv_states = torch.empty(0) - mgr._commit_gdn_cached_replay_history_layers = MagicMock() promote_calls = [] monkeypatch.setattr( @@ -1534,7 +1541,6 @@ def test_cpp_hybrid_replay_bookkeeping_is_fused_into_conv_promotion( state_indices=torch.tensor([0, 1], dtype=torch.int32), ) - mgr._commit_gdn_cached_replay_history_layers.assert_called_once() assert len(promote_calls) == 1 _, kwargs = promote_calls[0] assert kwargs["replay_pnat"] is mgr.prev_num_accepted_tokens @@ -2460,6 +2466,7 @@ def _build_v2_hybrid_with_mamba_layer( dtype=DataType.HALF, kv_cache_dtype="auto", conv_state_layout="x_b_c", + use_gdn_mtp_replay=True, mamba_d_conv=4, mamba_num_heads=4, mamba_n_groups=1, @@ -2527,6 +2534,7 @@ def _build_v2_hybrid_with_mamba_layer( use_replay_state_update=use_replay_state_update, dtype=dtype, conv_state_layout=conv_state_layout, + use_gdn_mtp_replay=use_gdn_mtp_replay, kda_replay_num_spec=kda_replay_num_spec, kda_draft_records=kda_draft_records, is_disagg=is_disagg, @@ -5260,127 +5268,85 @@ def test_hybrid_replay_buffers_size_by_tokens_per_gen_step(builder): @skip_no_cuda -def test_v2_gdn_replay_builds_affine_state_layout(): +@pytest.mark.parametrize( + "use_gdn_mtp_replay,ssm_dtype", + [(True, torch.bfloat16), (False, torch.bfloat16), (False, torch.float32)], +) +def test_v2_gdn_replay_allocates_single_history_buffer(use_gdn_mtp_replay, ssm_dtype): mgr = _build_v2_hybrid_with_mamba_layer( max_batch_size=16, num_mamba_layers=3, spec_config=MTPDecodingConfig(max_draft_len=3), use_replay_state_update=True, conv_state_layout="q_k_v", + mamba_ssm_cache_dtype=ssm_dtype, + use_gdn_mtp_replay=use_gdn_mtp_replay, ) try: - assert mgr.use_gdn_cached_replay_all_layer_commit - assert mgr._gdn_cached_replay_state_descriptors is None - state_strides = mgr._gdn_cached_replay_state_strides - assert state_strides is not None - states = mgr.all_ssm_states - first_state = states[0] - expected_layer_stride = ( - states[1].data_ptr() - first_state.data_ptr() - ) // first_state.element_size() - assert state_strides == ( - mgr.local_num_mamba_layers, - expected_layer_stride, - first_state.stride(0), - ) - finally: - mgr.shutdown() - - -@skip_no_cuda -def test_v2_gdn_replay_all_layer_commit_matches_contiguous_layout(): - from tensorrt_llm._torch.modules.fla.cached_replay import ( - commit_gdn_cached_replay_history_layers, - ) - - batch_size = 16 - mgr = _build_v2_hybrid_with_mamba_layer( - max_batch_size=batch_size, - num_mamba_layers=3, - spec_config=MTPDecodingConfig(max_draft_len=3), - use_replay_state_update=True, - conv_state_layout="q_k_v", - ) - try: - torch.manual_seed(1234) - for state in mgr.all_ssm_states: - state.normal_() - mgr.old_x.normal_() - mgr.old_B.normal_() - mgr.old_dt.uniform_(-0.2, -0.01) - - positions = torch.arange(batch_size, device="cuda", dtype=torch.int32) - replay_work_items = torch.stack( - ( - positions, - positions, - torch.full_like(positions, mgr.replay_history_size), - torch.zeros_like(positions), - ), - dim=1, - ) - n_writes = torch.tensor([batch_size], device="cuda", dtype=torch.int32) - expected = torch.stack(mgr.all_ssm_states) - commit_gdn_cached_replay_history_layers( - ssm_states=expected, - old_u=mgr.old_x, - old_k=mgr.old_B, - old_G=mgr.old_dt, - replay_work_items=replay_work_items, - n_writes=n_writes, - history_size=mgr.replay_history_size, - ) - - mgr._commit_gdn_cached_replay_history_layers( - SimpleNamespace( - mamba_metadata=SimpleNamespace( - replay_num_decodes=batch_size, - replay_work_items=replay_work_items, - replay_n_writes=n_writes, - ) - ), - batch_size, + assert mgr.use_gdn_mtp_replay is use_gdn_mtp_replay + assert mgr.use_gdn_cached_replay_all_layer_commit is not use_gdn_mtp_replay + if not use_gdn_mtp_replay: + assert mgr.ssm_state_dtype == ssm_dtype + assert mgr.cache_buf_idx is not None + assert mgr.old_dA_cumsum is not None + assert mgr._gdn_cached_replay_state_strides is not None + assert mgr.intermediate_conv_states.shape == (3, 16, 4, *mgr.conv_state_shape) + layer_cache = mgr.mamba_layer_cache(0) + _assert_replay_layer_cache_uses_history_size(layer_cache, mgr.replay_history_size) + return + assert mgr.cache_buf_idx is None + assert mgr.old_dA_cumsum is None + assert mgr._gdn_mtp_replay_state_view is not None + assert mgr._gdn_mtp_replay_conv_state_view is not None + assert mgr.intermediate_conv_states.shape == ( + 3, + 16, + 4, + mgr.conv_state_shape[0], ) - torch.testing.assert_close(torch.stack(mgr.all_ssm_states), expected, rtol=0, atol=0) + replay_metadata = mgr.get_replay_state_update_metadata() + assert replay_metadata is not None + assert replay_metadata.cache_buf_idx is None + expected_capacity = replay_metadata.replay_history_size + replay_metadata.replay_step_width + layer_cache = mgr.mamba_layer_cache(0) + assert layer_cache.old_x.shape[1] == expected_capacity + assert layer_cache.old_dA_cumsum is None + assert layer_cache.cache_buf_idx is None finally: mgr.shutdown() -def test_v2_gdn_replay_commits_before_advancing_bookkeeping(monkeypatch): +def test_v2_gdn_replay_uses_combined_commit_only(monkeypatch): mgr = object.__new__(MambaHybridCacheManagerV2) - batch_size = 16 mgr.local_num_mamba_layers = 1 mgr._use_replay_state_update = True - mgr._use_gdn_cached_replay_all_layer_commit = True - mgr.replay_step_width = 4 - mgr.replay_history_size = MIN_REPLAY_HISTORY_SIZE - mgr.prev_num_accepted_tokens = torch.zeros(batch_size, dtype=torch.int32) - mgr.cache_buf_idx = torch.zeros(batch_size, dtype=torch.int32) - mgr.intermediate_state_indices = torch.arange(batch_size, dtype=torch.int32) - mgr._dummy_request_mask_buffer = DeviceCopyBuffer(torch.zeros(batch_size, dtype=torch.bool)) - mgr.all_conv_states = [torch.empty(0)] - mgr._stacked_conv_states = None - mgr.intermediate_conv_states = torch.empty(0) + mgr._use_gdn_mtp_replay = True + mgr.intermediate_state_indices = torch.arange(2, dtype=torch.int32) + mgr._dummy_request_mask_buffer = DeviceCopyBuffer(torch.tensor([False, True, False])) - events = [] - mgr._commit_gdn_cached_replay_history_layers = lambda *_args, **_kwargs: events.append("commit") - monkeypatch.setattr( - "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager._advance_replay_state", - lambda *_args, **_kwargs: events.append("advance"), - ) + captured = {} + + def combined_commit(_metadata, num_decodes, accepted_tokens, is_dummy): + captured["num_decodes"] = num_decodes + captured["accepted_tokens"] = accepted_tokens.clone() + captured["is_dummy"] = is_dummy.clone() + + mgr._commit_gdn_mtp_replay = combined_commit monkeypatch.setattr( "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager._promote_mamba_state_triton", - lambda *_args, **_kwargs: None, + lambda *_args, **_kwargs: pytest.fail("separate promotion must be skipped"), ) mgr.update_mamba_states( - SimpleNamespace(num_seqs=batch_size, num_contexts=0), - torch.full((batch_size,), 3, dtype=torch.int32), - state_indices=torch.arange(batch_size, dtype=torch.int32), + SimpleNamespace(num_seqs=3, num_contexts=1), + torch.tensor([1, 2, 4], dtype=torch.int32), + state_indices=torch.tensor([7, 8, 9], dtype=torch.int32), ) - assert events == ["commit", "advance"] + assert captured["num_decodes"] == 2 + assert torch.equal(captured["accepted_tokens"], torch.tensor([2, 4], dtype=torch.int32)) + assert torch.equal(captured["is_dummy"], torch.tensor([True, False])) def test_v2_hybrid_replay_update_skips_dummy_and_padding_rows(monkeypatch): @@ -5401,6 +5367,7 @@ def test_v2_hybrid_replay_update_skips_dummy_and_padding_rows(monkeypatch): mgr._dummy_request_mask_buffer = DeviceCopyBuffer(torch.zeros(4, dtype=torch.bool)) mgr._dummy_request_mask_host = torch.zeros(4, dtype=torch.bool) mgr._use_replay_state_update = True + mgr._use_gdn_mtp_replay = False mgr.replay_step_width = 5 mgr.replay_history_size = 16 mgr.prev_num_accepted_tokens = torch.full((4,), 13, dtype=torch.int32) diff --git a/tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py b/tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py index 8113e5d8c582..3a3b0152f1c4 100644 --- a/tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py +++ b/tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py @@ -43,6 +43,63 @@ def test_gdn_replay_env_defaults_on_and_requires_one_when_set(monkeypatch, value assert is_gdn_replay_enabled() is expected +def test_gdn_replay_dispatch_skips_standalone_conv(monkeypatch): + """GDN replay consumes raw QKV and bypasses the standalone Conv4 path.""" + import tensorrt_llm._torch.modules.mamba.gdn_mixer as gdn_mixer + + layer = gdn_mixer.Qwen3NextGatedDeltaNet.__new__(gdn_mixer.Qwen3NextGatedDeltaNet) + torch.nn.Module.__init__(layer) + layer.num_v_heads_per_tp = 2 + layer.head_v_dim = 4 + + batch = 2 + replay_width = 4 + channels = 12 + mixed_qkv = torch.randn(batch * replay_width, channels) + a = torch.randn(batch * replay_width, layer.num_v_heads_per_tp) + b = torch.randn_like(a) + expected = torch.randn(batch, replay_width, layer.num_v_heads_per_tp, layer.head_v_dim) + captured = {} + + def replay(**kwargs): + captured.update(kwargs) + return expected + + layer._replay_verify = replay + monkeypatch.setattr( + gdn_mixer, + "causal_conv1d_update_triton", + lambda *args, **kwargs: pytest.fail("standalone Conv4 must be skipped"), + ) + + result = layer.forward_decode( + conv_states=torch.empty(1), + ssm_states=torch.empty(1), + query_start_loc_long=torch.empty(0, dtype=torch.int64), + spec_metadata=SimpleNamespace(runtime_draft_len=replay_width - 1), + intermediate_conv_states=torch.empty(batch, replay_width, channels), + is_target_verify=True, + mixed_qkv=mixed_qkv, + a=a, + b=b, + cache_indices=torch.arange(batch, dtype=torch.int32), + num_decodes=batch, + intermediate_state_indices=torch.arange(batch, dtype=torch.int32), + use_replay=True, + replay_metadata=object(), + layer_cache=object(), + is_dummy=torch.zeros(batch, dtype=torch.bool), + ) + + assert captured["raw_x"] is mixed_qkv + assert captured["packed_ba"] is b + assert captured["candidate_x"].shape == (batch, replay_width, channels) + torch.testing.assert_close( + result, + expected.view(1, batch * replay_width, layer.num_v_heads_per_tp, layer.head_v_dim), + ) + + # ---- Reference implementations ---- diff --git a/tests/unittest/_torch/modules/mamba/test_gdn_mtp_replay.py b/tests/unittest/_torch/modules/mamba/test_gdn_mtp_replay.py new file mode 100644 index 000000000000..54946de7fd11 --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/test_gdn_mtp_replay.py @@ -0,0 +1,417 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE +from tensorrt_llm._utils import is_sm_100f + +DIM = 128 +WINDOW = 16 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or not is_sm_100f() or not IS_CUTLASS_DSL_AVAILABLE, + reason="GDN MTP replay requires CUTLASS DSL on datacenter Blackwell", +) + + +def _randn(shape, scale=1.0, shift=0.0): + return torch.randn(shape, device="cuda", dtype=torch.float32) * scale + shift + + +def _conv4_silu(raw_x, conv_weight, conv_state, slots): + old = conv_state.index_select(0, slots).float() + sequence = torch.cat((old, raw_x.float().transpose(1, 2)), dim=-1) + terms = [ + sequence[:, :, tap : tap + raw_x.shape[1]] * conv_weight[:, tap].float()[None, :, None] + for tap in range(4) + ] + convolved = torch.stack(terms).sum(dim=0) + return (convolved * torch.sigmoid(convolved)).transpose(1, 2).to(torch.bfloat16) + + +@torch.no_grad() +def _replay_reference( + raw_x, + conv_weight, + conv_state, + packed_ba, + checkpoint, + state_indices, + is_dummy, + history_u, + history_k, + history_G, + history_len, + A_log, + dt_bias, +): + batch, width, _ = raw_x.shape + key_heads = history_k.shape[2] + value_heads = history_u.shape[2] + heads_per_key = value_heads // key_heads + slots = state_indices.long() + + packed_qkv = _conv4_silu(raw_x, conv_weight, conv_state, slots) + key_width = key_heads * DIM + value_width = value_heads * DIM + query = packed_qkv[..., :key_width].view(batch, width, key_heads, DIM) + key = packed_qkv[..., key_width : 2 * key_width].view_as(query) + value = packed_qkv[..., 2 * key_width : 2 * key_width + value_width].view( + batch, width, value_heads, DIM + ) + packed_ba = packed_ba.view(batch, width, 2 * value_heads) + beta_input = packed_ba[..., :value_heads] + decay_input = packed_ba[..., value_heads:] + + query_f32 = query.float() + key_f32 = key.float() + query_norm = ( + query_f32 + * (DIM**-0.5 / torch.sqrt(torch.sum(query_f32 * query_f32, dim=-1, keepdim=True) + 1e-6)) + ).to(torch.bfloat16) + key_norm = (key_f32 / torch.sqrt(torch.sum(key_f32 * key_f32, dim=-1, keepdim=True) + 1e-6)).to( + torch.bfloat16 + ) + + gate_input = decay_input.float() + dt_bias.float()[None, None, :] + softplus = torch.where( + gate_input <= 20.0, + torch.log1p(torch.exp(gate_input)), + gate_input, + ) + gate = -torch.exp(A_log.float())[None, None, :] * softplus + beta = torch.sigmoid(beta_input.float()) + + lengths = history_len.index_select(0, slots).long() + selected_u = history_u[:, :WINDOW].index_select(0, slots).float() + selected_k = history_k[:, :WINDOW].index_select(0, slots).float() + selected_G = history_G[:, :, :WINDOW].index_select(0, slots).float() + key_for_value_head = torch.arange(value_heads, device=raw_x.device) // heads_per_key + selected_k = selected_k[:, :, key_for_value_head, :] + + positions = torch.arange(WINDOW, device=raw_x.device)[None, None, :] + active = positions < lengths[:, None, None] + previous_G = torch.zeros(batch, value_heads, device=raw_x.device, dtype=torch.float32) + nonempty = lengths > 0 + if bool(nonempty.any()): + rows = torch.where(nonempty)[0] + previous_G[rows] = selected_G[rows, :, lengths[rows] - 1] + history_decay = torch.where( + active, + torch.exp(previous_G[:, :, None] - selected_G), + 0.0, + ) + state = checkpoint.index_select(0, slots).float() + state *= torch.exp(previous_G)[:, :, None, None] + state += torch.einsum("btiv,btik,bit->bivk", selected_u, selected_k, history_decay) + + output = torch.empty( + batch, + width, + value_heads, + DIM, + device=raw_x.device, + dtype=torch.bfloat16, + ) + current_u = torch.empty_like(output) + current_G = torch.empty(batch, width, value_heads, device=raw_x.device, dtype=torch.float32) + query_hv = query_norm[:, :, key_for_value_head].float() + key_hv = key_norm[:, :, key_for_value_head].float() + running_G = previous_G + for token in range(width): + running_G = running_G + gate[:, token] + current_G[:, token] = running_G + state *= torch.exp(gate[:, token])[:, :, None, None] + state_key = torch.einsum("bhvk,bhk->bhv", state, key_hv[:, token]) + update = beta[:, token, :, None] * (value[:, token].float() - state_key) + state += update[:, :, :, None] * key_hv[:, token, :, None, :] + output[:, token] = torch.einsum("bhvk,bhk->bhv", state, query_hv[:, token]).to( + torch.bfloat16 + ) + current_u[:, token] = update.to(torch.bfloat16) + + writable_rows = torch.where(~is_dummy)[0] + writable_slots = slots[writable_rows] + writable_lengths = lengths[writable_rows] + for token in range(width): + write_positions = writable_lengths + token + history_u[writable_slots, write_positions] = current_u[writable_rows, token] + history_k[writable_slots, write_positions] = key_norm[writable_rows, token] + history_G[writable_slots, :, write_positions] = current_G[writable_rows, token] + return output, raw_x.contiguous() + + +def test_gdn_mtp_replay_matches_reference(): + from tensorrt_llm._torch.custom_ops.cute_dsl_gdn_mtp_replay import gdn_mtp_replay + + torch.manual_seed(7) + batch, slots, width, configured_width = 2, 5, 4, 6 + key_heads, value_heads = 1, 2 + channels = (2 * key_heads + value_heads) * DIM + qkvz_width = channels + value_heads * DIM + + projected_qkvz = _randn((batch * width, qkvz_width), 0.09).to(torch.bfloat16) + raw_x_2d = projected_qkvz[:, :channels] + raw_x = raw_x_2d.as_strided( + (batch, width, channels), + (width * raw_x_2d.stride(0), raw_x_2d.stride(0), 1), + ) + packed_ba = _randn((batch * width, 2 * value_heads), 0.5, -0.1).to(torch.bfloat16) + packed_beta = packed_ba[:, :value_heads] + conv_weight = _randn((channels, 4), 0.18).to(torch.bfloat16) + conv_state = _randn((slots, channels, 3), 0.09).to(torch.bfloat16) + checkpoint = _randn((slots, value_heads, DIM, DIM), 0.012).to(torch.bfloat16) + configured_capacity = WINDOW + configured_width + capacity = WINDOW + width + history_u_storage = _randn((slots, configured_capacity, value_heads, DIM), 0.012).to( + torch.bfloat16 + ) + history_u = history_u_storage[:, :capacity] + raw_history_k = _randn((slots, configured_capacity, key_heads, DIM), 0.09) + history_k_storage = ( + raw_history_k / torch.sqrt(raw_history_k.square().sum(dim=-1, keepdim=True) + 1e-6) + ).to(torch.bfloat16) + history_k = history_k_storage[:, :capacity] + history_G_storage = torch.cumsum( + -(_randn((slots, value_heads, configured_capacity), 0.025).abs() + 0.002), + dim=-1, + ) + history_G = history_G_storage[:, :, :capacity] + state_indices = torch.tensor([3, 1], device="cuda", dtype=torch.int32) + history_len = torch.tensor([2, 7, 4, 11, 1], device="cuda", dtype=torch.int32) + is_dummy = torch.tensor([False, True], device="cuda") + A_log = _randn((value_heads,), 0.35, -0.25).clamp_(-1.0, 0.5) + dt_bias = _randn((value_heads,), 0.65, -2.5).clamp_(-5.0, 0.5) + + reference_u_storage = history_u_storage.clone() + reference_k_storage = history_k_storage.clone() + reference_G_storage = history_G_storage.clone() + reference_u = reference_u_storage[:, :capacity] + reference_k = reference_k_storage[:, :capacity] + reference_G = reference_G_storage[:, :, :capacity] + reference_output, reference_candidates = _replay_reference( + raw_x, + conv_weight, + conv_state, + packed_ba, + checkpoint, + state_indices, + is_dummy, + reference_u, + reference_k, + reference_G, + history_len, + A_log, + dt_bias, + ) + output = torch.empty_like(reference_output) + candidate_storage = torch.full( + (batch, configured_width, channels), + torch.nan, + device="cuda", + dtype=torch.bfloat16, + ) + candidate_x = candidate_storage[:, :width] + checkpoint_before = checkpoint.clone() + conv_before = conv_state.clone() + + gdn_mtp_replay( + raw_x, + conv_weight, + conv_state, + packed_beta, + checkpoint, + state_indices, + is_dummy, + history_u, + history_k, + history_G, + history_len, + A_log, + dt_bias, + output, + candidate_x, + ) + + live = ~is_dummy + torch.testing.assert_close(output[live], reference_output[live], rtol=2e-2, atol=5e-3) + torch.testing.assert_close(candidate_x[live], reference_candidates[live], rtol=0, atol=0) + torch.testing.assert_close(history_u_storage, reference_u_storage, rtol=2e-2, atol=5e-3) + torch.testing.assert_close(history_k_storage, reference_k_storage, rtol=1e-2, atol=1.5e-3) + torch.testing.assert_close(history_G_storage, reference_G_storage, rtol=1e-5, atol=1e-5) + assert torch.isnan(candidate_storage[:, width:]).all() + torch.testing.assert_close(checkpoint, checkpoint_before, rtol=0, atol=0) + torch.testing.assert_close(conv_state, conv_before, rtol=0, atol=0) + + +@torch.no_grad() +def _commit_reference( + state_views, + conv_views, + candidate_x, + history_u, + history_k, + history_G, + history_len, + replay_work_items, + accepted_tokens, + is_dummy, +): + value_heads = history_u.shape[-2] + key_heads = history_k.shape[-2] + key_for_value_head = torch.arange(value_heads, device=history_u.device) // ( + value_heads // key_heads + ) + for work_item in replay_work_items: + position = int(work_item[0].item()) + slot = int(work_item[1].item()) + start = int(work_item[2].item()) + if bool(is_dummy[position].item()): + continue + + accepted = int(accepted_tokens[position].item()) + old_conv = torch.stack([view[slot] for view in conv_views]) + accepted_x = candidate_x[:, position, :accepted].transpose(1, 2) + new_conv = torch.cat((old_conv, accepted_x), dim=-1)[..., -3:] + for layer, view in enumerate(conv_views): + view[slot].copy_(new_conv[layer]) + + total = start + accepted + if total <= WINDOW: + history_len[slot] = total + continue + + tail = total - WINDOW + end_G = history_G[:, slot, :, WINDOW - 1].float() + first_G = history_G[:, slot, :, :WINDOW].float() + first_u = history_u[:, slot, :WINDOW].float() + first_k = history_k[:, slot, :WINDOW].float()[:, :, key_for_value_head] + update = torch.einsum( + "ltiv,ltik,lit->livk", + first_u, + first_k, + torch.exp(end_G[:, :, None] - first_G), + ) + checkpoint = torch.stack([view[slot] for view in state_views]).float() + committed = checkpoint * torch.exp(end_G)[:, :, None, None] + update + for layer, view in enumerate(state_views): + view[slot].copy_(committed[layer].to(torch.bfloat16)) + + tail_u = history_u[:, slot, WINDOW : WINDOW + tail].clone() + tail_k = history_k[:, slot, WINDOW : WINDOW + tail].clone() + tail_G = history_G[:, slot, :, WINDOW : WINDOW + tail].clone() + history_u[:, slot, :tail] = tail_u + history_k[:, slot, :tail] = tail_k + history_G[:, slot, :, :tail] = tail_G - end_G[:, :, None] + history_len[slot] = tail + + +def _make_state_storage(indirect, layers, slots, value_heads, channels): + if indirect: + states = [ + _randn((slots, value_heads, DIM, DIM), 0.012).to(torch.bfloat16) for _ in range(layers) + ] + conv = [_randn((slots, channels, 3), 0.09).to(torch.bfloat16) for _ in range(layers)] + state_descriptors = torch.tensor( + [(state.data_ptr(), state.stride(0)) for state in states], + device="cuda", + dtype=torch.int64, + ) + conv_descriptors = torch.tensor( + [(state.data_ptr(), state.stride(0)) for state in conv], + device="cuda", + dtype=torch.int64, + ) + return states, conv, states[0], conv[0], state_descriptors, conv_descriptors + + state_storage = _randn((slots, layers, value_heads, DIM, DIM), 0.012).to(torch.bfloat16) + conv_storage = _randn((slots, layers, channels, 3), 0.09).to(torch.bfloat16) + states = [state_storage[:, layer] for layer in range(layers)] + conv = [conv_storage[:, layer] for layer in range(layers)] + return states, conv, state_storage, conv_storage, None, None + + +@pytest.mark.parametrize("indirect", [False, True], ids=["affine", "indirect"]) +def test_gdn_mtp_commit_matches_reference(indirect): + from tensorrt_llm._torch.custom_ops.cute_dsl_gdn_mtp_commit import gdn_mtp_commit + + torch.manual_seed(11) + layers, slots, candidate_capacity, width = 2, 5, 5, 4 + key_heads, value_heads = 1, 2 + channels = (2 * key_heads + value_heads) * DIM + capacity = WINDOW + width + + states, conv, state_arg, conv_arg, state_desc, conv_desc = _make_state_storage( + indirect, layers, slots, value_heads, channels + ) + reference_states = [state.clone() for state in states] + reference_conv = [state.clone() for state in conv] + candidate_x = _randn((layers, candidate_capacity, width, channels), 0.09).to(torch.bfloat16) + history_u = _randn((layers, slots, capacity, value_heads, DIM), 0.012).to(torch.bfloat16) + raw_history_k = _randn((layers, slots, capacity, key_heads, DIM), 0.09) + history_k = ( + raw_history_k / torch.sqrt(raw_history_k.square().sum(dim=-1, keepdim=True) + 1e-6) + ).to(torch.bfloat16) + history_G = torch.cumsum( + -(_randn((layers, slots, value_heads, capacity), 0.025).abs() + 0.002), + dim=-1, + ) + accepted_tokens = torch.tensor([3, 2, 1], device="cuda", dtype=torch.int32) + is_dummy = torch.tensor([False, False, True], device="cuda") + starts = torch.tensor([15, 10, 16], device="cuda", dtype=torch.int32) + positions = torch.tensor([2, 0, 1], device="cuda", dtype=torch.int32) + slot_for_position = torch.tensor([1, 3, 4], device="cuda", dtype=torch.int32) + selected_slots = slot_for_position[positions.long()] + replay_work_items = torch.stack( + ( + positions, + selected_slots, + starts[positions.long()], + torch.zeros_like(positions), + ), + dim=1, + ).contiguous() + history_len = torch.tensor([4, 15, 2, 10, 16], device="cuda", dtype=torch.int32) + + reference_u = history_u.clone() + reference_k = history_k.clone() + reference_G = history_G.clone() + reference_len = history_len.clone() + _commit_reference( + reference_states, + reference_conv, + candidate_x, + reference_u, + reference_k, + reference_G, + reference_len, + replay_work_items, + accepted_tokens, + is_dummy, + ) + + gdn_mtp_commit( + state_arg, + conv_arg, + candidate_x, + history_u, + history_k, + history_G, + history_len, + replay_work_items, + accepted_tokens, + is_dummy, + ssm_state_descriptors=state_desc, + conv_state_descriptors=conv_desc, + ) + + for actual, reference in zip(states, reference_states): + torch.testing.assert_close(actual, reference, rtol=2e-2, atol=8e-3) + for actual, reference in zip(conv, reference_conv): + torch.testing.assert_close(actual, reference, rtol=0, atol=0) + torch.testing.assert_close(history_u, reference_u, rtol=0, atol=0) + torch.testing.assert_close(history_k, reference_k, rtol=0, atol=0) + torch.testing.assert_close(history_G, reference_G, rtol=1e-6, atol=1e-6) + torch.testing.assert_close(history_len, reference_len, rtol=0, atol=0) diff --git a/tests/unittest/_torch/modules/mamba/test_gdn_replay_recurrent.py b/tests/unittest/_torch/modules/mamba/test_gdn_replay_recurrent.py deleted file mode 100644 index e9c8a7dfa2bb..000000000000 --- a/tests/unittest/_torch/modules/mamba/test_gdn_replay_recurrent.py +++ /dev/null @@ -1,669 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# 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 -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# 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. -"""Regression tests for GDN normalization and MTP replay verification. - -Checks recurrent Q/K normalization against the prefill epsilon convention. -Simulates several target-verify iterations with randomized acceptance and -checks cached replay plus PNAT/double-buffer bookkeeping mirroring -update_mamba_states against: - 1. the legacy path (fused_recurrent_gated_delta_rule_update with - intermediate_states_buffer + accepted-state copy), and - 2. a sequential fp32 reference. -""" - -from typing import Literal - -import pytest -import torch - -from tensorrt_llm._torch.modules.fla.cached_replay import ( - fused_recurrent_gated_delta_rule_cached_replay_update, -) -from tensorrt_llm._torch.modules.fla.fused_recurrent import fused_recurrent_gated_delta_rule_update - - -def _seq_ref_step(S, q, k, v, g, beta, scale): - """Sequential fp32 reference over T steps for one request. - - S: [HV, V, K] fp32 (modified out-of-place). Returns (o [T, HV, V], - per-step states list). - """ - T = q.shape[0] - HV = S.shape[0] - H = k.shape[1] - outs, states = [], [] - S = S.clone() - for t in range(T): - o_t = torch.empty(*S.shape[:1], S.shape[1], dtype=torch.float32, device=S.device) - for hv in range(HV): - h = hv // (HV // H) - qt = q[t, h].float() - kt = k[t, h].float() - vt = v[t, hv].float() - qt = qt / (qt.square().sum() + 1e-6).sqrt() * scale - kt = kt / (kt.square().sum() + 1e-6).sqrt() - St = S[hv] * torch.exp(g[t, hv].float()) - vt = (vt - (St * kt[None, :]).sum(-1)) * beta[t, hv].float() - St = St + kt[None, :] * vt[:, None] - o_t[hv] = (St * qt[None, :]).sum(-1) - S[hv] = St - outs.append(o_t) - states.append(S.clone()) - return torch.stack(outs), states - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("kernel", ["fused_gating", "recurrent", "cached_replay"]) -@pytest.mark.parametrize("amplitude", [0.0, 1e-4, 0.1]) -@pytest.mark.parametrize("nonzero_state", [False, True]) -@torch.inference_mode() -def test_gdn_recurrent_matches_l2_reference( - kernel: Literal["fused_gating", "recurrent", "cached_replay"], - amplitude: float, - nonzero_state: bool, -) -> None: - from tensorrt_llm._torch.modules.fla.fused_sigmoid_gating_recurrent import ( - fused_sigmoid_gating_delta_rule_update, - ) - - torch.manual_seed(2026) - batch, heads, value_heads, dim, slots = 2, 2, 4, 64, 4 - scale = dim**-0.5 - q = (torch.randn(batch, 1, heads, dim, device="cuda") * amplitude).bfloat16() - k = (torch.randn_like(q.float()) * amplitude).bfloat16() - v = torch.randn(batch, 1, value_heads, dim, device="cuda", dtype=torch.bfloat16) - state = torch.randn(slots, value_heads, dim, dim, device="cuda") * 0.01 - if not nonzero_state: - state.zero_() - original = state.clone() - indices = torch.tensor([2, 0], device="cuda", dtype=torch.int32) - a = torch.zeros(batch, 1, value_heads, device="cuda") - b = torch.zeros_like(a) - g = torch.full_like(a, -torch.log(torch.tensor(2.0)).item()) - beta = torch.full_like(a, 0.5) - - # Independent FP64 recurrence: epsilon is added to the squared norm, - # matching l2norm_fwd and fused_gdn_post_conv on the prefill path. - qr, kr = q.double(), k.double() - qr = qr / (qr.square().sum(-1, keepdim=True) + 1e-6).sqrt() - kr = kr / (kr.square().sum(-1, keepdim=True) + 1e-6).sqrt() - qr = qr[:, 0].repeat_interleave(value_heads // heads, dim=1) * scale - kr = kr[:, 0].repeat_interleave(value_heads // heads, dim=1) - sr = original[indices.long()].double() * 0.5 - residual = (v[:, 0].double() - (sr * kr[:, :, None, :]).sum(-1)) * 0.5 - sr += residual[:, :, :, None] * kr[:, :, None, :] - expected = (sr * qr[:, :, None, :]).sum(-1).unsqueeze(1) - - if kernel == "fused_gating": - actual = fused_sigmoid_gating_delta_rule_update( - A_log=torch.zeros(value_heads, device="cuda"), - a=a, - dt_bias=torch.zeros(value_heads, device="cuda"), - softplus_beta=1.0, - softplus_threshold=20.0, - q=q, - k=k, - v=v, - b=b, - initial_state_source=state, - initial_state_indices=indices, - use_qk_l2norm_in_kernel=True, - ) - elif kernel == "recurrent": - actual = fused_recurrent_gated_delta_rule_update( - q=q, - k=k, - v=v, - g=g, - beta=beta, - initial_state_source=state, - initial_state_indices=indices, - use_qk_l2norm_in_kernel=True, - ) - else: - history = 16 - actual = fused_recurrent_gated_delta_rule_cached_replay_update( - q, - k, - v, - g, - beta, - state, - indices, - torch.zeros(slots, 2, history, value_heads, dim, device="cuda", dtype=v.dtype), - torch.zeros(slots, 2, history, heads, dim, device="cuda", dtype=k.dtype), - torch.zeros(slots, 2, value_heads, history, device="cuda"), - torch.zeros(slots, 2, value_heads, history, device="cuda"), - torch.zeros(slots, device="cuda", dtype=torch.int32), - torch.zeros(slots, device="cuda", dtype=torch.int32), - history_size=history, - scale=scale, - use_qk_l2norm_in_kernel=True, - ) - - # BF16 outputs require room for final rounding. Cached replay additionally - # rounds normalized Q/K to BF16 before the dot products. - rtol = 0.015 if kernel == "cached_replay" else 0.005 - torch.testing.assert_close(actual.double(), expected, rtol=rtol, atol=3e-5) - if kernel != "cached_replay": - torch.testing.assert_close(state[indices.long()].double(), sr, rtol=3e-5, atol=2e-7) - torch.testing.assert_close(state[[1, 3]], original[[1, 3]], rtol=0, atol=0) - - -@pytest.mark.parametrize("fused_gating", [False, True], ids=["pre_gated", "fused_gating"]) -@pytest.mark.parametrize( - "pool_dtype", [torch.bfloat16, torch.float32], ids=["bf16_pool", "fp32_pool"] -) -@pytest.mark.parametrize( - "T,HIST,iters", - [(2, 14, 16), (4, 12, 10), (5, 11, 12), (4, 16, 10)], - ids=["T2_H14", "T4_H12", "T5_H11", "T4_H16"], -) -@pytest.mark.parametrize( - "H,HV,K,V", - [(4, 8, 128, 128), (2, 4, 64, 64), (4, 16, 128, 128)], - ids=["qwen3_like", "small", "qwen3_5_like"], -) -def test_gdn_replay_vs_legacy_and_ref(H, HV, K, V, T, HIST, iters, pool_dtype, fused_gating): - if not torch.cuda.is_available(): - pytest.skip("CUDA required") - torch.manual_seed(1234) - device = "cuda" - dtype = torch.bfloat16 - N = 4 - scale = K**-0.5 - slots = N + 2 # exercise non-trivial slot indices - - # Fused-gating mode: the replay kernel receives RAW a/b plus - # A_log/dt_bias and applies g = -exp(A_log)*softplus(a + dt_bias), - # beta = sigmoid(b) in-kernel; legacy + reference use the same gating - # computed on the host. - A_log_t = (torch.randn(HV, device=device, dtype=torch.float32) * 0.5) if fused_gating else None - dt_bias_t = ( - (torch.randn(HV, device=device, dtype=torch.float32) * 0.5) if fused_gating else None - ) - - pool_init = torch.randn(slots, HV, V, K, device=device, dtype=torch.float32) * 0.5 - # Exercise a pool-backed (non-contiguous, padded slot stride) layout: - # the kernel must commit through the real strides — a contiguous copy - # would silently discard checkpoint commits (Qwen3.5 looping root - # cause). Odd T uses the strided layout, even T stays dense. - strided_pool = T % 2 == 1 - # .clone() everywhere: for fp32 pools .to(pool_dtype)/.float() are no-ops - # returning the same tensor, which would alias pool/reference storage. - if strided_pool: - _pad = 1024 - _backing = torch.zeros(slots, HV * V * K + _pad, device=device, dtype=pool_dtype) - pool_replay = _backing[:, : HV * V * K].view(slots, HV, V, K) - pool_replay.copy_(pool_init.to(pool_dtype)) - assert not pool_replay.is_contiguous() - else: - pool_replay = pool_init.to(pool_dtype).clone() - pool_legacy = pool_init.to(pool_dtype).clone() - ref_S = pool_init.to(pool_dtype).float().clone() # follow quantized init - - # Replay buffers + bookkeeping (per-layer slice shapes) - old_v = torch.zeros(slots, 2, HIST, HV, V, device=device, dtype=dtype) - old_k = torch.zeros(slots, 2, HIST, H, K, device=device, dtype=dtype) - old_g = torch.zeros(slots, 2, HV, HIST, device=device, dtype=torch.float32) - old_beta = torch.zeros(slots, 2, HV, HIST, device=device, dtype=torch.float32) - buf_idx = torch.zeros(slots, dtype=torch.int32, device=device) - pnat = torch.zeros(slots, dtype=torch.int32, device=device) - # Non-identity slot mapping - state_indices = torch.arange(1, N + 1, dtype=torch.int32, device=device) - - # Legacy intermediate buffer, indexed by [0..N) - intermediate_ssm = torch.zeros(N, T, HV, V, K, device=device, dtype=pool_dtype) - arange_n = torch.arange(N, dtype=torch.int32, device=device) - - n_checkpoints = 0 - prev_ref_S = ref_S.clone() - for it in range(iters): - q = torch.randn(N, T, H, K, device=device, dtype=dtype) - k = torch.randn(N, T, H, K, device=device, dtype=dtype) - v = torch.randn(N, T, HV, V, device=device, dtype=dtype) * 0.5 - packed_qkv = None - if T == 4 and HIST == 16: - packed_qkv = torch.cat((q.flatten(2), k.flatten(2), v.flatten(2)), dim=-1).reshape( - N * T, -1 - ) - packed_qkv_3d = packed_qkv.view(N, T, -1) - key_width = H * K - q = packed_qkv_3d[..., :key_width].view(N, T, H, K) - k = packed_qkv_3d[..., key_width : 2 * key_width].view(N, T, H, K) - v = packed_qkv_3d[..., 2 * key_width :].view(N, T, HV, V) - assert packed_qkv.is_contiguous() - assert not q.is_contiguous() and not k.is_contiguous() and not v.is_contiguous() - if fused_gating: - packed_ba = torch.randn(N * T, 2 * HV, device=device, dtype=dtype) - packed_ba[:, HV:].mul_(0.5) - b_raw = packed_ba[:, :HV].view(N, T, HV) - a_raw = packed_ba[:, HV:].view(N, T, HV) - assert not a_raw.is_contiguous() and not b_raw.is_contiguous() - # Host reference of the in-kernel gating (fp32, same inputs). - g = -(A_log_t.exp() * torch.nn.functional.softplus(a_raw.float() + dt_bias_t)) - beta = torch.sigmoid(b_raw.float()) - else: - g = -torch.rand(N, T, HV, device=device, dtype=torch.float32) * 2.0 - beta = torch.rand(N, T, HV, device=device, dtype=torch.float32) - num_accepted = torch.randint(1, T + 1, (N,)) - - # --- replay path --- - replay_kwargs = {} - # Piggyback PDL coverage on the fused-gating arm: a PDL launch - # without a producer signal is legal (waits resolve at predecessor - # completion), so this exercises the gdc_wait code path. - replay_kwargs["launch_with_pdl"] = fused_gating - if packed_qkv is not None: - replay_kwargs["packed_qkv"] = packed_qkv - replay_output = torch.empty((N, T, HV, V), device=device, dtype=dtype) - replay_kwargs["output"] = replay_output - o_replay = fused_recurrent_gated_delta_rule_cached_replay_update( - q, - k, - v, - a_raw if fused_gating else g, - b_raw if fused_gating else beta, - pool_replay, - state_indices, - old_v, - old_k, - old_g, - old_beta, - buf_idx, - pnat, - history_size=HIST, - scale=scale, - use_qk_l2norm_in_kernel=True, - A_log=A_log_t, - dt_bias=dt_bias_t, - **replay_kwargs, - ) - assert o_replay is replay_output - - # --- legacy path --- - src = pool_legacy[state_indices] - o_legacy = fused_recurrent_gated_delta_rule_update( - q=q, - k=k, - v=v, - g=g.to(dtype), - beta=beta.to(dtype), - scale=scale, - initial_state_source=src, - initial_state_indices=arange_n, - use_qk_l2norm_in_kernel=True, - disable_state_update=True, - intermediate_states_buffer=intermediate_ssm, - cache_steps=T, - ) - - # --- reference outputs + acceptance --- - for n in range(N): - o_ref, states = _seq_ref_step( - ref_S[state_indices[n]], q[n], k[n], v[n], g[n], beta[n], scale - ) - torch.testing.assert_close(o_replay[n].float(), o_ref, atol=6e-3, rtol=2e-2) - torch.testing.assert_close(o_legacy[n].float(), o_ref, atol=6e-3, rtol=2e-2) - ref_S[state_indices[n]] = states[num_accepted[n] - 1] - - # replay vs legacy outputs directly (both bf16 outputs, tighter) - torch.testing.assert_close(o_replay.float(), o_legacy.float(), atol=4e-3, rtol=2e-2) - - # --- legacy acceptance: copy accepted state --- - pool_legacy[state_indices] = intermediate_ssm[ - arange_n.long(), (num_accepted - 1).to(device).long() - ] - - # --- replay bookkeeping (mirrors update_mamba_states) --- - pnat_h = pnat.cpu() - buf_h = buf_idx.cpu() - for n in range(N): - slot = state_indices[n].item() - if (pnat_h[slot] + T) > HIST: - # Kernel committed the pre-iteration state this launch - # (prev_ref_S = reference state before this iteration's - # acceptance advance). - err = (pool_replay[slot].float() - prev_ref_S[slot]).abs().max().item() - assert err < 6e-3, f"checkpoint state err {err}" - n_checkpoints += 1 - pnat_h[slot] = num_accepted[n] - buf_h[slot] = 1 - buf_h[slot] - else: - pnat_h[slot] = pnat_h[slot] + num_accepted[n] - pnat.copy_(pnat_h.to(device)) - buf_idx.copy_(buf_h.to(device)) - prev_ref_S = ref_S.clone() - - # Ensure the test exercised checkpoint commits and buffer flips. - assert n_checkpoints >= N, f"only {n_checkpoints} checkpoints hit; increase iters" - - -def test_gdn_replay_negative_slot_skipped(): - """Padding requests (slot < 0) must not touch pool or history.""" - if not torch.cuda.is_available(): - pytest.skip("CUDA required") - torch.manual_seed(0) - device = "cuda" - dtype = torch.bfloat16 - N, T, H, HV, K, V, HIST = 2, 4, 2, 4, 64, 64, 16 - - pool = torch.randn(N, HV, V, K, device=device, dtype=dtype) - pool_orig = pool.clone() - old_v = torch.zeros(N, 2, HIST, HV, V, device=device, dtype=dtype) - old_k = torch.zeros(N, 2, HIST, H, K, device=device, dtype=dtype) - old_g = torch.zeros(N, 2, HV, HIST, device=device, dtype=torch.float32) - old_beta = torch.zeros(N, 2, HV, HIST, device=device, dtype=torch.float32) - buf_idx = torch.zeros(N, dtype=torch.int32, device=device) - pnat = torch.full((N,), HIST, dtype=torch.int32, device=device) # force write - state_indices = torch.tensor([-1, -1], dtype=torch.int32, device=device) - - q = torch.randn(N, T, H, K, device=device, dtype=dtype) - k = torch.randn(N, T, H, K, device=device, dtype=dtype) - v = torch.randn(N, T, HV, V, device=device, dtype=dtype) - g = -torch.rand(N, T, HV, device=device, dtype=torch.float32) - beta = torch.rand(N, T, HV, device=device, dtype=torch.float32) - - fused_recurrent_gated_delta_rule_cached_replay_update( - q, - k, - v, - g, - beta, - pool, - state_indices, - old_v, - old_k, - old_g, - old_beta, - buf_idx, - pnat, - history_size=HIST, - use_qk_l2norm_in_kernel=True, - ) - - torch.testing.assert_close(pool, pool_orig) - assert old_v.abs().max().item() == 0 - assert old_k.abs().max().item() == 0 - - -def _cached_replay_commit_reference(initial_states, old_u, old_k, old_G, work_items, n_writes): - reference = initial_states.clone() - num_layers, _, HV, _, _ = reference.shape - H = old_k.shape[-2] - rows = work_items[: int(n_writes.item())].cpu().tolist() - for _, slot, pnat, active_buffer in rows: - for layer in range(num_layers): - for hv in range(HV): - h = hv // (HV // H) - history_k = old_k[layer, slot, active_buffer, :pnat, h] - history_u = old_u[layer, slot, active_buffer, :pnat, hv] - history_G = old_G[layer, slot, active_buffer, hv, :pnat] - g_start = history_G[-1] - commit_decay = torch.exp(g_start - history_G) - scaled_u = history_u.float() * commit_decay[:, None] - scaled_u_hi = scaled_u.to(history_u.dtype) - scaled_u_lo = (scaled_u - scaled_u_hi.float()).to(history_u.dtype) - delta = history_k.float().T @ scaled_u_hi.float() - delta += history_k.float().T @ scaled_u_lo.float() - committed = reference[layer, slot, hv].float() * torch.exp(g_start) - reference[layer, slot, hv] = (committed + delta.T).to(reference.dtype) - return reference - - -@pytest.mark.parametrize("state_layout", ["contiguous", "affine", "indirect"]) -def test_gdn_cached_replay_all_layer_commit_matches_reference(state_layout): - """The all-layer commit must match its direct mathematical reference.""" - if not torch.cuda.is_available(): - pytest.skip("CUDA required") - - from tensorrt_llm._torch.modules.fla.cached_replay import ( - commit_gdn_cached_replay_history_layers, - ) - - torch.manual_seed(1234) - device = "cuda" - dtype = torch.bfloat16 - num_layers, N, num_slots = 3, 16, 18 - H, HV, K, V, HIST = 2, 8, 128, 128, 16 - initial_states = torch.randn(num_layers, num_slots, HV, V, K, device=device, dtype=dtype) - old_u = torch.randn(num_layers, num_slots, 2, HIST, HV, V, device=device, dtype=dtype) - old_k = torch.randn(num_layers, num_slots, 2, HIST, H, K, device=device, dtype=dtype) - old_G = -torch.cumsum( - torch.rand( - num_layers, - num_slots, - 2, - HV, - HIST, - device=device, - dtype=torch.float32, - ) - * 0.1, - dim=-1, - ) - - positions = torch.arange(N, device=device, dtype=torch.int32) - cache_slots = positions + 1 - pnat = positions.remainder(HIST - 1) + 1 - active_buffers = positions.remainder(2) - work_items = torch.stack((positions, cache_slots, pnat, active_buffers), dim=1) - n_writes = torch.tensor([N // 2], device=device, dtype=torch.int32) - - expected_states = _cached_replay_commit_reference( - initial_states, - old_u, - old_k, - old_G, - work_items, - n_writes, - ) - state_descriptors = None - state_num_layers = None - state_layer_stride = None - state_slot_stride = None - if state_layout == "indirect": - state_numel = HV * V * K - state_views = [] - for layer in range(num_layers): - backing = torch.empty( - num_slots, - state_numel + layer + 1, - device=device, - dtype=dtype, - ) - state = backing[:, :state_numel].view(num_slots, HV, V, K) - state.copy_(initial_states[layer]) - state_views.append(state) - state_descriptors = torch.tensor( - [(state.data_ptr(), state.stride(0)) for state in state_views], - device=device, - dtype=torch.int64, - ) - commit_states = state_views[0] - elif state_layout == "affine": - backing = torch.empty( - num_slots, - num_layers, - HV, - V, - K, - device=device, - dtype=dtype, - ) - state_views = [backing[:, layer] for layer in range(num_layers)] - for layer, state in enumerate(state_views): - state.copy_(initial_states[layer]) - commit_states = state_views[0] - state_num_layers = num_layers - state_layer_stride = backing.stride(1) - state_slot_stride = backing.stride(0) - else: - actual_states = initial_states.clone() - commit_states = actual_states - commit_gdn_cached_replay_history_layers( - ssm_states=commit_states, - ssm_state_descriptors=state_descriptors, - ssm_state_num_layers=state_num_layers, - ssm_state_layer_stride=state_layer_stride, - ssm_state_slot_stride=state_slot_stride, - old_u=old_u, - old_k=old_k, - old_G=old_G, - replay_work_items=work_items, - n_writes=n_writes, - history_size=HIST, - ) - if state_layout in ("affine", "indirect"): - actual_states = torch.stack(state_views) - - torch.testing.assert_close(actual_states.float(), expected_states.float(), rtol=2e-2, atol=2e-2) - - -@pytest.mark.parametrize("batch_size", [8, 16], ids=["small_fused", "large_all_layer"]) -def test_gdn_cached_replay_dispatch_cuda_graph_matches_eager(batch_size): - """Both sides of the BS16 dispatch must be safe under CUDA graphs.""" - if not torch.cuda.is_available(): - pytest.skip("CUDA required") - - from tensorrt_llm._torch.modules.fla.cached_replay import ( - CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, - commit_gdn_cached_replay_history_layers, - ) - - torch.manual_seed(1234) - device = "cuda" - dtype = torch.bfloat16 - num_layers, T, HIST = 2, 4, 16 - H, HV, K, V = 2, 4, 64, 64 - num_slots = batch_size + 1 - use_all_layer_commit = batch_size >= CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE - - initial_states = torch.randn(num_layers, num_slots, HV, V, K, device=device, dtype=dtype) - initial_u = torch.randn(num_layers, num_slots, 2, HIST, HV, V, device=device, dtype=dtype) - initial_k = torch.randn(num_layers, num_slots, 2, HIST, H, K, device=device, dtype=dtype) - initial_G = -torch.cumsum( - torch.rand( - num_layers, - num_slots, - 2, - HV, - HIST, - device=device, - dtype=torch.float32, - ), - dim=-1, - ) - old_beta = torch.empty_like(initial_G) - state_indices = torch.arange(1, batch_size + 1, device=device, dtype=torch.int32) - cache_buf_idx = torch.zeros(num_slots, device=device, dtype=torch.int32) - pnat = torch.full((num_slots,), HIST, device=device, dtype=torch.int32) - q = torch.randn(num_layers, batch_size, T, H, K, device=device, dtype=dtype) - k = torch.randn_like(q) - v = torch.randn(num_layers, batch_size, T, HV, V, device=device, dtype=dtype) - g = -torch.rand(num_layers, batch_size, T, HV, device=device, dtype=torch.float32) - beta = torch.rand_like(g) - - replay_work_items = None - n_writes = None - if use_all_layer_commit: - positions = torch.arange(batch_size, device=device, dtype=torch.int32) - replay_work_items = torch.stack( - ( - positions, - state_indices, - torch.full_like(positions, HIST), - torch.zeros_like(positions), - ), - dim=1, - ) - n_writes = torch.tensor([batch_size], device=device, dtype=torch.int32) - - def run_replay(states, old_u, old_k, old_G): - outputs = [] - for layer in range(num_layers): - outputs.append( - fused_recurrent_gated_delta_rule_cached_replay_update( - q[layer], - k[layer], - v[layer], - g[layer], - beta[layer], - states[layer], - state_indices, - old_u[layer], - old_k[layer], - old_G[layer], - old_beta[layer], - cache_buf_idx, - pnat, - history_size=HIST, - use_qk_l2norm_in_kernel=True, - replay_work_items=replay_work_items, - n_writes=n_writes, - use_all_layer_commit=use_all_layer_commit, - ) - ) - return outputs - - def run_all_layer_commit(states, old_u, old_k, old_G): - if use_all_layer_commit: - commit_gdn_cached_replay_history_layers( - ssm_states=states, - old_u=old_u, - old_k=old_k, - old_G=old_G, - replay_work_items=replay_work_items, - n_writes=n_writes, - history_size=HIST, - ) - - warmup_states = initial_states.clone() - warmup_u = initial_u.clone() - warmup_k = initial_k.clone() - warmup_G = initial_G.clone() - run_replay(warmup_states, warmup_u, warmup_k, warmup_G) - run_all_layer_commit(warmup_states, warmup_u, warmup_k, warmup_G) - torch.cuda.synchronize() - - eager_states = initial_states.clone() - eager_u = initial_u.clone() - eager_k = initial_k.clone() - eager_G = initial_G.clone() - eager_outputs = run_replay(eager_states, eager_u, eager_k, eager_G) - run_all_layer_commit(eager_states, eager_u, eager_k, eager_G) - - graph_states = initial_states.clone() - graph_u = initial_u.clone() - graph_k = initial_k.clone() - graph_G = initial_G.clone() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - graph_outputs = run_replay(graph_states, graph_u, graph_k, graph_G) - graph_states.copy_(initial_states) - graph_u.copy_(initial_u) - graph_k.copy_(initial_k) - graph_G.copy_(initial_G) - torch.cuda.synchronize() - graph.replay() - run_all_layer_commit(graph_states, graph_u, graph_k, graph_G) - torch.cuda.synchronize() - - for eager_output, graph_output in zip(eager_outputs, graph_outputs): - torch.testing.assert_close(graph_output, eager_output, rtol=0, atol=0) - torch.testing.assert_close(graph_states, eager_states, rtol=0, atol=0) - torch.testing.assert_close(graph_u, eager_u, rtol=0, atol=0) - torch.testing.assert_close(graph_k, eager_k, rtol=0, atol=0) - torch.testing.assert_close(graph_G, eager_G, rtol=0, atol=0) diff --git a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py index 04f06439febd..79d7acd16834 100644 --- a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py +++ b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py @@ -20,15 +20,13 @@ import torch from tensorrt_llm._device_copy import DeviceCopyBuffer -from tensorrt_llm._torch.modules.mamba import mamba2_metadata from tensorrt_llm._torch.modules.mamba.mamba2_metadata import ( REPLAY_WORK_CACHE_BUF_IDX, REPLAY_WORK_CACHE_SLOT, REPLAY_WORK_PNAT, REPLAY_WORK_POSITION_IN_DECODE_BATCH, Mamba2Metadata, - _build_replay_work_items_torch, - _build_replay_work_items_triton, + _build_gdn_replay_work_items, cu_seqlens_to_chunk_indices_offsets, cu_seqlens_to_chunk_indices_offsets_triton, ) @@ -45,16 +43,15 @@ class _GdnReplayCacheManager: use_replay_state_update = True - use_gdn_cached_replay_all_layer_commit = True + use_gdn_mtp_replay = True - def __init__(self, prev_num_accepted_tokens, cache_buf_idx): - self.prev_num_accepted_tokens = prev_num_accepted_tokens - self.cache_buf_idx = cache_buf_idx + def __init__(self, history_len): + self.history_len = history_len def get_replay_state_update_metadata(self): return ReplayStateUpdateMetadata( - prev_num_accepted_tokens=self.prev_num_accepted_tokens, - cache_buf_idx=self.cache_buf_idx, + prev_num_accepted_tokens=self.history_len, + cache_buf_idx=None, replay_step_width=6, replay_history_size=MIN_REPLAY_HISTORY_SIZE, ) @@ -84,25 +81,17 @@ def _decode_attn_metadata(kv_cache_manager): ) -def _torch_reference_work_items(state_indices, prev_num_accepted_tokens, cache_buf_idx): - """Run the production ATen path into fresh buffers. - - The Triton kernel's contract is that it reproduces this path exactly, so - this path -- not a third hand-written copy -- is what it is compared to. - """ - num_decodes = state_indices.numel() - work_items = torch.zeros(num_decodes, 4, dtype=torch.int32, device="cuda") - n_writes = torch.zeros(1, dtype=torch.int32, device="cuda") - _build_replay_work_items_torch( - state_indices, - prev_num_accepted_tokens, - cache_buf_idx, - work_items, - n_writes, - 6, - MIN_REPLAY_HISTORY_SIZE, +def _gdn_reference_work_items(state_indices, history_len): + positions = torch.arange(state_indices.numel(), dtype=torch.int32, device=state_indices.device) + return torch.stack( + ( + positions, + state_indices, + history_len[state_indices.to(torch.long)], + torch.zeros_like(positions), + ), + dim=1, ) - return work_items, n_writes @skip_no_cuda @@ -322,104 +311,52 @@ def get_replay_state_update_metadata(self): assert actual[0, REPLAY_WORK_PNAT] == 11 assert actual[0, REPLAY_WORK_CACHE_BUF_IDX] == 1 - @pytest.mark.parametrize("num_decodes", [16, 17, 40, 255, 256]) - def test_replay_work_items_triton_matches_torch(self, num_decodes): - """The two builders must agree everywhere the dispatch may pick either. - - num_decodes 17 and 255 are not powers of two, so the kernel runs with - masked-off lanes -- those must contribute nothing to the write-first - prefix sum. - """ + @pytest.mark.parametrize("num_decodes", [1, 16, 17, 40, 255, 256, 257]) + def test_gdn_replay_work_items_match_reference(self, num_decodes): + """The GDN builder preserves decode order and snapshots history.""" num_slots = num_decodes + 7 - prev_num_accepted_tokens = torch.arange(num_slots, dtype=torch.int32, device="cuda") % 21 - cache_buf_idx = torch.arange(num_slots, dtype=torch.int32, device="cuda") % 2 + history_len = torch.arange(num_slots, dtype=torch.int32, device="cuda") % 21 state_indices = torch.randperm(num_slots, device="cuda")[:num_decodes].to(torch.int32) - triton_items = torch.zeros(num_decodes, 4, dtype=torch.int32, device="cuda") - triton_n_writes = torch.zeros(1, dtype=torch.int32, device="cuda") - _build_replay_work_items_triton( + work_items = torch.zeros(num_decodes, 4, dtype=torch.int32, device="cuda") + n_writes = torch.ones(1, dtype=torch.int32, device="cuda") + _build_gdn_replay_work_items( state_indices, - prev_num_accepted_tokens, - cache_buf_idx, - triton_items, - triton_n_writes, - 6, - MIN_REPLAY_HISTORY_SIZE, - ) - torch_items, torch_n_writes = _torch_reference_work_items( - state_indices, prev_num_accepted_tokens, cache_buf_idx + history_len, + work_items, + n_writes, ) - torch.testing.assert_close(triton_items, torch_items) - torch.testing.assert_close(triton_n_writes, torch_n_writes) + torch.testing.assert_close( + work_items, _gdn_reference_work_items(state_indices, history_len) + ) + torch.testing.assert_close(n_writes.cpu(), torch.zeros(1, dtype=torch.int32)) - def test_prepare_replay_work_items_uses_the_selected_builder(self): + def test_prepare_gdn_replay_work_items_uses_decode_slice(self): """The entry point must feed the builder the decode slice, not the whole batch.""" num_contexts, num_decodes = 3, 40 num_slots = num_contexts + num_decodes + 7 - prev_num_accepted_tokens = torch.arange(num_slots, dtype=torch.int32, device="cuda") % 21 - cache_buf_idx = torch.arange(num_slots, dtype=torch.int32, device="cuda") % 2 + history_len = torch.arange(num_slots, dtype=torch.int32, device="cuda") % 21 state_indices = torch.randperm(num_slots, device="cuda")[: num_contexts + num_decodes].to( torch.int32 ) - manager = _GdnReplayCacheManager(prev_num_accepted_tokens, cache_buf_idx) + manager = _GdnReplayCacheManager(history_len) metadata = Mamba2Metadata(max_batch_size=num_contexts + num_decodes, chunk_size=8) metadata.state_indices[: num_contexts + num_decodes].copy_(state_indices) metadata._prepare_replay_work_items(manager, num_contexts + num_decodes, num_contexts) - expected_items, expected_n_writes = _torch_reference_work_items( - state_indices[num_contexts:], prev_num_accepted_tokens, cache_buf_idx - ) + expected_items = _gdn_reference_work_items(state_indices[num_contexts:], history_len) torch.testing.assert_close(metadata.replay_work_items[:num_decodes], expected_items) - torch.testing.assert_close(metadata.replay_n_writes, expected_n_writes) - - @pytest.mark.parametrize( - ("num_decodes", "expect_fused"), - [ - pytest.param(8, False, id="below-partition-min"), - pytest.param(16, True, id="partition-min"), - pytest.param(256, True, id="fused-max"), - pytest.param(257, False, id="above-fused-max"), - ], - ) - def test_prepare_gdn_replay_work_items_dispatch_boundary( - self, monkeypatch, num_decodes, expect_fused - ): - """Pin which batch sizes reach the single-launch kernel. - - The equivalence tests above pass on either side of this boundary, so - without this the fused launch could silently stop being used. - """ - num_slots = num_decodes + 7 - prev_num_accepted_tokens = torch.arange(num_slots, dtype=torch.int32, device="cuda") % 21 - cache_buf_idx = torch.arange(num_slots, dtype=torch.int32, device="cuda") % 2 - manager = _GdnReplayCacheManager(prev_num_accepted_tokens, cache_buf_idx) - metadata = Mamba2Metadata(max_batch_size=num_decodes, chunk_size=8) - metadata.state_indices.copy_(torch.arange(num_decodes, dtype=torch.int32, device="cuda")) - - launches = [] - original_kernel = mamba2_metadata._prepare_gdn_replay_work_items_kernel - - class _CountingKernel: - def __getitem__(self, grid): - launches.append(grid) - return original_kernel[grid] - - monkeypatch.setattr( - mamba2_metadata, "_prepare_gdn_replay_work_items_kernel", _CountingKernel() + torch.testing.assert_close( + metadata.replay_n_writes.cpu(), torch.zeros(1, dtype=torch.int32) ) - metadata._prepare_replay_work_items(manager, num_decodes, 0) - - assert bool(launches) is expect_fused - def test_prepare_gdn_replay_work_items_cuda_graph_replay(self): num_decodes = 40 num_slots = num_decodes + 7 - prev_num_accepted_tokens = torch.arange(num_slots, dtype=torch.int32, device="cuda") % 21 - cache_buf_idx = torch.arange(num_slots, dtype=torch.int32, device="cuda") % 2 - manager = _GdnReplayCacheManager(prev_num_accepted_tokens, cache_buf_idx) + history_len = torch.arange(num_slots, dtype=torch.int32, device="cuda") % 21 + manager = _GdnReplayCacheManager(history_len) metadata = Mamba2Metadata(max_batch_size=num_decodes, chunk_size=8) metadata.state_indices.copy_(torch.arange(num_decodes, dtype=torch.int32, device="cuda")) @@ -437,18 +374,15 @@ def test_prepare_gdn_replay_work_items_cuda_graph_replay(self): device="cuda", ) metadata.state_indices.copy_(updated_state_indices) - prev_num_accepted_tokens.copy_( - (torch.arange(num_slots, dtype=torch.int32, device="cuda") * 7) % 21 - ) - cache_buf_idx.bitwise_xor_(1) + history_len.copy_((torch.arange(num_slots, dtype=torch.int32, device="cuda") * 7) % 21) graph.replay() torch.cuda.synchronize() - expected_items, expected_n_writes = _torch_reference_work_items( - updated_state_indices, prev_num_accepted_tokens, cache_buf_idx - ) + expected_items = _gdn_reference_work_items(updated_state_indices, history_len) torch.testing.assert_close(metadata.replay_work_items[:num_decodes], expected_items) - torch.testing.assert_close(metadata.replay_n_writes, expected_n_writes) + torch.testing.assert_close( + metadata.replay_n_writes.cpu(), torch.zeros(1, dtype=torch.int32) + ) def test_single_sequence_unaligned(self): """Test with a single sequence that doesn't align with chunk size.""" From 3b2ba13dbccdf8a8b34ef9d1b9e77e4841dd0724 Mon Sep 17 00:00:00 2001 From: jiant <107457950+JadoTu@users.noreply.github.com> Date: Wed, 14 Oct 2026 01:49:21 +0000 Subject: [PATCH 2/4] fix the conner case of draft token length=0 Signed-off-by: jiant <107457950+JadoTu@users.noreply.github.com> --- tensorrt_llm/_torch/speculative/interface.py | 13 +++++ .../modules/mamba/test_gdn_mtp_replay.py | 34 ++++++++--- .../speculative/hw_agnostic/test_mtp.py | 57 +++++++++++++++++++ 3 files changed, 97 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index eb9067c54f95..fde2500a6dc4 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -1815,6 +1815,19 @@ def skip_drafting( dtype=torch.int, device=logits.device) + if (getattr(attn_metadata.kv_cache_manager, "use_gdn_mtp_replay", False) + and batch_size > num_contexts): + state_indices = attn_metadata.mamba_metadata.state_indices + attn_metadata.kv_cache_manager.update_mamba_states( + attn_metadata=attn_metadata, + num_accepted_tokens=num_accepted_tokens, + state_indices=state_indices, + ) + if self._auxiliary_state_handlers: + self.commit_auxiliary_speculative_states( + num_accepted_tokens, state_indices[:batch_size], + num_contexts) + next_draft_tokens = torch.zeros((batch_size, 0), dtype=torch.int, device=logits.device) diff --git a/tests/unittest/_torch/modules/mamba/test_gdn_mtp_replay.py b/tests/unittest/_torch/modules/mamba/test_gdn_mtp_replay.py index 54946de7fd11..2a881ceb80ac 100644 --- a/tests/unittest/_torch/modules/mamba/test_gdn_mtp_replay.py +++ b/tests/unittest/_torch/modules/mamba/test_gdn_mtp_replay.py @@ -143,11 +143,18 @@ def _replay_reference( return output, raw_x.contiguous() -def test_gdn_mtp_replay_matches_reference(): +@pytest.mark.parametrize( + "width,configured_width,history_length", + [(4, 6, 11), (1, 4, 0), (1, 4, 15), (1, 4, 16)], + ids=["mtp3", "zero-draft-empty", "zero-draft-near-full", "zero-draft-full"], +) +def test_gdn_mtp_replay_matches_reference( + width: int, configured_width: int, history_length: int +) -> None: from tensorrt_llm._torch.custom_ops.cute_dsl_gdn_mtp_replay import gdn_mtp_replay torch.manual_seed(7) - batch, slots, width, configured_width = 2, 5, 4, 6 + batch, slots = 2, 5 key_heads, value_heads = 1, 2 channels = (2 * key_heads + value_heads) * DIM qkvz_width = channels + value_heads * DIM @@ -180,7 +187,7 @@ def test_gdn_mtp_replay_matches_reference(): ) history_G = history_G_storage[:, :, :capacity] state_indices = torch.tensor([3, 1], device="cuda", dtype=torch.int32) - history_len = torch.tensor([2, 7, 4, 11, 1], device="cuda", dtype=torch.int32) + history_len = torch.tensor([2, 7, 4, history_length, 1], device="cuda", dtype=torch.int32) is_dummy = torch.tensor([False, True], device="cuda") A_log = _randn((value_heads,), 0.35, -0.25).clamp_(-1.0, 0.5) dt_bias = _randn((value_heads,), 0.65, -2.5).clamp_(-5.0, 0.5) @@ -334,7 +341,16 @@ def _make_state_storage(indirect, layers, slots, value_heads, channels): @pytest.mark.parametrize("indirect", [False, True], ids=["affine", "indirect"]) -def test_gdn_mtp_commit_matches_reference(indirect): +@pytest.mark.parametrize( + "accepted_counts,starting_lengths", + [((3, 2, 1), (15, 10, 16)), ((1, 1, 1), (16, 15, 0))], + ids=["mtp3", "zero-draft"], +) +def test_gdn_mtp_commit_matches_reference( + indirect: bool, + accepted_counts: tuple[int, int, int], + starting_lengths: tuple[int, int, int], +) -> None: from tensorrt_llm._torch.custom_ops.cute_dsl_gdn_mtp_commit import gdn_mtp_commit torch.manual_seed(11) @@ -358,9 +374,9 @@ def test_gdn_mtp_commit_matches_reference(indirect): -(_randn((layers, slots, value_heads, capacity), 0.025).abs() + 0.002), dim=-1, ) - accepted_tokens = torch.tensor([3, 2, 1], device="cuda", dtype=torch.int32) + accepted_tokens = torch.tensor(accepted_counts, device="cuda", dtype=torch.int32) is_dummy = torch.tensor([False, False, True], device="cuda") - starts = torch.tensor([15, 10, 16], device="cuda", dtype=torch.int32) + starts = torch.tensor(starting_lengths, device="cuda", dtype=torch.int32) positions = torch.tensor([2, 0, 1], device="cuda", dtype=torch.int32) slot_for_position = torch.tensor([1, 3, 4], device="cuda", dtype=torch.int32) selected_slots = slot_for_position[positions.long()] @@ -373,7 +389,11 @@ def test_gdn_mtp_commit_matches_reference(indirect): ), dim=1, ).contiguous() - history_len = torch.tensor([4, 15, 2, 10, 16], device="cuda", dtype=torch.int32) + history_len = torch.tensor( + [4, starting_lengths[0], 2, starting_lengths[1], starting_lengths[2]], + device="cuda", + dtype=torch.int32, + ) reference_u = history_u.clone() reference_k = history_k.clone() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py index 7304212eeba0..efbafe3c7003 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py @@ -15,6 +15,7 @@ import unittest from types import SimpleNamespace +from unittest import mock import pytest import torch @@ -1870,6 +1871,62 @@ class ReplacementHeadTarget: assert spec_config._use_shared_kv_cache +@pytest.mark.cpu_only +@pytest.mark.parametrize("worker_cls", [MTPWorker, MTPEagleWorker], ids=["mtp", "mtp-eagle"]) +@pytest.mark.parametrize( + "use_gdn_replay,num_contexts", + [(True, 1), (True, 2), (False, 0), (None, 0)], + ids=["gdn-mixed", "context-only", "other-hybrid", "attention"], +) +def test_mtp_zero_draft_commits_gdn_replay( + worker_cls: type[MTPWorker] | type[MTPEagleWorker], + use_gdn_replay: bool | None, + num_contexts: int, +) -> None: + worker = worker_cls(MTPDecodingConfig(max_draft_len=3)) + worker._execute_guided_decoder_if_present = mock.Mock() + worker._sample_tokens_for_batch = mock.Mock(return_value=torch.tensor([3, 6])) + handler = mock.Mock() + worker.register_auxiliary_state_handler(handler) + manager = SimpleNamespace(update_mamba_states=mock.Mock()) + if use_gdn_replay is not None: + manager.use_gdn_mtp_replay = use_gdn_replay + state_indices = torch.tensor([3, 1], dtype=torch.int32) + attn_metadata = SimpleNamespace( + num_seqs=2, + num_contexts=num_contexts, + kv_cache_manager=manager, + mamba_metadata=SimpleNamespace(state_indices=state_indices), + ) + result = worker._forward_impl( + input_ids=mock.sentinel.input_ids, + position_ids=mock.sentinel.position_ids, + hidden_states=mock.sentinel.hidden_states, + logits=torch.zeros(2, 8), + attn_metadata=attn_metadata, + spec_metadata=SimpleNamespace(runtime_draft_len=0), + draft_model=mock.sentinel.draft_model, + ) + + assert result["next_draft_tokens"].shape == (2, 0) + torch.testing.assert_close(result["new_tokens"], torch.tensor([[3], [6]], dtype=torch.int)) + torch.testing.assert_close(result["new_tokens_lens"], torch.ones(2, dtype=torch.int)) + if use_gdn_replay and num_contexts < 2: + manager.update_mamba_states.assert_called_once_with( + attn_metadata=attn_metadata, + num_accepted_tokens=result["new_tokens_lens"], + state_indices=state_indices, + ) + handler.commit_speculative_states.assert_called_once() + accepted, indices, contexts = handler.commit_speculative_states.call_args.args + assert accepted is result["new_tokens_lens"] + torch.testing.assert_close(indices, state_indices) + assert contexts == num_contexts + else: + manager.update_mamba_states.assert_not_called() + handler.commit_speculative_states.assert_not_called() + + def test_mtp_shared_kv_draft_inputs(): spec_config = MTPDecodingConfig( max_draft_len=3, From 0c61cffb6ea873da67ee918fe25940bf6563b535 Mon Sep 17 00:00:00 2001 From: jiant <107457950+JadoTu@users.noreply.github.com> Date: Fri, 16 Oct 2026 02:48:56 +0000 Subject: [PATCH 3/4] support the V2's request suspension in GDN MTP replay Signed-off-by: jiant <107457950+JadoTu@users.noreply.github.com> --- .../custom_ops/cute_dsl_gdn_mtp_commit.py | 322 ++++++++++++++++++ .../kv_cache/mamba_cache_manager.py | 47 ++- .../kv_cache/test_mamba_cache_manager.py | 166 +++++++++ .../modules/mamba/test_gdn_mtp_replay.py | 90 +++++ 4 files changed, 624 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_gdn_mtp_commit.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_gdn_mtp_commit.py index 2bbcf83fdff2..bb7e73f37723 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_gdn_mtp_commit.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_gdn_mtp_commit.py @@ -816,3 +816,325 @@ def gdn_mtp_commit( _CACHE[key] = fn fn(*args, cutlass.Int32(grid), stream) + + +@cute.kernel +def _flush_kernel( + STATE_BASE: cutlass.Int64, + DESC_BASE: cutlass.Int64, + U_BASE: cutlass.Int64, + K_BASE: cutlass.Int64, + G_BASE: cutlass.Int64, + LEN_BASE: cutlass.Int64, + IDX_BASE: cutlass.Int64, + B: cutlass.Int32, + L: cutlass.Constexpr, + S: cutlass.Constexpr, + H: cutlass.Constexpr, + HV: cutlass.Constexpr, + CAP: cutlass.Constexpr, + W: cutlass.Constexpr, + INDIRECT: cutlass.Constexpr, + LAYER_STRIDE: cutlass.Constexpr, + SLOT_STRIDE: cutlass.Constexpr, +): + item, layer, hv = cute.arch.block_idx() + tid, _, _ = cute.arch.thread_idx() + f32, bf16 = cutlass.Float32, cutlass.BFloat16 + irow0, chunk = tid // NCH, tid % NCH + indices = cute.make_tensor( + cute.make_ptr(cutlass.Int64, IDX_BASE, AddressSpace.gmem, assumed_align=8), + cute.make_layout(B), + ) + lengths = cute.make_tensor( + cute.make_ptr(cutlass.Int32, LEN_BASE, AddressSpace.gmem, assumed_align=16), + cute.make_layout(S), + ) + slot = indices[item] + length = lengths[slot] + smem = utils.SmemAllocator() + acc_ptr = smem.allocate_tensor(bf16, cute.make_layout(DIM * PAD), byte_alignment=16) + u_ptr = smem.allocate_tensor(bf16, cute.make_layout(W * PAD), byte_alignment=16) + k_ptr = smem.allocate_tensor(bf16, cute.make_layout(W * PAD), byte_alignment=16) + sAcc = cute.make_tensor( + acc_ptr.iterator, cute.make_layout((NSTRIP, DIM, STRIPW), stride=(STRIPW, PAD, 1)) + ) + sAccVec = cute.make_tensor( + acc_ptr.iterator, cute.make_layout((DIM, NCH, VEC), stride=(PAD, VEC, 1)) + ) + sU = cute.make_tensor(u_ptr.iterator, cute.make_layout((DIM, W), stride=(1, PAD))) + sUVec = cute.make_tensor(u_ptr.iterator, cute.make_layout((W, NCH, VEC), stride=(PAD, VEC, 1))) + sK = cute.make_tensor( + k_ptr.iterator, cute.make_layout((NSTRIP, STRIPW, W), stride=(STRIPW, 1, PAD)) + ) + sKVec = cute.make_tensor(k_ptr.iterator, cute.make_layout((W, NCH, VEC), stride=(PAD, VEC, 1))) + + if length > 0: + state_addr = ( + STATE_BASE + + (cutlass.Int64(layer) * LAYER_STRIDE + cutlass.Int64(slot) * SLOT_STRIDE) * 2 + ) + if cutlass.const_expr(INDIRECT): + descriptors = cute.make_tensor( + cute.make_ptr(cutlass.Int64, DESC_BASE, AddressSpace.gmem, assumed_align=16), + cute.make_layout((L, 2), stride=(2, 1)), + ) + state_addr = descriptors[layer, 0] + slot * descriptors[layer, 1] * 2 + state_addr += cutlass.Int64(hv) * DIM * DIM * 2 + state = cute.make_tensor( + cute.make_ptr(bf16, state_addr, AddressSpace.gmem, assumed_align=16), + cute.make_layout((DIM, NCH, VEC), stride=(DIM, VEC, 1)), + ) + layer_slot = cutlass.Int64(layer) * S + slot + history_u = cute.make_tensor( + cute.make_ptr( + bf16, U_BASE + layer_slot * CAP * HV * DIM * 2, AddressSpace.gmem, assumed_align=16 + ), + cute.make_layout((CAP, HV, NCH, VEC), stride=(HV * DIM, DIM, VEC, 1)), + ) + history_k = cute.make_tensor( + cute.make_ptr( + bf16, K_BASE + layer_slot * CAP * H * DIM * 2, AddressSpace.gmem, assumed_align=16 + ), + cute.make_layout((CAP, H, NCH, VEC), stride=(H * DIM, DIM, VEC, 1)), + ) + history_G = cute.make_tensor( + cute.make_ptr( + f32, G_BASE + layer_slot * HV * CAP * 4, AddressSpace.gmem, assumed_align=16 + ), + cute.make_layout((HV, CAP), stride=(CAP, 1)), + ) + end_G = history_G[hv, length - 1] + saved = [cute.make_rmem_tensor((VEC,), bf16) for _ in range(NJ)] + for jj in cutlass.range_constexpr(NJ): + cute.autovec_copy(state[irow0 + MSTEP * jj, chunk, None], saved[jj]) + for wb in cutlass.range_constexpr(W // MSTEP): + row = irow0 + wb * MSTEP + key = cute.make_rmem_tensor((VEC,), bf16) + update = cute.make_rmem_tensor((VEC,), bf16) + key.fill(0.0) + update.fill(0.0) + if row < length: + cute.autovec_copy(history_k[row, hv // (HV // H), chunk, None], key) + cute.autovec_copy(history_u[row, hv, chunk, None], update) + decay = cute.exp2((end_G - history_G[hv, row]) * LOG2E, approx=True, ftz=True) + update.store((update.load().to(f32) * decay).to(bf16)) + cute.autovec_copy(key, sKVec[row, chunk, None]) + cute.autovec_copy(update, sUVec[row, chunk, None]) + scale = cute.exp2(end_G * LOG2E, approx=True, ftz=True) + cute.arch.barrier() + + mma = warp.MmaF16BF16Op(bf16, f32, (16, 8, 16)) + tiled_mma = cute.make_tiled_mma( + mma, cute.make_layout((MW, NW, 1)), permutation_mnk=(64, 32, 16) + ) + thread_mma = tiled_mma.get_slice(tid) + load_atom = cute.make_copy_atom(warp.LdMatrix8x8x16bOp(True, 4), bf16) + copy_a = cute.make_tiled_copy_A(load_atom, tiled_mma) + copy_b = cute.make_tiled_copy_B(load_atom, tiled_mma) + thread_a, thread_b = copy_a.get_slice(tid), copy_b.get_slice(tid) + store_atom = cute.make_copy_atom(warp.StMatrix8x8x16bOp(False, 4), bf16) + store_c = cute.make_tiled_copy_C_atom(store_atom, tiled_mma) + store_copy = cute.make_tiled_copy_S(store_atom, store_c) + thread_store = store_copy.get_slice(tid) + fragment_a = tiled_mma.make_fragment_A(thread_mma.partition_A(sU)) + cute.copy(copy_a, thread_a.partition_S(sU), thread_a.retile(fragment_a)) + for strip in cutlass.range_constexpr(NSTRIP): + operand_b = sK[strip, None, None] + fragment_b = tiled_mma.make_fragment_B(thread_mma.partition_B(operand_b)) + cute.copy(copy_b, thread_b.partition_S(operand_b), thread_b.retile(fragment_b)) + acc = cute.make_rmem_tensor(tiled_mma.partition_shape_C((DIM, STRIPW)), f32) + acc.fill(0.0) + for kb in cutlass.range_constexpr(cute.size(fragment_a, mode=[2])): + cute.gemm( + tiled_mma, acc, fragment_a[None, None, kb], fragment_b[None, None, kb], acc + ) + result_fragment = store_copy.retile(acc) + rounded = cute.make_rmem_tensor(result_fragment.shape, bf16) + rounded.store(result_fragment.load().to(bf16)) + cute.copy(store_copy, rounded, thread_store.partition_D(sAcc[strip, None, None])) + cute.arch.barrier() + for jj in cutlass.range_constexpr(NJ): + row = irow0 + MSTEP * jj + accumulated = cute.make_rmem_tensor((VEC,), bf16) + cute.autovec_copy(sAccVec[row, chunk, None], accumulated) + saved[jj].store( + (saved[jj].load().to(f32) * scale + accumulated.load().to(f32)).to(bf16) + ) + cute.autovec_copy(saved[jj], state[row, chunk, None]) + + +@cute.jit +def _launch_flush( + STATE_BASE: cutlass.Int64, + DESC_BASE: cutlass.Int64, + U_BASE: cutlass.Int64, + K_BASE: cutlass.Int64, + G_BASE: cutlass.Int64, + LEN_BASE: cutlass.Int64, + IDX_BASE: cutlass.Int64, + B: cutlass.Int32, + L: cutlass.Constexpr, + S: cutlass.Constexpr, + H: cutlass.Constexpr, + HV: cutlass.Constexpr, + CAP: cutlass.Constexpr, + W: cutlass.Constexpr, + INDIRECT: cutlass.Constexpr, + LAYER_STRIDE: cutlass.Constexpr, + SLOT_STRIDE: cutlass.Constexpr, + stream, +): + _flush_kernel( + STATE_BASE, + DESC_BASE, + U_BASE, + K_BASE, + G_BASE, + LEN_BASE, + IDX_BASE, + B, + L, + S, + H, + HV, + CAP, + W, + INDIRECT, + LAYER_STRIDE, + SLOT_STRIDE, + ).launch( + grid=[B, L, HV], + block=[NT, 1, 1], + stream=stream, + min_blocks_per_mp=CTAS_PER_SM, + preferred_smem_carveout=100, + ) + + +_FLUSH_CACHE = {} + + +@torch.no_grad() +def gdn_mtp_flush( + ssm_states: torch.Tensor, + history_u: torch.Tensor, + history_k: torch.Tensor, + history_G: torch.Tensor, + history_len: torch.Tensor, + state_indices: torch.Tensor, + *, + history_window: int, + ssm_state_descriptors: torch.Tensor | None = None, +) -> None: + """Flush selected slots' accepted histories into every recurrent checkpoint. + + U/K have shapes [L,S,CAP,HV,128] and [L,S,CAP,H,128]; G is FP32 + [L,S,HV,CAP]. The BF16 checkpoint uses the normal commit's affine or + indirect V2 layout, with descriptor strides in BF16 elements. + Distinct int64 state_indices select slots with accepted lengths at most + history_window. Conv and rejected candidates are not accessed. + """ + if state_indices.numel() == 0: + return + if history_u.ndim != 5 or history_k.ndim != 5: + raise ValueError("GDN flush K/U histories must have five dimensions") + layers, slots, capacity, value_heads, value_dim = history_u.shape + key_heads = history_k.shape[3] + if value_dim != DIM or history_k.shape != (layers, slots, capacity, key_heads, DIM): + raise ValueError("GDN flush histories have inconsistent dimensions") + if key_heads <= 0 or value_heads <= 0 or value_heads % key_heads != 0: + raise ValueError("value_heads must be divisible by key_heads") + if history_window not in (16, 32) or not 1 <= capacity - history_window <= 8: + raise ValueError("GDN flush requires history 16 or 32 plus 1..8 candidates") + if not 1 <= layers <= 128: + raise ValueError("GDN flush supports 1..128 layers") + if history_G.shape != (layers, slots, value_heads, capacity): + raise ValueError("GDN flush cumulative-G history has an unexpected shape") + if history_u.dtype != torch.bfloat16 or history_k.dtype != torch.bfloat16: + raise TypeError("GDN flush K/U histories must be BF16") + if history_G.dtype != torch.float32: + raise TypeError("GDN flush cumulative-G history must be FP32") + if history_len.shape != (slots,) or history_len.dtype != torch.int32: + raise ValueError("GDN flush history_len must be int32 with shape [slots]") + if state_indices.ndim != 1 or state_indices.dtype != torch.int64: + raise ValueError("GDN flush state_indices must be a rank-1 int64 tensor") + if ssm_states.dtype != torch.bfloat16: + raise TypeError("GDN flush recurrent checkpoints must be BF16") + if ssm_states.stride()[-3:] != (DIM * DIM, DIM, 1): + raise ValueError("GDN flush checkpoint inner dimensions must be dense") + indirect = ssm_state_descriptors is not None + if indirect: + if ssm_states.shape != (slots, value_heads, DIM, DIM): + raise ValueError("Indirect GDN state must be the first rank-4 layer view") + if ( + ssm_state_descriptors.shape != (layers, 2) + or ssm_state_descriptors.dtype != torch.int64 + or not ssm_state_descriptors.is_contiguous() + ): + raise ValueError("SSM descriptors must be contiguous int64 [layers,2]") + layer_stride = slot_stride = 0 + descriptors = ssm_state_descriptors + else: + if ssm_states.shape != (slots, layers, value_heads, DIM, DIM): + raise ValueError("Affine GDN state must have shape [S,L,HV,128,128]") + layer_stride = ssm_states.stride(1) if layers > 1 else 0 + slot_stride = ssm_states.stride(0) + descriptors = _dummy_descriptors(ssm_states.device) + tensors = (history_u, history_k, history_G, history_len, state_indices, descriptors) + if any(not tensor.is_contiguous() for tensor in tensors): + raise ValueError("GDN flush histories, indices, and descriptors must be contiguous") + if any( + not tensor.is_cuda or tensor.device != ssm_states.device + for tensor in (*tensors, ssm_states) + ): + raise ValueError("GDN flush tensors must be on the same CUDA device") + args = tuple( + cutlass.Int64(tensor.data_ptr()) + for tensor in ( + ssm_states, + descriptors, + history_u, + history_k, + history_G, + history_len, + state_indices, + ) + ) + from cuda.bindings import driver as _drv + + key = ( + ssm_states.device.index, + layers, + slots, + key_heads, + value_heads, + capacity, + history_window, + indirect, + layer_stride, + slot_stride, + ) + with torch.cuda.device(ssm_states.device): + stream = _drv.CUstream(torch.cuda.current_stream().cuda_stream) + fn = _FLUSH_CACHE.get(key) + if fn is None: + fn = cute.compile( + _launch_flush, + *args, + cutlass.Int32(state_indices.numel()), + layers, + slots, + key_heads, + value_heads, + capacity, + history_window, + indirect, + layer_stride, + slot_stride, + stream, + ) + _FLUSH_CACHE[key] = fn + fn(*args, cutlass.Int32(state_indices.numel()), stream) + # Every layer/head must finish reading the length before it is reset. + history_len.index_fill_(0, state_indices, 0) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index 94a39d8efa99..eaff2a639be2 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -4679,7 +4679,7 @@ def _reset_kda_replay_slots( self.prev_num_accepted_tokens.index_fill_(0, slots, 0) def suspend_request(self, req: LlmRequest) -> None: - """Suspend the request's KV cache with its whole KDA state in V2. + """Flush replay state before suspending the request's V2 pages. The drafts the last verify accepted are pending in the slot's KDA replay buffers, and so is the conv window, which the fused verify does @@ -4688,6 +4688,10 @@ def suspend_request(self, req: LlmRequest) -> None: pages are released, so the suspended state is self-contained in the recurrent state; the request's next generation step re-seeds the replay buffers of whichever slot it resumes in from them. + + GDN replay commits Conv each round, but its accepted K/U/G history + must likewise be flushed into the recurrent checkpoint before the + slot can be released. """ request_id = req.py_request_id kv_cache = self.kv_cache_map.get(request_id) @@ -4709,6 +4713,12 @@ def suspend_request(self, req: LlmRequest) -> None: self._commit_kda_pending_drafts(slot) del self._request_id_to_state_index[request_id] self._kda_replay_suspended.add(request_id) + elif (self.use_gdn_mtp_replay and slot >= 0 and kv_cache is not None + and kv_cache.is_active): + self._flush_gdn_mtp_replay(slot) + # A resumed request uses its restored V2 state with no history, + # even if another request has overwritten the old slot meanwhile. + del self._request_id_to_state_index[request_id] super().suspend_request(req) @contextmanager @@ -4766,6 +4776,41 @@ def _commit_kda_pending_drafts(self, slot: int) -> None: ) self.prev_num_accepted_tokens[slot] = 0 + def _flush_gdn_mtp_replay(self, slot: int) -> None: + """Flush accepted recurrent updates without advancing committed Conv. + + The post-verification accepted-prefix commit must already be enqueued. + V2 page migration is ordered against the KV-cache stream. + """ + from tensorrt_llm._torch.custom_ops.cute_dsl_gdn_mtp_commit import \ + gdn_mtp_flush + + if (not self.all_ssm_states + or (self._gdn_mtp_replay_state_descriptors is None + and self._gdn_mtp_replay_state_view is None) + or self.old_x is None or self.old_B is None + or self.old_dt is None or self.prev_num_accepted_tokens is None + or self.replay_history_size is None): + raise RuntimeError( + "GDN flush requires V2 state views and replay buffers") + with self._on_kv_cache_stream(): + slots = torch.full((1, ), + slot, + dtype=torch.long, + device=self.prev_num_accepted_tokens.device) + state_view = self._gdn_mtp_replay_state_view + gdn_mtp_flush( + ssm_states=(state_view if state_view is not None else + self.all_ssm_states[0]), + history_u=self.old_x, + history_k=self.old_B, + history_G=self.old_dt, + history_len=self.prev_num_accepted_tokens, + state_indices=slots, + history_window=self.replay_history_size, + ssm_state_descriptors=self._gdn_mtp_replay_state_descriptors, + ) + def _reset_gdn_mtp_replay_slots( self, slots: Union[List[int], torch.Tensor], diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 5f163e217b43..af3ddebd9346 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -5349,6 +5349,172 @@ def combined_commit(_metadata, num_decodes, accepted_tokens, is_dummy): assert torch.equal(captured["is_dummy"], torch.tensor([True, False])) +def _v2_gdn_mtp_replay_stub() -> MambaHybridCacheManagerV2: + mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.local_num_mamba_layers = 1 + mgr._use_kda_replay_update = False + mgr._use_gdn_mtp_replay = True + mgr._request_id_to_state_index = {} + mgr._request_id_to_is_dummy = {} + mgr.kv_cache_map = {} + mgr.ssm_layer_group_id = 0 + mgr._host_state_indices = torch.zeros(4, dtype=torch.int32) + mgr.cuda_state_indices = torch.zeros(4, dtype=torch.int32) + mgr._dummy_request_mask = None + mgr.prev_num_accepted_tokens = torch.zeros(4, dtype=torch.int32) + mgr.old_x, mgr.old_B, mgr.old_dt = [torch.zeros(1, 4, 2, 4) for _ in range(3)] + mgr.all_conv_states = [torch.zeros(4, 6, 3)] + mgr.flushed_lengths = [] + + def flush(slot: int) -> None: + mgr.flushed_lengths.append(int(mgr.prev_num_accepted_tokens[slot])) + mgr.prev_num_accepted_tokens[slot] = 0 + + mgr._flush_gdn_mtp_replay = MagicMock(side_effect=flush) + return mgr + + +@pytest.mark.parametrize("resume_slot", [0, 2], ids=["same_slot", "different_slot"]) +def test_v2_gdn_replay_suspend_resume_starts_with_empty_history(resume_slot: int) -> None: + mgr = _v2_gdn_mtp_replay_stub() + request = SimpleNamespace(py_request_id=1, is_dummy=False) + other = SimpleNamespace(py_request_id=2, is_dummy=False) + mgr.kv_cache_map[1] = _SlotCache(0) + mgr._setup_state_indices([request]) + mgr.prev_num_accepted_tokens[0] = 7 + mgr.suspend_request(request) + assert mgr.flushed_lengths == [7] + assert 1 not in mgr._request_id_to_state_index + assert not mgr.kv_cache_map[1].is_active + mgr.suspend_request(request) + mgr._flush_gdn_mtp_replay.assert_called_once_with(0) + + mgr.kv_cache_map[2] = _SlotCache(0) + mgr._setup_state_indices([other]) + mgr.prev_num_accepted_tokens[0] = 3 + for history in (mgr.old_x, mgr.old_B, mgr.old_dt): + history[:, 0] = 9 + if resume_slot == 0: + del mgr.kv_cache_map[2] + del mgr._request_id_to_state_index[2] + restored_conv = torch.arange(18.0).view(6, 3) + mgr.all_conv_states[0][resume_slot] = restored_conv + mgr.kv_cache_map[1].slot = resume_slot + mgr.kv_cache_map[1].is_active = True + mgr._setup_state_indices([request] if resume_slot == 0 else [request, other]) + assert mgr.prev_num_accepted_tokens[resume_slot] == 0 + torch.testing.assert_close(mgr.all_conv_states[0][resume_slot], restored_conv, rtol=0, atol=0) + if resume_slot != 0: + assert mgr.prev_num_accepted_tokens[0] == 3 + for history in (mgr.old_x, mgr.old_B, mgr.old_dt): + assert (history[:, 0] == 9).all() + + +@skip_no_cuda +@pytest.mark.parametrize("indirect", [False, True], ids=["affine", "indirect"]) +@pytest.mark.parametrize("resume_slot", [1, 2], ids=["same_slot", "different_slot"]) +def test_v2_gdn_replay_flush_precedes_slot_release(indirect: bool, resume_slot: int) -> None: + from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE + from tensorrt_llm._utils import is_sm_100f + + if not IS_CUTLASS_DSL_AVAILABLE or not is_sm_100f(): + pytest.skip("GDN replay requires datacenter Blackwell and CUTLASS DSL") + torch.manual_seed(23) + mgr = _v2_gdn_mtp_replay_stub() + del mgr._flush_gdn_mtp_replay + layers, slots, heads, value_heads, dim, window = 2, 4, 2, 4, 128, 16 + capacity = window + 4 + if indirect: + mgr.all_ssm_states = [ + torch.randn(slots, value_heads + layer + 1, dim, dim, device="cuda").bfloat16()[ + :, :value_heads + ] + for layer in range(layers) + ] + mgr._gdn_mtp_replay_state_view = None + mgr._gdn_mtp_replay_state_descriptors = torch.tensor( + [(state.data_ptr(), state.stride(0)) for state in mgr.all_ssm_states], + device="cuda", + dtype=torch.int64, + ) + else: + state_storage = torch.randn(slots, layers, value_heads, dim, dim, device="cuda").bfloat16() + mgr.all_ssm_states = [state_storage[:, layer] for layer in range(layers)] + mgr._gdn_mtp_replay_state_view = state_storage + mgr._gdn_mtp_replay_state_descriptors = None + mgr.all_conv_states = [ + torch.randn(slots, 8, 3, device="cuda").bfloat16() for _ in range(layers) + ] + mgr.old_x = ( + torch.randn(layers, slots, capacity, value_heads, dim, device="cuda") * 0.03 + ).bfloat16() + mgr.old_B = (torch.randn(layers, slots, capacity, heads, dim, device="cuda") * 0.09).bfloat16() + mgr.old_dt = torch.cumsum( + -torch.rand(layers, slots, value_heads, capacity, device="cuda") * 0.1, dim=-1 + ) + mgr.prev_num_accepted_tokens = torch.zeros(slots, device="cuda", dtype=torch.int32) + mgr.cuda_state_indices = torch.zeros(slots, device="cuda", dtype=torch.int32) + mgr.replay_history_size = window + mgr._stream = torch.cuda.Stream() + mgr.local_num_mamba_layers = layers + request = SimpleNamespace(py_request_id=1, is_dummy=False) + cache = _SlotCache(1) + mgr.kv_cache_map[1] = cache + mgr._setup_state_indices([request]) + mgr.prev_num_accepted_tokens[1] = 7 + state_before = torch.stack([state[1].clone() for state in mgr.all_ssm_states]).float() + conv_before = [conv[1].clone() for conv in mgr.all_conv_states] + end_G = mgr.old_dt[:, 1, :, 6] + key_heads = torch.arange(value_heads, device="cuda") // (value_heads // heads) + update = torch.einsum( + "ltiv,ltik,lit->livk", + mgr.old_x[:, 1, :7].float(), + mgr.old_B[:, 1, :7][:, :, key_heads].float(), + torch.exp(end_G[:, :, None] - mgr.old_dt[:, 1, :, :7]), + ) + expected = (state_before * torch.exp(end_G)[:, :, None, None] + update).bfloat16() + saved_states, saved_conv = [], [] + + def suspend() -> None: + saved_states.extend(state[1].clone() for state in mgr.all_ssm_states) + saved_conv.extend(conv[1].clone() for conv in mgr.all_conv_states) + cache.is_active = False + + cache.suspend = suspend + mgr.suspend_request(request) + assert not cache.is_active and 1 not in mgr._request_id_to_state_index + assert mgr.prev_num_accepted_tokens[1] == 0 + for actual, reference in zip(saved_states, expected): + torch.testing.assert_close(actual, reference, rtol=1e-2, atol=8e-3) + for actual, reference in zip(saved_conv, conv_before): + torch.testing.assert_close(actual, reference, rtol=0, atol=0) + for layer, (state, conv) in enumerate(zip(mgr.all_ssm_states, mgr.all_conv_states)): + state[1].fill_(torch.nan) + conv[1].fill_(torch.nan) + state[resume_slot].copy_(saved_states[layer]) + conv[resume_slot].copy_(saved_conv[layer]) + for history in (mgr.old_x, mgr.old_B, mgr.old_dt): + history[:, 1].fill_(torch.nan) + mgr.prev_num_accepted_tokens.fill_(window) + cache.slot, cache.is_active = resume_slot, True + mgr._setup_state_indices([request]) + assert mgr.prev_num_accepted_tokens[resume_slot] == 0 + for actual, reference in zip(mgr.all_conv_states, conv_before): + torch.testing.assert_close(actual[resume_slot], reference, rtol=0, atol=0) + + +def test_v2_gdn_replay_flush_failure_does_not_release_the_slot() -> None: + mgr = _v2_gdn_mtp_replay_stub() + request = SimpleNamespace(py_request_id=1, is_dummy=False) + mgr.kv_cache_map[1] = _SlotCache(0) + mgr._setup_state_indices([request]) + mgr._flush_gdn_mtp_replay.side_effect = RuntimeError("flush failed") + with pytest.raises(RuntimeError, match="flush failed"): + mgr.suspend_request(request) + assert mgr.kv_cache_map[1].is_active + assert mgr._request_id_to_state_index[1] == 0 + + def test_v2_hybrid_replay_update_skips_dummy_and_padding_rows(monkeypatch): mgr = object.__new__(MambaHybridCacheManagerV2) mgr.local_num_mamba_layers = 1 diff --git a/tests/unittest/_torch/modules/mamba/test_gdn_mtp_replay.py b/tests/unittest/_torch/modules/mamba/test_gdn_mtp_replay.py index 2a881ceb80ac..1b3b254be99d 100644 --- a/tests/unittest/_torch/modules/mamba/test_gdn_mtp_replay.py +++ b/tests/unittest/_torch/modules/mamba/test_gdn_mtp_replay.py @@ -435,3 +435,93 @@ def test_gdn_mtp_commit_matches_reference( torch.testing.assert_close(history_k, reference_k, rtol=0, atol=0) torch.testing.assert_close(history_G, reference_G, rtol=1e-6, atol=1e-6) torch.testing.assert_close(history_len, reference_len, rtol=0, atol=0) + + +@pytest.mark.parametrize("window", [16, 32]) +@pytest.mark.parametrize("indirect", [False, True], ids=["affine", "indirect"]) +@pytest.mark.parametrize("head_ratio", [1, 2, 4, 8]) +@torch.no_grad() +def test_gdn_mtp_flush_matches_reference(window: int, indirect: bool, head_ratio: int) -> None: + from tensorrt_llm._torch.custom_ops.cute_dsl_gdn_mtp_commit import gdn_mtp_flush + + torch.manual_seed(19) + layers, slots, key_heads, width = 2, 6, 2, 4 + value_heads = key_heads * head_ratio + channels = (2 * key_heads + value_heads) * DIM + capacity = window + width + states, conv, state_arg, _conv_arg, state_desc, _conv_desc = _make_state_storage( + indirect, layers, slots, value_heads, channels + ) + if indirect: + # The layers need not have equal slot strides in V2's indirect layout. + states = [ + _randn((slots, value_heads + layer + 1, DIM, DIM), 0.012).bfloat16()[:, :value_heads] + for layer in range(layers) + ] + state_arg = states[0] + state_desc = torch.tensor( + [(state.data_ptr(), state.stride(0)) for state in states], + device="cuda", + dtype=torch.int64, + ) + expected_states = [state.clone() for state in states] + conv_before = [value.clone() for value in conv] + history_u = _randn((layers, slots, capacity, value_heads, DIM), 0.012).bfloat16() + history_k = _randn((layers, slots, capacity, key_heads, DIM), 0.09).bfloat16() + history_G = torch.cumsum( + -(_randn((layers, slots, value_heads, capacity), 0.025).abs() + 0.002), dim=-1 + ) + history_len = torch.tensor([3, 1, window, 5, 0, 7], device="cuda", dtype=torch.int32) + selected = torch.tensor([4, 1, 5, 2], device="cuda", dtype=torch.int64) + keys_for_value = torch.arange(value_heads, device="cuda") // head_ratio + for slot in selected.tolist(): + length = int(history_len[slot].item()) + if length: + end_G = history_G[:, slot, :, length - 1] + update = torch.einsum( + "ltiv,ltik,lit->livk", + history_u[:, slot, :length].float(), + history_k[:, slot, :length][:, :, keys_for_value].float(), + torch.exp(end_G[:, :, None] - history_G[:, slot, :, :length]), + ) + checkpoint = torch.stack([state[slot] for state in states]).float() + committed = checkpoint * torch.exp(end_G)[:, :, None, None] + update + for layer, state in enumerate(expected_states): + state[slot] = committed[layer].bfloat16() + history_u[:, slot, length:] = torch.nan + history_k[:, slot, length:] = torch.nan + history_G[:, slot, :, length:] = torch.nan + histories_before = [value.clone() for value in (history_u, history_k, history_G)] + expected_len = history_len.clone() + expected_len[selected] = 0 + gdn_mtp_flush( + state_arg, + history_u, + history_k, + history_G, + history_len, + selected, + history_window=window, + ssm_state_descriptors=state_desc, + ) + for actual, expected in zip(states, expected_states): + torch.testing.assert_close(actual, expected, rtol=1e-2, atol=5e-4) + for actual, expected in zip(conv, conv_before): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + for actual, expected in zip((history_u, history_k, history_G), histories_before): + torch.testing.assert_close(actual, expected, rtol=0, atol=0, equal_nan=True) + torch.testing.assert_close(history_len, expected_len, rtol=0, atol=0) + # Empty histories are a no-op when the entry point is reused with another B. + saved_states = [state.clone() for state in states] + gdn_mtp_flush( + state_arg, + history_u, + history_k, + history_G, + history_len, + selected[:1], + history_window=window, + ssm_state_descriptors=state_desc, + ) + for actual, expected in zip(states, saved_states): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) From cd11d732a036bce0b47603301727ab6ab9f4acb7 Mon Sep 17 00:00:00 2001 From: jiant <107457950+JadoTu@users.noreply.github.com> Date: Tue, 20 Oct 2026 09:48:05 +0000 Subject: [PATCH 4/4] CI fixes Signed-off-by: jiant <107457950+JadoTu@users.noreply.github.com> --- .../kv_cache/test_mamba_cache_manager.py | 8 +- .../hw_agnostic/test_guided_hostfunc_drain.py | 2 + .../speculative/hw_agnostic/test_mtp.py | 57 ------------- .../hw_agnostic/test_mtp_zero_draft.py | 80 +++++++++++++++++++ 4 files changed, 87 insertions(+), 60 deletions(-) create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_mtp_zero_draft.py diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index af3ddebd9346..a692a70af297 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -5359,8 +5359,8 @@ def _v2_gdn_mtp_replay_stub() -> MambaHybridCacheManagerV2: mgr.kv_cache_map = {} mgr.ssm_layer_group_id = 0 mgr._host_state_indices = torch.zeros(4, dtype=torch.int32) - mgr.cuda_state_indices = torch.zeros(4, dtype=torch.int32) - mgr._dummy_request_mask = None + mgr._state_indices_buffer = DeviceCopyBuffer(torch.zeros(4, dtype=torch.int32)) + mgr._dummy_request_mask_buffer = None mgr.prev_num_accepted_tokens = torch.zeros(4, dtype=torch.int32) mgr.old_x, mgr.old_B, mgr.old_dt = [torch.zeros(1, 4, 2, 4) for _ in range(3)] mgr.all_conv_states = [torch.zeros(4, 6, 3)] @@ -5453,7 +5453,9 @@ def test_v2_gdn_replay_flush_precedes_slot_release(indirect: bool, resume_slot: -torch.rand(layers, slots, value_heads, capacity, device="cuda") * 0.1, dim=-1 ) mgr.prev_num_accepted_tokens = torch.zeros(slots, device="cuda", dtype=torch.int32) - mgr.cuda_state_indices = torch.zeros(slots, device="cuda", dtype=torch.int32) + mgr._state_indices_buffer = DeviceCopyBuffer( + torch.zeros(slots, device="cuda", dtype=torch.int32) + ) mgr.replay_history_size = window mgr._stream = torch.cuda.Stream() mgr.local_num_mamba_layers = layers diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_guided_hostfunc_drain.py b/tests/unittest/_torch/speculative/hw_agnostic/test_guided_hostfunc_drain.py index 7437308c915a..7fa24b28f8b8 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_guided_hostfunc_drain.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_guided_hostfunc_drain.py @@ -175,6 +175,7 @@ def test_skip_drafting_routes_through_the_guided_helper(): worker = _Worker(MagicMock(name="guided_decoder")) batch_size = 3 attn_metadata = MagicMock(num_seqs=batch_size, num_contexts=1) + attn_metadata.kv_cache_manager.use_gdn_mtp_replay = False logits = torch.zeros((batch_size, 8)) with ( patch.object(_Worker, "_execute_guided_decoder_if_present") as helper, @@ -195,4 +196,5 @@ def test_skip_drafting_routes_through_the_guided_helper(): ) helper.assert_called_once_with(logits) + attn_metadata.kv_cache_manager.update_mamba_states.assert_not_called() assert out["new_tokens"].tolist() == [[0], [1], [2]] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py index efbafe3c7003..7304212eeba0 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py @@ -15,7 +15,6 @@ import unittest from types import SimpleNamespace -from unittest import mock import pytest import torch @@ -1871,62 +1870,6 @@ class ReplacementHeadTarget: assert spec_config._use_shared_kv_cache -@pytest.mark.cpu_only -@pytest.mark.parametrize("worker_cls", [MTPWorker, MTPEagleWorker], ids=["mtp", "mtp-eagle"]) -@pytest.mark.parametrize( - "use_gdn_replay,num_contexts", - [(True, 1), (True, 2), (False, 0), (None, 0)], - ids=["gdn-mixed", "context-only", "other-hybrid", "attention"], -) -def test_mtp_zero_draft_commits_gdn_replay( - worker_cls: type[MTPWorker] | type[MTPEagleWorker], - use_gdn_replay: bool | None, - num_contexts: int, -) -> None: - worker = worker_cls(MTPDecodingConfig(max_draft_len=3)) - worker._execute_guided_decoder_if_present = mock.Mock() - worker._sample_tokens_for_batch = mock.Mock(return_value=torch.tensor([3, 6])) - handler = mock.Mock() - worker.register_auxiliary_state_handler(handler) - manager = SimpleNamespace(update_mamba_states=mock.Mock()) - if use_gdn_replay is not None: - manager.use_gdn_mtp_replay = use_gdn_replay - state_indices = torch.tensor([3, 1], dtype=torch.int32) - attn_metadata = SimpleNamespace( - num_seqs=2, - num_contexts=num_contexts, - kv_cache_manager=manager, - mamba_metadata=SimpleNamespace(state_indices=state_indices), - ) - result = worker._forward_impl( - input_ids=mock.sentinel.input_ids, - position_ids=mock.sentinel.position_ids, - hidden_states=mock.sentinel.hidden_states, - logits=torch.zeros(2, 8), - attn_metadata=attn_metadata, - spec_metadata=SimpleNamespace(runtime_draft_len=0), - draft_model=mock.sentinel.draft_model, - ) - - assert result["next_draft_tokens"].shape == (2, 0) - torch.testing.assert_close(result["new_tokens"], torch.tensor([[3], [6]], dtype=torch.int)) - torch.testing.assert_close(result["new_tokens_lens"], torch.ones(2, dtype=torch.int)) - if use_gdn_replay and num_contexts < 2: - manager.update_mamba_states.assert_called_once_with( - attn_metadata=attn_metadata, - num_accepted_tokens=result["new_tokens_lens"], - state_indices=state_indices, - ) - handler.commit_speculative_states.assert_called_once() - accepted, indices, contexts = handler.commit_speculative_states.call_args.args - assert accepted is result["new_tokens_lens"] - torch.testing.assert_close(indices, state_indices) - assert contexts == num_contexts - else: - manager.update_mamba_states.assert_not_called() - handler.commit_speculative_states.assert_not_called() - - def test_mtp_shared_kv_draft_inputs(): spec_config = MTPDecodingConfig( max_draft_len=3, diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp_zero_draft.py b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp_zero_draft.py new file mode 100644 index 000000000000..c100c8f00ac5 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp_zero_draft.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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 +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# 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. + +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch + +from tensorrt_llm._torch.speculative.eagle3 import MTPEagleWorker +from tensorrt_llm._torch.speculative.mtp import MTPWorker +from tensorrt_llm.llmapi import MTPDecodingConfig + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("worker_cls", [MTPWorker, MTPEagleWorker], ids=["mtp", "mtp-eagle"]) +@pytest.mark.parametrize( + "use_gdn_replay,num_contexts", + [(True, 1), (True, 2), (False, 0), (None, 0)], + ids=["gdn-mixed", "context-only", "other-hybrid", "attention"], +) +def test_mtp_zero_draft_commits_gdn_replay( + worker_cls: type[MTPWorker] | type[MTPEagleWorker], + use_gdn_replay: bool | None, + num_contexts: int, +) -> None: + worker = worker_cls(MTPDecodingConfig(max_draft_len=3)) + worker._execute_guided_decoder_if_present = mock.Mock() + worker._sample_tokens_for_batch = mock.Mock(return_value=torch.tensor([3, 6])) + handler = mock.Mock() + worker.register_auxiliary_state_handler(handler) + manager = SimpleNamespace(update_mamba_states=mock.Mock()) + if use_gdn_replay is not None: + manager.use_gdn_mtp_replay = use_gdn_replay + state_indices = torch.tensor([3, 1], dtype=torch.int32) + attn_metadata = SimpleNamespace( + num_seqs=2, + num_contexts=num_contexts, + kv_cache_manager=manager, + mamba_metadata=SimpleNamespace(state_indices=state_indices), + ) + result = worker._forward_impl( + input_ids=mock.sentinel.input_ids, + position_ids=mock.sentinel.position_ids, + hidden_states=mock.sentinel.hidden_states, + logits=torch.zeros(2, 8), + attn_metadata=attn_metadata, + spec_metadata=SimpleNamespace(runtime_draft_len=0), + draft_model=mock.sentinel.draft_model, + ) + + assert result["next_draft_tokens"].shape == (2, 0) + torch.testing.assert_close(result["new_tokens"], torch.tensor([[3], [6]], dtype=torch.int)) + torch.testing.assert_close(result["new_tokens_lens"], torch.ones(2, dtype=torch.int)) + if use_gdn_replay and num_contexts < 2: + manager.update_mamba_states.assert_called_once_with( + attn_metadata=attn_metadata, + num_accepted_tokens=result["new_tokens_lens"], + state_indices=state_indices, + ) + handler.commit_speculative_states.assert_called_once() + accepted, indices, contexts = handler.commit_speculative_states.call_args.args + assert accepted is result["new_tokens_lens"] + torch.testing.assert_close(indices, state_indices) + assert contexts == num_contexts + else: + manager.update_mamba_states.assert_not_called() + handler.commit_speculative_states.assert_not_called()