STree: Speculative Tree Decoding for Hybrid State Space Models

We thanks the reviewer for their reviews. We address the concerns and questions raised below:

W1: See Q2.

W2: See Q1.

W3: See Q2.

Q1. Dynamic tree: We note that optimizing the tree structure is orthogonal to our work. Our work focuses on enabling efficient computation of token tree with SSMs. With a baseline static tree, we are already able to obtain runtime improvement, as shown in Table 2 in paper. We also note that the beam search tree we used in Table 1 is not static, as the expansion of the tree depends on the joint probability and will vary from sample to sample. Here, we additionally show the end-to-end result for drafting with a beam search tree. We use greedy decoding with the M = 3 and N = 4 (M, N explained in table 1). The results are as follow:

We can see that with a beam search tree, we are able to get a better acceptance length and thus a faster end-to-end generation speed. We believe that more advance tree construction technique would bring even more benefits.

Q2: improvement not obvious: At temperature=0, our STree is always better than Vanilla SD, though by a smaller margin. However, at temperature=1, our method achieves an average 12% gain, which is an obvious difference. Averaging over all 6 experiment, there is still a 7% gain, which is non-trivial. Meanwhile, we show the runtime for forward passes for bigger Mamba2 models. STree and Vanilla SD both pay an overhead over autoregressive runs to compute multiple tokens. As model size increase, the ratio between STree overhead and Vanilla SD overhead decreases from 1.48 to 1.19. This means that when we have a bigger model，it is much easier to achieve an end-to-end performance gain. Given that we already achieved end-to-end performance gain in 8B models as show in Table 2 in the paper, there is reason to believe that performance with bigger models will further improve with similar acceptance rate.

Q3: Generation Temperature=2,3,4: In practice, no previous works in speculative decoding has used generation with temperature beyond 1. Using too high a temperature can lead to high perplexity in generated output and therefore hinders downstream tasks. We have the ablation study for temperature between 0 and 1 in figure 4b and 4c in the paper. Here, we additionally provide experiment with even higher temperature=1.1 and 1.2. We can see that the acceptance length and hence generation speed continues to decrease with such high temperature.

Q4: More benchmarks: The three benchmarks we provide in Table 2 (MT_Bench, GSM8K, Human_Eval) is extensively used by previous works to benchmark speculative decoding performance with text completion as the task. We note that MMLU, HellaSwag, SWAG, BBH, BoolQ are all datasets on common sense reasoning and Natural Language inference. Some of them feature multiple choice questions, requiring only a single word as answer, which can lead to premature termination in generation. This could lead to abnormal token distribution if we continue to generate beyond the eos token, affecting the acceptance length of both the draft and target model. This is probably why these datasets are often not used in speculative decoding works. Nonetheless, we additionally provide BoolQ's reults here. We note that as we employ a lossless verification strategy with temperature=0, the results we generated will be identical to the result of the target model and we only focus on the runtime. We use the question in BoolQ as a prompt and let it generate for 1024 tokens to measure the runtime.

Q5: Larger GPUs: While we do not have access to H100/A100 GPUs, we provide results obtained using A6000 GPU, which is the best we have. The results are presented below. We note that inference for SSM and hybrid models are likely still memory bound and computation speed is not the bottleneck. Therefore, faster computation units from better GPUs should not invalidate our improvements.

Q5: Draft model Ablation: For MambaInLlama-8B model, there are no currently viable open-source drafting models that uses the same tokenizer. Therefore we distilled our own drafting model. If the reviewer would like to suggest any suitable model for drafting, we are happy do experiments. Here, we additionally provide an ablation study on our drafting model distilled to different steps.

We can see that as the distillation goes on, the quality of the drafting model gets better, which is reflected by the increase in acceptance length. The result also suggest that with better drafting model and technique, the end-to-end runtime of STree will become better.

L1: Novelty: To the best of our knowledge, we are the first to apply speculative tree decoding to SSMs and hybrid models, which makes our work novel. While tree based speculative decoding exists for transformer based models, it heavily relies on the transformer's property that computing a tree of tokens doesn't incur much overhead with the use of attention masks. Without this property, it is much slower to compute a tree of tokens in SSMs due to their recurrent state update. Our work addresses this critical challenge by improving the efficiency of computing a tree of tokens with SSMs, which is not done in previous works.

