torch_bindings.cpp 40.7 KB
Newer Older
1
2
3
#include "cache.h"
#include "cuda_utils.h"
#include "ops.h"
4
#include "core/registration.h"
5
6

#include <torch/library.h>
7
#include <torch/version.h>
8
9
10
11
12
13
14
15
16
17
18
19
20

// Note on op signatures:
// The X_meta signatures are for the meta functions corresponding to op X.
// They must be kept in sync with the signature for X. Generally, only
// functions that return Tensors require a meta function.
//
// See the following links for detailed docs on op registration and function
// schemas.
// https://docs.google.com/document/d/1_W62p8WJOQQUzPsJYa7s701JXt0qf2OfLub2sbkHOaU/edit#heading=h.ptttacy8y1u9
// https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/README.md#annotations

TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  // vLLM custom ops
21
22
23
24
25
26
27
28
29
  //

  // The default behavior in PyTorch 2.6 is "requires_contiguous", so we need
  // to override this for many GEMMs with the following tag. Otherwise,
  // torch.compile will force all input tensors to be contiguous(), which
  // will break many custom ops that require column-major weight matrices.
  // TODO: remove this for PyTorch 2.8, when the default is planned to switch
  // to match exact eager-mode strides.
  at::Tag stride_tag = at::Tag::needs_fixed_stride_order;
30

31
32
33
  ops.def("weak_ref_tensor(Tensor input) -> Tensor");
  ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);

34
35
36
37
  ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
  ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU,
           &get_cuda_view_from_cpu_tensor);

38
39
40
41
42
43
44
45
46
  // Attention ops
  // Compute the attention between an input query and the cached
  // keys/values using PagedAttention.
  ops.def(
      "paged_attention_v1("
      "    Tensor! out, Tensor query, Tensor key_cache,"
      "    Tensor value_cache, int num_kv_heads, float scale,"
      "    Tensor block_tables, Tensor seq_lens, int block_size,"
      "    int max_seq_len, Tensor? alibi_slopes,"
47
      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
48
      "    int tp_rank, int blocksparse_local_blocks,"
49
      "    int blocksparse_vert_stride, int blocksparse_block_size,"
50
51
52
53
54
55
56
57
58
59
60
      "    int blocksparse_head_sliding_step) -> ()");
  ops.impl("paged_attention_v1", torch::kCUDA, &paged_attention_v1);

  // PagedAttention V2.
  ops.def(
      "paged_attention_v2("
      "    Tensor! out, Tensor! exp_sums, Tensor! max_logits,"
      "    Tensor! tmp_out, Tensor query, Tensor key_cache,"
      "    Tensor value_cache, int num_kv_heads, float scale,"
      "    Tensor block_tables, Tensor seq_lens, int block_size,"
      "    int max_seq_len, Tensor? alibi_slopes,"
61
      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
62
63
64
65
66
67
68
69
70
71
72
73
74
      "    int tp_rank, int blocksparse_local_blocks,"
      "    int blocksparse_vert_stride, int blocksparse_block_size,"
      "    int blocksparse_head_sliding_step) -> ()");
  ops.impl("paged_attention_v2", torch::kCUDA, &paged_attention_v2);

  // Compute the attention between an input query and the cached
  // keys/values using PagedAttention. (opt)
  ops.def(
      "paged_attention_v1_opt("
      "    Tensor! out, Tensor query, Tensor key_cache,"
      "    Tensor value_cache, int num_kv_heads, float scale,"
      "    Tensor block_tables, Tensor seq_lens, int block_size,"
      "    int max_seq_len, Tensor? alibi_slopes,"
75
      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
76
77
78
79
80
81
82
83
84
85
86
87
88
      "    int tp_rank, int blocksparse_local_blocks,"
      "    int blocksparse_vert_stride, int blocksparse_block_size,"
      "    int blocksparse_head_sliding_step) -> ()");
  ops.impl("paged_attention_v1_opt", torch::kCUDA, &paged_attention_v1_opt);

  // PagedAttention V2 (opt). 
  ops.def(
      "paged_attention_v2_opt("
      "    Tensor! out, Tensor exp_sums, Tensor max_logits,"
      "    Tensor tmp_out, Tensor query, Tensor key_cache,"
      "    Tensor value_cache, int num_kv_heads, float scale,"
      "    Tensor block_tables, Tensor seq_lens, int block_size,"
      "    int max_seq_len, Tensor? alibi_slopes,"
89
      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
90
91
92
93
94
95
96
97
98
99
100
101
102
      "    int tp_rank, int blocksparse_local_blocks,"
      "    int blocksparse_vert_stride, int blocksparse_block_size,"
      "    int blocksparse_head_sliding_step) -> ()");
  ops.impl("paged_attention_v2_opt", torch::kCUDA, &paged_attention_v2_opt);

  // Compute the attention between an input query and the cached
  // keys/values using PagedAttention. (opt)
  ops.def(
      "paged_attention_v1_opt_tc("
      "    Tensor! out, Tensor query, Tensor key_cache,"
      "    Tensor value_cache, int num_kv_heads, float scale,"
      "    Tensor block_tables, Tensor seq_lens, int block_size,"
      "    int max_seq_len, Tensor? alibi_slopes,"
103
      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
104
105
106
107
108
109
110
111
112
113
114
115
116
      "    int tp_rank, int blocksparse_local_blocks,"
      "    int blocksparse_vert_stride, int blocksparse_block_size,"
      "    int blocksparse_head_sliding_step) -> ()");
  ops.impl("paged_attention_v1_opt_tc", torch::kCUDA, &paged_attention_v1_opt_tc);

  // PagedAttention V2 (opt). 
  ops.def(
      "paged_attention_v2_opt_tc("
      "    Tensor! out, Tensor exp_sums, Tensor max_logits,"
      "    Tensor tmp_out, Tensor query, Tensor key_cache,"
      "    Tensor value_cache, int num_kv_heads, float scale,"
      "    Tensor block_tables, Tensor seq_lens, int block_size,"
      "    int max_seq_len, Tensor? alibi_slopes,"
117
      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
118
119
120
121
122
123
124
125
126
127
128
129
130
      "    int tp_rank, int blocksparse_local_blocks,"
      "    int blocksparse_vert_stride, int blocksparse_block_size,"
      "    int blocksparse_head_sliding_step) -> ()");
  ops.impl("paged_attention_v2_opt_tc", torch::kCUDA, &paged_attention_v2_opt_tc);


  // paged_attention with atth_masks
  ops.def(
      "paged_attention_v1_with_mask("
      "    Tensor! out, Tensor query, Tensor key_cache,"
      "    Tensor value_cache, int num_kv_heads, float scale,"
      "    Tensor block_tables, Tensor seq_lens, int block_size,"
      "    int max_seq_len, Tensor? alibi_slopes,"
131
      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
132
133
      "    int tp_rank, int blocksparse_local_blocks,"
      "    int blocksparse_vert_stride, int blocksparse_block_size,"
134
135
136
      "    int blocksparse_head_sliding_step,"
      "    Tensor? attn_masks,"
      "    int attn_masks_stride) -> ()");
