Skip to content

vllm.lora.ops.xla_ops

Modules:

Name Description
lora_ops

__all__ module-attribute 

__all__ = [
    "bgmv_expand",
    "bgmv_expand_slice",
    "bgmv_shrink",
]

bgmv_expand 

bgmv_expand(
    inputs: Tensor,
    lora_b_weights: Tensor,
    output_tensor: Tensor,
    lora_indices_tensor: Tensor,
    add_inputs: bool = True,
)

Parameters:

Name Type Description Default
inputs Tensor

Input tensor of shape [num_tokens, hidden_size].

 required
lora_b_weights Tensor

LoRA weights of shape [num_loras, lora_rank, hidden_size].

 required
output_tensor Tensor

output tensor of shape [num_tokens, hidden_size * num_slices].

 required
lora_indices_tensor Tensor

Tensor of shape [num_tokens] indicating which LoRA matrix to use for each token.

 required
add_inputs bool

Whether or not to add the input tensor to the output tensor.

 True
Source code in vllm/lora/ops/xla_ops/lora_ops.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
def bgmv_expand(
    inputs: torch.Tensor,
    lora_b_weights: torch.Tensor,
    output_tensor: torch.Tensor,
    lora_indices_tensor: torch.Tensor,
    add_inputs: bool = True,
):
    """
    Args:
        inputs (torch.Tensor): Input tensor of shape [num_tokens, hidden_size].

        lora_b_weights (torch.Tensor): LoRA weights of shape
            [num_loras, lora_rank, hidden_size].

        output_tensor (torch.Tensor): output tensor of shape
            [num_tokens, hidden_size * num_slices].

        lora_indices_tensor (torch.Tensor): Tensor of shape [num_tokens]
            indicating which LoRA matrix to use for each token.
        add_inputs (bool): Whether or not to add the input tensor to the output
            tensor.
    """

    outputs = torch.ops.xla.bgmv(inputs, lora_b_weights, lora_indices_tensor)

    limit = output_tensor.shape[0]
    if outputs.shape[0] == 1 and output_tensor.shape[0] != 1:
        limit = 1

    if output_tensor.shape[1] > outputs.shape[1]:
        outputs = F.pad(outputs, (0, output_tensor.shape[1] - outputs.shape[1], 0, 0))

    if add_inputs:
        return output_tensor + outputs[:limit, : output_tensor.shape[1]]
    else:
        return outputs[:limit, : output_tensor.shape[1]]

bgmv_expand_slice 

bgmv_expand_slice(
    inputs: Tensor,
    lora_b_weights: Tensor,
    output_tensor: Tensor,
    lora_indices_tensor: Tensor,
    slice_offset: int,
    slice_size: int,
    add_inputs: bool = True,
)

Parameters:

Name Type Description Default
inputs Tensor

Input tensor of shape [num_tokens, hidden_size].

 required
lora_b_weights Tensor

LoRA weights of shape [num_loras, lora_rank, hidden_size].

 required
output_tensor Tensor

output tensor of shape [num_tokens, hidden_size * num_slices].

 required
lora_indices_tensor Tensor

Tensor of shape [num_tokens] indicating which LoRA matrix to use for each token.

 required
add_inputs bool

Whether or not to add the input tensor to the output tensor.

 True
Source code in vllm/lora/ops/xla_ops/lora_ops.py 
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
def bgmv_expand_slice(
    inputs: torch.Tensor,
    lora_b_weights: torch.Tensor,
    output_tensor: torch.Tensor,
    lora_indices_tensor: torch.Tensor,
    slice_offset: int,
    slice_size: int,
    add_inputs: bool = True,
):
    """
    Args:
        inputs (torch.Tensor): Input tensor of shape [num_tokens, hidden_size].

        lora_b_weights (torch.Tensor): LoRA weights of shape
            [num_loras, lora_rank, hidden_size].

        output_tensor (torch.Tensor): output tensor of shape
            [num_tokens, hidden_size * num_slices].

        lora_indices_tensor (torch.Tensor): Tensor of shape [num_tokens]
            indicating which LoRA matrix to use for each token.
        add_inputs (bool): Whether or not to add the input tensor to the output
            tensor.
    """
    outputs = torch.ops.xla.bgmv(inputs, lora_b_weights, lora_indices_tensor)

    outputs = F.pad(
        outputs,
        (
            slice_offset,
            output_tensor.shape[1] - (slice_offset + slice_size),
            0,
            0,
        ),
    )

    if add_inputs:
        return output_tensor + outputs
    else:
        return outputs

bgmv_shrink 

bgmv_shrink(
    inputs: Tensor,
    lora_b_weights: Tensor,
    lora_indices_tensor: Tensor,
    scaling: float = 1.0,
)

Parameters:

Name Type Description Default
inputs Tensor

Input tensor of shape [num_tokens, hidden_size].

 required
lora_b_weights Tensor

LoRA weights of shape [num_loras, lora_rank, hidden_size].

 required
lora_indices_tensor Tensor

Tensor of shape [num_tokens] indicating which LoRA matrix to use for each token.

 required
scaling float

Scalar multiplier applied to the output.

 1.0
Source code in vllm/lora/ops/xla_ops/lora_ops.py 
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
def bgmv_shrink(
    inputs: torch.Tensor,
    lora_b_weights: torch.Tensor,
    lora_indices_tensor: torch.Tensor,
    scaling: float = 1.0,
):
    """
    Args:
        inputs (torch.Tensor): Input tensor of shape [num_tokens, hidden_size].
        lora_b_weights (torch.Tensor): LoRA weights of shape
            [num_loras, lora_rank, hidden_size].
        lora_indices_tensor (torch.Tensor): Tensor of shape [num_tokens]
            indicating which LoRA matrix to use for each token.
        scaling (float, optional): Scalar multiplier applied to the output.
    """

    return scaling * torch.ops.xla.bgmv(inputs, lora_b_weights, lora_indices_tensor)