| 26 | |
| 27 | @trtllm_plugin("TritonLookUp") |
| 28 | class LookUpPlugin(PluginBase): |
| 29 | |
| 30 | def __init__(self, use_torch_tensor, fp32_output): |
| 31 | super().__init__() |
| 32 | self.use_torch_tensor = use_torch_tensor |
| 33 | self.fp32_output = fp32_output |
| 34 | |
| 35 | def shape_dtype_inference(self, inputs: Sequence[SymTensor]) -> SymTensor: |
| 36 | shape = inputs[1].shape |
| 37 | shape[0] = inputs[0].shape[0] + inputs[1].shape[0] - inputs[1].shape[0] |
| 38 | return SymTensor( |
| 39 | inputs[1].dtype if not self.fp32_output else torch.float32, shape) |
| 40 | |
| 41 | def forward(self, inputs: Sequence[TensorWrapper], |
| 42 | outputs: Sequence[TensorWrapper]): |
| 43 | assert len(inputs) == 2 |
| 44 | assert inputs[0].dtype in [torch.int32 or torch.int64] |
| 45 | assert inputs[1].dtype in [torch.float32, torch.float16, torch.bfloat16] |
| 46 | assert (self.fp32_output and outputs[0].dtype |
| 47 | == torch.float32) or outputs[0].dtype == inputs[1].dtype |
| 48 | |
| 49 | x = inputs[0] |
| 50 | y = inputs[1] |
| 51 | z = outputs[0] |
| 52 | if self.use_torch_tensor: |
| 53 | x = convert_to_torch_tensor(x) |
| 54 | y = convert_to_torch_tensor(y) |
| 55 | z = convert_to_torch_tensor(z) |
| 56 | MAX_BLOCK_NUM = 65536 |
| 57 | MAX_BLOCK_SIZE = 512 |
| 58 | grid = lambda meta: (min(MAX_BLOCK_NUM, x.shape[0]) * min( |
| 59 | MAX_BLOCK_SIZE, y.shape[1]), ) |
| 60 | lookup_kernel[grid](x, y, z, y.shape[0], y.shape[1], x.shape[0]) |
no outgoing calls