fixed predicts dim

Signed-off-by: xintin <[email protected]>

Author: xintin

iree/turbine/kernel/wave/templates/speculative_decoding.py

@@ -120,10 +120,10 @@ def tree_speculative_sampling(
         sum_relu = tkw.sum(relu_diff, dim=VOCAB_SIZE)
         cdf = tkw.cumsum(relu_diff, dim=VOCAB_SIZE)
 
-        threshold_u = tkw.broadcast(
+        threshold_dist_u = tkw.broadcast(
             coin * sum_relu, target_shape=[BATCH_SIZE, NUM_DRAFT_TOKENS, VOCAB_SIZE]
         )
-        greater_than_u = cdf > threshold_u
+        greater_than_u = cdf > threshold_dist_u
         pad_token = tkl.Register[BATCH_SIZE, NUM_DRAFT_TOKENS, VOCAB_SIZE, tkl.i32](1e6)
         token_idx = tkl.Register[BATCH_SIZE, NUM_DRAFT_TOKENS, VOCAB_SIZE, tkl.i32](
             THREAD_0
@@ -146,25 +146,28 @@ def get_speculative_sampling_kernel(
     threshold_single: float,
     num_draft_tokens: int,
     vocab_size: int,
+    seq_len: int,
 ):
     CUR_INDEX = sympy.Symbol("CUR_INDEX")
     J = sympy.Symbol("J")
     BATCH_SIZE = tkl.sym.BATCH_SIZE
     NUM_DRAFT_TOKENS = tkl.sym.NUM_DRAFT_TOKENS
     VOCAB_SIZE = tkl.sym.VOCAB_SIZE
+    SEQ_LEN = tkl.sym.SEQ_LEN
     BLOCK_BATCH_SIZE = tkl.sym.BLOCK_BATCH_SIZE
     BLOCK_NUM_DRAFT_TOK = tkl.sym.BLOCK_NUM_DRAFT_TOK
     ADDRESS_SPACE = tkl.sym.ADDRESS_SPACE
-    ADDRESS_SPACE_0 = tkl.sym.ADDRESS_SPACE_0
+    GLOBAL_ADDRESS_SPACE = tkl.sym.GLOBAL_ADDRESS_SPACE
 
     hyperparams = {
         BLOCK_NUM_DRAFT_TOK: 1,
         NUM_DRAFT_TOKENS: num_draft_tokens,
         ADDRESS_SPACE: SHARED_ADDRESS_SPACE,
-        ADDRESS_SPACE_0: GLOBAL_ADDRESS_SPACE,
+        GLOBAL_ADDRESS_SPACE: GLOBAL_ADDRESS_SPACE,
         BATCH_SIZE: batch_size,
         BLOCK_BATCH_SIZE: 1,
         VOCAB_SIZE: vocab_size,
+        SEQ_LEN: seq_len,
     }
 
     dynamic_symbols = []
@@ -233,7 +236,7 @@ def get_speculative_sampling_kernel(
     write_mapping_1d = tkw.IndexMapping(
         num_iterators=2,
         inputs={BATCH_SIZE: i, NUM_DRAFT_TOKENS: j},
-        outputs={NUM_DRAFT_TOKENS: LAST_ACCEPTED_RETRIEVE_IDX},
+        outputs={SEQ_LEN: LAST_ACCEPTED_RETRIEVE_IDX},
     )
 
     write_mapping_3d = tkw.IndexMapping(
@@ -285,59 +288,63 @@ def write_with_zero_offset(x, y):
     accept_index_layout = tkl.MemoryLayout(shape=[batch_size, num_speculative_tokens])
     cur_prob_offset_vec_layout = tkl.MemoryLayout(shape=[batch_size, 1, 1])
     last_accepted_retrieve_idx_vec_layout = tkl.MemoryLayout(shape=[batch_size, 1, 1])
-    predict_layout = tkl.MemoryLayout(shape=[batch_size * num_draft_tokens])
+    predict_layout = tkl.MemoryLayout(shape=[seq_len])
 
     # Kernel.
     # =================================================================================
     @tkw.wave(constraints)
     def speculative_sampling(
         uniform_samples: tkl.Memory[
-            BATCH_SIZE, NUM_DRAFT_TOKENS, ADDRESS_SPACE_0, tkl.f32
+            BATCH_SIZE, NUM_DRAFT_TOKENS, GLOBAL_ADDRESS_SPACE, tkl.f32
         ],
         target_probs: tkl.Memory[
-            BATCH_SIZE, NUM_DRAFT_TOKENS, VOCAB_SIZE, ADDRESS_SPACE_0, tkl.f32
+            BATCH_SIZE, NUM_DRAFT_TOKENS, VOCAB_SIZE, GLOBAL_ADDRESS_SPACE, tkl.f32
         ],
         draft_probs: tkl.Memory[
-            BATCH_SIZE, NUM_DRAFT_TOKENS, VOCAB_SIZE, ADDRESS_SPACE_0, tkl.f32
+            BATCH_SIZE, NUM_DRAFT_TOKENS, VOCAB_SIZE, GLOBAL_ADDRESS_SPACE, tkl.f32
+        ],
+        candidates: tkl.Memory[
+            BATCH_SIZE, NUM_DRAFT_TOKENS, GLOBAL_ADDRESS_SPACE, tkl.i32
         ],
-        candidates: tkl.Memory[BATCH_SIZE, NUM_DRAFT_TOKENS, ADDRESS_SPACE_0, tkl.i32],
         retrieve_index: tkl.Memory[
-            BATCH_SIZE, NUM_DRAFT_TOKENS, ADDRESS_SPACE_0, tkl.i32
+            BATCH_SIZE, NUM_DRAFT_TOKENS, GLOBAL_ADDRESS_SPACE, tkl.i32
         ],
         retrieve_next_token: tkl.Memory[
-            BATCH_SIZE, NUM_DRAFT_TOKENS, ADDRESS_SPACE_0, tkl.i32
+            BATCH_SIZE, NUM_DRAFT_TOKENS, GLOBAL_ADDRESS_SPACE, tkl.i32
         ],
         retrieve_next_sibling: tkl.Memory[
-            BATCH_SIZE, NUM_DRAFT_TOKENS, ADDRESS_SPACE_0, tkl.i32
+            BATCH_SIZE, NUM_DRAFT_TOKENS, GLOBAL_ADDRESS_SPACE, tkl.i32
         ],
         # Outputs
-        predicts: tkl.Memory[
-            NUM_DRAFT_TOKENS, ADDRESS_SPACE_0, tkl.i32, predict_layout
-        ],
+        predicts: tkl.Memory[SEQ_LEN, GLOBAL_ADDRESS_SPACE, tkl.i32, predict_layout],
         accept_token_num: tkl.Memory[
             BATCH_SIZE,
             NUM_DRAFT_TOKENS,
             VOCAB_SIZE,
-            ADDRESS_SPACE_0,
+            GLOBAL_ADDRESS_SPACE,
             tkl.i32,
             accept_token_num_layout,
         ],
         accept_index: tkl.Memory[
-            BATCH_SIZE, NUM_DRAFT_TOKENS, ADDRESS_SPACE_0, tkl.i32, accept_index_layout
+            BATCH_SIZE,
+            NUM_DRAFT_TOKENS,
+            GLOBAL_ADDRESS_SPACE,
+            tkl.i32,
+            accept_index_layout,
         ],
         cur_prob_offset_vec: tkl.Memory[
             BATCH_SIZE,
             NUM_DRAFT_TOKENS,
             VOCAB_SIZE,
-            ADDRESS_SPACE_0,
+            GLOBAL_ADDRESS_SPACE,
             tkl.i32,
             cur_prob_offset_vec_layout,
         ],
         last_accepted_retrieve_idx_vec: tkl.Memory[
             BATCH_SIZE,
             NUM_DRAFT_TOKENS,
             VOCAB_SIZE,
-            ADDRESS_SPACE_0,
+            GLOBAL_ADDRESS_SPACE,
             tkl.i32,
             last_accepted_retrieve_idx_vec_layout,
         ],

lit_tests/kernel/wave/speculative_decoding.py

@@ -25,6 +25,7 @@ def test_speculative_decoding():
         threshold_acc=0.01,
         num_draft_tokens=6,
         vocab_size=20,
+        seq_len=12,
     )
 
     # Create the kernel with the hyperparameters

tests/kernel/wave/speculative_decode_test.py

@@ -61,6 +61,7 @@ def get_wave_speculative_sampling_kernel(
     threshold_single,
     num_draft_tokens,
     vocab_size,
+    seq_len,
 ):
     speculative_sampling, symbols, _, _ = get_speculative_sampling_kernel(
         batch_size,
@@ -69,6 +70,7 @@ def get_wave_speculative_sampling_kernel(
         threshold_single,
         num_draft_tokens,
         vocab_size,
+        seq_len,
     )
     symbols.update(get_default_scheduling_params())
 
@@ -188,6 +190,7 @@ def tree_speculative_sampling_target_only(
         threshold_single,
         num_draft_tokens,
         vocab_size,
+        seq_len,
     )
     sampling_kernel(
         uniform_samples,