@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 130 lines 4.02 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from dataclasses import dataclass
5
6 import torch
7
8 from vllm.model_executor.layers.mamba.checkpoint import (
9 MambaPrefillCheckpointExporter,
10 MambaPrefillCheckpointMetadata,
11 )
12 from vllm.triton_utils import tl, triton
13 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
14
15
16 def kda_prefill_checkpoint_alignment(backend: str) -> int | None:
17 return 16 if backend == "flashkda" else None
18
19
20 @dataclass(frozen=True)
21 class FlashKDAPrefillCheckpointExporter(MambaPrefillCheckpointExporter):
22 """Store FlashKDA recurrent and convolution checkpoint states."""
23
24 state_len: int | None = None
25
26 def export(
27 self,
28 checkpoint: MambaPrefillCheckpointMetadata,
29 *,
30 raw_qkv: torch.Tensor,
31 conv_state: torch.Tensor,
32 recurrent_checkpoint: torch.Tensor,
33 recurrent_state: torch.Tensor,
34 cu_seqlens: torch.Tensor,
35 ) -> None:
36 state_len = (
37 self.state_len if self.state_len is not None else conv_state.shape[-1]
38 )
39 width = raw_qkv.shape[-1]
40 recurrent_row_size = recurrent_checkpoint[0].numel()
41 block_size = 256
42 store_cache_checkpoints_kernel[
43 (
44 checkpoint.checkpoint_offsets.numel(),
45 triton.cdiv(max(width * state_len, recurrent_row_size), block_size),
46 )
47 ](
48 raw_qkv,
49 conv_state,
50 recurrent_checkpoint,
51 recurrent_state,
52 cu_seqlens,
53 checkpoint.checkpoint_offsets,
54 checkpoint.state_indices,
55 raw_qkv.stride(0),
56 raw_qkv.stride(1),
57 conv_state.stride(0),
58 conv_state.stride(1),
59 conv_state.stride(2),
60 recurrent_checkpoint.stride(0),
61 recurrent_state.stride(0),
62 checkpoint.checkpoint_offsets.stride(0),
63 state_len,
64 width,
65 recurrent_row_size,
66 NULL_BLOCK_ID,
67 block_size,
68 )
69
70
71 @triton.jit
72 def store_cache_checkpoints_kernel(
73 x_ptr,
74 conv_state_ptr,
75 recurrent_checkpoint_ptr,
76 recurrent_state_ptr,
77 query_start_loc_ptr,
78 checkpoint_offsets_ptr,
79 checkpoint_state_indices_ptr,
80 x_stride_0: tl.constexpr,
81 x_stride_1: tl.constexpr,
82 state_stride_0: tl.constexpr,
83 state_stride_1: tl.constexpr,
84 state_stride_2: tl.constexpr,
85 checkpoint_stride_0: tl.constexpr,
86 recurrent_state_stride_0: tl.constexpr,
87 checkpoint_offset_stride: tl.constexpr,
88 STATE_LEN: tl.constexpr,
89 WIDTH: tl.constexpr,
90 RECURRENT_ROW_SIZE: tl.constexpr,
91 NULL_STATE_IDX: tl.constexpr,
92 BLOCK_SIZE: tl.constexpr,
93 ):
94 seq_idx = tl.program_id(0)
95 cols = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
96 state_idx = tl.load(checkpoint_state_indices_ptr + seq_idx).to(tl.int64)
97 checkpoint_offset = tl.load(
98 checkpoint_offsets_ptr + seq_idx * checkpoint_offset_stride
99 )
100 valid_checkpoint = (state_idx != NULL_STATE_IDX) & (checkpoint_offset > 0)
101 valid_conv = (
102 (cols < WIDTH * STATE_LEN) & valid_checkpoint & (checkpoint_offset >= STATE_LEN)
103 )
104 width_idx = cols // STATE_LEN
105 history_idx = cols % STATE_LEN
106 checkpoint_end = tl.load(query_start_loc_ptr + seq_idx) + checkpoint_offset
107 token_idx = checkpoint_end - STATE_LEN + history_idx
108 values = tl.load(
109 x_ptr + token_idx * x_stride_0 + width_idx * x_stride_1,
110 mask=valid_conv,
111 )
112 tl.store(
113 conv_state_ptr
114 + state_idx * state_stride_0
115 + width_idx * state_stride_1
116 + history_idx * state_stride_2,
117 values,
118 mask=valid_conv,
119 )
120
121 valid_recurrent = (cols < RECURRENT_ROW_SIZE) & valid_checkpoint
122 recurrent = tl.load(
123 recurrent_checkpoint_ptr + seq_idx * checkpoint_stride_0 + cols,
124 mask=valid_recurrent,
125 )
126 tl.store(
127 recurrent_state_ptr + state_idx * recurrent_state_stride_0 + cols,
128 recurrent,
129 mask=valid_recurrent,
130 )