Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Kernel] LoRA - Enable CUDAGraphs for V1 #14626

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
+27 −14
Open

[Kernel] LoRA - Enable CUDAGraphs for V1 #14626

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f92f3e2
add cudagraph support
Mar 10, 2025
a085845
fix input shape inference
Mar 12, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 14 additions & 8 deletions vllm/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2282,9 +2282,14 @@ def compute_hash(self) -> str:
excluding anything before input ids/embeddings and after
the final hidden states.
"""
# no factors to consider.
# LoRA is not compatible with `torch.compile` .
factors: list[Any] = []
factors.append(self.max_lora_rank)
factors.append(self.max_loras)
factors.append(self.fully_sharded_loras)
factors.append(self.lora_dtype)
factors.append(self.lora_extra_vocab_size)
factors.append(self.long_lora_scaling_factors)
factors.append(self.bias_enabled)
hash_str = hashlib.md5(str(factors).encode()).hexdigest()
return hash_str

Expand Down Expand Up @@ -3290,6 +3295,13 @@ def compute_hash(self) -> str:
vllm_factors.append("None")
if self.lora_config:
vllm_factors.append(self.lora_config.compute_hash())
# LoRA creates static buffers based on max_num_batched_tokens.
# The tensor sizes and strides get captured in the torch.compile
# graph explicitly.
vllm_factors.append(
hashlib.md5(
str(self.scheduler_config.max_num_batched_tokens).encode()
).hexdigest())
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

During torch.compile, LoRA static buffers like in

vllm/vllm/lora/punica_wrapper/punica_base.py

Line 133 in 5305673

self._token_lora_indices = torch.empty(max_num_batched_tokens,
and

vllm/vllm/lora/ops/triton_ops/v1/v1_kernel_metadata.py

Line 24 in 5305673

token_lora_mapping = torch.empty(max_num_tokens,
get captured along with their sizes and strides (they aren't dynamic)

When max_num_batched_tokens changes, and when the captured graph is executed, we hit assert_size_stride asserts on these tensors. As a solution, we simply recompile when max_num_batched_tokens change.

else:
vllm_factors.append("None")
if self.speculative_config:
Expand Down Expand Up @@ -3440,12 +3452,6 @@ def __post_init__(self):
" Disabling `torch.compile`.")
self.compilation_config.level = CompilationLevel.NO_COMPILATION

if self.lora_config is not None and self.compilation_config.level !=\
CompilationLevel.NO_COMPILATION:
logger.warning("LoRA is not supported with `torch.compile` yet. "
"Disabling `torch.compile`.")
self.compilation_config.level = CompilationLevel.NO_COMPILATION

if self.model_config and self.model_config.use_mla and \
not current_platform.is_cuda():
logger.info(
Expand Down
11 changes: 7 additions & 4 deletions vllm/lora/layers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -237,16 +237,19 @@ def set_lora(
self.embeddings_weights[:embeddings.shape[0]].copy_(embeddings)

def forward(self, x: torch.Tensor) -> torch.Tensor:
added_tokens_mask = x > self.base_layer.org_vocab_size - 1
embeddings_indices = self.punica_wrapper.embeddings_indices
added_tokens_mask = torch.where(x > self.base_layer.org_vocab_size - 1,
1, 0)
embeddings_indices = torch.narrow(
self.punica_wrapper._embeddings_indices, 1, 0, x.size(0))

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

^ changes are to avoid errors such as,

  raise ConstraintViolationError(
torch.fx.experimental.symbolic_shapes.ConstraintViolationError: Constraints violated (L['input_ids'].size()[0], L['positions'].size()[0])! For more information, run with TORCH_LOGS="+dynamic".
  - Not all values of RelaxedUnspecConstraint(L['input_ids'].size()[0]) are valid because L['input_ids'].size()[0] was inferred to be a constant (8192).
  - Not all values of RelaxedUnspecConstraint(L['positions'].size()[0]) are valid because L['positions'].size()[0] was inferred to be a constant (8192).

indices = embeddings_indices[1].view_as(x)
full_lora_a_embeddings = F.embedding(
x + indices,
self.lora_a_stacked_2d,
)
indices = embeddings_indices[0].view_as(x)
full_output = self.base_layer.forward(
x.add_(indices * added_tokens_mask))
full_output = self.base_layer.forward(x +
(indices * added_tokens_mask))
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

x here is the input_ids. In V1, we don't zero out the cuda graph pad region.
Avoid the in-place update here to prevent accumulating garbage into the input buffer.


full_output_org = full_output
if full_output.ndim == 3:
Expand Down