@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d93c976a0d

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 09:45 UTC

[Kernel] Have rotary embeddings support tensors (#18046)

Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>

Lucas Wilkinson committed May 14, 2025 at 18:43 UTC d93c976a0d78639d0ea9074a9e01607f0d5c5670

4 files changed +59 -31

csrc/pos_encoding_kernels.cu

+28 -12

@@ -44,7 +44,8 @@ inline __device__ void apply_rotary_embedding(

44 // head_size]

45 const scalar_t* cache_ptr, const int head_size, const int num_heads,

46 const int num_kv_heads, const int rot_dim, const int token_idx,

47 - const int64_t query_stride, const int64_t key_stride) {

47 + const int64_t query_stride, const int64_t key_stride,

48 + const int64_t head_stride) {

49 const int embed_dim = rot_dim / 2;

50 const scalar_t* cos_ptr = cache_ptr;

51 const scalar_t* sin_ptr = cache_ptr + embed_dim;

@@ -52,7 +53,8 @@ inline __device__ void apply_rotary_embedding(

53 const int nq = num_heads * embed_dim;

54 for (int i = threadIdx.x; i < nq; i += blockDim.x) {

55 const int head_idx = i / embed_dim;

55 - const int64_t token_head = token_idx * query_stride + head_idx * head_size;

56 + const int64_t token_head =

57 + token_idx * query_stride + head_idx * head_stride;

58 const int rot_offset = i % embed_dim;

59 apply_token_rotary_embedding<scalar_t, IS_NEOX>(

60 query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);

@@ -62,7 +64,8 @@ inline __device__ void apply_rotary_embedding(

64 const int nk = num_kv_heads * embed_dim;

65 for (int i = threadIdx.x; i < nk; i += blockDim.x) {

66 const int head_idx = i / embed_dim;

65 - const int64_t token_head = token_idx * key_stride + head_idx * head_size;

67 + const int64_t token_head =

68 + token_idx * key_stride + head_idx * head_stride;

69 const int rot_offset = i % embed_dim;

70 apply_token_rotary_embedding<scalar_t, IS_NEOX>(

71 key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);

@@ -84,7 +87,8 @@ __global__ void rotary_embedding_kernel(

87 const scalar_t* __restrict__ cos_sin_cache, // [max_position, 2, rot_dim //

88 // 2]

89 const int rot_dim, const int64_t query_stride, const int64_t key_stride,

87 - const int num_heads, const int num_kv_heads, const int head_size) {

90 + const int64_t head_stride, const int num_heads, const int num_kv_heads,

91 + const int head_size) {

92 // Each thread block is responsible for one token.

93 const int token_idx = blockIdx.x;

94 int64_t pos = positions[token_idx];

@@ -92,7 +96,7 @@ __global__ void rotary_embedding_kernel(

97 apply_rotary_embedding<scalar_t, IS_NEOX>(

98 query, key, cache_ptr, head_size, num_heads, num_kv_heads, rot_dim,

95 - token_idx, query_stride, key_stride);

99 + token_idx, query_stride, key_stride, head_stride);

100 }

101

102 template <typename scalar_t, bool IS_NEOX>

@@ -109,9 +113,9 @@ __global__ void batched_rotary_embedding_kernel(

113 const scalar_t* __restrict__ cos_sin_cache, // [max_position, 2, rot_dim //

114 // 2]

115 const int64_t* __restrict__ cos_sin_cache_offsets, // [batch_size, seq_len]

112 - // or [num_tokens]

116 const int rot_dim, const int64_t query_stride, const int64_t key_stride,

114 - const int num_heads, const int num_kv_heads, const int head_size) {

117 + const int64_t head_stride, const int num_heads, const int num_kv_heads,

118 + const int head_size) {

119 // Each thread block is responsible for one token.

120 const int token_idx = blockIdx.x;

121 int64_t pos = positions[token_idx];

@@ -121,7 +125,7 @@ __global__ void batched_rotary_embedding_kernel(

125

126 apply_rotary_embedding<scalar_t, IS_NEOX>(

127 query, key, cache_ptr, head_size, num_heads, num_kv_heads, rot_dim,

124 - token_idx, query_stride, key_stride);

128 + token_idx, query_stride, key_stride, head_stride);

129 }

130

131 } // namespace vllm

@@ -179,6 +183,12 @@ void rotary_embedding(

183 int seq_dim_idx = positions_ndim - 1;

184 int64_t query_stride = query.stride(seq_dim_idx);

185 int64_t key_stride = key.has_value() ? key->stride(seq_dim_idx) : 0;

186 + // Determine head stride: for [*, heads, head_size] use stride of last dim;

187 + // for flat [*, heads*head_size], heads blocks are contiguous of size

188 + // head_size

189 + int query_ndim = query.dim();

190 + int64_t head_stride =

191 + (query_ndim == positions_ndim + 2) ? query.stride(-2) : head_size;

192

193 dim3 grid(num_tokens);

194 dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));

@@ -190,14 +200,14 @@ void rotary_embedding(

200 positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),

201 key.has_value() ? key->data_ptr<scalar_t>() : nullptr,

202 cos_sin_cache.data_ptr<scalar_t>(), rot_dim, query_stride, key_stride,

193 - num_heads, num_kv_heads, head_size);

203 + head_stride, num_heads, num_kv_heads, head_size);

204 } else {

205 vllm::rotary_embedding_kernel<scalar_t, false>

206 <<<grid, block, 0, stream>>>(

207 positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),

208 key.has_value() ? key->data_ptr<scalar_t>() : nullptr,

209 cos_sin_cache.data_ptr<scalar_t>(), rot_dim, query_stride,

200 - key_stride, num_heads, num_kv_heads, head_size);

210 + key_stride, head_stride, num_heads, num_kv_heads, head_size);