137
  ops.impl("paged_attention_v1_with_mask", torch::kCUDA, &paged_attention_v1_with_mask);
138
139
140

  // PagedAttention V2.
  ops.def(
141
      "paged_attention_v2_with_mask("
142
143
      "    Tensor! out, Tensor! exp_sums, Tensor! max_logits,"
      "    Tensor! tmp_out, Tensor query, Tensor key_cache,"
144
145
146
      "    Tensor value_cache, int num_kv_heads, float scale,"
      "    Tensor block_tables, Tensor seq_lens, int block_size,"
      "    int max_seq_len, Tensor? alibi_slopes,"
147
      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
148
      "    int tp_rank, int blocksparse_local_blocks,"
149
      "    int blocksparse_vert_stride, int blocksparse_block_size,"
150
151
152
      "    int blocksparse_head_sliding_step,"
      "    Tensor? attn_masks,"
      "    int attn_masks_stride) -> ()");
153
  ops.impl("paged_attention_v2_with_mask", torch::kCUDA, &paged_attention_v2_with_mask);
154

zhuwenwen's avatar
zhuwenwen committed
155
156
157
  // Compute the attention between an input query and the cached
  // keys/values using PagedAttention. (opt)
  ops.def(
158
      "paged_attention_v1_opt_with_mask("
zhuwenwen's avatar
zhuwenwen committed
159
160
161
162
      "    Tensor! out, Tensor query, Tensor key_cache,"
      "    Tensor value_cache, int num_kv_heads, float scale,"
      "    Tensor block_tables, Tensor seq_lens, int block_size,"
      "    int max_seq_len, Tensor? alibi_slopes,"
163
      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
zhuwenwen's avatar
zhuwenwen committed
164
165
      "    int tp_rank, int blocksparse_local_blocks,"
      "    int blocksparse_vert_stride, int blocksparse_block_size,"
166
167
168
      "    int blocksparse_head_sliding_step,"
      "    Tensor? attn_masks,"
      "    int attn_masks_stride) -> ()");
169
  ops.impl("paged_attention_v1_opt_with_mask", torch::kCUDA, &paged_attention_v1_opt_with_mask);
zhuwenwen's avatar
zhuwenwen committed
170
171
172

  // PagedAttention V2 (opt). 
  ops.def(
173
      "paged_attention_v2_opt_with_mask("
zhuwenwen's avatar
zhuwenwen committed
174
175
176
177
178
      "    Tensor! out, Tensor exp_sums, Tensor max_logits,"
      "    Tensor tmp_out, Tensor query, Tensor key_cache,"
      "    Tensor value_cache, int num_kv_heads, float scale,"
      "    Tensor block_tables, Tensor seq_lens, int block_size,"
      "    int max_seq_len, Tensor? alibi_slopes,"
179
      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
zhuwenwen's avatar
zhuwenwen committed
180
181
      "    int tp_rank, int blocksparse_local_blocks,"
      "    int blocksparse_vert_stride, int blocksparse_block_size,"
182
183
184
      "    int blocksparse_head_sliding_step,"
      "    Tensor? attn_masks,"
      "    int attn_masks_stride) -> ()");
185
  ops.impl("paged_attention_v2_opt_with_mask", torch::kCUDA, &paged_attention_v2_opt_with_mask);
zhuwenwen's avatar
zhuwenwen committed
186

187
  // Compute the attention between an input query and the cached
zhuwenwen's avatar
zhuwenwen committed
188
  // keys/values using PagedAttention. (opt)
189
  ops.def(
190
      "paged_attention_v1_opt_tc_with_mask("
191
192
193
194
      "    Tensor! out, Tensor query, Tensor key_cache,"
      "    Tensor value_cache, int num_kv_heads, float scale,"
      "    Tensor block_tables, Tensor seq_lens, int block_size,"
      "    int max_seq_len, Tensor? alibi_slopes,"
195
      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
zhuwenwen's avatar
zhuwenwen committed
196
      "    int tp_rank, int blocksparse_local_blocks,"
197
      "    int blocksparse_vert_stride, int blocksparse_block_size,"
198
199
200
      "    int blocksparse_head_sliding_step,"
      "    Tensor? attn_masks,"
      "    int attn_masks_stride) -> ()");
201
  ops.impl("paged_attention_v1_opt_tc_with_mask", torch::kCUDA, &paged_attention_v1_opt_tc_with_mask);
202

zhuwenwen's avatar
zhuwenwen committed
203
  // PagedAttention V2 (opt). 
204
  ops.def(
205
      "paged_attention_v2_opt_tc_with_mask("
206
207
208
209
210
      "    Tensor! out, Tensor exp_sums, Tensor max_logits,"
      "    Tensor tmp_out, Tensor query, Tensor key_cache,"
      "    Tensor value_cache, int num_kv_heads, float scale,"
      "    Tensor block_tables, Tensor seq_lens, int block_size,"
      "    int max_seq_len, Tensor? alibi_slopes,"
211
      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
zhuwenwen's avatar
zhuwenwen committed
212
      "    int tp_rank, int blocksparse_local_blocks,"
213
      "    int blocksparse_vert_stride, int blocksparse_block_size,"
214
215
216
      "    int blocksparse_head_sliding_step,"
      "    Tensor? attn_masks,"
      "    int attn_masks_stride) -> ()");
217
218
  ops.impl("paged_attention_v2_opt_tc_with_mask", torch::kCUDA, &paged_attention_v2_opt_tc_with_mask);

