| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # Copyright (c) 2024, Tri Dao. |
| 5 | # Adapted from https://github.com/Dao-AILab/causal-conv1d/blob/main/causal_conv1d/causal_conv1d_interface.py |
| 6 | |
| 7 | |
| 8 | import numpy as np |
| 9 | import torch |
| 10 | |
| 11 | from vllm.platforms import current_platform |
| 12 | from vllm.triton_utils import tl, triton |
| 13 | from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID |
| 14 | |
| 15 | |
| 16 | @triton.jit(do_not_specialize_on_alignment=["num_cache_lines"]) |
| 17 | def _causal_conv1d_fwd_kernel( # continuous batching |
| 18 | # Pointers to matrices |
| 19 | x_ptr, # (dim, cu_seqlen) holding `batch` of actual sequences + padded sequences |
| 20 | w_ptr, # (dim, width) |
| 21 | bias_ptr, |
| 22 | initial_states_ptr, # conv_states_ptr |
| 23 | cache_indices_ptr, # (batch, n_blocks + padding) The second dimension contains |
| 24 | # the block indices relevant for each sequence |
| 25 | # plus potential 0-padding at the beginning and at the end |
| 26 | has_initial_states_ptr, |
| 27 | query_start_loc_ptr, |
| 28 | batch_ptr, |
| 29 | token_chunk_offset_ptr, |
| 30 | block_idx_first_scheduled_token, # (batch,) |
| 31 | block_idx_last_scheduled_token, # (batch,) |
| 32 | initial_state_idx, # (batch,) |
| 33 | num_computed_tokens, # (batch,) |
| 34 | o_ptr, # (dim, seqlen) - actually pointing to x_ptr |
| 35 | # Matrix dimensions |
| 36 | dim: tl.constexpr, |
| 37 | num_cache_lines, # added to support vLLM larger cache lines |
| 38 | # Strides |
| 39 | stride_x_dim: tl.constexpr, # stride to get to next feature-value, |
| 40 | stride_x_token: tl.int64, # stride to get to next token (same feature-index, same sequence-index) |
| 41 | stride_w_dim: tl.constexpr, # stride to get to next dim-axis value |
| 42 | stride_w_width: tl.constexpr, # stride to get to next width-axis value |
| 43 | stride_istate_seq: tl.constexpr, |
| 44 | stride_istate_dim: tl.constexpr, |
| 45 | stride_istate_token: tl.constexpr, |
| 46 | stride_cache_indices: tl.constexpr, |
| 47 | stride_o_dim: tl.constexpr, |
| 48 | stride_o_token: tl.int64, |
| 49 | stride_block_m: tl.constexpr, # Stride block to align divided by BLOCK_M |
| 50 | # others |
| 51 | pad_slot_id: tl.constexpr, |
| 52 | null_block_id: tl.constexpr, |
| 53 | # Meta-parameters |
| 54 | HAS_BIAS: tl.constexpr, |
| 55 | KERNEL_WIDTH: tl.constexpr, |
| 56 | SILU_ACTIVATION: tl.constexpr, |
| 57 | IS_APC_ENABLED: tl.constexpr, |
| 58 | HAS_NULL_BLOCK: tl.constexpr, |
| 59 | NP2_STATELEN: tl.constexpr, |
| 60 | BLOCK_M: tl.constexpr, |
| 61 | BLOCK_N: tl.constexpr, |
| 62 | launch_pdl: tl.constexpr, |
| 63 | ): |
| 64 | conv_states_ptr = initial_states_ptr |
| 65 | conv_state_indices_ptr = cache_indices_ptr |
| 66 | stride_conv_state_seq = stride_istate_seq |
| 67 | stride_conv_state_dim = stride_istate_dim |
| 68 | stride_conv_state_tok = stride_istate_token |
| 69 | state_len = ( |
| 70 | KERNEL_WIDTH - 1 |
| 71 | ) # can be passed via argument if it's not the same as this value |
| 72 | |
| 73 | if launch_pdl: |
| 74 | tl.extra.cuda.gdc_wait() |
| 75 | |
| 76 | # one program handles one chunk in a single sequence |
| 77 | # rather than mixing sequences - to make updating initial_states across sequences efficiently |
| 78 | |
| 79 | # single-sequence id |
| 80 | idx_seq = tl.load(batch_ptr + tl.program_id(0)).to(tl.int64) |
| 81 | chunk_offset = tl.load(token_chunk_offset_ptr + tl.program_id(0)) |
| 82 | |
| 83 | # BLOCK_N elements along the feature-dimension (channel) |
| 84 | idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) |
| 85 | |
| 86 | if idx_seq == pad_slot_id: |
| 87 | if launch_pdl: |
| 88 | tl.extra.cuda.gdc_launch_dependents() |
| 89 | return |
| 90 | |
| 91 | sequence_start_index = tl.load(query_start_loc_ptr + idx_seq) |
| 92 | sequence_end_index = tl.load(query_start_loc_ptr + idx_seq + 1) |
| 93 | # find the actual sequence length |
| 94 | seqlen = sequence_end_index - sequence_start_index |
| 95 | |
| 96 | B_size: tl.constexpr = stride_block_m * BLOCK_M |
| 97 | |
| 98 | if IS_APC_ENABLED: |
| 99 | # Handle the case if prefix caching is enabled. |
| 100 | # In particular, if prefix caching is enabled, the program write additional cache states to "cache_indices_ptr" |
| 101 | |
| 102 | # Get the length of the completed sequence so far and compute the offset. |
| 103 | current_first_index = tl.load(block_idx_first_scheduled_token + idx_seq) |
| 104 | current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq) |
| 105 | sequence_completed_index = tl.load(num_computed_tokens + idx_seq) |
| 106 | |
| 107 | # Compute the offset where the first stride_block_m-aligned first full block is |
| 108 | # Value in "token-space" |
| 109 | sequence_completed_offset_token = sequence_completed_index % B_size |
| 110 | seq_completed_offset = B_size - sequence_completed_offset_token |
| 111 | seq_end_offset = (seqlen - seq_completed_offset) % B_size |
| 112 | last_full_block_token_index = sequence_end_index - seq_end_offset |
| 113 | # If the sequence without the sequence_offset_index is stride_cache_chunk-aligned, then the last full chunk is the second-to-last one |
| 114 | if seq_end_offset == 0: |
| 115 | last_full_block_token_index = last_full_block_token_index - B_size |
| 116 | |
| 117 | # Get the number of blocks to be filled for the current sequence |
| 118 | # If n_block_to_fill = 0, then only the state at the sequence end is stored |
| 119 | n_block_to_fill = current_last_index - current_first_index |
| 120 | |
| 121 | # Get the index of the init block |
| 122 | conv_state_init_index = tl.load(initial_state_idx + idx_seq) |
| 123 | else: |
| 124 | n_block_to_fill = 0 |
| 125 | current_last_index = 0 |
| 126 | conv_state_init_index = 0 |
| 127 | current_first_index = 0 |
| 128 | last_full_block_token_index = 0 |
| 129 | |
| 130 | token_offset = BLOCK_M * chunk_offset |
| 131 | segment_len = min(BLOCK_M, seqlen - token_offset) |
| 132 | |
| 133 | # base of the sequence |
| 134 | x_base = ( |
| 135 | x_ptr + sequence_start_index * stride_x_token + idx_feats * stride_x_dim |
| 136 | ) # [BLOCK_N,] |
| 137 | |
| 138 | # cache_idx |
| 139 | conv_states_input_coord = tl.load( |
| 140 | conv_state_indices_ptr + idx_seq * stride_cache_indices + conv_state_init_index |
| 141 | ).to(tl.int64) |
| 142 | |
| 143 | if HAS_NULL_BLOCK: # noqa |
| 144 | if conv_states_input_coord == null_block_id: |
| 145 | # not processing as this is a null block (padding) |
| 146 | if launch_pdl: |
| 147 | tl.extra.cuda.gdc_launch_dependents() |
| 148 | return |
| 149 | conv_states_base = ( |
| 150 | conv_states_ptr |
| 151 | + (conv_states_input_coord * stride_conv_state_seq) |
| 152 | + (idx_feats * stride_conv_state_dim) |
| 153 | ) # [BLOCK_N,] |
| 154 | |
| 155 | w_base = w_ptr + (idx_feats * stride_w_dim) # [BLOCK_N,] |
| 156 | |
| 157 | # Does 2 things: |
| 158 | # 1. READ prior-block init-state data - [done by every Triton programs] |
| 159 | # 2. update conv_state with new data [only by the Triton program handles chunk_offset=0] |
| 160 | if chunk_offset == 0: |
| 161 | # read from conv_states |
| 162 | load_init_state = tl.load(has_initial_states_ptr + idx_seq).to(tl.int1) |
| 163 | if load_init_state: |
| 164 | # load from conv_states |
| 165 | prior_tokens = conv_states_base + (state_len - 1) * stride_conv_state_tok |
| 166 | mask_w = idx_feats < dim |
| 167 | if KERNEL_WIDTH == 2: |
| 168 | conv_states_ptrs = prior_tokens # [BLOCK_N] |
| 169 | col0 = tl.load(conv_states_ptrs, mask_w, 0.0) |
| 170 | if KERNEL_WIDTH == 3: |
| 171 | conv_states_ptrs = prior_tokens # [BLOCK_N] |
| 172 | col1 = tl.load(conv_states_ptrs, mask_w, 0.0) |
| 173 | conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N] |
| 174 | col0 = tl.load(conv_states_ptrs, mask_w, 0.0) |
| 175 | if KERNEL_WIDTH == 4: |
| 176 | conv_states_ptrs = prior_tokens # [BLOCK_N] |
| 177 | col2 = tl.load(conv_states_ptrs, mask_w, 0.0) |
| 178 | conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N] |
| 179 | col1 = tl.load(conv_states_ptrs, mask_w, 0.0) |
| 180 | conv_states_ptrs = prior_tokens - 2 * stride_conv_state_tok # [BLOCK_N] |
| 181 | col0 = tl.load(conv_states_ptrs, mask_w, 0.0) |
| 182 | if KERNEL_WIDTH == 5: |
| 183 | conv_states_ptrs = prior_tokens # [BLOCK_N] |
| 184 | col3 = tl.load(conv_states_ptrs, mask_w, 0.0) |
| 185 | conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N] |
| 186 | col2 = tl.load(conv_states_ptrs, mask_w, 0.0) |
| 187 | conv_states_ptrs = prior_tokens - 2 * stride_conv_state_tok # [BLOCK_N] |
| 188 | col1 = tl.load(conv_states_ptrs, mask_w, 0.0) |
| 189 | conv_states_ptrs = prior_tokens - 3 * stride_conv_state_tok # [BLOCK_N] |
| 190 | col0 = tl.load(conv_states_ptrs, mask_w, 0.0) |
| 191 | else: |
| 192 | # prior-tokens are zeros |
| 193 | if KERNEL_WIDTH >= 2: # STRATEGY1 |
| 194 | # first chunk and does not have prior-token, so just set to 0 |
| 195 | col0 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) |
| 196 | if KERNEL_WIDTH >= 3: # STRATEGY1 |
| 197 | col1 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) |
| 198 | if KERNEL_WIDTH >= 4: # STRATEGY1 |
| 199 | col2 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) |
| 200 | if KERNEL_WIDTH >= 5: # STRATEGY1 |
| 201 | col3 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) |
| 202 | |
| 203 | # STEP 2: |
| 204 | # here prepare data for updating conv_state |
| 205 | if ( |
| 206 | state_len <= seqlen |
| 207 | ): # SMALL_CACHE=True (only move part of 'x' into conv_state cache) |
| 208 | # just read from 'x' |
| 209 | # copy 'x' data to conv_state |
| 210 | # load only 'x' data (and set 0 before 'x' if seqlen < state_len) |
| 211 | idx_tokens_last = (seqlen - state_len) + tl.arange( |
| 212 | 0, NP2_STATELEN |
| 213 | ) # [BLOCK_M] |
| 214 | x_ptrs = ( |
| 215 | x_ptr |
| 216 | + ((sequence_start_index + idx_tokens_last) * stride_x_token)[:, None] |
| 217 | + (idx_feats * stride_x_dim)[None, :] |
| 218 | ) # [BLOCK_M,BLOCK_N,] |
| 219 | mask_x = ( |
| 220 | (idx_tokens_last >= 0)[:, None] |
| 221 | & (idx_tokens_last < seqlen)[:, None] |
| 222 | & (idx_feats < dim)[None, :] |
| 223 | ) # token-index # token-index # feature-index |
| 224 | loaded_x = tl.load(x_ptrs, mask_x, 0.0) |
| 225 | idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] |
| 226 | |
| 227 | # Compute the offset where the last block should be written in the conv_states |
| 228 | conv_states_output_coord = tl.load( |
| 229 | conv_state_indices_ptr |
| 230 | + idx_seq * stride_cache_indices |
| 231 | + current_last_index |
| 232 | ).to(tl.int64) |
| 233 | |
| 234 | conv_states_ptrs_target = ( |
| 235 | conv_states_ptr |
| 236 | + (conv_states_output_coord * stride_conv_state_seq) # Offset from seq |
| 237 | + (idx_feats * stride_conv_state_dim) |
| 238 | )[None, :] + ( # [BLOCK_N,] |
| 239 | idx_tokens_conv * stride_conv_state_tok |
| 240 | )[:, None] |
| 241 | |
| 242 | mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :] |
| 243 | tl.debug_barrier() # NOTE: use this due to bug in Triton compiler |
| 244 | tl.store(conv_states_ptrs_target, loaded_x, mask) |
| 245 | |
| 246 | else: |
| 247 | if load_init_state: |
| 248 | # update conv_state by shifting left, i.e. take last few cols from conv_state + cols from 'x' |
| 249 | idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] |
| 250 | |
| 251 | conv_states_ptrs_source = ( |
| 252 | conv_states_ptr |
| 253 | + (conv_states_input_coord * stride_conv_state_seq) |
| 254 | + (idx_feats * stride_conv_state_dim)[None, :] |
| 255 | + ((idx_tokens_conv + seqlen) * stride_conv_state_tok)[:, None] |
| 256 | ) # [BLOCK_M, BLOCK_N] |
| 257 | mask = ( |
| 258 | (conv_states_input_coord < num_cache_lines) |
| 259 | & ((idx_tokens_conv + seqlen) < state_len)[:, None] |
| 260 | & (idx_feats < dim)[None, :] |
| 261 | ) |
| 262 | conv_state = tl.load(conv_states_ptrs_source, mask, other=0.0) |
| 263 | |
| 264 | VAL = state_len - seqlen |
| 265 | |
| 266 | x_ptrs = ( |
| 267 | x_base[None, :] |
| 268 | + ((idx_tokens_conv - VAL) * stride_x_token)[:, None] |
| 269 | ) # [BLOCK_M, BLOCK_N] |
| 270 | |
| 271 | mask_x = ( |
| 272 | (idx_tokens_conv - VAL >= 0)[:, None] |
| 273 | & (idx_tokens_conv - VAL < seqlen)[:, None] |
| 274 | & (idx_feats < dim)[None, :] |
| 275 | ) # token-index # token-index # feature-index |
| 276 | loaded_x = tl.load(x_ptrs, mask_x, 0.0) |
| 277 | |
| 278 | tl.debug_barrier() # need this due to the bug in tl.where not enforcing this when data is the result of another tl.load |
| 279 | new_conv_state = tl.where( |
| 280 | mask, conv_state, loaded_x |
| 281 | ) # BUG in 'tl.where' which requires a barrier before this |
| 282 | conv_states_ptrs_target = ( |
| 283 | conv_states_base |
| 284 | + (idx_tokens_conv * stride_conv_state_tok)[:, None] |
| 285 | ) # [BLOCK_M, BLOCK_N] |
| 286 | mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[ |
| 287 | None, : |
| 288 | ] |
| 289 | tl.store(conv_states_ptrs_target, new_conv_state, mask) |
| 290 | else: # load_init_state == False |
| 291 | # update conv_state by shifting left, BUT |
| 292 | # set cols prior to 'x' as zeros + cols from 'x' |
| 293 | idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] |
| 294 | |
| 295 | VAL = state_len - seqlen |
| 296 | |
| 297 | x_ptrs = ( |
| 298 | x_base[None, :] |
| 299 | + ((idx_tokens_conv - VAL) * stride_x_token)[:, None] |
| 300 | ) # [BLOCK_M, BLOCK_N] |
| 301 | |
| 302 | mask_x = ( |
| 303 | (idx_tokens_conv - VAL >= 0)[:, None] |
| 304 | & (idx_tokens_conv - VAL < seqlen)[:, None] |
| 305 | & (idx_feats < dim)[None, :] |
| 306 | ) # token-index # token-index # feature-index |
| 307 | new_conv_state = tl.load(x_ptrs, mask_x, 0.0) |
| 308 | |
| 309 | conv_states_ptrs_target = ( |
| 310 | conv_states_base |
| 311 | + (idx_tokens_conv * stride_conv_state_tok)[:, None] |
| 312 | ) # [BLOCK_M, BLOCK_N] |
| 313 | mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[ |
| 314 | None, : |
| 315 | ] |
| 316 | tl.store(conv_states_ptrs_target, new_conv_state, mask) |
| 317 | |
| 318 | else: # chunk_offset > 0 |
| 319 | # read prior-token data from `x` |
| 320 | load_init_state = True |
| 321 | prior_tokens = x_base + (token_offset - 1) * stride_x_token |
| 322 | mask_w = idx_feats < dim |
| 323 | if KERNEL_WIDTH == 2: |
| 324 | conv_states_ptrs = prior_tokens # [BLOCK_N] |
| 325 | col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") |
| 326 | if KERNEL_WIDTH == 3: |
| 327 | conv_states_ptrs = prior_tokens # [BLOCK_N] |
| 328 | col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") |
| 329 | conv_states_ptrs = prior_tokens - 1 * stride_x_token # [BLOCK_N] |
| 330 | col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") |
| 331 | if KERNEL_WIDTH == 4: |
| 332 | conv_states_ptrs = prior_tokens # [BLOCK_N] |
| 333 | col2 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") |
| 334 | conv_states_ptrs = prior_tokens - 1 * stride_x_token # [BLOCK_N] |
| 335 | col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") |
| 336 | conv_states_ptrs = prior_tokens - 2 * stride_x_token # [BLOCK_N] |
| 337 | col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") |
| 338 | if KERNEL_WIDTH == 5: |
| 339 | # ruff: noqa: F841 |
| 340 | conv_states_ptrs = prior_tokens # [BLOCK_N] |
| 341 | col3 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") |
| 342 | conv_states_ptrs = prior_tokens - 1 * stride_x_token # [BLOCK_N] |
| 343 | col2 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") |
| 344 | conv_states_ptrs = prior_tokens - 2 * stride_x_token # [BLOCK_N] |
| 345 | col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") |
| 346 | conv_states_ptrs = prior_tokens - 3 * stride_x_token # [BLOCK_N] |
| 347 | col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") |
| 348 | |
| 349 | # Store intermediate states aligned with stride_block_m |
| 350 | # The additional states are cached starting from the last stride_block_m. |
| 351 | # For example: |
| 352 | # If n_block_to_fill = 0, then only the state at the sequence end is cached and the process below is not involved. |
| 353 | # If n_block_to_fill > 0, then the states at the sequence end and at the n_block_to_fill-last |
| 354 | # stride_block_m are cached. |
| 355 | # For example chunk_offset = n_block_to_fill stores the state at last_full_block |
| 356 | if (chunk_offset - 1) < n_block_to_fill: |
| 357 | # Store the states at the chunk boundaries from the start of the sequence |
| 358 | idx_tokens_last = ( |
| 359 | last_full_block_token_index |
| 360 | - (n_block_to_fill - chunk_offset) * B_size |
| 361 | - state_len |
| 362 | ) + tl.arange(0, NP2_STATELEN) # [BLOCK_M] |
| 363 | x_ptrs = ( |
| 364 | x_ptr |
| 365 | + (idx_tokens_last * stride_x_token)[:, None] |
| 366 | + (idx_feats * stride_x_dim)[None, :] |
| 367 | ) # [BLOCK_M,BLOCK_N,] |
| 368 | |
| 369 | mask_x = (idx_tokens_last >= 0)[:, None] & (idx_feats < dim)[ |
| 370 | None, : |
| 371 | ] # token-index # token-index # feature-index |
| 372 | loaded_x = tl.load(x_ptrs, mask_x, 0.0) |
| 373 | idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] |
| 374 | |
| 375 | # cache_idx |
| 376 | conv_states_output_coord = tl.load( |
| 377 | conv_state_indices_ptr |
| 378 | + idx_seq * stride_cache_indices |
| 379 | + current_first_index |
| 380 | + (chunk_offset - 1) |
| 381 | ).to(tl.int64) |
| 382 | |
| 383 | conv_states_ptrs_target = ( |
| 384 | conv_states_ptr |
| 385 | + (conv_states_output_coord * stride_conv_state_seq) # Offset from seq |
| 386 | + (idx_feats * stride_conv_state_dim) |
| 387 | )[None, :] + ( # [BLOCK_N,] |
| 388 | idx_tokens_conv * stride_conv_state_tok |
| 389 | )[:, None] |
| 390 | |
| 391 | mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :] |
| 392 | tl.debug_barrier() # NOTE: use this due to bug in Triton compiler |
| 393 | tl.store(conv_states_ptrs_target, loaded_x, mask) |
| 394 | |
| 395 | if HAS_BIAS: |
| 396 | bias = bias_ptr + idx_feats |
| 397 | mask_bias = idx_feats < dim |
| 398 | acc_preload = tl.load(bias, mask=mask_bias, other=0.0).to( |
| 399 | tl.float32 |
| 400 | ) # [BLOCK_N] |
| 401 | else: |
| 402 | acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32) |
| 403 | |
| 404 | x_base_1d = x_base + token_offset * stride_x_token # starting of chunk |
| 405 | |
| 406 | # PRE-LOAD WEIGHTS |
| 407 | mask_w = idx_feats < dim |
| 408 | if KERNEL_WIDTH >= 2: |
| 409 | w_ptrs = w_base + (0 * stride_w_width) # [BLOCK_N] tensor |
| 410 | w_col0 = tl.load(w_ptrs, mask_w, other=0.0) |
| 411 | w_ptrs = w_base + (1 * stride_w_width) # [BLOCK_N] tensor |
| 412 | w_col1 = tl.load(w_ptrs, mask_w, other=0.0) |
| 413 | if KERNEL_WIDTH >= 3: |
| 414 | w_ptrs = w_base + (2 * stride_w_width) # [BLOCK_N] tensor |
| 415 | w_col2 = tl.load(w_ptrs, mask_w, other=0.0) |
| 416 | if KERNEL_WIDTH >= 4: |
| 417 | w_ptrs = w_base + (3 * stride_w_width) # [BLOCK_N] tensor |
| 418 | w_col3 = tl.load(w_ptrs, mask_w, other=0.0) |
| 419 | mask_x_1d = idx_feats < dim |
| 420 | |
| 421 | if launch_pdl: |
| 422 | tl.extra.cuda.gdc_launch_dependents() |
| 423 | |
| 424 | for idx_token in range(segment_len): |
| 425 | acc = acc_preload |
| 426 | |
| 427 | matrix_w = w_col0 |
| 428 | matrix_x = col0 |
| 429 | for j in tl.static_range(KERNEL_WIDTH): |
| 430 | if KERNEL_WIDTH == 2: |
| 431 | if j == 1: # KERNEL_WIDTH-1: |
| 432 | matrix_w = w_col1 |
| 433 | x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] |
| 434 | matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) |
| 435 | elif KERNEL_WIDTH == 3: |
| 436 | if j == 1: |
| 437 | matrix_w = w_col1 |
| 438 | matrix_x = col1 |
| 439 | elif j == 2: |
| 440 | matrix_w = w_col2 |
| 441 | x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] |
| 442 | matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) |
| 443 | elif KERNEL_WIDTH == 4: |
| 444 | if j == 1: |
| 445 | matrix_w = w_col1 |
| 446 | matrix_x = col1 |
| 447 | elif j == 2: |
| 448 | matrix_w = w_col2 |
| 449 | matrix_x = col2 |
| 450 | elif j == 3: |
| 451 | matrix_w = w_col3 |
| 452 | x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] |
| 453 | matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) |
| 454 | |
| 455 | acc += matrix_x * matrix_w # [BLOCK_N] |
| 456 | |
| 457 | if KERNEL_WIDTH == 2: |
| 458 | col0 = matrix_x |
| 459 | elif KERNEL_WIDTH == 3: |
| 460 | col0 = col1 |
| 461 | col1 = matrix_x |
| 462 | elif KERNEL_WIDTH == 4: |
| 463 | col0 = col1 |
| 464 | col1 = col2 |
| 465 | col2 = matrix_x |
| 466 | |
| 467 | if SILU_ACTIVATION: |
| 468 | acc = acc / (1 + tl.exp(-acc)) |
| 469 | mask_1d = (idx_token < segment_len) & ( |
| 470 | idx_feats < dim |
| 471 | ) # token-index # feature-index |
| 472 | o_ptrs = ( |
| 473 | o_ptr |
| 474 | + (sequence_start_index + token_offset + idx_token) * stride_o_token |
| 475 | + (idx_feats * stride_o_dim) |
| 476 | ) |
| 477 | |
| 478 | tl.store(o_ptrs, acc, mask=mask_1d) |
| 479 | |
| 480 | |
| 481 | def causal_conv1d_fn( |
| 482 | x: torch.Tensor, |
| 483 | weight: torch.Tensor, |
| 484 | bias: torch.Tensor | None, |
| 485 | conv_states: torch.Tensor, |
| 486 | query_start_loc: torch.Tensor, |
| 487 | cache_indices: torch.Tensor | None = None, |
| 488 | has_initial_state: torch.Tensor | None = None, |
| 489 | activation: str | None = "silu", |
| 490 | pad_slot_id: int = PAD_SLOT_ID, |
| 491 | null_block_id: int = NULL_BLOCK_ID, |
| 492 | block_idx_first_scheduled_token: torch.Tensor | None = None, |
| 493 | block_idx_last_scheduled_token: torch.Tensor | None = None, |
| 494 | initial_state_idx: torch.Tensor | None = None, |
| 495 | num_computed_tokens: torch.Tensor | None = None, |
| 496 | block_size_to_align=0, |
| 497 | metadata=None, |
| 498 | validate_data=False, |
| 499 | ): |
| 500 | """Support varlen + continuous batching when x is 2D tensor. |
| 501 | |
| 502 | x: (dim,cu_seq_len) |
| 503 | cu_seq_len = total tokens of all seqs in that batch |
| 504 | sequences are concatenated from left to right for varlen |
| 505 | weight: (dim, width) |
| 506 | conv_states: (...,dim,width - 1) itype |
| 507 | updated inplace if cache_indices are not provided |
| 508 | [it use `cache_indices` to get the index to the cache of conv_state for that sequence |
| 509 | |
| 510 | conv_state[cache_indices[i]] for seq-i - to be used as initial_state when has_initial_state[i] = True |
| 511 | and after that conv_state[cache_indices[i]] need to be shift-left and updated with values from 'x' |
| 512 | ] |
| 513 | query_start_loc: (batch + 1) int32 |
| 514 | The cumulative sequence lengths of the sequences in |
| 515 | the batch, used to index into sequence. prepended by 0. |
| 516 | if |
| 517 | x = [5, 1, 1, 1] <- continuous batching (batch=4) |
| 518 | then |
| 519 | query_start_loc = [0, 5, 6, 7, 8] <- the starting index of the next sequence; while the last value is |
| 520 | the ending index of the last sequence |
| 521 | [length(query_start_loc)-1 == batch] |
| 522 | for example: query_start_loc = torch.Tensor([0,10,16,17]), |
| 523 | x.shape=(dim,17) |
| 524 | cache_indices: (batch) int32 |
| 525 | indicates the corresponding state index, |
| 526 | like so: conv_state = conv_states[cache_indices[batch_id]] |
| 527 | has_initial_state: (batch) bool |
| 528 | indicates whether should the kernel take the current state as initial |
| 529 | state for the calculations |
| 530 | [single boolean for each sequence in the batch: True or False] |
| 531 | bias: (dim,) |
| 532 | activation: either None or "silu" or "swish" or True |
| 533 | pad_slot_id: int |
| 534 | if cache_indices is passed, lets the kernel identify padded |
| 535 | entries that will not be processed, |
| 536 | for example: cache_indices = [pad_slot_id, 1, 20, pad_slot_id] |
| 537 | in this case, the kernel will not process entries at |
| 538 | indices 0 and 3 |
| 539 | block_idx_first_scheduled_token: (batch,), dtype int32 |
| 540 | The pointer into cache_indices, where the first cache block to be filled is located. |
| 541 | block_idx_last_scheduled_token: (batch,), dtype int32 |
| 542 | The pointer into cache_indices, where the last cache block to be filled is located. |
| 543 | initial_state_idx: (batch,), dtype int32 |
| 544 | The pointer into cache_indices, where the cache block containing the initial state is located. |
| 545 | num_computed_tokens: (batch,), dtype int32 |
| 546 | The number of tokens already completed for each sequence |
| 547 | block_size_to_align: int |
| 548 | The block size to align the cached states to |
| 549 | out: same shape as `x` |
| 550 | """ |
| 551 | if isinstance(activation, bool) and activation: |
| 552 | activation = "silu" |
| 553 | |
| 554 | args = None |
| 555 | # Store original dtype to cast back at the end |
| 556 | original_x_dtype = x.dtype |
| 557 | x = x.to(conv_states.dtype) |
| 558 | out = torch.empty_like(x) |
| 559 | if metadata is not None: |
| 560 | nums_dict = metadata.nums_dict |
| 561 | args = nums_dict |
| 562 | batch_ptr = metadata.batch_ptr |
| 563 | token_chunk_offset_ptr = metadata.token_chunk_offset_ptr |
| 564 | else: |
| 565 | seqlens = query_start_loc.diff().to("cpu") |
| 566 | args = seqlens |
| 567 | MAX_NUM_PROGRAMS = 1024 |
| 568 | |
| 569 | batch_ptr = torch.full( |
| 570 | (MAX_NUM_PROGRAMS,), PAD_SLOT_ID, dtype=torch.int32, device=x.device |
| 571 | ) # tracking which seq-idx the Triton program is handling |
| 572 | token_chunk_offset_ptr = torch.full( |
| 573 | (MAX_NUM_PROGRAMS,), PAD_SLOT_ID, dtype=torch.int32, device=x.device |
| 574 | ) # tracking BLOCK_M-based index in the sequence the Triton program is handling |
| 575 | |
| 576 | is_channel_last = (x.stride(0) == 1) & (x.stride(1) > 1) |
| 577 | dim, cu_seqlen = x.shape |
| 578 | _, width = weight.shape |
| 579 | state_len = width - 1 |
| 580 | np2_statelen = triton.next_power_of_2(state_len) |
| 581 | |
| 582 | padded_batch = query_start_loc.size(0) - 1 |
| 583 | stride_x_dim = x.stride(0) |
| 584 | stride_x_token = x.stride(1) |
| 585 | stride_w_dim = weight.stride(0) |
| 586 | stride_w_width = weight.stride(1) |
| 587 | stride_istate_seq = 0 |
| 588 | stride_istate_dim = 0 |
| 589 | stride_istate_token = 0 |
| 590 | num_cache_lines = 0 |
| 591 | BLOCK_M = 8 |
| 592 | if conv_states is not None: |
| 593 | # extensions to support vLLM: |
| 594 | # 1. conv_states is used to replaced initial_states |
| 595 | # 2. conv_states serve as a cache with num cache lines can be larger than batch size |
| 596 | # 3. mapping from sequence x[idx] to a cache line at index as specified via cache_indices[idx] |
| 597 | # 4. computation can be skipped if cache_indices[idx] == pad_slot_id |
| 598 | num_cache_lines = conv_states.size(0) |
| 599 | assert ( |
| 600 | num_cache_lines == conv_states.shape[0] |
| 601 | and dim == conv_states.shape[1] |
| 602 | and width - 1 <= conv_states.shape[2] |
| 603 | ) |
| 604 | stride_istate_seq = conv_states.stride(0) |
| 605 | stride_istate_dim = conv_states.stride(1) |
| 606 | stride_istate_token = conv_states.stride(2) |
| 607 | if out.dim() == 2: |
| 608 | stride_o_dim = out.stride(0) |
| 609 | stride_o_token = out.stride(1) |
| 610 | else: |
| 611 | stride_o_dim = out.stride(1) |
| 612 | stride_o_token = out.stride(2) |
| 613 | stride_cache_indices = cache_indices.stride(0) if cache_indices is not None else 0 |
| 614 | |
| 615 | if validate_data: |
| 616 | assert x.dim() == 2 |
| 617 | assert query_start_loc is not None |
| 618 | assert query_start_loc.dim() == 1 |
| 619 | assert x.stride(0) == 1 or x.stride(1) == 1 |
| 620 | if bias is not None: |
| 621 | assert bias.dim() == 1 |
| 622 | assert dim == bias.size(0) |
| 623 | if cache_indices is not None: |
| 624 | assert cache_indices.dim() == 1 |
| 625 | assert padded_batch == cache_indices.size(0) |
| 626 | if has_initial_state is not None: |
| 627 | assert has_initial_state.size() == (padded_batch,) |
| 628 | assert conv_states is not None, ( |
| 629 | "ERROR: `has_initial_state` is used, which needs also `conv_states`" |
| 630 | ) |
| 631 | assert weight.stride(1) == 1 |
| 632 | assert (dim, width) == weight.shape |
| 633 | assert is_channel_last, "Need to run in channel-last layout" |
| 634 | if block_size_to_align is not None and block_size_to_align > 0: |
| 635 | assert (block_size_to_align % BLOCK_M) == 0, ( |
| 636 | "The mamba block size needs to be divisible by the BLOCK_M" |
| 637 | ) |
| 638 | else: |
| 639 | block_size_to_align = BLOCK_M |
| 640 | |
| 641 | if metadata is None: |
| 642 | |
| 643 | def num_program(META, seqlens): |
| 644 | tot = 0 |
| 645 | |
| 646 | mlist = [] |
| 647 | offsetlist = [] # type: ignore |
| 648 | |
| 649 | nums = -(-seqlens // META["BLOCK_M"]) |
| 650 | |
| 651 | tot = nums.sum().item() |
| 652 | mlist = np.repeat(np.arange(len(nums)), nums) |
| 653 | for idx, num in enumerate(nums): |
| 654 | offsetlist.extend( |
| 655 | range(num) |
| 656 | ) # chunk-idx if a sequence is split into multiple chunks |
| 657 | |
| 658 | if META["batch_ptr"].nelement() < len(mlist): |
| 659 | newlen = len(mlist) + 1 |
| 660 | META["batch_ptr"].resize_(newlen).fill_(PAD_SLOT_ID) |
| 661 | META["token_chunk_offset_ptr"].resize_(newlen).fill_(PAD_SLOT_ID) |
| 662 | |
| 663 | if META["batch_ptr"].nelement() >= len(mlist): |
| 664 | META["batch_ptr"][0 : len(mlist)].copy_( |
| 665 | torch.from_numpy(np.array(mlist)) |
| 666 | ) |
| 667 | META["token_chunk_offset_ptr"][0 : len(mlist)].copy_( |
| 668 | torch.from_numpy(np.array(offsetlist)) |
| 669 | ) |
| 670 | |
| 671 | META["batch_ptr"] = META["batch_ptr"].to(META["x_ptr"].device) |
| 672 | META["token_chunk_offset_ptr"] = META["token_chunk_offset_ptr"].to( |
| 673 | META["x_ptr"].device |
| 674 | ) |
| 675 | return tot |
| 676 | else: |
| 677 | |
| 678 | def num_program(META, nums_dict): |
| 679 | tot = nums_dict[META["BLOCK_M"]]["tot"] |
| 680 | |
| 681 | mlist = nums_dict[META["BLOCK_M"]]["mlist"] |
| 682 | mlist_len = nums_dict[META["BLOCK_M"]]["mlist_len"] |
| 683 | |
| 684 | offsetlist = nums_dict[META["BLOCK_M"]]["offsetlist"] |
| 685 | |
| 686 | if nums_dict[META["BLOCK_M"]]["batch_ptr"] is not None: |
| 687 | META["batch_ptr"] = nums_dict[META["BLOCK_M"]]["batch_ptr"] |
| 688 | META["token_chunk_offset_ptr"] = nums_dict[META["BLOCK_M"]][ |
| 689 | "token_chunk_offset_ptr" |
| 690 | ] |
| 691 | else: |
| 692 | if META["batch_ptr"].nelement() < mlist_len: |
| 693 | newlen = mlist_len + 1 |
| 694 | META["batch_ptr"].resize_(newlen).fill_(PAD_SLOT_ID) |
| 695 | META["token_chunk_offset_ptr"].resize_(newlen).fill_(PAD_SLOT_ID) |
| 696 | |
| 697 | if META["batch_ptr"].nelement() >= mlist_len: |
| 698 | META["batch_ptr"][0:mlist_len].copy_(mlist) |
| 699 | META["token_chunk_offset_ptr"][0:mlist_len].copy_(offsetlist) |
| 700 | return tot |
| 701 | |
| 702 | def grid(META): |
| 703 | return ( |
| 704 | num_program(META, args), |
| 705 | triton.cdiv(dim, META["BLOCK_N"]), |
| 706 | ) |
| 707 | |
| 708 | if batch_ptr.device != x.device: |
| 709 | batch_ptr = batch_ptr.to(x.device) |
| 710 | token_chunk_offset_ptr = token_chunk_offset_ptr.to(x.device) |
| 711 | |
| 712 | _causal_conv1d_fwd_kernel[grid]( |
| 713 | # Pointers to matrices |
| 714 | x, |
| 715 | weight, |
| 716 | bias, |
| 717 | conv_states, |
| 718 | cache_indices, |
| 719 | has_initial_state, |
| 720 | query_start_loc, |
| 721 | batch_ptr, |
| 722 | token_chunk_offset_ptr, |
| 723 | block_idx_first_scheduled_token, |
| 724 | block_idx_last_scheduled_token, |
| 725 | initial_state_idx, |
| 726 | num_computed_tokens, |
| 727 | out, |
| 728 | # Matrix dimensions |
| 729 | dim, |
| 730 | num_cache_lines, |
| 731 | # stride |
| 732 | stride_x_dim, |
| 733 | stride_x_token, |
| 734 | stride_w_dim, |
| 735 | stride_w_width, |
| 736 | stride_istate_seq, |
| 737 | stride_istate_dim, |
| 738 | stride_istate_token, |
| 739 | stride_cache_indices, |
| 740 | stride_o_dim, |
| 741 | stride_o_token, |
| 742 | block_size_to_align // BLOCK_M, |
| 743 | # others |
| 744 | pad_slot_id, |
| 745 | null_block_id, |
| 746 | # META |
| 747 | HAS_BIAS=bias is not None, |
| 748 | KERNEL_WIDTH=width, |
| 749 | SILU_ACTIVATION=activation in ["silu", "swish"], |
| 750 | IS_APC_ENABLED=block_idx_last_scheduled_token is not None, |
| 751 | HAS_NULL_BLOCK=null_block_id is not None, |
| 752 | NP2_STATELEN=np2_statelen, |
| 753 | # launch_cooperative_grid=True |
| 754 | BLOCK_M=BLOCK_M, |
| 755 | BLOCK_N=256, |
| 756 | num_stages=2, |
| 757 | launch_pdl=current_platform.is_arch_support_pdl(), |
| 758 | ) |
| 759 | return out.to(original_x_dtype) |
| 760 | |
| 761 | |
| 762 | @triton.jit(do_not_specialize_on_alignment=["num_cache_lines"]) |
| 763 | def _causal_conv1d_update_kernel( |
| 764 | # Pointers to matrices |
| 765 | x_ptr, # (batch, dim, seqlen) |
| 766 | w_ptr, # (dim, width) |
| 767 | bias_ptr, |
| 768 | conv_state_ptr, |
| 769 | conv_state_indices_ptr, |
| 770 | num_accepted_tokens_ptr, |
| 771 | query_start_loc_ptr, # (batch + 1) |
| 772 | block_idx_last_scheduled_token, # (batch,) |
| 773 | initial_state_idx, # (batch,) |
| 774 | o_ptr, # (batch, dim, seqlen) |
| 775 | # Matrix dimensions |
| 776 | batch: int, |
| 777 | dim: tl.constexpr, |
| 778 | seqlen: tl.constexpr, |
| 779 | state_len: tl.constexpr, |
| 780 | num_cache_lines, # added to support vLLM larger cache lines |
| 781 | # Strides |
| 782 | stride_x_seq: tl.constexpr, |
| 783 | stride_x_dim: tl.constexpr, |
| 784 | stride_x_token: tl.int64, |
| 785 | stride_w_dim: tl.constexpr, |
| 786 | stride_w_width: tl.constexpr, |
| 787 | stride_conv_state_seq: tl.constexpr, |
| 788 | stride_conv_state_dim: tl.constexpr, |
| 789 | stride_conv_state_tok: tl.constexpr, |
| 790 | stride_state_indices: tl.constexpr, |
| 791 | stride_o_seq: tl.constexpr, |
| 792 | stride_o_dim: tl.constexpr, |
| 793 | stride_o_token: tl.int64, |
| 794 | # others |
| 795 | null_block_id: tl.constexpr, |
| 796 | # Meta-parameters |
| 797 | HAS_BIAS: tl.constexpr, |
| 798 | KERNEL_WIDTH: tl.constexpr, |
| 799 | SILU_ACTIVATION: tl.constexpr, |
| 800 | IS_VARLEN: tl.constexpr, |
| 801 | IS_APC_ENABLED: tl.constexpr, |
| 802 | IS_SPEC_DECODING: tl.constexpr, |
| 803 | NP2_STATELEN: tl.constexpr, |
| 804 | HAS_NULL_BLOCK: tl.constexpr, |
| 805 | BLOCK_N: tl.constexpr, |
| 806 | launch_pdl: tl.constexpr, |
| 807 | ): |
| 808 | if launch_pdl: |
| 809 | tl.extra.cuda.gdc_wait() |
| 810 | |
| 811 | # ruff: noqa: E501 |
| 812 | idx_seq = tl.program_id(0) |
| 813 | if idx_seq >= batch: |
| 814 | if launch_pdl: |
| 815 | tl.extra.cuda.gdc_launch_dependents() |
| 816 | return |
| 817 | |
| 818 | # [BLOCK_N,] elements along the feature-dimension (channel) |
| 819 | idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) |
| 820 | |
| 821 | if IS_APC_ENABLED: |
| 822 | # Get the state from the initial_state_idx |
| 823 | conv_state_init = tl.load(initial_state_idx + idx_seq) |
| 824 | current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq) |
| 825 | else: |
| 826 | conv_state_init = 0 |
| 827 | current_last_index = 0 |
| 828 | |
| 829 | # cache_idx |
| 830 | conv_states_input_coord = tl.load( |
| 831 | conv_state_indices_ptr + idx_seq * stride_state_indices + conv_state_init |
| 832 | ).to(tl.int64) |
| 833 | |
| 834 | if HAS_NULL_BLOCK: # noqa |
| 835 | if conv_states_input_coord == null_block_id: |
| 836 | # not processing as this is not the actual sequence |
| 837 | if launch_pdl: |
| 838 | tl.extra.cuda.gdc_launch_dependents() |
| 839 | return |
| 840 | |
| 841 | if IS_VARLEN: |
| 842 | query_start_index = tl.load(query_start_loc_ptr + idx_seq).to(tl.int64) |
| 843 | query_end_index = tl.load(query_start_loc_ptr + (idx_seq + 1)).to(tl.int64) |
| 844 | # revise state_len and seqlen |
| 845 | state_len = state_len - (seqlen - (query_end_index - query_start_index)) |
| 846 | seqlen = query_end_index - query_start_index |
| 847 | x_offset = query_start_index * stride_x_token |
| 848 | o_offset = query_start_index * stride_o_token |
| 849 | else: |
| 850 | query_start_index = idx_seq * seqlen |
| 851 | query_end_index = query_start_index + seqlen |
| 852 | x_offset = idx_seq * stride_x_seq |
| 853 | o_offset = idx_seq * stride_o_seq |
| 854 | |
| 855 | if query_start_index == query_end_index: |
| 856 | if launch_pdl: |
| 857 | tl.extra.cuda.gdc_launch_dependents() |
| 858 | return |
| 859 | |
| 860 | if IS_SPEC_DECODING: |
| 861 | # The rolling of conv state: |
| 862 | # |
| 863 | # Before forward, the conv_state is: |
| 864 | # [history1, history2, ..., historyM]. |
| 865 | # |
| 866 | # After forward, the conv_state becomes: |
| 867 | # [history2, ..., historyM, draft1, draft2, ..., draftN]. |
| 868 | # |
| 869 | # After acceptance, it becomes: |
| 870 | # |
| 871 | # - accept 1 tokens: [history2, ..., historyM, draft1] |
| 872 | # - accept 2 tokens: [history3, ..., historyM, draft1, draft2] |
| 873 | # - and so on. |
| 874 | conv_state_token_offset = ( |
| 875 | tl.load(num_accepted_tokens_ptr + idx_seq).to(tl.int64) - 1 |
| 876 | ) |
| 877 | else: |
| 878 | conv_state_token_offset = 0 |
| 879 | |
| 880 | # STEP 1: READ init_state data |
| 881 | conv_states_base = ( |
| 882 | conv_state_ptr |
| 883 | + (conv_states_input_coord * stride_conv_state_seq) |
| 884 | + (idx_feats * stride_conv_state_dim) |
| 885 | ) |
| 886 | mask_w = idx_feats < dim |
| 887 | |
| 888 | prior_tokens = conv_states_base + conv_state_token_offset * stride_conv_state_tok |
| 889 | if KERNEL_WIDTH >= 2: |
| 890 | conv_states_ptrs = prior_tokens # [BLOCK_N] |
| 891 | col0 = tl.load(conv_states_ptrs, mask_w, 0.0) |
| 892 | if KERNEL_WIDTH >= 3: |
| 893 | conv_states_ptrs = prior_tokens + 1 * stride_conv_state_tok # [BLOCK_N] |
| 894 | col1 = tl.load(conv_states_ptrs, mask_w, 0.0) |
| 895 | if KERNEL_WIDTH >= 4: |
| 896 | conv_states_ptrs = prior_tokens + 2 * stride_conv_state_tok # [BLOCK_N] |
| 897 | col2 = tl.load(conv_states_ptrs, mask_w, 0.0) |
| 898 | if KERNEL_WIDTH >= 5: |
| 899 | conv_states_ptrs = prior_tokens + 3 * stride_conv_state_tok # [BLOCK_N] |
| 900 | col3 = tl.load(conv_states_ptrs, mask_w, 0.0) |
| 901 | if KERNEL_WIDTH >= 6: |
| 902 | conv_states_ptrs = prior_tokens + 4 * stride_conv_state_tok # [BLOCK_N] |
| 903 | col4 = tl.load(conv_states_ptrs, mask_w, 0.0) |
| 904 | |
| 905 | # STEP 2: assume state_len > seqlen |
| 906 | idx_tokens = tl.arange(0, NP2_STATELEN) # [BLOCK_M] |
| 907 | |
| 908 | # With speculative decoding, the conv_state updates works in a sliding |
| 909 | # window manner, at each forward pass, the tokens are shift by 1, so we |
| 910 | # load since idx_tokens + 1. |
| 911 | conv_state_ptrs_source = ( |
| 912 | conv_state_ptr |
| 913 | + (conv_states_input_coord * stride_conv_state_seq) |
| 914 | + conv_state_token_offset * stride_conv_state_tok |
| 915 | + (idx_feats * stride_conv_state_dim)[None, :] |
| 916 | + ((idx_tokens + (1 if IS_SPEC_DECODING else seqlen)) * stride_conv_state_tok)[ |
| 917 | :, None |
| 918 | ] |
| 919 | ) # [BLOCK_M, BLOCK_N] |
| 920 | mask = ( |
| 921 | (conv_states_input_coord < num_cache_lines) |
| 922 | & ((idx_tokens + seqlen) < state_len)[:, None] |
| 923 | & (idx_feats < dim)[None, :] |
| 924 | ) |
| 925 | conv_state = tl.load(conv_state_ptrs_source, mask, other=0.0) |
| 926 | |
| 927 | VAL = state_len - seqlen |
| 928 | x_base = x_ptr + x_offset + (idx_feats * stride_x_dim) # [BLOCK_N] |
| 929 | |
| 930 | x_ptrs = ( |
| 931 | x_base[None, :] + ((idx_tokens - VAL) * stride_x_token)[:, None] |
| 932 | ) # [BLOCK_M, BLOCK_N] |
| 933 | |
| 934 | mask_x = ( |
| 935 | (idx_tokens - VAL >= 0)[:, None] |
| 936 | & (idx_tokens - VAL < seqlen)[:, None] |
| 937 | & (idx_feats < dim)[None, :] |
| 938 | ) # token-index # token-index # feature-index |
| 939 | loaded_x = tl.load(x_ptrs, mask_x, 0.0) |
| 940 | tl.debug_barrier() |
| 941 | |
| 942 | new_conv_state = tl.where(mask, conv_state, loaded_x) |
| 943 | |
| 944 | # Get the state from the initial_state_idx |
| 945 | # cache_idx |
| 946 | conv_states_offset = tl.load( |
| 947 | conv_state_indices_ptr + idx_seq * stride_state_indices + current_last_index |
| 948 | ).to(tl.int64) |
| 949 | conv_state_ptrs_target = ( |
| 950 | conv_state_ptr |
| 951 | + (conv_states_offset * stride_conv_state_seq) # Offset from seq |
| 952 | + (idx_feats * stride_conv_state_dim) |
| 953 | )[None, :] + ( # [BLOCK_N,] |
| 954 | idx_tokens * stride_conv_state_tok |
| 955 | )[:, None] |
| 956 | mask = (idx_tokens < state_len)[:, None] & (idx_feats < dim)[None, :] |
| 957 | tl.store(conv_state_ptrs_target, new_conv_state, mask) |
| 958 | |
| 959 | # STEP 3: init accumulator |
| 960 | if HAS_BIAS: |
| 961 | bias = bias_ptr + idx_feats |
| 962 | mask_bias = idx_feats < dim |
| 963 | acc_preload = tl.load(bias, mask=mask_bias, other=0.0).to( |
| 964 | tl.float32 |
| 965 | ) # [BLOCK_N] |
| 966 | else: |
| 967 | acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32) |
| 968 | |
| 969 | # STEP 4: |
| 970 | # PRE-LOAD WEIGHTS |
| 971 | # first kernel column, configured for weights to handle BLOCK_N features in range |
| 972 | w_base = w_ptr + (idx_feats * stride_w_dim) # [BLOCK_N,] |
| 973 | mask_w = idx_feats < dim |
| 974 | if KERNEL_WIDTH >= 2: |
| 975 | w_ptrs = w_base + (0 * stride_w_width) # [BLOCK_N] tensor |
| 976 | w_col0 = tl.load(w_ptrs, mask_w, other=0.0) |
| 977 | w_ptrs = w_base + (1 * stride_w_width) # [BLOCK_N] tensor |
| 978 | w_col1 = tl.load(w_ptrs, mask_w, other=0.0) |
| 979 | if KERNEL_WIDTH >= 3: |
| 980 | w_ptrs = w_base + (2 * stride_w_width) # [BLOCK_N] tensor |
| 981 | w_col2 = tl.load(w_ptrs, mask_w, other=0.0) |
| 982 | if KERNEL_WIDTH >= 4: |
| 983 | w_ptrs = w_base + (3 * stride_w_width) # [BLOCK_N] tensor |
| 984 | w_col3 = tl.load(w_ptrs, mask_w, other=0.0) |
| 985 | if KERNEL_WIDTH >= 5: |
| 986 | w_ptrs = w_base + (4 * stride_w_width) # [BLOCK_N] tensor |
| 987 | w_col4 = tl.load(w_ptrs, mask_w, other=0.0) |
| 988 | if KERNEL_WIDTH >= 6: |
| 989 | w_ptrs = w_base + (5 * stride_w_width) # [BLOCK_N] tensor |
| 990 | w_col5 = tl.load(w_ptrs, mask_w, other=0.0) |
| 991 | |
| 992 | x_base_1d = x_base # starting of chunk [BLOCK_N] |
| 993 | mask_x_1d = idx_feats < dim |
| 994 | |
| 995 | # STEP 5: compute each token |
| 996 | if launch_pdl: |
| 997 | tl.extra.cuda.gdc_launch_dependents() |
| 998 | |
| 999 | for idx_token in tl.range(seqlen): |
| 1000 | acc = acc_preload |
| 1001 | |
| 1002 | matrix_w = w_col0 |
| 1003 | matrix_x = col0 |
| 1004 | for j in tl.static_range(KERNEL_WIDTH): |
| 1005 | if KERNEL_WIDTH == 2: |
| 1006 | if j == 1: # KERNEL_WIDTH-1: |
| 1007 | matrix_w = w_col1 |
| 1008 | x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] |
| 1009 | matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) |
| 1010 | elif KERNEL_WIDTH == 3: |
| 1011 | if j == 1: |
| 1012 | matrix_w = w_col1 |
| 1013 | matrix_x = col1 |
| 1014 | elif j == 2: |
| 1015 | matrix_w = w_col2 |
| 1016 | x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] |
| 1017 | matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) |
| 1018 | elif KERNEL_WIDTH == 4: |
| 1019 | if j == 1: |
| 1020 | matrix_w = w_col1 |
| 1021 | matrix_x = col1 |
| 1022 | elif j == 2: |
| 1023 | matrix_w = w_col2 |
| 1024 | matrix_x = col2 |
| 1025 | elif j == 3: |
| 1026 | matrix_w = w_col3 |
| 1027 | x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] |
| 1028 | matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) |
| 1029 | elif KERNEL_WIDTH == 5: |
| 1030 | if j == 1: |
| 1031 | matrix_w = w_col1 |
| 1032 | matrix_x = col1 |
| 1033 | elif j == 2: |
| 1034 | matrix_w = w_col2 |
| 1035 | matrix_x = col2 |
| 1036 | elif j == 3: |
| 1037 | matrix_w = w_col3 |
| 1038 | matrix_x = col3 |
| 1039 | elif j == 4: |
| 1040 | matrix_w = w_col4 |
| 1041 | x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] |
| 1042 | matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) |
| 1043 | elif KERNEL_WIDTH == 6: |
| 1044 | if j == 1: |
| 1045 | matrix_w = w_col1 |
| 1046 | matrix_x = col1 |
| 1047 | elif j == 2: |
| 1048 | matrix_w = w_col2 |
| 1049 | matrix_x = col2 |
| 1050 | elif j == 3: |
| 1051 | matrix_w = w_col3 |
| 1052 | matrix_x = col3 |
| 1053 | elif j == 4: |
| 1054 | matrix_w = w_col4 |
| 1055 | matrix_x = col4 |
| 1056 | elif j == 5: |
| 1057 | matrix_w = w_col5 |
| 1058 | x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] |
| 1059 | matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) |
| 1060 | |
| 1061 | acc += matrix_x * matrix_w # [BLOCK_N] |
| 1062 | |
| 1063 | if KERNEL_WIDTH == 2: |
| 1064 | col0 = matrix_x |
| 1065 | elif KERNEL_WIDTH == 3: |
| 1066 | col0 = col1 |
| 1067 | col1 = matrix_x |
| 1068 | elif KERNEL_WIDTH == 4: |
| 1069 | col0 = col1 |
| 1070 | col1 = col2 |
| 1071 | col2 = matrix_x |
| 1072 | elif KERNEL_WIDTH == 5: |
| 1073 | col0 = col1 |
| 1074 | col1 = col2 |
| 1075 | col2 = col3 |
| 1076 | col3 = matrix_x |
| 1077 | elif KERNEL_WIDTH == 6: |
| 1078 | col0 = col1 |
| 1079 | col1 = col2 |
| 1080 | col2 = col3 |
| 1081 | col3 = col4 |
| 1082 | col4 = matrix_x |
| 1083 | |
| 1084 | if SILU_ACTIVATION: |
| 1085 | acc = acc / (1 + tl.exp(-acc)) |
| 1086 | mask_1d = (idx_token < seqlen) & ( |
| 1087 | idx_feats < dim |
| 1088 | ) # token-index # feature-index |
| 1089 | o_ptrs = ( |
| 1090 | o_ptr + o_offset + idx_token * stride_o_token + (idx_feats * stride_o_dim) |
| 1091 | ) |
| 1092 | |
| 1093 | tl.store(o_ptrs, acc, mask=mask_1d) |
| 1094 | |
| 1095 | |
| 1096 | def causal_conv1d_update( |
| 1097 | x: torch.Tensor, |
| 1098 | conv_state: torch.Tensor, |
| 1099 | weight: torch.Tensor, |
| 1100 | bias: torch.Tensor | None = None, |
| 1101 | activation: bool | str | None = None, |
| 1102 | conv_state_indices: torch.Tensor | None = None, |
| 1103 | num_accepted_tokens: torch.Tensor | None = None, |
| 1104 | query_start_loc: torch.Tensor | None = None, |
| 1105 | max_query_len: int = -1, |
| 1106 | null_block_id: int = NULL_BLOCK_ID, |
| 1107 | block_idx_last_scheduled_token: torch.Tensor | None = None, |
| 1108 | initial_state_idx: torch.Tensor | None = None, |
| 1109 | validate_data=False, |
| 1110 | out: torch.Tensor | None = None, |
| 1111 | ): |
| 1112 | """x: Input tensor which can take the following shapes: |
| 1113 | |
| 1114 | - `[batch, dim]` - single token prediction |
| 1115 | - `[batch, dim, seqlen]` - single or multiple tokens prediction |
| 1116 | - `[num_tokens, dim]` - continuous batching, where num_tokens is |
| 1117 | the total tokens of all sequences in that batch |
| 1118 | |
| 1119 | conv_state: (..., dim, state_len), where state_len >= width - 1 |
| 1120 | weight: (dim, width) |
| 1121 | bias: (dim,) |
| 1122 | conv_state_indices: (batch,), dtype int32 |
| 1123 | If not None, the conv_state is a larger tensor along the batch dim, |
| 1124 | and we are selecting the batch coords specified by conv_state_indices. |
| 1125 | Useful for a continuous batching scenario. |
| 1126 | block_idx_last_scheduled_token: (batch,), dtype int32 |
| 1127 | The pointer into conv_state_indices, where the last cache block to be filled is located. |
| 1128 | initial_state_idx: (batch,), dtype int32 |
| 1129 | The pointer into conv_state_indices, where the cache block containing the initial state is located. |
| 1130 | num_accepted_tokens: (batch,), dtype int32 |
| 1131 | If not None, it indicates the number of accepted tokens for each |
| 1132 | sequence in the batch. |
| 1133 | This is used in speculative decoding, where the conv_state is updated |
| 1134 | in a sliding window manner. |
| 1135 | query_start_loc: (batch + 1,) int32 |
| 1136 | If not None, the inputs is given in a varlen fashion and this indicates |
| 1137 | the starting index of each sequence in the batch. |
| 1138 | max_query_len: int |
| 1139 | If query_start_loc is not None, this indicates the maximum query |
| 1140 | length in the batch. |
| 1141 | null_block_id: int |
| 1142 | Block ID used to identify padded entries in |
| 1143 | conv_state_indices. Block 0 is the null block. |
| 1144 | for example: conv_state_indices = [null_block_id, 1, 20, null_block_id] |
| 1145 | in this case, the kernel will not process entries at |
| 1146 | indices 0 and 3 |
| 1147 | out: optional output tensor with the same shape as `x`. When omitted, |
| 1148 | the input is overwritten. |
| 1149 | """ |
| 1150 | if validate_data: |
| 1151 | assert null_block_id is not None |
| 1152 | assert x.stride(1) == 1 |
| 1153 | if isinstance(activation, bool): |
| 1154 | activation = "silu" if activation is True else None |
| 1155 | elif activation is not None: |
| 1156 | assert activation in ["silu", "swish"] |
| 1157 | |
| 1158 | original_x_dtype = x.dtype |
| 1159 | x = x.to(conv_state.dtype) |
| 1160 | if out is None: |
| 1161 | out = x |
| 1162 | else: |
| 1163 | if out.shape != x.shape: |
| 1164 | raise ValueError( |
| 1165 | f"`out` shape {tuple(out.shape)} must match `x` shape {tuple(x.shape)}." |
| 1166 | ) |
| 1167 | if out.dtype != original_x_dtype or out.device != x.device: |
| 1168 | raise ValueError( |
| 1169 | "`out` must have the same dtype and device as the input `x`." |
| 1170 | ) |
| 1171 | unsqueeze = query_start_loc is None and x.dim() == 2 |
| 1172 | if unsqueeze: |
| 1173 | # make it (batch, dim, seqlen) with seqlen == 1 |
| 1174 | x = x.unsqueeze(-1) |
| 1175 | out = out.unsqueeze(-1) |
| 1176 | if query_start_loc is None: |
| 1177 | batch, dim, seqlen = x.shape |
| 1178 | else: |
| 1179 | assert conv_state_indices is not None |
| 1180 | batch = conv_state_indices.size(0) |
| 1181 | dim = x.size(1) |
| 1182 | seqlen = max_query_len |
| 1183 | _, width = weight.shape |
| 1184 | # conv_state: (..., dim, state_len), where state_len >= width - 1 |
| 1185 | num_cache_lines, _, state_len = conv_state.size() |
| 1186 | |
| 1187 | if validate_data: |
| 1188 | assert dim == weight.size(0) |
| 1189 | assert state_len >= width - 1 |
| 1190 | # when above happens, we don't shift-left to keep any records in conv_state |
| 1191 | assert dim == conv_state.size(1) |
| 1192 | if conv_state_indices is None: |
| 1193 | assert conv_state.size(0) >= batch |
| 1194 | else: |
| 1195 | assert batch == conv_state_indices.shape[0], ( |
| 1196 | f"ERROR: conv_state_indices should have shape ({batch},*) but got {conv_state_indices.shape}" |
| 1197 | ) |
| 1198 | |
| 1199 | assert num_cache_lines >= batch |
| 1200 | assert weight.stride(1) == 1 # Need this |
| 1201 | |
| 1202 | stride_w_dim, stride_w_width = weight.stride() |
| 1203 | |
| 1204 | if query_start_loc is None: |
| 1205 | # X (batch, dim, seqlen) |
| 1206 | stride_x_seq, stride_x_dim, stride_x_token = x.stride() |
| 1207 | stride_o_seq, stride_o_dim, stride_o_token = out.stride() |
| 1208 | else: |
| 1209 | # X (dim, cu_seqlen) |
| 1210 | stride_x_token, stride_x_dim = x.stride() |
| 1211 | stride_x_seq = 0 |
| 1212 | stride_o_token, stride_o_dim = out.stride() |
| 1213 | stride_o_seq = 0 |
| 1214 | |
| 1215 | stride_istate_seq, stride_istate_dim, stride_istate_token = conv_state.stride() |
| 1216 | stride_state_indices = ( |
| 1217 | conv_state_indices.stride(0) if conv_state_indices is not None else 0 |
| 1218 | ) |
| 1219 | if num_accepted_tokens is not None: |
| 1220 | state_len = width - 1 + (seqlen - 1) # effective state_len needed |
| 1221 | else: |
| 1222 | state_len = width - 1 |
| 1223 | np2_statelen = triton.next_power_of_2(state_len) |
| 1224 | |
| 1225 | def grid(META): |
| 1226 | return ( |
| 1227 | batch, |
| 1228 | triton.cdiv(dim, META["BLOCK_N"]), |
| 1229 | ) |
| 1230 | |
| 1231 | _causal_conv1d_update_kernel[grid]( |
| 1232 | # Pointers to matrices |
| 1233 | x, |
| 1234 | weight, |
| 1235 | bias, |
| 1236 | conv_state, |
| 1237 | conv_state_indices, |
| 1238 | num_accepted_tokens, |
| 1239 | query_start_loc, |
| 1240 | block_idx_last_scheduled_token, |
| 1241 | initial_state_idx, |
| 1242 | out, |
| 1243 | # Matrix dimensions |
| 1244 | batch, |
| 1245 | dim, |
| 1246 | seqlen, |
| 1247 | state_len, |
| 1248 | num_cache_lines, |
| 1249 | # stride |
| 1250 | stride_x_seq, |
| 1251 | stride_x_dim, |
| 1252 | stride_x_token, |
| 1253 | stride_w_dim, |
| 1254 | stride_w_width, |
| 1255 | stride_istate_seq, |
| 1256 | stride_istate_dim, |
| 1257 | stride_istate_token, |
| 1258 | stride_state_indices, |
| 1259 | stride_o_seq, |
| 1260 | stride_o_dim, |
| 1261 | stride_o_token, |
| 1262 | # others |
| 1263 | null_block_id, |
| 1264 | # META |
| 1265 | HAS_BIAS=bias is not None, |
| 1266 | KERNEL_WIDTH=width, |
| 1267 | SILU_ACTIVATION=activation in ["silu", "swish"], |
| 1268 | IS_VARLEN=query_start_loc is not None, |
| 1269 | IS_APC_ENABLED=block_idx_last_scheduled_token is not None, |
| 1270 | IS_SPEC_DECODING=num_accepted_tokens is not None, |
| 1271 | NP2_STATELEN=np2_statelen, |
| 1272 | HAS_NULL_BLOCK=null_block_id is not None, |
| 1273 | BLOCK_N=256, |
| 1274 | launch_pdl=current_platform.is_arch_support_pdl(), |
| 1275 | ) |
| 1276 | if unsqueeze: |
| 1277 | out = out.squeeze(-1) |
| 1278 | return out.to(original_x_dtype) |
| 1279 | |
| 1280 | |
| 1281 | if current_platform.is_cpu(): |
| 1282 | from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( |
| 1283 | causal_conv1d_fn_cpu, |
| 1284 | causal_conv1d_update_cpu, |
| 1285 | ) |
| 1286 | |
| 1287 | causal_conv1d_fn = causal_conv1d_fn_cpu # type: ignore |
| 1288 | causal_conv1d_update = causal_conv1d_update_cpu # type: ignore |