[None][feat] Checkpointing variant of replay for MTP for mamba models

[hnover-nv]

Description

Summary of Changes
Quick summary, see ideas below for more detail.

We replace old per-step replay with a "checkpointing" replay, where we don't write out mamba state every step. Main changes:

• Families of tunable kernels for the actual state update, tuned for Nemo v3 Super on B200.
• Framework adjustments to support the new conditions on mamba cache management reflecting when we write state. Unfortunate switch to double-buffered old_x as Triton can't tell two derived memory addresses are actually the same, leading to a read hazard.
• New fields on MambaMetadata that track the number of requests in a batch that will want to write checkpoints this step, plus a tensor that stores per-request data for the batch, sorted into writes and nowrites.
• Removal of benchmark script. It go insanely complicated to support efficient tuning runs. A version may come back some day.

Timing results
These are the time from the start of conv1d to the end of the last replay kernel. Replay state update is broken up into 2 or 3 kernels, and conv1d PDLs into the first one. It's a kernel microbenchmark that tries to simulate hot and cold inputs based on what is coming out of in_proj in the mamba2_mixer and what is not.

batch	kernel	fp32/RN	fp16/SR	fp8/SR	int8/SR
1	replay	5.85	6.43	6.21	6.79
	old replay	7.00	7.32
		+16.4%	+12.2%
2	replay	6.80	6.93	7.00	7.22
	old replay	7.35	7.36
		+7.5%	+5.8%
4	replay	7.23	7.02	7.08	7.72
	old replay	7.82	7.69
		+7.5%	+8.6%
8	replay	8.09	7.76	7.56	8.29
	old replay	9.12	8.88
		+11.3%	+12.6%
16	replay	9.59	9.17	8.71	10.02
	old replay	12.67	11.78
		+24.3%	+22.2%
32	replay	13.39	12.50	10.83	13.03
	old replay	22.15	17.59
		+39.5%	+28.9%
64	replay	19.16	17.43	15.76	17.84
	old replay	26.69	25.20
		+28.2%	+30.8%
128	replay	29.65	25.50	23.67	27.58
	old replay	44.49	40.80
		+33.3%	+37.5%
256	replay	49.64	42.44	38.77	45.11
	old replay	84.85	76.35
		+41.5%	+44.4%
512	replay	88.46	72.10	68.62	78.06
	old replay	162.78	146.12
		+45.7%	+50.7%
1024	replay	168.68	131.89	123.81	142.73
	old replay	320.07	285.84
		+47.3%	+53.9%

Ideas

Building on #13453, here we take the replay idea a step further, to a "checkpointing" replay. This replaces the old replay version in our code.

This section walks through the ideas, skip below to the tl;dr summary.

To summarize "old" replay, we would save a mamba state every step. It just lagged what would be the true state by one, because we don't know how many to accept. Compared to the prior method of saving every possible state at each step, we saved substantially on memory traffic in the kernel and at step-end when we'd otherwise copy the "winning" state over. And needing to only materialize one state let us save further by skipping a lot of computation, and accelerating what remained with tensor cores.

But there's no reason we have to write even one state every step. We can have an ssm state checkpoint arbitrarily far back, store all the needed inputs since, and reconstruct the state. When we don't accept tokens, we just overwrite the unaccepted ones with the next step's new ones. Of course, in the limit, that just becomes linear-ish attention where you kept the whole KV cache.

How often should we checkpoint? One could try to balance the memory traffic of the replay window with the state. But we're actually not just memory bound. The easiest choice is that our tensor core usage requires us to pad all our dimension T (draft length + 1) operations to at least 16, so that's a natural history size. In the prior PR you can see that the old kernels' runtimes were quite flat with respect to T up to 16, and that had larger T for both state replay and output generation. Here we replay up to 16 from history but only output T.

The algorithm, then is every step, if the number of previously accepted tokens (PNAT)+ T > window size, we can't store the T new tokens so we need to write a checkpoint and restart the history buffer. Otherwise we have room to append.

