wehub-resource-sync
24 lines
700 B
C
24 lines
700 B
C
// Copyright (c) Microsoft Corporation.
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
// DeepSpeed Team
|
|
|
|
#pragma once
|
|
|
|
#include <c10/cuda/CUDAStream.h>
|
|
#include <torch/extension.h>
|
|
#include "embed.cuh"
|
|
|
|
/*
|
|
Embeddings kernel aware of ragged batch structure.
|
|
*/
|
|
void ragged_embed(torch::Tensor& embedded_tokens,
|
|
torch::Tensor& input_ids,
|
|
torch::Tensor& embedding_weight,
|
|
c10::optional<torch::Tensor>& position_weight,
|
|
int32_t position_embed_offset,
|
|
torch::Tensor& batch_metadata,
|
|
torch::Tensor& seq_metadata,
|
|
torch::Tensor& tokens_to_seq,
|
|
torch::Tensor& kv_ptrs);
|