update

Author: preminstrel

index.html

@@ -327,11 +327,11 @@ <h2 class="title is-3"><img src="static/images/Telescope.png" style="height: 50p
   <section class="section" id="BibTeX">
     <div class="container is-max-desktop content">
       <h2 class="title">BibTeX</h2>
-      <pre><code>@article{sun2024TriForce,
-  title={TriForce: Lossless Acceleration of Long Sequence Generation with Hierarchical Speculative Decoding},
-  author={Sun, Hanshi and Chen, Zhuoming and Yang, Xinyu and Tian, Yuandong and Chen, Beidi},
-  journal={arXiv preprint arXiv:2404.11912},
-  year={2024}
+      <pre><code>@article{sun2024triforce,
+    title={Triforce: Lossless acceleration of long sequence generation with hierarchical speculative decoding},
+    author={Sun, Hanshi and Chen, Zhuoming and Yang, Xinyu and Tian, Yuandong and Chen, Beidi},
+    journal={arXiv preprint arXiv:2404.11912},
+    year={2024}
 }</code></pre>
     </div>
 </section>

models/modeling_llama_68m.py

@@ -156,7 +156,7 @@ def forward(
             query_states = query_states.transpose(1, 2)
             key_states = key_states.transpose(1, 2)
 
-            position_ids = torch.arange(graph_cache.real_budget-graph_cache.gamma-1, graph_cache.real_budget-graph_cache.gamma+gamma_offset, device=position_ids.device).unsqueeze(0)
+            position_ids = torch.arange(graph_cache.real_budget-graph_cache.gamma-3, graph_cache.real_budget-graph_cache.gamma+gamma_offset-2, device=position_ids.device).unsqueeze(0)
             query_states = apply_rotary_pos_emb_single(query_states, cos, sin, position_ids)
             key_position_ids = torch.arange(kv_seq_len, device=position_ids.device).unsqueeze(0)
             key_states = apply_rotary_pos_emb_single(key_states, cos, sin, key_position_ids)