219
220
221
222
223
224
225
226
227
228
229
230
  // Merge attn states
  // Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
  // can be used to combine partial attention results (in the split-KV case)
  ops.def(
      "merge_attn_states("
      "    Tensor! output,"
      "    Tensor!? output_lse,"
      "    Tensor prefix_output,"
      "    Tensor prefix_lse,"
      "    Tensor suffix_output,"
      "    Tensor suffix_lse) -> ()");
  ops.impl("merge_attn_states", torch::kCUDA, &merge_attn_states);
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253

  ops.def(
      "convert_vertical_slash_indexes("
      "   Tensor! block_count, Tensor! block_offset, "
      "   Tensor! column_count, Tensor! column_index, "
      "   Tensor q_seqlens, Tensor q_seqlens, "
      "   Tensor vertical_indexes, Tensor slash_indexes, "
      "   int context_size, int block_size_M, int block_size_N, "
      "   bool causal) -> ()");
  ops.impl("convert_vertical_slash_indexes", torch::kCUDA,
           &convert_vertical_slash_indexes);

  ops.def(
      "convert_vertical_slash_indexes_mergehead("
      "   Tensor! block_count, Tensor! block_offset, "
      "   Tensor! column_count, Tensor! column_index, "
      "   Tensor q_seqlens, Tensor q_seqlens, "
      "   Tensor vertical_indexes, Tensor slash_indexes, "
      "   Tensor vertical_indices_count, Tensor slash_indices_count, "
      "   int context_size, int block_size_M, int block_size_N, "
      "   bool causal) -> ()");
  ops.impl("convert_vertical_slash_indexes_mergehead", torch::kCUDA,
           &convert_vertical_slash_indexes_mergehead);
254

255
256
  // Activation ops
  // Activation function used in SwiGLU.
257
  ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
258
259
  ops.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);

zhuwenwen's avatar
zhuwenwen committed
260
261
262
//   ops.def(
//       "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
//   ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
263

264
265
266
  ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()");
  ops.impl("mul_and_silu", torch::kCUDA, &mul_and_silu);

267
268
269
270
271
272
273
274
  // Activation function used in GeGLU with `none` approximation.
  ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
  ops.impl("gelu_and_mul", torch::kCUDA, &gelu_and_mul);

  // Activation function used in GeGLU with `tanh` approximation.
  ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()");
  ops.impl("gelu_tanh_and_mul", torch::kCUDA, &gelu_tanh_and_mul);

zhuwenwen's avatar
zhuwenwen committed
275
276
  // Activation function used in SwiGLU. (opt)
  ops.def("silu_and_mul_opt(Tensor! out, Tensor input) -> ()");
zhuwenwen's avatar
zhuwenwen committed
277
  ops.impl("silu_and_mul_opt", torch::kCUDA, &silu_and_mul_opt);
zhuwenwen's avatar
zhuwenwen committed
278
279
280

  // Activation function used in GeGLU with `none` approximation. (opt)
  ops.def("gelu_and_mul_opt(Tensor! out, Tensor input) -> ()");
zhuwenwen's avatar
zhuwenwen committed
281
  ops.impl("gelu_and_mul_opt", torch::kCUDA, &gelu_and_mul_opt);
zhuwenwen's avatar
zhuwenwen committed
282
283
284

  // Activation function used in GeGLU with `tanh` approximation. (opt)
  ops.def("gelu_tanh_and_mul_opt(Tensor! out, Tensor input) -> ()");
zhuwenwen's avatar
zhuwenwen committed
285
  ops.impl("gelu_tanh_and_mul_opt", torch::kCUDA, &gelu_tanh_and_mul_opt);
zhuwenwen's avatar
zhuwenwen committed
286

287
288
289
290
  // FATReLU implementation.
  ops.def("fatrelu_and_mul(Tensor! out, Tensor input, float threshold) -> ()");
  ops.impl("fatrelu_and_mul", torch::kCUDA, &fatrelu_and_mul);

291
292
293
294
295
296
297
298
  // GELU implementation used in GPT-2.
  ops.def("gelu_new(Tensor! out, Tensor input) -> ()");
  ops.impl("gelu_new", torch::kCUDA, &gelu_new);

  // Approximate GELU implementation.
  ops.def("gelu_fast(Tensor! out, Tensor input) -> ()");
  ops.impl("gelu_fast", torch::kCUDA, &gelu_fast);

299
300
301
302
  // Quick GELU implementation.
  ops.def("gelu_quick(Tensor! out, Tensor input) -> ()");
  ops.impl("gelu_quick", torch::kCUDA, &gelu_quick);

303
  // prepare_inputs advance_step
304
  ops.def(
305
      "advance_step_flashattn(int num_seqs, int num_queries, int block_size, "
306
307
308
      "Tensor! input_tokens, Tensor sampled_token_ids, "
      "Tensor! input_positions, Tensor! seq_lens, Tensor! slot_mapping, "
      "Tensor block_tables) -> ()");
309
310
311
312
313
314
315
316
317
318
319
320
  ops.impl("advance_step_flashattn", torch::kCUDA, &advance_step_flashattn);

  ops.def(
      "advance_step_flashinfer("
      "    int num_seqs, int num_queries, int block_size,"
      "    Tensor! input_tokens, Tensor sampled_token_ids,"
      "    Tensor! input_positions, Tensor! seq_lens, Tensor! slot_mapping,"
      "    Tensor block_tables, Tensor! paged_kv_indices,"
      "    Tensor! paged_kv_indptr, Tensor! paged_kv_last_page_len,"
      "    Tensor! block_table_bounds"
      ") -> ()");
  ops.impl("advance_step_flashinfer", torch::kCUDA, &advance_step_flashinfer);
321

322
323
324
  // Layernorm
  // Apply Root Mean Square (RMS) Normalization to the input tensor.
  ops.def(
325
      "rms_norm(Tensor! result, Tensor input, Tensor weight, float epsilon) -> "
326
327
328
329
330
331
332
333
334
      "()");
  ops.impl("rms_norm", torch::kCUDA, &rms_norm);

  // In-place fused Add and RMS Normalization.
  ops.def(
      "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor weight, "
      "float epsilon) -> ()");
  ops.impl("fused_add_rms_norm", torch::kCUDA, &fused_add_rms_norm);

335
336
337
338
339
340
341
  // Apply repetition penalties to logits in-place
  ops.def(
      "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
      "Tensor output_mask, Tensor repetition_penalties) -> ()");
  ops.impl("apply_repetition_penalties_", torch::kCUDA,
           &apply_repetition_penalties_);

342
343
  // Layernorm-quant
  // Apply Root Mean Square (RMS) Normalization to the input tensor.
zhuwenwen's avatar
zhuwenwen committed
344
345
346
347
348
349
350
351
352
353
354
  ops.def(
      "rms_norm_opt(Tensor! out, Tensor input, Tensor weight, float epsilon) -> "
      "()");
  ops.impl("rms_norm_opt", torch::kCUDA, &rms_norm_opt);

  // In-place fused Add and RMS Normalization. (opt)
  ops.def(
      "fused_add_rms_norm_opt(Tensor! input, Tensor! residual, Tensor weight, "
      "float epsilon) -> ()");
  ops.impl("fused_add_rms_norm_opt", torch::kCUDA, &fused_add_rms_norm_opt);