Already this approach saves time. We don't issue instructions for saving the state, and for quantized, stochastically rounded, or block scaled SSM dstates, we also save on that pre-writing computation. Also, by quantizing less often, we accumulate less quantization error, which means we could quantize more aggressively. For example, perhaps int8 w/ stochastic rounding is viable if it only happens every 8 tokens, but not every token. Speculative decoding already gives some of that benefit, but it depends on acceptance rate. Here, we guarantee 16 - T tokens between checkpoints, and it can be as high as 16.

This approach also unlocks a number of additional speedups:

The Rectangle
Ignoring dA/dt scaling factors, the original replay algorithm is:
last_step_state = very_old_state + (old_X^t @ old_B)
outputs = new_C @ one_back_state^t + (new_C @ new_B^t) @ new_X. # bypass all the draft states, straight to output tokens. That C @ B^t we call the CB matrix, it's TxT square lower triangular, padded into 16x16.
Here old_* are last step's new_*, except we only take the first k of them, the acceptance rate.

All the matrix multiplies have some dimension at T, padded to 16.
But if we aren't writing, then just as we go straight from one_back_state to output tokens, we can actually go straight from very_old_state. If you just plug last_step_state in and expand, you get:
outputs = new_C @ very_old_state^t + ( new_C @ (old_B cat new_B)^t) @ (old_X + new_X)
Now the CB matrix is a rectangular lower triangular (PNAT + T) x T, and this fits in our window size (16) exactly when PNAT + T <= window_size. Oh look, our non-checkpointing condition!

The old state computation is one of the largest matmuls, so this is a substantial savings. But due to implementation details it is not always faster at small batch sizes.

constexpr-specialized Persistent Kernels
For a given request, sometimes we now write a checkpoint and sometimes don't. And when we don't we may even use a different algorithm. We can specialize nowrite algorithm choice as a constexpr we tune based on batch/dtype/rounding mode, but not write vs nowrite. Or can we? If we do specialize them, the compiler can take a lot of advantage. Nowrite should need less live state, even more so when it uses rectangle or writing uses stochastic rounding or quantizing. The compiler generates very different code. We also build on our prior PR where we tune parameters like M (slice of a head_dim done by a block), number of warps, etc. Separate kernels lets us tune these knobs for each case.

Ideally we'd launch one correctly sized grid for all the write requests, and one for the nowrite. But we're using cuda graphs in pytorch, so we can't vary that. Instead, we need to launch enough CTAs to cover all requests as write, and all as nowrite. Although the kernel can early out on requests that aren't its type, that overhead made specialized kernels uncompetative.

Except, we can use persistent kernels. As long as we have enough work to fill all the SMs a few times over, then there's no real overhead. And if we add some secondary data structures to present the requests partitioned by write/nowrite, the kernels can just loop over their portion. If we pack other needed data into that structure, it's not even added latency.

This makes specialized kernels a huge win at medium and large batches. At smaller batches a "dynamic"(i.e. unspecialized, non-constexpr) kernel still wins, probably because we can't easily optimally PDL in PyTorch between the precomputation kernel and the two main kernels. And we didn't have time to explore emulating this with atomics. But the dynamic kernel also wins by being persistent. Similarly constrained by time, the precompute kernel is non-persistent dynamic, whereas at least one of persistent or specialized would surely help.

Other tweaks

• Optionally support TMA loads and stores of mamba state, as a specializable knob per kernel and scenario.
• Use of tensor-ization in triton to replace some for loops, notably in the heads per block part of precompute kernel. This both removed some undefined behavior and improves performance.
• Sadly, we had to double-buffer old_x. In the main kernel, we read old x from the history buffer and new x from inputs from conv1d. Then on nowrite requests we append new x to old x, but on write requests it puts new x at the start of the history buffer, on top of old x. This is the same block and should be even threads doing the work, so should be safe, but triton does not realize they're the same, so we have to double buffer. In replay, where old_x and new_x are the same size, it was presumably more clear to the compiler. Or a latent bug, hard to know! If space is an issue, we could replace this with a circular buffer of size window + T, instead of today's 2 * window.