L2: See Q1

L3: See Q4

We thank the reviewer for recognizing the presentation of method and evaluation. We address the concerns raised by the reviewer below:

W1: Multiple Runs: Thank you for you suggestion. Results in table 1 and table 2 (temperature=0) is performed with greedy decoding so re-running the experiment for multiple runs will yield the same results modulo hardware fluctuations. For Table 2 (temperature=1), we additionally provide the average and standard deviation for 4 runs here. We will include that in the updated version of the paper.

W2: Better GPU: While we do not have access to H100/A100 GPUs, we provide results obtained using A6000 GPU, which is the best we have. The results are presented below. We note that inference for SSM and hybrid models are still memory bound and computation speed is not the bottleneck. Therefore, faster computation units from better GPUs will not invalidate our improvements.

Per the request of the reviewer, we acquire a server with a h100 GPU and obtain some results using our current implementation as it is on h100. The setting is the same as table 2 mentioned in the paper and the A6000 results provided in the rebuttal. We can see that without any further optimization for h100 GPUs, we are already seeing a end-to-end speedup improvement for STree against Vanilla SD. This shows the potential of STree on better hardware. We will include more result with H100 GPUs in the future version of our paper.

We thank the reviewer for recognizing the motivation and methods of our work. We hereby address the concerns and questions raised by the reviewer:

W1: reduced novelty: To the best of our knowledge, we are the first to apply speculative tree decoding to SSMs and hybrid models, which makes our work novel. While tree based speculative decoding exists for transformer based models, it heavily relies on the transformer's property that computing a tree of tokens doesn't incur much overhead with the use of attention masks. Without this property, it is much slower to compute a tree of tokens in SSMs due to their recurrent state update. Our work addresses this critical challenge by improving the efficiency of computing a tree of tokens with SSMs, which is not done in previous works.

W2: AR transformer clarification: Yes, the reviewer is right that AR transformers retain the key and value of past tokens in memory (a structure called KV cache) to avoid recomputing them again for future generation. With speculative decoding, certain tokens computed might not be verified, which then needs to be removed from the KV cache to avoid contaminating future predictions. This can be done easily with KV cache, since it is partitioned on token boundaries and removing unverified tokens is equivalent to zeroing out the corresponding KV cache entries.

W3: Code availability: We thank the reviewer for their interest in our work. Here we provide code segments where we apply the tree mask in the kernel to compute SSM output. We will publish the complete code base upon paper acceptance.

    ...
    # We don't need to iterate K since tree scan sequence length is probably short. 
    for k in range(0, K_MAX, BLOCK_SIZE_K):
        cb = tl.load(cb_ptrs, mask=(offs_m[:, None] < chunk_size) & (offs_k[None, :] < chunk_size - k), other=0.0).to(tl.float32)
        dA_cs_k = tl.load(dA_cumsum_ptrs, mask=offs_k < chunk_size - k, other=0.0).to(tl.float32)
        # If there's seq_idx, we already set cb[i, j] = 0 for seq_idx[i] != seq_idx[j].
        # So we don't need masking wrt seq_idx here.
        cb *= tl.exp((dA_cs_m[:, None] - dA_cs_k[None, :]))
        dt_k = tl.load(dt_ptrs, mask=offs_k < chunk_size - k, other=0.0).to(tl.float32)
        cb *= dt_k

        # apply tree masking here
        tree_mask = tl.load(tree_mask_ptrs, mask=(offs_m[:, None] < chunk_size) & (offs_k[None, :] < chunk_size - k), other=0.0).to(tl.int1)
        cb = tl.where(tree_mask, cb, 0)

        cb = cb.to(x_ptr.dtype.element_ty)
        x = tl.load(x_ptrs, mask=(offs_k[:, None] < chunk_size_limit - k) & (offs_n[None, :] < hdim), other=0.0)
        acc += tl.dot(cb, x)
        cb_ptrs += BLOCK_SIZE_K * stride_cb_csize_k
        x_ptrs += BLOCK_SIZE_K * stride_x_seqlen
        dt_ptrs += BLOCK_SIZE_K * stride_dt_csize
        dA_cumsum_ptrs += BLOCK_SIZE_K * stride_dA_cs_csize
        tree_mask_ptrs += BLOCK_SIZE_K * stride_tm_csize2
        ...