zhuwenwen's avatar
zhuwenwen committed
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
  // Layernorm-quant
  // Apply Root Mean Square (RMS) Normalization to the input tensor.
//   ops.def(
//       "rms_norm_static_fp8_quant(Tensor! result, Tensor input, Tensor weight, "
//       "Tensor scale, float epsilon) -> "
//       "()");
//   ops.impl("rms_norm_static_fp8_quant", torch::kCUDA,
//            &rms_norm_static_fp8_quant);

  // In-place fused Add and RMS Normalization.
//   ops.def(
//       "fused_add_rms_norm_static_fp8_quant(Tensor! result, Tensor input, "
//       "Tensor! residual, Tensor weight, "
//       "Tensor scale, float epsilon) -> ()");
//   ops.impl("fused_add_rms_norm_static_fp8_quant", torch::kCUDA,
//            &fused_add_rms_norm_static_fp8_quant);
371

372
  // Fused Layernorm + Quant kernels
373
374
375
376
377
378
  ops.def(
      "rms_norm_dynamic_per_token_quant(Tensor! result, Tensor input, "
      "Tensor weight, Tensor! scale, float epsilon, "
      "Tensor? scale_ub, Tensor!? residual) -> ()");
  ops.impl("rms_norm_dynamic_per_token_quant", torch::kCUDA,
           &rms_norm_dynamic_per_token_quant);
379

380
381
382
383
  // Rotary embedding
  // Apply GPT-NeoX or GPT-J style rotary embedding to query and key.
  ops.def(
      "rotary_embedding(Tensor positions, Tensor! query,"
384
      "                 Tensor!? key, int head_size,"
385
386
387
      "                 Tensor cos_sin_cache, bool is_neox) -> ()");
  ops.impl("rotary_embedding", torch::kCUDA, &rotary_embedding);

huangwb's avatar
huangwb committed
388
389
390
391
392
393
394
395
396
  // Rotary embedding TGI for TGI
  // Apply GPT-NeoX or GPT-J style rotary embedding to query and key.
  ops.def(
      "rotary_embedding_tgi(Tensor! query, Tensor! key,"
      "                 int head_size, Tensor cos_cache,"
      "                 Tensor sin_cache, bool is_neox) -> ()");
//   ops.def("rotary_embedding_tgi",&rotary_embedding_tgi);
  ops.impl("rotary_embedding_tgi", torch::kCUDA, &rotary_embedding_tgi);

397
398
399
400
  // Apply GPT-NeoX or GPT-J style rotary embedding to query and key
  // (supports multiple loras).
  ops.def(
      "batched_rotary_embedding(Tensor positions, Tensor! query,"
401
      "                         Tensor!? key, int head_size,"
402
403
404
405
406
      "                         Tensor cos_sin_cache, bool is_neox,"
      "                         int rot_dim,"
      "                         Tensor cos_sin_cache_offsets) -> ()");
  ops.impl("batched_rotary_embedding", torch::kCUDA, &batched_rotary_embedding);

zhuwenwen's avatar
zhuwenwen committed
407
408
409
410
  // trans w16
  ops.def("trans_w16_gemm(Tensor! dst, Tensor src, int row, int col) -> ()");
  ops.impl("trans_w16_gemm", torch::kCUDA, &trans_w16_gemm);

411
412
413
  // Quantization ops
#ifndef USE_ROCM
  // Quantized GEMM for AQLM.
414
415
416
  ops.def(
      "aqlm_gemm(Tensor input, Tensor codes, Tensor codebooks, "
      "Tensor scales, int[] codebook_partition_sizes, Tensor? bias) "
417
418
      "-> Tensor",
      {stride_tag});
419
420
421
  ops.impl("aqlm_gemm", torch::kCUDA, &aqlm_gemm);

  // Decompression method for AQLM.
422
423
  ops.def(
      "aqlm_dequant(Tensor codes, Tensor codebooks, "
424
425
      "int[] codebook_partition_sizes) -> Tensor",
      {stride_tag});
426
427
428
  ops.impl("aqlm_dequant", torch::kCUDA, &aqlm_dequant);

  // Quantized GEMM for AWQ.
429
430
  ops.def(
      "awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, "
431
432
      "Tensor _zeros, SymInt split_k_iters) -> Tensor",
      {stride_tag});
433
434
435
  ops.impl("awq_gemm", torch::kCUDA, &awq_gemm);

  // Dequantization for AWQ.
436
437
  ops.def(
      "awq_dequantize(Tensor _kernel, Tensor _scaling_factors, "
438
439
      "Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor",
      {stride_tag});
440
441
  ops.impl("awq_dequantize", torch::kCUDA, &awq_dequantize);

442
443
444
445
446
447
448
449
450
451
452
453
454
455
  // Note about marlin kernel 'workspace' arguments:
  // Technically these should be mutable since they are modified by the kernel.
  // But since they are set back to zero once the kernel is finished we can
  // hand wave and say that they have no net effect.
  //
  // The reason to mark 'workspace' as immutable is so that they don't interfere
  // with using ScalarType arguments in the ops. If they are marked as mutable,
  // pytorch throws an assert in
  // 'torch._higher_order_ops._register_effectful_op' that prevents these
  // kernels from being torch.compile'd.
  // See the following document for more info on custom types and ops that use
  // custom types:
  // https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA

456
  // Marlin (Dense) Optimized Quantized GEMM for GPTQ.
457
458
  ops.def(
      "marlin_gemm(Tensor a, Tensor b_q_weight, Tensor b_scales, "
459
      "Tensor! workspace, SymInt size_m, SymInt size_n, SymInt size_k) -> "
460
461
      "Tensor",
      {stride_tag});
462
  // conditionally compiled so impl in source file
463
464

  // Marlin_24 (Sparse) Optimized Quantized GEMM for GPTQ.
465
466
467
  ops.def(
      "gptq_marlin_24_gemm(Tensor a, Tensor b_q_weight, Tensor b_meta, "
      "Tensor b_scales, Tensor workspace, "
468
      "int b_q_type, "
469
470
      "SymInt size_m, SymInt size_n, SymInt size_k) -> Tensor",
      {stride_tag});
471
  //  conditionally compiled so impl in source file
472