211 }

212 });

213 }

@@ -263,6 +273,12 @@ void batched_rotary_embedding(

273 int seq_dim_idx = positions_ndim - 1;

274 int64_t query_stride = query.stride(seq_dim_idx);

275 int64_t key_stride = key.has_value() ? key->stride(seq_dim_idx) : 0;

276 + // Determine head stride: for [*, heads, head_size] use stride of last dim;

277 + // for flat [*, heads*head_size], heads blocks are contiguous of size

278 + // head_size

279 + int query_ndim = query.dim();

280 + int64_t head_stride =

281 + (query_ndim == positions_ndim + 2) ? query.stride(-2) : head_size;

282

283 dim3 grid(num_tokens);

284 dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));

@@ -276,7 +292,7 @@ void batched_rotary_embedding(

292 key.has_value() ? key->data_ptr<scalar_t>() : nullptr,

293 cos_sin_cache.data_ptr<scalar_t>(),

294 cos_sin_cache_offsets.data_ptr<int64_t>(), rot_dim, query_stride,

279 - key_stride, num_heads, num_kv_heads, head_size);

295 + key_stride, head_stride, num_heads, num_kv_heads, head_size);

296 } else {

297 vllm::batched_rotary_embedding_kernel<scalar_t, false>

298 <<<grid, block, 0, stream>>>(

@@ -284,7 +300,7 @@ void batched_rotary_embedding(

300 key.has_value() ? key->data_ptr<scalar_t>() : nullptr,

301 cos_sin_cache.data_ptr<scalar_t>(),

302 cos_sin_cache_offsets.data_ptr<int64_t>(), rot_dim, query_stride,

287 - key_stride, num_heads, num_kv_heads, head_size);

303 + key_stride, head_stride, num_heads, num_kv_heads, head_size);

304 }

305 });

306 }

tests/kernels/core/test_pos_encoding.py

+13 -1

