@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 153 lines 5.59 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import torch
5
6 from vllm.triton_utils import tl, triton
7 from vllm.utils.torch_utils import async_tensor_h2d
8
9
10 # num_tokens/num_grids vary per request (incl. the ==1 class) and grids_ptr
11 # is an offset slice of the packed H2D buffer, so its 16B alignment varies
12 # with grid count; without these the kernel recompiles mid-serving.
13 @triton.jit(
14 do_not_specialize=["num_tokens", "num_grids"],
15 do_not_specialize_on_alignment=["grids_ptr"],
16 )
17 def _vit_mrope_setup_kernel(
18 inv_freq_t_ptr,
19 inv_freq_h_ptr,
20 inv_freq_w_ptr,
21 cos_out_ptr,
22 sin_out_ptr,
23 cu_seqlens_ptr, # (G + 1,) int32 token boundaries per grid
24 grids_ptr, # (G, 3) int32: t, h, w per grid
25 num_tokens,
26 num_grids,
27 half_t: tl.constexpr,
28 half_h: tl.constexpr,
29 half_w: tl.constexpr,
30 spatial_merge: tl.constexpr,
31 BLOCK_N: tl.constexpr,
32 BLOCK_D: tl.constexpr,
33 ):
34 """Build triton_mrope's (3, N, half_rot) cos/sin t/h/w planes in one
35 launch: compute each token's (t, h, w) position arithmetically and the
36 rotation in fp32 from the per-axis inv_freq. Only each axis's mrope
37 section columns are written; the rest is left uninitialized (the
38 consumer's per-axis masked loads read only the section columns)."""
39 pid = tl.program_id(0)
40 tok = pid * BLOCK_N + tl.arange(0, BLOCK_N)
41 tok_mask = tok < num_tokens
42
43 # Grid (segment) index per token; cu_seqlens is nondecreasing.
44 seg = tl.zeros([BLOCK_N], dtype=tl.int32)
45 for g in range(num_grids):
46 end = tl.load(cu_seqlens_ptr + g + 1)
47 seg += (tok >= end).to(tl.int32)
48
49 seg_start = tl.load(cu_seqlens_ptr + seg, mask=tok_mask, other=0)
50 grid_h = tl.load(grids_ptr + seg * 3 + 1, mask=tok_mask, other=1)
51 grid_w = tl.load(grids_ptr + seg * 3 + 2, mask=tok_mask, other=1)
52
53 local = tok - seg_start
54 rem = local % (grid_h * grid_w)
55 tpos = local // (grid_h * grid_w)
56 # Within a frame tokens are ordered by spatial-merge block (bh, bw, mh,
57 # mw), matching the reference ViT position-id permute.
58 bh = rem // (spatial_merge * grid_w)
59 r = rem % (spatial_merge * grid_w)
60 bw = r // (spatial_merge * spatial_merge)
61 r2 = r % (spatial_merge * spatial_merge)
62 hpos = bh * spatial_merge + r2 // spatial_merge
63 wpos = bw * spatial_merge + r2 % spatial_merge
64
65 d = tl.arange(0, BLOCK_D)
66 half_rot = half_t + half_h + half_w
67 row = (tok * half_rot)[:, None] + d[None, :]
68 plane = num_tokens * half_rot
69
70 # Per-axis compute+store to cap live registers; the fp32 trig sequence
71 # wants few elements per thread (hence small BLOCK_N, many warps).
72 in_t = d < half_t
73 inv_t = tl.load(inv_freq_t_ptr + d, mask=in_t, other=0.0)
74 freq = tpos.to(tl.float32)[:, None] * inv_t[None, :]
75 mask = tok_mask[:, None] & in_t[None, :]
76 tl.store(cos_out_ptr + row, tl.cos(freq), mask=mask)
77 tl.store(sin_out_ptr + row, tl.sin(freq), mask=mask)
78
79 in_h = (d >= half_t) & (d < half_t + half_h)
80 inv_h = tl.load(inv_freq_h_ptr + d - half_t, mask=in_h, other=0.0)
81 freq = hpos.to(tl.float32)[:, None] * inv_h[None, :]
82 mask = tok_mask[:, None] & in_h[None, :]
83 tl.store(cos_out_ptr + plane + row, tl.cos(freq), mask=mask)
84 tl.store(sin_out_ptr + plane + row, tl.sin(freq), mask=mask)
85
86 in_w = (d >= half_t + half_h) & (d < half_rot)
87 inv_w = tl.load(inv_freq_w_ptr + d - half_t - half_h, mask=in_w, other=0.0)
88 freq = wpos.to(tl.float32)[:, None] * inv_w[None, :]
89 mask = tok_mask[:, None] & in_w[None, :]
90 tl.store(cos_out_ptr + 2 * plane + row, tl.cos(freq), mask=mask)
91 tl.store(sin_out_ptr + 2 * plane + row, tl.sin(freq), mask=mask)
92
93
94 def vit_mrope_setup(
95 inv_freq_t: torch.Tensor,
96 inv_freq_h: torch.Tensor,
97 inv_freq_w: torch.Tensor,
98 grid_thw: list[list[int]],
99 spatial_merge_size: int,
100 ) -> tuple[torch.Tensor, torch.Tensor]:
101 """(3, total_N, half_rot_dim) fp32 cos/sin t/h/w planes for triton_mrope.
102
103 One fused kernel: positions are computed arithmetically from the grids
104 and the rotation comes from the per-axis inv_freq, so there is no
105 cos/sin cache to size or grow.
106
107 Args:
108 inv_freq_t: t-axis inverse frequencies, (t_dim // 2,).
109 inv_freq_h: h-axis inverse frequencies, (h_dim // 2,).
110 inv_freq_w: w-axis inverse frequencies, (w_dim // 2,).
111 grid_thw: per-grid [t, h, w] patch counts (host metadata).
112 spatial_merge_size: spatial merge factor of the ViT.
113
114 """
115 device = inv_freq_t.device
116 cu = [0]
117 for t, h, w in grid_thw:
118 cu.append(cu[-1] + t * h * w)
119 n = cu[-1]
120 g = len(grid_thw)
121 # grid_thw is per-request host metadata, so one H2D is unavoidable;
122 # pack cu_seqlens and grids into a single pinned async copy.
123 buf = async_tensor_h2d(
124 cu + [d for grid in grid_thw for d in grid], device=device, dtype=torch.int32
125 )
126 cu_seqlens = buf[: g + 1]
127 grids = buf[g + 1 :].view(g, 3)
128
129 half_t = inv_freq_t.numel()
130 half_h = inv_freq_h.numel()
131 half_w = inv_freq_w.numel()
132 half_rot = half_t + half_h + half_w
133 cos = torch.empty(3, n, half_rot, device=device, dtype=torch.float32)
134 sin = torch.empty(3, n, half_rot, device=device, dtype=torch.float32)
135 _vit_mrope_setup_kernel[(triton.cdiv(n, 32),)](
136 inv_freq_t,
137 inv_freq_h,
138 inv_freq_w,
139 cos,
140 sin,
141 cu_seqlens,
142 grids,
143 n,
144 g,
145 half_t=half_t,
146 half_h=half_h,
147 half_w=half_w,
148 spatial_merge=spatial_merge_size,
149 BLOCK_N=32,
150 BLOCK_D=triton.next_power_of_2(half_rot),
151 num_warps=8,
152 )
153 return cos, sin