@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 715681c127
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 5, 06:00 UTC

[LoRA] Support dual CUDA streams-Linear Layer (#35721)

Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>

Jee Jee Li committed Apr 13, 2026 at 10:57 UTC 715681c12710edc2b6bfb63236b09ac93cade793
10 files changed +321 -104
tests/lora/conftest.py
+7
@@ -43,6 +43,13 @@ def cleanup_fixture(should_do_global_cleanup_after_test: bool):
43 cleanup_dist_env_and_memory(shutdown_ray=True)
44
45
46 +@pytest.fixture
47 +def maybe_enable_lora_dual_stream(monkeypatch: pytest.MonkeyPatch):
48 + if current_platform.is_cuda():
49 + monkeypatch.setenv("VLLM_LORA_ENABLE_DUAL_STREAM", "1")
50 + yield
51 +
52 +
53 @pytest.fixture
54 def dist_init():
55 from tests.utils import ensure_current_vllm_config
tests/lora/test_layers.py
+43 -15
@@ -521,8 +521,10 @@ def test_linear_replicated(
521 punica_wrapper = get_punica_wrapper(8192, 256, device, lora_config=lora_config)
522 assert check_punica_wrapper(punica_wrapper)
523
524 - def create_random_linear_replicated_layer():
525 - linear = ReplicatedLinear(4096, 4096, bias=False, params_dtype=torch.float16)
524 + def create_random_linear_replicated_layer(idx: int = 0):
525 + linear = ReplicatedLinear(
526 + 4096, 4096, bias=False, params_dtype=torch.float16, prefix=f"layer_{idx}"
527 + )
528 linear.weight.data = torch.rand_like(linear.weight.data)
529 lora_linear = ReplicatedLinearWithLoRA(linear)
530
@@ -539,7 +541,7 @@ def test_linear_replicated(
541 set_random_seed(i)
542
543 id_to_index = get_random_id_to_index(num_loras, max_loras)
542 - linear, lora_linear = create_random_linear_replicated_layer()
544 + linear, lora_linear = create_random_linear_replicated_layer(i)
545 assert torch.equal(linear.weight, lora_linear.weight)
546 lora_linear.set_mapping(punica_wrapper)
547 lora_dict, _ = populate_loras(
@@ -629,10 +631,14 @@ def test_linear_parallel(
631 punica_wrapper = get_punica_wrapper(8192, 256, device, lora_config=lora_config)
632 assert check_punica_wrapper(punica_wrapper)
633
632 - def create_random_linear_parallel_layer():
634 + def create_random_linear_parallel_layer(idx: int = 0):
635 if orientation == "row":
636 linear = RowParallelLinear(
635 - 4096, 4096, bias=False, params_dtype=torch.float16
637 + 4096,
638 + 4096,
639 + bias=False,
640 + params_dtype=torch.float16,
641 + prefix=f"layer_{idx}",
642 )
643 linear.weight.data = torch.rand_like(linear.weight.data)
644 lora_linear = (
@@ -642,7 +648,11 @@ def test_linear_parallel(
648 )
649 else:
650 linear = ColumnParallelLinear(
645 - 4096, 4096, bias=False, params_dtype=torch.float16
651 + 4096,
652 + 4096,
653 + bias=False,
654 + params_dtype=torch.float16,
655 + prefix=f"layer_{idx}",
656 )
657 linear.weight.data = torch.rand_like(linear.weight.data)
658 lora_linear = (
@@ -664,7 +674,7 @@ def test_linear_parallel(
674 set_random_seed(i)
675
676 id_to_index = get_random_id_to_index(num_loras, max_loras)
667 - linear, lora_linear = create_random_linear_parallel_layer()
677 + linear, lora_linear = create_random_linear_parallel_layer(i)
678 assert torch.equal(linear.weight, lora_linear.weight)
679 lora_linear.set_mapping(punica_wrapper)
680 lora_dict, _ = populate_loras(
@@ -754,10 +764,14 @@ def test_column_parallel_packed(
764 punica_wrapper = get_punica_wrapper(8192, 256, device, lora_config=lora_config)
765 assert check_punica_wrapper(punica_wrapper)
766
757 - def create_column_parallel_packed_layer():
767 + def create_column_parallel_packed_layer(idx: int = 0):
768 if repeats == 2:
769 linear = MergedColumnParallelLinear(
760 - 4096, [4096] * repeats, bias=False, params_dtype=torch.float16
770 + 4096,
771 + [4096] * repeats,
772 + bias=False,
773 + params_dtype=torch.float16,
774 + prefix=f"layer_{idx}",
775 )
776 linear.weight.data = torch.rand_like(linear.weight.data)
777 lora_linear = (
@@ -767,7 +781,12 @@ def test_column_parallel_packed(
781 )
782 elif repeats == 3:
783 linear = QKVParallelLinear(
770 - 4096, 64, 32, bias=False, params_dtype=torch.float16
784 + 4096,
785 + 64,
786 + 32,
787 + bias=False,
788 + params_dtype=torch.float16,
789 + prefix=f"layer_{idx}",
790 )
791 linear.weight.data = torch.rand_like(linear.weight.data)
792 lora_linear = (
@@ -777,7 +796,12 @@ def test_column_parallel_packed(
796 )
797 else:
798 linear = QKVParallelLinear(
780 - 4096, 64, 32, bias=False, params_dtype=torch.float16
799 + 4096,
800 + 64,
801 + 32,
802 + bias=False,
803 + params_dtype=torch.float16,
804 + prefix=f"layer_{idx}",
805 )
806 linear.weight.data = torch.rand_like(linear.weight.data)
807 lora_linear = (
@@ -810,7 +834,7 @@ def test_column_parallel_packed(
834
835 id_to_index = get_random_id_to_index(num_loras, max_loras)
836
813 - linear, lora_linear = create_column_parallel_packed_layer()
837 + linear, lora_linear = create_column_parallel_packed_layer(i)
838 assert torch.equal(linear.weight, lora_linear.weight)
839 lora_linear.set_mapping(punica_wrapper)
840 lora_dict, sublora_dict = populate_loras(
@@ -902,10 +926,14 @@ def test_merged_column_parallel_variable_slice(
926 output_sizes = [1024 + i * 256 for i in range(num_slices)]
927 total_output = sum(output_sizes)
928
905 - def create_layer():
929 + def create_layer(idx: int = 0):
930 # Create linear layer
931 linear = MergedColumnParallelLinear(
908 - 4096, output_sizes, bias=False, params_dtype=torch.float16
932 + 4096,
933 + output_sizes,
934 + bias=False,
935 + params_dtype=torch.float16,
936 + prefix=f"layer_{idx}",
937 )
938 linear.weight.data = torch.rand_like(linear.weight.data)
939
@@ -917,7 +945,7 @@ def test_merged_column_parallel_variable_slice(
945 for i in range(NUM_RANDOM_SEEDS):
946 set_random_seed(i)
947 id_to_index = get_random_id_to_index(num_loras, max_loras)
920 - linear, lora_linear = create_layer()
948 + linear, lora_linear = create_layer(i)
949 lora_linear.set_mapping(punica_wrapper)
950
951 # Populate LoRA weights
tests/lora/test_olmoe_tp.py
+4 -2
@@ -110,7 +110,7 @@ def generate_and_test(
110 )
111
112
113 -def test_olmoe_lora(olmoe_lora_files):
113 +def test_olmoe_lora(olmoe_lora_files, maybe_enable_lora_dual_stream):
114 # We enable enforce_eager=True here to reduce VRAM usage for lora-test CI,
115 # Otherwise, the lora-test will fail due to CUDA OOM.
116 llm = vllm.LLM(
@@ -141,7 +141,9 @@ def test_olmoe_lora_mixed(olmoe_lora_files):
141 generate_and_test(llm, olmoe_lora_files, lora_id=[1, None, 3, None])
142
143
144 -def test_olmoe_lora_mixed_random(olmoe_lora_files, tmp_path):
144 +def test_olmoe_lora_mixed_random(
145 + olmoe_lora_files, tmp_path, maybe_enable_lora_dual_stream
146 +):
147 # Create a dummy LoRA with random weights based on the real one
148 random_lora_path = tmp_path / "random_lora"
149 shutil.copytree(olmoe_lora_files, random_lora_path)
tests/lora/test_qwen35_densemodel_lora.py
+6 -2
@@ -312,7 +312,9 @@ def _assert_qwen35_text_vl_and_mixed_lora(
312
313
314 @create_new_process_for_each_test()
315 -def test_qwen35_text_lora(qwen35_text_lora_files, qwen35_vl_lora_files):
315 +def test_qwen35_text_lora(
316 + qwen35_text_lora_files, qwen35_vl_lora_files, maybe_enable_lora_dual_stream
317 +):
318 llm = vllm.LLM(
319 model=MODEL_PATH,
320 max_model_len=4096,
@@ -335,7 +337,9 @@ def test_qwen35_text_lora(qwen35_text_lora_files, qwen35_vl_lora_files):
337
338
339 @multi_gpu_test(num_gpus=4)
338 -def test_qwen35_text_lora_tp4(qwen35_text_lora_files, qwen35_vl_lora_files):
340 +def test_qwen35_text_lora_tp4(
341 + qwen35_text_lora_files, qwen35_vl_lora_files, maybe_enable_lora_dual_stream
342 +):
343 llm = vllm.LLM(
344 model=MODEL_PATH,
345 max_model_len=4096,
vllm/config/lora.py
+10 -1
@@ -7,8 +7,10 @@ import torch
7 from pydantic import ConfigDict, Field, model_validator
8 from typing_extensions import Self
9
10 +from vllm import envs
11 from vllm.config.utils import config
12 from vllm.logger import init_logger
13 +from vllm.platforms import current_platform
14 from vllm.utils.hashing import safe_hash
15
16 if TYPE_CHECKING:
@@ -105,7 +107,14 @@ class LoRAConfig:
107 f"max_cpu_loras ({self.max_cpu_loras}) must be >= "
108 f"max_loras ({self.max_loras})."
109 )
108 -
110 + if envs.VLLM_LORA_ENABLE_DUAL_STREAM and not current_platform.is_cuda_alike():
111 + raise ValueError("Dual CUDA streams are only supported on CUDA platforms.")
112 + if envs.VLLM_LORA_ENABLE_DUAL_STREAM and self.fully_sharded_loras:
113 + logger.warning_once(
114 + "fully_sharded_loras isn't compatible with "
115 + "VLLM_LORA_ENABLE_DUAL_STREAM, set VLLM_LORA_ENABLE_DUAL_STREAM=False"
116 + )
117 + envs.VLLM_LORA_ENABLE_DUAL_STREAM = False
118 return self
119
120 def verify_with_model_config(self, model_config: ModelConfig):
vllm/envs.py
+78 -71
@@ -258,6 +258,7 @@ if TYPE_CHECKING:
258 VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS: bool = False
259 VLLM_NIXL_EP_MAX_NUM_RANKS: int = 32
260 VLLM_XPU_ENABLE_XPU_GRAPH: bool = False
261 + VLLM_LORA_ENABLE_DUAL_STREAM: bool = False
262
263
264 def get_default_cache_root():
@@ -496,8 +497,9 @@ environment_variables: dict[str, Callable[[], Any]] = {
497 # rocm, cpu]
498 "VLLM_TARGET_DEVICE": lambda: os.getenv("VLLM_TARGET_DEVICE", "cuda").lower(),
499 # Main CUDA version of vLLM. This follows PyTorch but can be overridden.
499 - "VLLM_MAIN_CUDA_VERSION": lambda: os.getenv("VLLM_MAIN_CUDA_VERSION", "").lower()
500 - or "12.9",
500 + "VLLM_MAIN_CUDA_VERSION": lambda: (
501 + os.getenv("VLLM_MAIN_CUDA_VERSION", "").lower() or "12.9"
502 + ),
503 # Controls PyTorch float32 matmul precision mode within vLLM workers.
504 # Valid options mirror torch.set_float32_matmul_precision
505 "VLLM_FLOAT32_MATMUL_PRECISION": env_with_choices(
@@ -517,21 +519,19 @@ environment_variables: dict[str, Callable[[], Any]] = {
519 # If set, `MAX_JOBS` will be reduced to avoid oversubscribing the CPU.
520 "NVCC_THREADS": lambda: os.getenv("NVCC_THREADS", None),
521 # If set, vllm will use precompiled binaries (*.so)
520 - "VLLM_USE_PRECOMPILED": lambda: os.environ.get("VLLM_USE_PRECOMPILED", "")
521 - .strip()
522 - .lower()
523 - in ("1", "true")
524 - or bool(os.environ.get("VLLM_PRECOMPILED_WHEEL_LOCATION")),
522 + "VLLM_USE_PRECOMPILED": lambda: (
523 + os.environ.get("VLLM_USE_PRECOMPILED", "").strip().lower() in ("1", "true")
524 + or bool(os.environ.get("VLLM_PRECOMPILED_WHEEL_LOCATION"))
525 + ),
526 # If set, skip adding +precompiled suffix to version string
527 "VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX": lambda: bool(
528 int(os.environ.get("VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX", "0"))
529 ),
530 # Used to mark that setup.py is running in a Docker build context,
531 # in order to force the use of precompiled binaries.
531 - "VLLM_DOCKER_BUILD_CONTEXT": lambda: os.environ.get("VLLM_DOCKER_BUILD_CONTEXT", "")
532 - .strip()
533 - .lower()
534 - in ("1", "true"),
532 + "VLLM_DOCKER_BUILD_CONTEXT": lambda: (
533 + os.environ.get("VLLM_DOCKER_BUILD_CONTEXT", "").strip().lower() in ("1", "true")
534 + ),
535 # CMake build type
536 # If not set, defaults to "Debug" or "RelWithDebInfo"
537 # Available options: "Debug", "Release", "RelWithDebInfo"
@@ -577,10 +577,9 @@ environment_variables: dict[str, Callable[[], Any]] = {
577 ),
578 # If true, will load models from ModelScope instead of Hugging Face Hub.
579 # note that the value is true or false, not numbers
580 - "VLLM_USE_MODELSCOPE": lambda: os.environ.get(
581 - "VLLM_USE_MODELSCOPE", "False"
582 - ).lower()
583 - == "true",
580 + "VLLM_USE_MODELSCOPE": lambda: (
581 + os.environ.get("VLLM_USE_MODELSCOPE", "False").lower() == "true"
582 + ),
583 # Interval in seconds to log a warning message when the ring buffer is full
584 "VLLM_RINGBUFFER_WARNING_INTERVAL": lambda: int(
585 os.environ.get("VLLM_RINGBUFFER_WARNING_INTERVAL", "60")
@@ -601,19 +600,17 @@ environment_variables: dict[str, Callable[[], Any]] = {
600 # Feature flag to enable/disable Inductor standalone compile.
601 # In torch <= 2.7 we ignore this flag; in torch >= 2.9 this is
602 # enabled by default.
604 - "VLLM_USE_STANDALONE_COMPILE": lambda: os.environ.get(
605 - "VLLM_USE_STANDALONE_COMPILE", "1"
606 - )
607 - == "1",
603 + "VLLM_USE_STANDALONE_COMPILE": lambda: (
604 + os.environ.get("VLLM_USE_STANDALONE_COMPILE", "1") == "1"
605 + ),
606 # Inductor's pre-grad passes don't do anything for vLLM.
607 # The pre-grad passes get run even on cache-hit and negatively impact
608 # vllm cold compile times by O(1s)
609 # Can remove this after the following issue gets fixed
610 # https://github.com/pytorch/pytorch/issues/174502
613 - "VLLM_ENABLE_PREGRAD_PASSES": lambda: os.environ.get(
614 - "VLLM_ENABLE_PREGRAD_PASSES", "0"
615 - )
616 - == "1",
611 + "VLLM_ENABLE_PREGRAD_PASSES": lambda: (
612 + os.environ.get("VLLM_ENABLE_PREGRAD_PASSES", "0") == "1"
613 + ),
614 # Debug pattern matching inside custom passes.
615 # Should be set to the fx.Node name (e.g. 'getitem_34' or 'scaled_mm_3').
616 "VLLM_PATTERN_MATCH_DEBUG": lambda: os.environ.get(
@@ -656,10 +653,9 @@ environment_variables: dict[str, Callable[[], Any]] = {
653 # API key for vLLM API server
654 "VLLM_API_KEY": lambda: os.environ.get("VLLM_API_KEY", None),
655 # Whether to log responses from API Server for debugging
659 - "VLLM_DEBUG_LOG_API_SERVER_RESPONSE": lambda: os.environ.get(
660 - "VLLM_DEBUG_LOG_API_SERVER_RESPONSE", "False"
661 - ).lower()
662 - == "true",
656 + "VLLM_DEBUG_LOG_API_SERVER_RESPONSE": lambda: (
657 + os.environ.get("VLLM_DEBUG_LOG_API_SERVER_RESPONSE", "False").lower() == "true"
658 + ),
659 # S3 access information, used for tensorizer to load model from S3
660 "S3_ACCESS_KEY_ID": lambda: os.environ.get("S3_ACCESS_KEY_ID", None),
661 "S3_SECRET_ACCESS_KEY": lambda: os.environ.get("S3_SECRET_ACCESS_KEY", None),
@@ -670,11 +666,13 @@ environment_variables: dict[str, Callable[[], Any]] = {
666 ),
667 "VLLM_NO_USAGE_STATS": lambda: os.environ.get("VLLM_NO_USAGE_STATS", "0") == "1",
668 "VLLM_DO_NOT_TRACK": lambda: (
673 - os.environ.get("VLLM_DO_NOT_TRACK", None)
674 - or os.environ.get("DO_NOT_TRACK", None)
675 - or "0"
676 - )
677 - == "1",
669 + (
670 + os.environ.get("VLLM_DO_NOT_TRACK", None)
671 + or os.environ.get("DO_NOT_TRACK", None)
672 + or "0"
673 + )
674 + == "1"
675 + ),
676 "VLLM_USAGE_SOURCE": lambda: os.environ.get("VLLM_USAGE_SOURCE", "production"),
677 # Logging configuration
678 # If set to 0, vllm will not configure logging
@@ -697,36 +695,40 @@ environment_variables: dict[str, Callable[[], Any]] = {
695 "NO_COLOR": lambda: os.getenv("NO_COLOR", "0") != "0",
696 # If set, vllm will log stats at this interval in seconds
697 # If not set, vllm will log stats every 10 seconds.
700 - "VLLM_LOG_STATS_INTERVAL": lambda: val
701 - if (val := float(os.getenv("VLLM_LOG_STATS_INTERVAL", "10."))) > 0.0
702 - else 10.0,
698 + "VLLM_LOG_STATS_INTERVAL": lambda: (
699 + val
700 + if (val := float(os.getenv("VLLM_LOG_STATS_INTERVAL", "10."))) > 0.0
701 + else 10.0
702 + ),
703 # Trace function calls
704 # If set to 1, vllm will trace function calls
705 # Useful for debugging
706 "VLLM_TRACE_FUNCTION": lambda: int(os.getenv("VLLM_TRACE_FUNCTION", "0")),
707 # If set, vllm will use flashinfer sampler
708 - "VLLM_USE_FLASHINFER_SAMPLER": lambda: bool(
709 - int(os.environ["VLLM_USE_FLASHINFER_SAMPLER"])
710 - )
711 - if "VLLM_USE_FLASHINFER_SAMPLER" in os.environ
712 - else None,
708 + "VLLM_USE_FLASHINFER_SAMPLER": lambda: (
709 + bool(int(os.environ["VLLM_USE_FLASHINFER_SAMPLER"]))
710 + if "VLLM_USE_FLASHINFER_SAMPLER" in os.environ
711 + else None
712 + ),
713 # Pipeline stage partition strategy
714 "VLLM_PP_LAYER_PARTITION": lambda: os.getenv("VLLM_PP_LAYER_PARTITION", None),
715 # (CPU backend only) CPU key-value cache space.
716 # default is None and will be set as 4 GB
717 - "VLLM_CPU_KVCACHE_SPACE": lambda: int(os.getenv("VLLM_CPU_KVCACHE_SPACE", "0"))
718 - if "VLLM_CPU_KVCACHE_SPACE" in os.environ
719 - else None,
717 + "VLLM_CPU_KVCACHE_SPACE": lambda: (
718 + int(os.getenv("VLLM_CPU_KVCACHE_SPACE", "0"))
719 + if "VLLM_CPU_KVCACHE_SPACE" in os.environ
720 + else None
721 + ),
722 # (CPU backend only) CPU core ids bound by OpenMP threads, e.g., "0-31",
723 # "0,1,2", "0-31,33". CPU cores of different ranks are separated by '|'.
724 "VLLM_CPU_OMP_THREADS_BIND": lambda: os.getenv("VLLM_CPU_OMP_THREADS_BIND", "auto"),
725 # (CPU backend only) CPU cores not used by OMP threads .
726 # Those CPU cores will not be used by OMP threads of a rank.
725 - "VLLM_CPU_NUM_OF_RESERVED_CPU": lambda: int(
726 - os.getenv("VLLM_CPU_NUM_OF_RESERVED_CPU", "0")
727 - )
728 - if "VLLM_CPU_NUM_OF_RESERVED_CPU" in os.environ
729 - else None,
727 + "VLLM_CPU_NUM_OF_RESERVED_CPU": lambda: (
728 + int(os.getenv("VLLM_CPU_NUM_OF_RESERVED_CPU", "0"))
729 + if "VLLM_CPU_NUM_OF_RESERVED_CPU" in os.environ
730 + else None
731 + ),
732 # (CPU backend only) whether to use SGL kernels, optimized for small batch.
733 "VLLM_CPU_SGL_KERNEL": lambda: bool(int(os.getenv("VLLM_CPU_SGL_KERNEL", "0"))),
734 # (CPU backend only) whether to enable attention spilt KV.
@@ -920,9 +922,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
922 # a list of plugin names to load, separated by commas.
923 # if this is not set, it means all plugins will be loaded
924 # if this is set to an empty string, no plugins will be loaded
923 - "VLLM_PLUGINS": lambda: None
924 - if "VLLM_PLUGINS" not in os.environ
925 - else os.environ["VLLM_PLUGINS"].split(","),
925 + "VLLM_PLUGINS": lambda: (
926 + None
927 + if "VLLM_PLUGINS" not in os.environ
928 + else os.environ["VLLM_PLUGINS"].split(",")
929 + ),
930 # a local directory to look in for unrecognized LoRA adapters.
931 # only works if plugins are enabled and
932 # VLLM_ALLOW_RUNTIME_LORA_UPDATING is enabled.
@@ -954,9 +958,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
958 # and performance comparisons. Currently only affects MPLinearKernel
959 # selection
960 # (kernels: MacheteLinearKernel, MarlinLinearKernel, ExllamaLinearKernel)
957 - "VLLM_DISABLED_KERNELS": lambda: []
958 - if "VLLM_DISABLED_KERNELS" not in os.environ
959 - else os.environ["VLLM_DISABLED_KERNELS"].split(","),
961 + "VLLM_DISABLED_KERNELS": lambda: (
962 + []
963 + if "VLLM_DISABLED_KERNELS" not in os.environ
964 + else os.environ["VLLM_DISABLED_KERNELS"].split(",")
965 + ),
966 "VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE": lambda: bool(
967 int(os.getenv("VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE", "1"))
968 ),
@@ -1147,10 +1153,9 @@ environment_variables: dict[str, Callable[[], Any]] = {
1153 int(os.getenv("VLLM_ENABLE_MOE_DP_CHUNK", "1"))
1154 ),
1155 # Randomize inputs during dummy runs when using Data Parallel
1150 - "VLLM_RANDOMIZE_DP_DUMMY_INPUTS": lambda: os.environ.get(
1151 - "VLLM_RANDOMIZE_DP_DUMMY_INPUTS", "0"
1152 - )
1153 - == "1",
1156 + "VLLM_RANDOMIZE_DP_DUMMY_INPUTS": lambda: (
1157 + os.environ.get("VLLM_RANDOMIZE_DP_DUMMY_INPUTS", "0") == "1"
1158 + ),
1159 # Strategy to pack the data parallel ranks for Ray.
1160 # Available options:
1161 # - "fill":
@@ -1190,10 +1195,9 @@ environment_variables: dict[str, Callable[[], Any]] = {
1195 "VLLM_MODEL_REDIRECT_PATH", None
1196 ),
1197 # Whether to use atomicAdd reduce in gptq/awq marlin kernel.
1193 - "VLLM_MARLIN_USE_ATOMIC_ADD": lambda: os.environ.get(
1194 - "VLLM_MARLIN_USE_ATOMIC_ADD", "0"
1195 - )
1196 - == "1",
1198 + "VLLM_MARLIN_USE_ATOMIC_ADD": lambda: (
1199 + os.environ.get("VLLM_MARLIN_USE_ATOMIC_ADD", "0") == "1"
1200 + ),
1201 # Whether to use marlin kernel in mxfp4 quantization method
1202 "VLLM_MXFP4_USE_MARLIN": lambda: maybe_convert_bool(
1203 os.environ.get("VLLM_MXFP4_USE_MARLIN", None)
@@ -1211,17 +1215,16 @@ environment_variables: dict[str, Callable[[], Any]] = {
1215 # Whether to turn on the outlines cache for V1
1216 # This cache is unbounded and on disk, so it's not safe to use in
1217 # an environment with potentially malicious users.
1214 - "VLLM_V1_USE_OUTLINES_CACHE": lambda: os.environ.get(
1215 - "VLLM_V1_USE_OUTLINES_CACHE", "0"
1216 - )
1217 - == "1",
1218 + "VLLM_V1_USE_OUTLINES_CACHE": lambda: (
1219 + os.environ.get("VLLM_V1_USE_OUTLINES_CACHE", "0") == "1"
1220 + ),
1221 # Gap between padding buckets for the forward pass. So we have
1222 # 8, we will run forward pass with [16, 24, 32, ...].
1220 - "VLLM_TPU_BUCKET_PADDING_GAP": lambda: int(
1221 - os.environ["VLLM_TPU_BUCKET_PADDING_GAP"]
1222 - )
1223 - if "VLLM_TPU_BUCKET_PADDING_GAP" in os.environ
1224 - else 0,
1223 + "VLLM_TPU_BUCKET_PADDING_GAP": lambda: (
1224 + int(os.environ["VLLM_TPU_BUCKET_PADDING_GAP"])
1225 + if "VLLM_TPU_BUCKET_PADDING_GAP" in os.environ
1226 + else 0
1227 + ),
1228 "VLLM_TPU_MOST_MODEL_LEN": lambda: maybe_convert_int(
1229 os.environ.get("VLLM_TPU_MOST_MODEL_LEN", None)
1230 ),
@@ -1714,6 +1717,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
1717 "VLLM_USE_SIMPLE_KV_OFFLOAD": lambda: bool(
1718 int(os.getenv("VLLM_USE_SIMPLE_KV_OFFLOAD", "0"))
1719 ),
1720 + # Whether to enable dual cuda streams for LoRA computation
1721 + "VLLM_LORA_ENABLE_DUAL_STREAM": lambda: bool(
1722 + int(os.getenv("VLLM_LORA_ENABLE_DUAL_STREAM", "0"))
1723 + ),
1724 }
1725
1726
vllm/lora/layers/base_linear.py
+149 -1
@@ -5,8 +5,15 @@
5 import torch
6 from transformers import PretrainedConfig
7
8 +from vllm import envs
9 +from vllm.config import get_current_vllm_config
10 from vllm.config.lora import LoRAConfig
11 from vllm.distributed.utils import divide
12 +from vllm.forward_context import (
13 + ForwardContext,
14 + get_forward_context,
15 + is_forward_context_available,
16 +)
17 from vllm.model_executor.layers.linear import (
18 ColumnParallelLinear,
19 LinearBase,
@@ -14,24 +21,88 @@ from vllm.model_executor.layers.linear import (
21 RowParallelLinear,
22 )
23 from vllm.platforms import current_platform
24 +from vllm.utils.multi_stream_utils import maybe_execute_in_parallel
25 +from vllm.utils.torch_utils import direct_register_custom_op
26
27 from .base import BaseLayerWithLoRA
28 from .utils import _get_lora_device
29
30 +if envs.VLLM_LORA_ENABLE_DUAL_STREAM:
31 + _lora_aux_cuda_stream: torch.cuda.Stream | None = None
32 +
33 + def _get_lora_aux_cuda_stream() -> torch.cuda.Stream | None:
34 + global _lora_aux_cuda_stream
35 + if _lora_aux_cuda_stream is None and current_platform.is_cuda_alike():
36 + _lora_aux_cuda_stream = torch.cuda.Stream()
37 + return _lora_aux_cuda_stream
38 +
39 + def lora_linear_async(
40 + layer_name: str,
41 + output_size: int,
42 + x: torch.Tensor,
43 + bias: torch.Tensor | None = None,
44 + ) -> torch.Tensor:
45 + forward_context: ForwardContext = get_forward_context()
46 + self = forward_context.no_compile_layers[layer_name]
47 + return self._apply_async_impl(x, bias)
48 +
49 + def lora_linear_async_fake(
50 + layer_name: str,
51 + output_size: int,
52 + x: torch.Tensor,
53 + bias: torch.Tensor | None = None,
54 + ) -> torch.Tensor:
55 + # The real function reshapes output back to the original 3D shape
56 + # when the input has an extra batch dimension (transformers backend).
57 + if x.ndim == 3:
58 + return torch.empty(
59 + (x.size(0), x.size(1), output_size),
60 + device=x.device,
61 + dtype=x.dtype,
62 + )
63 + return torch.empty(
64 + (x.size(0), output_size),
65 + device=x.device,
66 + dtype=x.dtype,
67 + )
68 +
69 + direct_register_custom_op(
70 + op_name="lora_linear_async",
71 + op_func=lora_linear_async,
72 + fake_impl=lora_linear_async_fake,
73 + )
74 +
75
76 class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):
77 def __init__(self, base_layer: LinearBase):
78 super().__init__()
79 +
80 + self._enable_aux_cuda_stream = envs.VLLM_LORA_ENABLE_DUAL_STREAM
81 self.base_layer = base_layer
82 self.input_size = self.base_layer.input_size
83 # Ensure tp_size and tp_rank consistency with the base_layer.
84 self.tp_size = self.base_layer.tp_size
85 self.tp_rank = self.base_layer.tp_rank
86 self.device = _get_lora_device(self.base_layer)
87 + self._init_lora_stream_context()
88 self.output_slices: tuple[int, ...]
89 self.output_size: int
90 self.n_slices: int
91
92 + def _init_lora_stream_context(self) -> None:
93 + if not self._enable_aux_cuda_stream:
94 + return
95 + vllm_config = get_current_vllm_config()
96 + self._lora_stream = _get_lora_aux_cuda_stream()
97 + assert current_platform.is_cuda_alike()
98 + self._events = [torch.cuda.Event(), torch.cuda.Event()]
99 + # lora_linear avoids prefix conflicts with the base layer
100 + self.layer_name = self.base_layer.prefix + ".lora_linear_async"
101 + compilation_config = vllm_config.compilation_config
102 + if self.layer_name in compilation_config.static_forward_context:
103 + raise ValueError("Duplicate layer name: {}".format(self.layer_name))
104 + compilation_config.static_forward_context[self.layer_name] = self
105 +
106 def create_lora_weights(
107 self,
108 max_loras: int,
@@ -39,7 +110,6 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):
110 model_config: PretrainedConfig | None = None,
111 ) -> None:
112 self.lora_config = lora_config
42 - #
113 if isinstance(self.base_layer, ReplicatedLinear):
114 lora_a_out_size = lora_config.max_lora_rank
115 lora_b_out_size = self.output_size
@@ -120,6 +190,18 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):
190 )
191
192 def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
193 + # is_forward_context_available for tower modules
194 + if self._enable_aux_cuda_stream and is_forward_context_available():
195 + output_size = sum(self.output_slices)
196 + return torch.ops.vllm.lora_linear_async(
197 + self.layer_name, output_size, x, bias
198 + )
199 + else:
200 + return self._apply_sync(x, bias)
201 +
202 + def _apply_sync(
203 + self, x: torch.Tensor, bias: torch.Tensor | None = None
204 + ) -> torch.Tensor:
205 output = self.base_layer.quant_method.apply(self.base_layer, x, bias)
206
207 original_shape = output.shape if output.ndim == 3 else None
@@ -144,6 +226,72 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):
226
227 return output
228
229 + def _apply_async_impl(
230 + self, x: torch.Tensor, bias: torch.Tensor | None = None
231 + ) -> torch.Tensor:
232 + """
233 + Forward pass with base linear and LoRA on separate CUDA streams
234 + for overlap, using maybe_execute_in_parallel.
235 + Base layer runs on default stream; LoRA runs on aux stream.
236 + """
237 + assert envs.VLLM_LORA_ENABLE_DUAL_STREAM
238 + assert x.ndim in (2, 3)
239 + num_tokens = x.size(0) if x.ndim == 2 else x.size(1)
240 + output_size = sum(self.output_slices)
241 +
242 + def base_fn() -> torch.Tensor:
243 + return self.base_layer.quant_method.apply(self.base_layer, x, bias)
244 +
245 + def lora_fn() -> torch.Tensor:
246 + # Must be zeros, not empty: _lora_expand_kernel exits early (without
247 + # writing) when lora_id == -1 (no active LoRA). If uninitialized,
248 + # output.add_(lora_result) below would corrupt the base output.
249 + lora_output = torch.zeros(
250 + (num_tokens, output_size),
251 + device=self.device,
252 + dtype=x.dtype,
253 + )
254 +
255 + # Flatten the batch dimension for the transformers backend
256 + # (which uses shape (1, seq_len, hidden)), matching _apply_sync.
257 + x_2d = x.flatten(0, 1) if x.ndim == 3 else x
258 + self.punica_wrapper.add_lora_linear(
259 + lora_output,
260 + x_2d,
261 + self.lora_a_stacked,
262 + self.lora_b_stacked,
263 + 1.0,
264 + self.output_slices,
265 + add_inputs=False,
266 + )
267 + return lora_output
268 +
269 + output, lora_result = maybe_execute_in_parallel(
270 + base_fn,
271 + lora_fn,
272 + self._events[0],
273 + self._events[1],
274 + self._lora_stream,
275 + )
276 +
277 + original_shape = output.shape if output.ndim == 3 else None
278 +
279 + # In transformers backend, x and output have extra batch dimension like
280 + # (1, seq_len, hidden_dim), while punica expects (seq_len, hidden_dim),
281 + # therefore we need to flatten the batch dimensions.
282 + if x.ndim == 3 and output.ndim == 3:
283 + output = output.flatten(0, 1)
284 + x = x.flatten(0, 1)
285 +
286 + output.add_(lora_result)
287 +
288 + # Reshape the flattened output back to its original shape,
289 + # as some MM encoders cannot handle flattened inputs.
290 + if original_shape is not None:
291 + output = output.reshape(original_shape)
292 +
293 + return output
294 +
295 @property
296 def weight(self) -> torch.Tensor:
297 # unquantizedLinear
vllm/lora/ops/triton_ops/lora_expand_op.py
+9 -4
@@ -9,8 +9,13 @@ https://arxiv.org/abs/2310.18547
9
10 import torch
11
12 +from vllm import envs
13 from vllm.lora.ops.triton_ops.kernel_utils import do_expand_kernel
13 -from vllm.lora.ops.triton_ops.utils import _get_lora_b_ptr, get_lora_op_configs
14 +from vllm.lora.ops.triton_ops.utils import (
15 + _get_lora_b_ptr,
16 + get_lora_op_configs,
17 + supports_pdl,
18 +)
19 from vllm.triton_utils import tl, triton
20 from vllm.utils.torch_utils import direct_register_custom_op
21
@@ -237,9 +242,9 @@ def _lora_expand(
242 NUM_SLICES,
243 num_active_loras.item(),
244 )
240 - # We disable PDL temporarily because LoRA kernels are not launching back-to-back,
241 - # making PDL invalid and affecting the kernel performance.
242 - use_gdc = False # supports_pdl(inputs.device)
245 +
246 + # PDL only works when dual-stream is being used.
247 + use_gdc = supports_pdl(inputs.device) and envs.VLLM_LORA_ENABLE_DUAL_STREAM
248 _lora_expand_kernel[grid](
249 inputs,
250 lora_ptr_tensor,
vllm/lora/ops/triton_ops/lora_shrink_op.py
+9 -4
@@ -9,8 +9,13 @@ https://arxiv.org/abs/2310.18547
9
10 import torch
11
12 +from vllm import envs
13 from vllm.lora.ops.triton_ops.kernel_utils import do_shrink_kernel
13 -from vllm.lora.ops.triton_ops.utils import _get_lora_a_ptr, get_lora_op_configs
14 +from vllm.lora.ops.triton_ops.utils import (
15 + _get_lora_a_ptr,
16 + get_lora_op_configs,
17 + supports_pdl,
18 +)
19 from vllm.triton_utils import tl, triton
20 from vllm.utils.torch_utils import direct_register_custom_op
21
@@ -220,9 +225,9 @@ def _lora_shrink(
225 NUM_SLICES,
226 num_active_loras.item(),
227 )
223 - # We disable PDL temporarily because LoRA kernels are not launching back-to-back,
224 - # making PDL invalid and affecting the kernel performance.
225 - use_gdc = False # supports_pdl(inputs.device)
228 +
229 + # PDL only works when dual-stream is being used.
230 + use_gdc = supports_pdl(inputs.device) and envs.VLLM_LORA_ENABLE_DUAL_STREAM
231 _lora_shrink_kernel[grid](
232 inputs,
233 lora_ptr_tensor,
vllm/lora/punica_wrapper/punica_gpu.py
+6 -4
@@ -144,7 +144,9 @@ class PunicaWrapperGPU(PunicaWrapperBase):
144 x (torch.Tensor): Input tensors
145 lora_b_stacked (tuple[torch.Tensor, ...]): lora_b's weight
146 output_slices (tuple[int, ...]): Every slice's size
147 - add_inputs (bool): Defaults to True.
147 + add_inputs (bool): If True, add LoRA output to y; if False, write
148 + LoRA-only output to y (used for dual-stream when base and LoRA
149 + run on different CUDA streams). Defaults to True.
150 """
151 y_org = y
152 y = y.view(-1, y.shape[-1])
@@ -161,7 +163,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):
163 num_tokens, self.lora_config.specialize_active_lora
164 ),
165 offset_start=offset_start,
164 - add_inputs=True,
166 + add_inputs=add_inputs,
167 )
168
169 y = y.view_as(y_org)
@@ -244,7 +246,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):
246 buffer = torch.empty(
247 (len(output_slices), x.size(0), r), dtype=torch.float32, device=x.device
248 )
247 -
249 + add_inputs = kwargs.pop("add_inputs", True)
250 self.add_shrink(
251 buffer, # type: ignore
252 x,
@@ -257,7 +259,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):
259 buffer, # type: ignore
260 lora_b_stacked,
261 output_slices,
260 - add_inputs=True,
262 + add_inputs=add_inputs,
263 **kwargs,
264 )
265