@@ -29,12 +29,20 @@ def _get_flat_tensor_shape(batch_size: int, seq_len: int, num_heads: int,

29 return (batch_size, seq_len, num_heads * head_size)

32 +# For testing sliced tensors

33 +def _get_padded_tensor_shape(batch_size: int, seq_len: int, num_heads: int,

34 + head_size: int) -> tuple[int, ...]:

35 + return (batch_size, seq_len, num_heads, head_size + 64)

36 +

37 +

38 def _get_batch_tensor_shape(batch_size: int, seq_len: int, num_heads: int,

39 head_size: int) -> tuple[int, ...]:

40 return (batch_size, seq_len, num_heads, head_size)

37 -TENSORS_SHAPES_FN = [_get_batch_tensor_shape, _get_flat_tensor_shape]

43 +TENSORS_SHAPES_FN = [

44 + _get_batch_tensor_shape, _get_flat_tensor_shape, _get_padded_tensor_shape

45 +]

48 @pytest.mark.parametrize("is_neox_style", IS_NEOX_STYLE)

@@ -79,6 +87,10 @@ def test_rotary_embedding(

87 query = torch.randn(query_shape, dtype=dtype)

88 key = torch.randn_like(query) if use_key else None

90 + # slice tensor if required, noop otherwise

91 + query = query[..., :head_size]

92 + key = key[..., :head_size] if use_key else None

93 +

94 # NOTE(woosuk): The reference implementation should be executed first

95 # because the custom kernel is in-place.

96 ref_query, ref_key = rope.forward_native(positions, query, key)

tests/kernels/core/test_rotary_embedding.py

+15 -2

@@ -38,9 +38,10 @@ def rotary_embedding_opcheck(rot,

38 @pytest.mark.parametrize("head_size", [32, 108])

39 @pytest.mark.parametrize("seq_len", [11, 1024])

40 @pytest.mark.parametrize("use_key", [True, False])

41 +@pytest.mark.parametrize("head_stride_is_contingous", [True, False])

42 def test_rotary_embedding_opcheck(dist_init, device, max_position,

43 is_neox_style, rotary_dim, head_size,

43 - seq_len, use_key):

44 + seq_len, use_key, head_stride_is_contingous):

45 batch_size = 1

46 base = 10000

47 num_heads = 7

@@ -50,15 +51,27 @@ def test_rotary_embedding_opcheck(dist_init, device, max_position,

51 positions = torch.randint(0,

52 max_position, (batch_size, seq_len),

53 device=device)

54 + head_stride = head_size + (64 if head_stride_is_contingous else 0)

55 +

56 query = torch.randn(batch_size,

57 seq_len,

55 - num_heads * head_size,

58 + num_heads,

59 + head_stride,

60 dtype=torch.float32,

61 device=device)

62 key = torch.randn_like(query) if use_key else None

63 + query = query[..., :head_size]

64 + key = key[..., :head_size] if use_key else None

66 rotary_embedding_opcheck(rot, positions, query, key)

67 offsets = torch.zeros(batch_size * seq_len,

68 device=device,

69 dtype=torch.long)

70 rotary_embedding_opcheck(rot, positions, query, key, offsets)

71 +

72 + # if we have a contiguous head stride, test the alternate

73 + # [..., num_heads * head_dim] shape/layout

74 + if head_stride_is_contingous:

75 + rotary_embedding_opcheck(

76 + rot, positions, query.flatten(start_dim=-2),

77 + key.flatten(start_dim=-2) if use_key else None)

vllm/_custom_ops.py

+3 -16

@@ -254,14 +254,8 @@ def rotary_embedding(

254 cos_sin_cache: torch.Tensor,

255 is_neox: bool,

256 ) -> None:

257 - # TODO: Remove this contiguous call when the kernel is updated to support tensor slices

258 - query_contiguous = query.contiguous()

259 - key_contiguous = key.contiguous() if key is not None else None

260 - torch.ops._C.rotary_embedding(positions, query_contiguous, key_contiguous,

261 - head_size, cos_sin_cache, is_neox)

262 - query.copy_(query_contiguous)

263 - if key is not None:

264 - key.copy_(key_contiguous)

257 + torch.ops._C.rotary_embedding(positions, query, key, head_size,

258 + cos_sin_cache, is_neox)

259

260

261 def batched_rotary_embedding(positions: torch.Tensor, query: torch.Tensor,

@@ -269,16 +263,9 @@ def batched_rotary_embedding(positions: torch.Tensor, query: torch.Tensor,

263 cos_sin_cache: torch.Tensor, is_neox: bool,

264 rot_dim: int,

265 cos_sin_cache_offsets: torch.Tensor) -> None:

272 - # TODO: Remove this contiguous call when the kernel is updated to support tensor slices

273 - query_contiguous = query.contiguous()

274 - key_contiguous = key.contiguous() if key is not None else None

275 - torch.ops._C.batched_rotary_embedding(positions, query_contiguous,

276 - key_contiguous, head_size,

266 + torch.ops._C.batched_rotary_embedding(positions, query, key, head_size,

267 cos_sin_cache, is_neox, rot_dim,

268 cos_sin_cache_offsets)

279 - query.copy_(query_contiguous)

280 - if key is not None:

281 - key.copy_(key_contiguous)

269

270

271 # layer norm ops