473
474
  // Machete (Dense) Optimized Mixed Precision GEMM for Hopper.
  ops.def(
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
      "machete_supported_schedules("
      "   ScalarType a_type,"
      "   int b_type,"
      "   ScalarType? maybe_group_scales_type,"
      "   ScalarType? maybe_group_zeros_type,"
      "   ScalarType? maybe_channel_scales_type,"
      "   ScalarType? maybe_token_scales_type,"
      "   ScalarType? maybe_out_type"
      ") -> str[]");
  ops.def(
      "machete_mm("
      "   Tensor A,"
      "   Tensor B,"
      "   int b_type,"
      "   ScalarType? out_type,"
      "   Tensor? group_scales,"
      "   Tensor? group_zeros,"
      "   int?    group_size,"
      "   Tensor? channel_scales,"
      "   Tensor? token_scales,"
      "   str?    schedule"
496
497
      ") -> Tensor",
      {stride_tag});
498
499
500
501
502
503
504
  ops.def(
      "machete_prepack_B("
      "   Tensor B,"
      "   ScalarType a_type,"
      "   int b_type,"
      "   ScalarType? group_scales_type"
      ") -> Tensor");
505
  // conditionally compiled so impl registration is in source file
506

507
508
509
  ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
  ops.impl("permute_cols", torch::kCUDA, &permute_cols);

510
  // gptq_marlin Optimized Quantized GEMM for GPTQ.
511
  ops.def(
512
      "gptq_marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, "
513
514
      "Tensor b_scales, Tensor? global_scale, Tensor? b_zeros_or_none, Tensor? "
      "g_idx_or_none, Tensor? perm_or_none, Tensor workspace, int b_q_type, "
515
      "SymInt size_m, SymInt size_n, SymInt size_k, bool is_k_full, "
516
      "bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor",
517
      {stride_tag});
518
  // conditionally compiled so impl registration is in source file
519
520

  // gptq_marlin repack from GPTQ.
521
522
523
  ops.def(
      "gptq_marlin_repack(Tensor b_q_weight, Tensor perm, "
      "SymInt size_k, SymInt size_n, int num_bits) -> Tensor");
524
  // conditionally compiled so impl registrations are in source file
525

526
  // awq_marlin repack from AWQ.
527
528
529
  ops.def(
      "awq_marlin_repack(Tensor b_q_weight, SymInt size_k, "
      "SymInt size_n, int num_bits) -> Tensor");
530
  // conditionally compiled so impl registrations are in source file
531
#endif
532

533
  // Dequantization for GGML.
534
535
536
  ops.def(
      "ggml_dequantize(Tensor W, int type, SymInt m, SymInt n, ScalarType? "
      "dtype) -> Tensor");
537
538
539
  ops.impl("ggml_dequantize", torch::kCUDA, &ggml_dequantize);

  // mmvq kernel for GGML.
540
  ops.def(
541
      "ggml_mul_mat_vec_a8(Tensor W, Tensor X, int type, SymInt row) "
542
      "-> Tensor");
543
544
545
  ops.impl("ggml_mul_mat_vec_a8", torch::kCUDA, &ggml_mul_mat_vec_a8);

  // mmq kernel for GGML.
546
547
  ops.def(
      "ggml_mul_mat_a8(Tensor W, Tensor X, int type, SymInt row) -> Tensor");
548
549
  ops.impl("ggml_mul_mat_a8", torch::kCUDA, &ggml_mul_mat_a8);

550
551
552
553
554
555
556
557
  // moe kernel for GGML.
  ops.def(
      "ggml_moe_a8(Tensor X, Tensor W, "
      "Tensor sorted_token_ids, Tensor expert_ids, Tensor "
      "num_tokens_post_padded, "
      "int type, SymInt row, SymInt top_k, SymInt tokens) -> Tensor");
  ops.impl("ggml_moe_a8", torch::kCUDA, &ggml_moe_a8);

558
559
560
561
562
563
  ops.def(
      "ggml_moe_a8_vec(Tensor X, Tensor W, "
      "Tensor topk_ids, int top_k, "
      "int type, SymInt row, SymInt tokens) -> Tensor");
  ops.impl("ggml_moe_a8_vec", torch::kCUDA, &ggml_moe_a8_vec);

564
565
  ops.def("ggml_moe_get_block_size", &ggml_moe_get_block_size);

566
#ifndef USE_ROCM
567
  // marlin_qqq_gemm for QQQ.
568
569
570
  ops.def(
      "marlin_qqq_gemm(Tensor a, Tensor b_q_weight, "
      "Tensor s_tok, Tensor s_ch, Tensor s_group, "
571
      "Tensor! workspace, SymInt size_m, SymInt size_n, "
572
573
      "SymInt size_k) -> Tensor",
      {stride_tag});
574
  // conditionally compiled so impl registration is in source file
575

576
577
578
579
  // CUTLASS nvfp4 block scaled GEMM
  ops.def(
      "cutlass_scaled_fp4_mm(Tensor! out, Tensor a, Tensor b,"
      "                      Tensor block_scale_a, Tensor block_scale_b,"
580
581
      "                      Tensor alpha) -> ()",
      {stride_tag});
582
583
  ops.impl("cutlass_scaled_fp4_mm", torch::kCUDA, &cutlass_scaled_fp4_mm);

584
585
586
587
588
589
  // cutlass blockwise scaledgroup GEMM
  ops.def(
      "cutlass_blockwise_scaled_grouped_mm(Tensor! output, Tensor a, Tensor b, "
      "Tensor scales_a, Tensor scales_b, "
      "Tensor problem_sizes, Tensor expert_offsets) -> ()",
      {stride_tag});
590
  // conditionally compiled so impl registration is in source file
591

592
593
594
595
596
597
598
599
  // cutlass nvfp4 block scaled group GEMM
  ops.def(
      "cutlass_fp4_group_mm(Tensor! out, Tensor a, Tensor b,"
      " Tensor a_blockscale, Tensor b_blockscales, Tensor alphas,"
      " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()",
      {stride_tag});
  ops.impl("cutlass_fp4_group_mm", torch::kCUDA, &cutlass_fp4_group_mm);

600
  // CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column
601
  // quantization, as well as bias
602
  ops.def(
603
604
      "cutlass_scaled_mm(Tensor! out, Tensor a,"
      "                  Tensor b, Tensor a_scales,"
605
606
      "                  Tensor b_scales, Tensor? bias) -> ()",
      {stride_tag});
607
  ops.impl("cutlass_scaled_mm", torch::kCUDA, &cutlass_scaled_mm);
608