Dear Reviewer 7Ar2,

Thanks for your review! The authors have provided their rebuttal. Please respond to the authors and update your review as appropriate.

Thank you!
AC

We thank the reviewer for recognizing the novelty in tree-base speculative decoding for SSMs. We address the concerns raised by the reviewer below:

W1: Larger LLMs: It is challenging to find open-source, pretrained hybrid models suitable for speculative decoding that are bigger than 30B parameters. We explored several options:

• Nemotron-H-47B: A hybrid model released by Nvidia in April. The smaller Nemotron-H-8B model can potentially act as the draft model. However, the model does not support generation with KV cache/SSM states, making it impossible for us to use.
• Jamba-mini-50B: A hybrid model from AI21. We cannot find any suitable draft model that is small in size and uses the same tokenizer as the target model.

If the reviewer has any specific open-source models that are suitable we would be happy to run experiments with them.

Absent end-to-end results on larger models, we analyze the runtime (in ms) of a forward pass through Mamba2 models using different methods below.

Theoretically, the effectiveness of speculative decoding is determined by the runtime overhead of forward pass using different methods as well as the acceptance length. Both STree and Vanilla SD are paying an overhead (having a longer runtime than autoregressive forward pass) to compute multiple tokens. From the table below, we can see that the overhead of an STree forward pass is decreasing from 2.04x to 1.26x as we scale the model from 2.7B to 23B. The ratio between STree and Vanilla SD is also shrinking (from 1.48 to 1.19) as model size increases. This means that, with a larger model, it is much easier to achieve an end-to-end performance gain.

Given that we already achieved end-to-end performance gain in 8B models as shown in Table 2 in the paper, there is a reason to believe that our performance with bigger models would be even better, assuming that the acceptance length is about the same.

W2: Dynamic tree: Our method does not rely on static trees and is applicable to any tree generation strategy. We note that optimizing the tree structure is orthogonal to our work. Our work focuses on enabling efficient computation of token tree with SSMs. With a baseline static tree, we are already able to obtain runtime improvement, as shown in Table 2 in the paper. We also note that the beam search tree we used in Table 1 is not static, as the expansion of the tree depends on the joint probability and will vary from sample to sample. Here, we additionally show the end-to-end result for drafting with a beam search tree. We use greedy decoding with M = 3 and N = 4 (M, N explained in table 1). The results are as follow:

We can see that with a beam-search tree, we are able to get a better acceptance length and thus a faster end-to-end generation speed. We believe that more advance tree construction technique would bring even more benefits.

We thank the reviewer for the reply. I am afraid there might be some misunderstanding with our additional experiments provided in the rebuttal, resulting in an opposite conclusion with the reviewer. As we go from left to right across columns in the first table provide above, we can see that as model size increases, the overhead of STree relative to an autoregressive run decreases from 2.04x to 1.26x. This is a drastic drop in overhead. Meanwhile, if we compare it with Vanilla SD, the ratio between the overhead of STree and Vanilla SD also drops from 1.48 to 1.19. This means that given the same acceptance length achieved with draft models, a 23B model is able to have a higher inference speed-up than a 7B model using STree.

We wonder if it is possible for the reviewer to elaborate more on why they think the speed up of STree on larger models seems to be limited by overhead. We would be happy to provide any further explanation.

Thanks for the authors' reply. The speedup performance of STree on larger models seems to be limited by its overhead. I choose to maintain my score.

Dear reviewer wh4v,

As the discussion period is coming to an end, we kindly ask you to review the Table above once more.

The conclusion drawn from the Table is that our method has a reduced overhead at larger model size. It is not limited by the overhead, and our method does scale well to larger model sizes.

We hope that this message gets to the Reviewer in time so that we can discuss this point of confusion.

Thank you!