| 1 | #include "core/registration.h" |
| 2 | #include "rocm/ops.h" |
| 3 | |
| 4 | // Note on op signatures: |
| 5 | // The X_meta signatures are for the meta functions corresponding to op X. |
| 6 | // They must be kept in sync with the signature for X. Generally, only |
| 7 | // functions that return Tensors require a meta function. |
| 8 | // |
| 9 | // See the following links for detailed docs on op registration and function |
| 10 | // schemas. |
| 11 | // https://docs.google.com/document/d/1_W62p8WJOQQUzPsJYa7s701JXt0qf2OfLub2sbkHOaU/edit#heading=h.ptttacy8y1u9 |
| 12 | // https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/README.md#annotations |
| 13 | |
| 14 | TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) { |
| 15 | // vLLM custom ops for rocm |
| 16 | |
| 17 | // skinny_gemms.cu (LLMM1/wvSplitK/wvSplitKrc/wvSplitKQ) is excluded on gfx1250 |
| 18 | // (gfx9/gfx11 ISA, unsupported there); skip these registrations to avoid |
| 19 | // undefined symbols. vLLM uses default/Triton GEMM for these ops on gfx1250. |
| 20 | #ifndef VLLM_SKIP_SKINNY_GEMMS |
| 21 | // Custom gemm op for matrix-vector multiplication |
| 22 | rocm_ops.def( |
| 23 | "LLMM1(Tensor in_a, Tensor in_b, int rows_per_block) -> " |
| 24 | "Tensor"); |
| 25 | rocm_ops.impl("LLMM1", torch::kCUDA, &LLMM1); |
| 26 | |
| 27 | // Custom gemm op for skinny matrix-matrix multiplication |
| 28 | rocm_ops.def( |
| 29 | "wvSplitK(Tensor in_a, Tensor in_b, Tensor? in_bias, int CuCount) -> " |
| 30 | "Tensor"); |
| 31 | rocm_ops.impl("wvSplitK", torch::kCUDA, &wvSplitK); |
| 32 | |
| 33 | // W4A16 grouped skinny GEMM: packed int4 weights, per-group scales, |
| 34 | // optional zero points [M/8, K/group_size] int32 for asymmetric |
| 35 | // quantization |
| 36 | rocm_ops.def( |
| 37 | "wvSplitK_int4_g(Tensor in_a, Tensor in_b, Tensor in_scale, " |
| 38 | "Tensor? in_zero_points, Tensor? in_bias, int CuCount, " |
| 39 | "int group_size) -> Tensor"); |
| 40 | rocm_ops.impl("wvSplitK_int4_g", torch::kCUDA, &wvSplitK_int4_g); |
| 41 | |
| 42 | // Custom gemm op for skinny matrix-matrix multiplication |
| 43 | rocm_ops.def( |
| 44 | "wvSplitKrc(Tensor in_a, Tensor in_b, Tensor? in_bias, int CuCount) -> " |
| 45 | "Tensor"); |
| 46 | rocm_ops.impl("wvSplitKrc", torch::kCUDA, &wvSplitKrc); |
| 47 | |
| 48 | // wvSplitK for fp8 |
| 49 | rocm_ops.def( |
| 50 | "wvSplitKQ(Tensor in_a, Tensor in_b, Tensor? in_bias, Tensor! out_c, " |
| 51 | "Tensor scale_a, " |
| 52 | " Tensor scale_b, int CuCount) -> ()"); |
| 53 | rocm_ops.impl("wvSplitKQ", torch::kCUDA, &wvSplitKQ); |
| 54 | #endif // VLLM_SKIP_SKINNY_GEMMS |
| 55 | |
| 56 | #ifdef VLLM_ROCM_GFX1100 |
| 57 | // W4A16 GPTQ kernels for AMD RDNA3 (gfx1100). |
| 58 | rocm_ops.def( |
| 59 | "gptq_gemm_rdna3(Tensor a, Tensor b_q_weight, Tensor b_qzeros, " |
| 60 | "Tensor b_scales, bool use_v2_format) -> Tensor"); |
| 61 | rocm_ops.impl("gptq_gemm_rdna3", torch::kCUDA, &gptq_gemm_rdna3); |
| 62 | |
| 63 | rocm_ops.def( |
| 64 | "gptq_gemm_rdna3_wmma(Tensor a, Tensor b_q_weight, Tensor b_qzeros, " |
| 65 | "Tensor b_scales, bool use_v2_format) -> Tensor"); |
| 66 | rocm_ops.impl("gptq_gemm_rdna3_wmma", torch::kCUDA, &gptq_gemm_rdna3_wmma); |
| 67 | |
| 68 | rocm_ops.def( |
| 69 | "moe_gptq_gemm_rdna3(Tensor a, Tensor! c, Tensor b_q_weight, " |
| 70 | "Tensor b_scales, Tensor b_qzeros, Tensor topk_weights, " |
| 71 | "Tensor sorted_token_ids, Tensor expert_ids, " |
| 72 | "Tensor num_tokens_post_padded, " |
| 73 | "int top_k, int block_size_m, bool mul_topk_weight, " |
| 74 | "int output_topk) -> ()"); |
| 75 | rocm_ops.impl("moe_gptq_gemm_rdna3", torch::kCUDA, &moe_gptq_gemm_rdna3); |
| 76 | #endif |
| 77 | |
| 78 | // Custom attention op |
| 79 | // Compute the attention between an input query and the cached |
| 80 | // keys/values using PagedAttention. |
| 81 | rocm_ops.def( |
| 82 | "paged_attention(Tensor! out, Tensor exp_sums," |
| 83 | " Tensor max_logits, Tensor tmp_out," |
| 84 | " Tensor query, Tensor key_cache," |
| 85 | " Tensor value_cache, int num_kv_heads," |
| 86 | " float scale, Tensor block_tables," |
| 87 | " Tensor seq_lens," |
| 88 | " Tensor? query_start_loc," |
| 89 | " int block_size," |
| 90 | " int max_seq_len," |
| 91 | " Tensor? alibi_slopes," |
| 92 | " str kv_cache_dtype," |
| 93 | " Tensor k_scale, Tensor v_scale," |
| 94 | " Tensor? fp8_out_scale," |
| 95 | " str mfma_type) -> ()"); |
| 96 | rocm_ops.impl("paged_attention", torch::kCUDA, &paged_attention); |
| 97 | } |
| 98 | |
| 99 | REGISTER_EXTENSION(TORCH_EXTENSION_NAME) |