609
610
611
612
613
614
  // CUTLASS w8a8 GEMM, supporting asymmetric per-tensor or per-row/column
  // quantization.
  ops.def(
      "cutlass_scaled_mm_azp(Tensor! out, Tensor a,"
      "                  Tensor b, Tensor a_scales,"
      "                  Tensor b_scales, Tensor azp_adj,"
615
616
      "                  Tensor? azp, Tensor? bias) -> ()",
      {stride_tag});
617
618
  ops.impl("cutlass_scaled_mm_azp", torch::kCUDA, &cutlass_scaled_mm_azp);

619
620
  // Check if cutlass scaled_mm is supported for CUDA devices of the given
  // capability
621
622
623
  ops.def("cutlass_scaled_mm_supports_fp8(int cuda_device_capability) -> bool");
  ops.impl("cutlass_scaled_mm_supports_fp8", &cutlass_scaled_mm_supports_fp8);

624
625
626
627
628
629
630
631
632
633
  // Check if cutlass grouped gemm is supported for CUDA devices of the given
  // capability
  ops.def("cutlass_group_gemm_supported(int cuda_device_capability) -> bool");
  ops.impl("cutlass_group_gemm_supported", &cutlass_group_gemm_supported);

  // CUTLASS w8a8 grouped GEMM
  ops.def(
      "cutlass_moe_mm(Tensor! out_tensors, Tensor a_tensors, Tensor b_tensors, "
      "               Tensor a_scales, Tensor b_scales, Tensor expert_offsets, "
      "               Tensor problem_sizes, Tensor a_strides, "
634
635
      "               Tensor b_strides, Tensor c_strides, bool per_act_token, "
      "               bool per_out_ch) -> ()",
636
637
638
639
640
641
642
643
644
645
646
647
648
649
      {stride_tag});
  ops.impl("cutlass_moe_mm", torch::kCUDA, &cutlass_moe_mm);

  // A function that computes data required to run fused MoE with w8a8 grouped
  // GEMM. It takes topk_ids as an input, and computes expert_offsets
  // (token start indices of each expert). In addition to this, it computes
  // problem sizes for each expert's multiplication used by the two mms called
  // from fused MoE operation, and arrays with permutations required to shuffle
  // and de-shuffle the input/output of the fused operation.
  ops.def(
      "get_cutlass_moe_mm_data(Tensor topk_ids, Tensor! expert_offsets, "
      "                        Tensor! problem_sizes1, Tensor! problem_sizes2, "
      "                        Tensor! input_permutation, "
      "                        Tensor! output_permutation, int num_experts, "
650
      "                        int n, int k, Tensor? blockscale_offsets) -> ()",
651
652
653
      {stride_tag});
  ops.impl("get_cutlass_moe_mm_data", torch::kCUDA, &get_cutlass_moe_mm_data);

654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
  // A function that computes data required to run fused MoE with w8a8 grouped
  // GEMM and PPLX. It takes expert_num_tokens and non_zero_expert_idxs
  // as an input, and computes expert_offsets (token start indices of each
  // expert). In addition to this, it computes problem sizes for each expert's
  // multiplication used by the two mms called from fused MoE operation.
  ops.def(
      "get_cutlass_pplx_moe_mm_data(Tensor! expert_offsets, "
      "                             Tensor! problem_sizes1, "
      "                             Tensor! problem_sizes2, "
      "                             Tensor expert_num_tokens, "
      "                             int num_local_experts, int padded_m, "
      "                             int n, int k) -> ()",
      {stride_tag});
  ops.impl("get_cutlass_pplx_moe_mm_data", torch::kCUDA,
           &get_cutlass_pplx_moe_mm_data);

670
671
672
673
674
  // Check if cutlass scaled_mm supports block quantization (used by DeepSeekV3)
  ops.def(
      "cutlass_scaled_mm_supports_block_fp8(int cuda_device_capability) -> "
      "bool");
  ops.impl("cutlass_scaled_mm_supports_block_fp8",
675
           &cutlass_scaled_mm_supports_block_fp8);
676

677
678
679
680
681
682
683
  // Check if cutlass sparse scaled_mm is supported for CUDA devices of the
  // given capability
  ops.def(
      "cutlass_sparse_scaled_mm_supported(int cuda_device_capability) -> bool");
  ops.impl("cutlass_sparse_scaled_mm_supported",
           &cutlass_sparse_scaled_mm_supported);

684
685
686
687
688
689
  // CUTLASS sparse GEMM, supporting symmetric per-tensor or per-row/column
  // quantization, as well as bias
  ops.def(
      "cutlass_scaled_sparse_mm(Tensor! out, Tensor a,"
      "                         Tensor bt_nzs,"
      "                         Tensor bt_meta, Tensor a_scales,"
690
691
      "                         Tensor b_scales, Tensor? bias) -> ()",
      {stride_tag});
692
693
694
  ops.impl("cutlass_scaled_sparse_mm", torch::kCUDA, &cutlass_scaled_sparse_mm);

  // CUTLASS sparse matrix compressor
695
696
  ops.def("cutlass_sparse_compress(Tensor a) -> Tensor[]");
  ops.impl("cutlass_sparse_compress", &cutlass_sparse_compress);
697

698
  // CUTLASS MLA decode
699
  ops.def(
700
701
702
703
      "cutlass_mla_decode(Tensor! out, Tensor q_nope, Tensor q_pe,"
      "                   Tensor kv_c_and_k_pe_cache, Tensor seq_lens,"
      "                   Tensor page_table, float scale) -> ()");
  ops.impl("cutlass_mla_decode", torch::kCUDA, &cutlass_mla_decode);
704
705
706
707
708
709
710

  // Compute NVFP4 block quantized tensor.
  ops.def(
      "scaled_fp4_quant(Tensor! output, Tensor input,"
      "                 Tensor! output_scale, Tensor input_scale) -> ()");
  ops.impl("scaled_fp4_quant", torch::kCUDA, &scaled_fp4_quant);

711
712
713
714
715
716
717
  // Compute NVFP4 experts quantization.
  ops.def(
      "scaled_fp4_experts_quant(Tensor! output, Tensor! output_scale,"
      "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts,"
      "Tensor output_scale_offset_by_experts) -> ()");
  ops.impl("scaled_fp4_experts_quant", torch::kCUDA, &scaled_fp4_experts_quant);

718
719
720
721
  // Check if cutlass_scaled_mm_fp4 is supported for CUDA devices
  // of the given capability
  ops.def("cutlass_scaled_mm_supports_fp4(int cuda_device_capability) -> bool");
  ops.impl("cutlass_scaled_mm_supports_fp4", &cutlass_scaled_mm_supports_fp4);
722
723
724
#endif

  // Quantized GEMM for GPTQ.
725
726
  // Note: even though the C++ inferred schema is correct for this op, it seems
  // to prevent the meta function registry.
zhuwenwen's avatar
zhuwenwen committed
727

zhuwenwen's avatar
zhuwenwen committed
728
729
730
//   ops.def(
//       "gptq_gemm(Tensor a, Tensor b_q_weight, Tensor b_gptq_qzeros, "
//       "Tensor b_gptq_scales, Tensor b_g_idx, bool use_exllama, int bit) "
zhuwenwen's avatar
zhuwenwen committed
731
732
//       "-> Tensor",
//       {stride_tag});
733
//   ops.impl("gptq_gemm", torch::kCUDA, &gptq_gemm);
734
735

  // Post processing for GPTQ.
736
737
//   ops.def("gptq_shuffle(Tensor! q_weight, Tensor q_perm, int bit) -> ()");
//   ops.impl("gptq_shuffle", torch::kCUDA, &gptq_shuffle);
738
739

  // Compute FP8 quantized tensor for given scaling factor.
zhuwenwen's avatar
zhuwenwen committed
740
//   ops.def(
zhuwenwen's avatar
zhuwenwen committed
741
742
//       "static_scaled_fp8_quant(Tensor! result, Tensor input, Tensor scale) -> "
//       "()");
zhuwenwen's avatar
zhuwenwen committed
743
//   ops.impl("static_scaled_fp8_quant", torch::kCUDA, &static_scaled_fp8_quant);
744

zhuwenwen's avatar
zhuwenwen committed
745
//   // Compute dynamic-per-tensor FP8 quantized tensor and scaling factor.
zhuwenwen's avatar
zhuwenwen committed
746
//   ops.def(
zhuwenwen's avatar
zhuwenwen committed
747
748
//       "dynamic_scaled_fp8_quant(Tensor! result, Tensor input, Tensor! scale) "
//       "-> "
zhuwenwen's avatar
zhuwenwen committed
749
750
//       "()");
//   ops.impl("dynamic_scaled_fp8_quant", torch::kCUDA, &dynamic_scaled_fp8_quant);
751

zhuwenwen's avatar
zhuwenwen committed
752
//   // Compute dynamic-per-token FP8 quantized tensor and scaling factor.
753
//   ops.def(
zhuwenwen's avatar
zhuwenwen committed
754
//       "dynamic_per_token_scaled_fp8_quant(Tensor! result, Tensor input, "
zhuwenwen's avatar
zhuwenwen committed
755
//       "Tensor! scale, Tensor? scale_ub) -> "
756
757
758
//       "()");
//   ops.impl("dynamic_per_token_scaled_fp8_quant", torch::kCUDA,
//            &dynamic_per_token_scaled_fp8_quant);
759

760
761
  // Compute int8 quantized tensor for given scaling factor.
  ops.def(
762
      "static_scaled_int8_quant(Tensor! result, Tensor input, Tensor scale,"
763
      "Tensor? azp) -> ()");
764
765
766
767
  ops.impl("static_scaled_int8_quant", torch::kCUDA, &static_scaled_int8_quant);

  // Compute int8 quantized tensor and scaling factor
  ops.def(
768
      "dynamic_scaled_int8_quant(Tensor! result, Tensor input, Tensor! scale, "
769
      "Tensor!? azp) -> ()");
770
771
  ops.impl("dynamic_scaled_int8_quant", torch::kCUDA,
           &dynamic_scaled_int8_quant);
772

773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
  // Mamba selective scan kernel
  ops.def(
      "selective_scan_fwd(Tensor! u, Tensor! delta,"
      "Tensor! A, Tensor! B, Tensor! C,"
      "Tensor? D_, Tensor!? z_, Tensor? delta_bias_,"
      "bool delta_softplus,"
      "Tensor? query_start_loc,"
      "Tensor? cache_indices,"
      "Tensor? has_initial_state,"
      "Tensor! ssm_states,"
      "int pad_slot_id) -> ()");
  ops.impl("selective_scan_fwd", torch::kCUDA, &selective_scan_fwd);

  ops.def(
      "causal_conv1d_update(Tensor! x,"
      "Tensor! conv_state,"
      "Tensor! weight,"
      "Tensor? bias_,"
      "bool silu_activation,"
      "Tensor? cache_seqlens_,"
      "Tensor? conv_state_indices,"
      "int pad_slot_id) -> ()");
  ops.impl("causal_conv1d_update", torch::kCUDA, &causal_conv1d_update);

  ops.def(
      "causal_conv1d_fwd(Tensor! x, Tensor! weight,"
      "Tensor? bias_,"
      "Tensor!? conv_states,"
      "Tensor? query_start_loc,"
      "Tensor? cache_indices,"
      "Tensor? has_initial_state,"
      "bool silu_activation,"
      "int pad_slot_id) -> ()");
  ops.impl("causal_conv1d_fwd", torch::kCUDA, &causal_conv1d_fwd);

808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
#ifndef USE_ROCM
  // reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel
  ops.def(
      "rearrange_kn_weight_as_n32k16_order(Tensor b_qweight, Tensor b_scales, "
      "Tensor? b_zeros, "
      "bool has_zp, Tensor! b_qweight_reorder, Tensor! b_scales_reorder, "
      "Tensor!? b_zeros_reorder, "
      "int K, int N, int N_32align) -> ()");
  //  conditionally compiled so impl in source file

  // AllSpark quantization ops
  ops.def(
      "allspark_w8a16_gemm(Tensor a, Tensor b_qweight, Tensor b_scales, "
      "Tensor? b_qzeros, "
      "SymInt n, SymInt group_size, SymInt sm_count, SymInt sm_version, SymInt "
      "CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) -> Tensor");
  //  conditionally compiled so impl in source file
#endif
826
827
828
829
830
831
832
833
834
835
836
}

TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
  // Cache ops
  // Swap in (out) the cache blocks from src to dst.
  cache_ops.def(
      "swap_blocks(Tensor src, Tensor! dst, Tensor block_mapping) -> ()");
  cache_ops.impl("swap_blocks", torch::kCUDA, &swap_blocks);

  // Copy the cache blocks from src to dst.
  cache_ops.def(
837
838
      "copy_blocks(Tensor(a!)[] key_caches, Tensor[](b!) value_caches, "
      "Tensor block_mapping) -> ()");
839
840
  cache_ops.impl("copy_blocks", torch::kCUDA, &copy_blocks);

841
842
843
844
  cache_ops.def(
      "copy_blocks_mla(Tensor(a!)[] kv_caches, Tensor block_mapping) -> ()");
  cache_ops.impl("copy_blocks_mla", torch::kCUDA, &copy_blocks_mla);

845
846
847
848
849
850
  // Reshape the key and value tensors and cache them.
  cache_ops.def(
      "reshape_and_cache(Tensor key, Tensor value,"
      "                  Tensor! key_cache, Tensor! value_cache,"
      "                  Tensor slot_mapping,"
      "                  str kv_cache_dtype,"
851
      "                  Tensor k_scale, Tensor v_scale) -> ()");
852
853
  cache_ops.impl("reshape_and_cache", torch::kCUDA, &reshape_and_cache);

zhuwenwen's avatar
zhuwenwen committed
854
855
856
857
858
859
860
861
862
  // Reshape the key(new) and value tensors and cache them. 
  cache_ops.def(
       "reshape_and_cache_cuda(Tensor key, Tensor value, "
       "Tensor! key_cache, Tensor! value_cache, Tensor slot_mapping, "
       "str kv_cache_dtype, Tensor k_scale, Tensor v_scale) -> ()");
  cache_ops.impl("reshape_and_cache_cuda",
                  torch::kCUDA,
                  &reshape_and_cache_cuda);

863
864
865
866
867
868
  // Reshape the key and value tensors and cache them.
  cache_ops.def(
      "reshape_and_cache_flash(Tensor key, Tensor value,"
      "                        Tensor! key_cache,"
      "                        Tensor! value_cache,"
      "                        Tensor slot_mapping,"
869
      "                        str kv_cache_dtype,"
870
      "                        Tensor k_scale, Tensor v_scale) -> ()");
871
872
873
  cache_ops.impl("reshape_and_cache_flash", torch::kCUDA,
                 &reshape_and_cache_flash);

874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
  // read key and value form kv cache
  cache_ops.def(
      "read_cache(Tensor keys, Tensor values,"
      "                  Tensor[]! key_caches, Tensor[]! value_caches,"
      "                  Tensor slot_mapping,"
      "                  str kv_cache_dtype) -> ()");
  cache_ops.impl("read_cache", torch::kCUDA, &read_cache);

  // write multi-layers key and value to kv cache
  cache_ops.def(
      "write_cache_multi_layers(Tensor keys, Tensor values,"
      "                  Tensor[]! key_caches, Tensor[]! value_caches,"
      "                  Tensor slot_mapping,"
      "                  str kv_cache_dtype) -> ()");
  cache_ops.impl("write_cache_multi_layers", torch::kCUDA, &write_cache_multi_layers);

890
891
892
893
894
895
896
897
898
  // Concat kv_c and k_pe and cache them.
  cache_ops.def(
      "concat_and_cache_mla(Tensor kv_c, Tensor k_pe,"
      "                     Tensor! kv_cache,"
      "                     Tensor slot_mapping,"
      "                     str kv_cache_dtype,"
      "                     Tensor scale) -> ()");
  cache_ops.impl("concat_and_cache_mla", torch::kCUDA, &concat_and_cache_mla);

899
900
  // Convert the key and value cache to fp8 data type.
  cache_ops.def(
901
902
      "convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, "
      "str kv_cache_dtype) -> ()");
903
  cache_ops.impl("convert_fp8", torch::kCUDA, &convert_fp8);
904
905
906
907
908
909

  // Gather cache blocks from src_cache to dst.
  cache_ops.def(
      "gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, "
      "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()");
  cache_ops.impl("gather_cache", torch::kCUDA, &gather_cache);
910
911
912
913
914
915
}

TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
  // Cuda utils

  // Gets the specified device attribute.
916
917
  cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int");
  cuda_utils.impl("get_device_attribute", &get_device_attribute);
918
919

  // Gets the maximum shared memory per block device attribute.
920
921
  cuda_utils.def(
      "get_max_shared_memory_per_block_device_attribute(int device_id) -> int");
922
923
924
925
926
927
  cuda_utils.impl("get_max_shared_memory_per_block_device_attribute",
                  &get_max_shared_memory_per_block_device_attribute);
}

TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) {
  // Custom all-reduce kernels
928
  custom_ar.def(
929
      "init_custom_ar(int[] ipc_tensors, Tensor rank_data, "
930
      "int rank, bool fully_connected) -> int");
931
932
  custom_ar.impl("init_custom_ar", torch::kCUDA, &init_custom_ar);
  custom_ar.def(
933
934
935
      "all_reduce(int fa, Tensor inp, Tensor! out, int reg_buffer, "
      "int reg_buffer_sz_bytes) -> ()");
  custom_ar.impl("all_reduce", torch::kCUDA, &all_reduce);
936
937
938
939

  custom_ar.def("dispose", &dispose);
  custom_ar.def("meta_size", &meta_size);

940
  custom_ar.def("register_buffer", &register_buffer);
941
942
  custom_ar.def("get_graph_buffer_ipc_meta", &get_graph_buffer_ipc_meta);
  custom_ar.def("register_graph_buffers", &register_graph_buffers);
943

zhuwenwen's avatar
zhuwenwen committed
944
  custom_ar.def("allocate_shared_buffer_and_handle",
945
                &allocate_shared_buffer_and_handle);
zhuwenwen's avatar
zhuwenwen committed
946
947
948
949
  custom_ar.def("open_mem_handle(Tensor mem_handle) -> int", &open_mem_handle);
  custom_ar.impl("open_mem_handle", torch::kCPU, &open_mem_handle);

  custom_ar.def("free_shared_buffer", &free_shared_buffer);
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
#ifdef USE_ROCM
  // Quick Reduce all-reduce kernels
  custom_ar.def(
      "qr_all_reduce(int fa, Tensor inp, Tensor out, int quant_level, bool "
      "cast_bf2half) -> ()");
  custom_ar.impl("qr_all_reduce", torch::kCUDA, &qr_all_reduce);

  custom_ar.def("init_custom_qr", &init_custom_qr);
  custom_ar.def("qr_destroy", &qr_destroy);

  custom_ar.def("qr_get_handle", &qr_get_handle);

  custom_ar.def("qr_open_handles(int _fa, Tensor[](b!) handles) -> ()");
  custom_ar.impl("qr_open_handles", torch::kCPU, &qr_open_handles);

  // Max input size in bytes
  custom_ar.def("qr_max_size", &qr_max_size);
#endif
968
969
970
}

REGISTER_EXTENSION(TORCH_EXTENSION_NAME)