@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 146 lines 4.96 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 # Copyright (c) 2024, Tri Dao, Albert Gu.
5 # Adapted from https://github.com/state-spaces/mamba/blob/v2.2.4/mamba_ssm/ops/triton/ssd_state_passing.py
6
7 # ruff: noqa: E501
8
9 import torch
10
11 from vllm.model_executor.layers.mamba.ops.triton_helpers import fast_exp
12 from vllm.triton_utils import tl, triton
13
14
15 @triton.autotune(
16 configs=[
17 triton.Config({"BLOCK_SIZE": 64}),
18 triton.Config({"BLOCK_SIZE": 128}),
19 triton.Config({"BLOCK_SIZE": 256}),
20 triton.Config({"BLOCK_SIZE": 512}),
21 triton.Config({"BLOCK_SIZE": 1024}),
22 triton.Config({"BLOCK_SIZE": 2048}),
23 ],
24 key=["dim"],
25 )
26 @triton.jit
27 def _state_passing_fwd_kernel(
28 # Pointers to matrices
29 states_ptr,
30 out_ptr,
31 dA_cs_ptr,
32 initstates_ptr,
33 last_chunk_indices_ptr,
34 # Matrix dimensions
35 dim: tl.constexpr,
36 chunk_size: tl.constexpr,
37 # Strides
38 stride_states_chunk: tl.int64,
39 stride_states_head: tl.int64,
40 stride_states_dim: tl.constexpr,
41 stride_out_chunk: tl.int64,
42 stride_out_head: tl.int64,
43 stride_out_dim: tl.constexpr,
44 stride_dA_cs_head: tl.int64,
45 stride_dA_cs_chunk: tl.int64,
46 stride_dA_cs_csize: tl.constexpr,
47 stride_initstates_batch: tl.int64,
48 stride_initstates_head: tl.int64,
49 stride_initstates_dim: tl.constexpr,
50 # Meta-parameters
51 HAS_INITSTATES: tl.constexpr,
52 BLOCK_SIZE: tl.constexpr,
53 ):
54 pid_m = tl.program_id(axis=0)
55 pid_b = tl.program_id(axis=1)
56 pid_h = tl.program_id(axis=2)
57
58 # Derive this sequence's chunk range from last_chunk_indices
59 chunk_end = tl.load(last_chunk_indices_ptr + pid_b) + 1
60 chunk_start = (
61 tl.load(last_chunk_indices_ptr + pid_b - 1, mask=pid_b > 0, other=-1) + 1
62 )
63
64 # Offset pointers to this sequence's first chunk
65 states_ptr += chunk_start * stride_states_chunk + pid_h * stride_states_head
66 dA_cs_ptr += (
67 pid_h * stride_dA_cs_head
68 + chunk_start * stride_dA_cs_chunk
69 + (chunk_size - 1) * stride_dA_cs_csize
70 )
71 out_ptr += chunk_start * stride_out_chunk + pid_h * stride_out_head
72
73 offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
74 states_ptrs = states_ptr + offs_m * stride_states_dim
75 out_ptrs = out_ptr + offs_m * stride_out_dim
76
77 # Load initial state once — no per-chunk branching needed
78 if HAS_INITSTATES:
79 initstates_ptrs = (
80 initstates_ptr
81 + pid_b * stride_initstates_batch
82 + pid_h * stride_initstates_head
83 + offs_m * stride_initstates_dim
84 )
85 states = tl.load(initstates_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)
86 else:
87 states = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
88
89 # Loop over only this sequence's chunks — branchless
90 nchunks_this_seq = chunk_end - chunk_start
91 for _ in range(nchunks_this_seq):
92 new_states = tl.load(states_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)
93 dA_cs = tl.load(dA_cs_ptr).to(tl.float32)
94 states = fast_exp(dA_cs) * states + new_states
95 tl.store(out_ptrs, states, mask=offs_m < dim)
96
97 states_ptrs += stride_states_chunk
98 dA_cs_ptr += stride_dA_cs_chunk
99 out_ptrs += stride_out_chunk
100
101
102 def _state_passing_fwd(
103 states,
104 dA_cumsum,
105 last_chunk_indices,
106 initial_states=None,
107 out_dtype=None,
108 ):
109 nchunks, nheads, dim = states.shape
110 chunk_size = dA_cumsum.shape[-1]
111 batch = last_chunk_indices.shape[0]
112 assert dA_cumsum.shape == (nheads, nchunks, chunk_size)
113 out_dtype = states.dtype if out_dtype is None else out_dtype
114 out = torch.empty((nchunks, nheads, dim), device=states.device, dtype=out_dtype)
115
116 initial_states_strides = (
117 (initial_states.stride(0), initial_states.stride(1), initial_states.stride(2))
118 if initial_states is not None
119 else (0, 0, 0)
120 )
121
122 grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE"]), batch, nheads)
123 with torch.accelerator.device_index(states.device.index):
124 _state_passing_fwd_kernel[grid](
125 states_ptr=states,
126 out_ptr=out,
127 dA_cs_ptr=dA_cumsum,
128 initstates_ptr=initial_states,
129 last_chunk_indices_ptr=last_chunk_indices,
130 dim=dim,
131 chunk_size=chunk_size,
132 stride_states_chunk=states.stride(0),
133 stride_states_head=states.stride(1),
134 stride_states_dim=states.stride(2),
135 stride_out_chunk=out.stride(0),
136 stride_out_head=out.stride(1),
137 stride_out_dim=out.stride(2),
138 stride_dA_cs_head=dA_cumsum.stride(0),
139 stride_dA_cs_chunk=dA_cumsum.stride(1),
140 stride_dA_cs_csize=dA_cumsum.stride(2),
141 stride_initstates_batch=initial_states_strides[0],
142 stride_initstates_head=initial_states_strides[1],
143 stride_initstates_dim=initial_states_strides[2],
144 HAS_INITSTATES=initial_states is not None,
145 )
146 return out