Tunings

We implemented a huge number of tunable knobs, and can pick which knobs to use based on batch size, dtype, and whether we are doing stochastic rounding. But how to pick the best? We have a mode choice of a single dynamic or a pair of specialized main kernels. Then each of these offers many tuning knobs. When we pick the specialized kernels, ther are a total of of 17 active knobs (more are in the code, but a few either don't compile or weren't tuned as they seemed dead), a mixture of boolean and integer. This resulted in over 61 billion possible combinations. While we'd hoped to be able to drop a lot of them, exhaustive searching of select subspaces showed none of the 17 were always dead. Obviously we can't exhaustively search the whole space for even one input cofiguration, let alone for all supported batch/dtype/rounding combinations. So we implemented a multi-process block coordinate descent optimization, with some tweaks to deal with cross-GPU performance drift. We also had to invest considerable effort into the actual benchmark runner: in-process CUPTI timing vs cuda events, multi-process for compilation and cupti analysis, optimization of cuda graph instantiation vs lauch overhead due to GIL contention.

These are honestly giant vibe coded messes, so we're removing the benchmarker from the repo and not attempting to submit the optimizer. No other human should have to review such things. But there's no arguing with the actual timing results.

Test Coverage

I've updated the replay unit tests to the new paradigm and added some more coverage scenarios, including better coverage of the rounding. Also added unit tests to cache manager for modified PNAT/double buffer behavior, and mamba metadata to test its new functionality of geneating the replay work items, needed for efficient persistent specialized kernel pairs.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

Summary by CodeRabbit

Release Notes

• New Features

  • Enhanced replay state update support for improved speculative decoding performance.
  • Added replay work indexing and metadata tracking for state management.

• Tests

  • Expanded test coverage for replay state updates and quantization-aware validation.
  • Added comprehensive testing for stochastic rounding correctness and multi-head scenarios.

• Chores

  • Removed benchmark utility script.

[hnover-nv]

/bot run

[tensorrt-cicd]

PR_Github #48673 [ run ] triggered by Bot. Commit: 32ebf3f Link to invocation

[tensorrt-cicd]

PR_Github #48673 [ run ] completed with state FAILURE. Commit: 32ebf3f
/LLM/main/L0_MergeRequest_PR pipeline #38453 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[hnover-nv]

/bot run

[tensorrt-cicd]

PR_Github #48787 [ run ] triggered by Bot. Commit: bab26f2 Link to invocation

[tensorrt-cicd]

PR_Github #48787 [ run ] completed with state SUCCESS. Commit: bab26f2
/LLM/main/L0_MergeRequest_PR pipeline #38551 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

[hnover-nv]

/bot run

[tensorrt-cicd]

PR_Github #49541 [ run ] triggered by Bot. Commit: 24f3dda Link to invocation

[tensorrt-cicd]

PR_Github #49541 [ run ] completed with state SUCCESS. Commit: 24f3dda
/LLM/main/L0_MergeRequest_PR pipeline #39171 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[hnover-nv]

/bot run

[tensorrt-cicd]

PR_Github #49775 [ run ] triggered by Bot. Commit: 5c44835 Link to invocation

[tensorrt-cicd]

PR_Github #49775 [ run ] completed with state SUCCESS. Commit: 5c44835
/LLM/main/L0_MergeRequest_PR pipeline #39373 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[hnover-nv]

/bot run

[hnover-nv]

/bot run

[tensorrt-cicd]

PR_Github #50880 [ run ] triggered by Bot. Commit: 3a55712 Link to invocation

[tensorrt-cicd]

PR_Github #50880 [ run ] completed with state SUCCESS. Commit: 3a55712
/LLM/main/L0_MergeRequest_PR pipeline #40347 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation