Add vLLM v0.18.1 source tree with KV transfer abort fix

third_party/vllm/ now tracked in git for direct patch management.
Based on vLLM v0.18.1 release with one patch applied:

  vllm/v1/core/sched/scheduler.py:
    Replace fatal assert with graceful skip when KV transfer callback
    arrives for an already-aborted request during PD disaggregated serving.

Future vLLM modifications should be made directly in third_party/vllm/
and committed normally. The patches/ directory is kept as documentation
of what changed from upstream.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

third_party/vllm/tests/lora/conftest.py

@@ -0,0 +1,306 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import tempfile
+from collections import OrderedDict
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+import torch.nn as nn
+from huggingface_hub import snapshot_download
+
+from vllm.distributed import (
+    cleanup_dist_env_and_memory,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    MergedColumnParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
+from vllm.model_executor.models.interfaces import SupportsLoRA
+from vllm.platforms import current_platform
+
+
+@pytest.fixture()
+def should_do_global_cleanup_after_test(request) -> bool:
+    """Allow subdirectories to skip global cleanup by overriding this fixture.
+    This can provide a ~10x speedup for non-GPU unit tests since they don't need
+    to initialize torch.
+    """
+
+    return not request.node.get_closest_marker("skip_global_cleanup")
+
+
+@pytest.fixture(autouse=True)
+def cleanup_fixture(should_do_global_cleanup_after_test: bool):
+    yield
+    if should_do_global_cleanup_after_test:
+        cleanup_dist_env_and_memory(shutdown_ray=True)
+
+
+@pytest.fixture
+def dist_init():
+    from tests.utils import ensure_current_vllm_config
+
+    temp_file = tempfile.mkstemp()[1]
+
+    backend = "nccl"
+    if current_platform.is_cpu() or current_platform.is_tpu():
+        backend = "gloo"
+
+    with ensure_current_vllm_config():
+        init_distributed_environment(
+            world_size=1,
+            rank=0,
+            distributed_init_method=f"file://{temp_file}",
+            local_rank=0,
+            backend=backend,
+        )
+        initialize_model_parallel(1, 1)
+        yield
+    cleanup_dist_env_and_memory(shutdown_ray=True)
+
+
+@pytest.fixture
+def dist_init_torch_only():
+    if torch.distributed.is_initialized():
+        return
+    backend = "nccl"
+    if current_platform.is_cpu():
+        backend = "gloo"
+
+    temp_file = tempfile.mkstemp()[1]
+    torch.distributed.init_process_group(
+        world_size=1, rank=0, init_method=f"file://{temp_file}", backend=backend
+    )
+
+
+class DummyLoRAModel(nn.Sequential, SupportsLoRA):
+    pass
+
+
+@pytest.fixture
+def dummy_model(default_vllm_config) -> nn.Module:
+    model = DummyLoRAModel(
+        OrderedDict(
+            [
+                ("dense1", ColumnParallelLinear(764, 100)),
+                ("dense2", RowParallelLinear(100, 50)),
+                (
+                    "layer1",
+                    nn.Sequential(
+                        OrderedDict(
+                            [
+                                ("dense1", ColumnParallelLinear(100, 10)),
+                                ("dense2", RowParallelLinear(10, 50)),
+                            ]
+                        )
+                    ),
+                ),
+                ("act2", nn.ReLU()),
+                ("output", ColumnParallelLinear(50, 10)),
+                ("outact", nn.Sigmoid()),
+                # Special handling for lm_head & sampler
+                ("lm_head", ParallelLMHead(32064, 10)),
+                ("logits_processor", LogitsProcessor(32064)),
+            ]
+        )
+    )
+    model.config = MagicMock()
+    model.embedding_modules = {"lm_head": "lm_head"}
+    model.unpadded_vocab_size = 32064
+    return model
+
+
+@pytest.fixture
+def dummy_model_gate_up(default_vllm_config) -> nn.Module:
+    model = DummyLoRAModel(
+        OrderedDict(
+            [
+                ("dense1", ColumnParallelLinear(764, 100)),
+                ("dense2", RowParallelLinear(100, 50)),
+                (
+                    "layer1",
+                    nn.Sequential(
+                        OrderedDict(
+                            [
+                                ("dense1", ColumnParallelLinear(100, 10)),
+                                ("dense2", RowParallelLinear(10, 50)),
+                            ]
+                        )
+                    ),
+                ),
+                ("act2", nn.ReLU()),
+                ("gate_up_proj", MergedColumnParallelLinear(50, [5, 5])),
+                ("outact", nn.Sigmoid()),
+                # Special handling for lm_head & sampler
+                ("lm_head", ParallelLMHead(32064, 10)),
+                ("logits_processor", LogitsProcessor(32064)),
+            ]
+        )
+    )
+    model.config = MagicMock()
+    model.packed_modules_mapping = {
+        "gate_up_proj": [
+            "gate_proj",
+            "up_proj",
+        ],
+    }
+    model.embedding_modules = {"lm_head": "lm_head"}
+    model.unpadded_vocab_size = 32064
+
+    return model
+
+
+@pytest.fixture(scope="session")
+def mixtral_lora_files():
+    # Note: this module has incorrect adapter_config.json to test
+    # https://github.com/vllm-project/vllm/pull/5909/files.
+    return snapshot_download(repo_id="SangBinCho/mixtral-lora")
+
+
+@pytest.fixture(scope="session")
+def chatglm3_lora_files():
+    return snapshot_download(repo_id="jeeejeee/chatglm3-text2sql-spider")
+
+
+@pytest.fixture(scope="session")
+def baichuan_lora_files():
+    return snapshot_download(repo_id="jeeejeee/baichuan7b-text2sql-spider")
+
+
+@pytest.fixture(scope="session")
+def baichuan_zero_lora_files():
+    # all the lora_B weights are initialized to zero.
+    return snapshot_download(repo_id="jeeejeee/baichuan7b-zero-init")
+
+
+@pytest.fixture(scope="session")
+def baichuan_regex_lora_files():
+    return snapshot_download(repo_id="jeeejeee/baichuan-7b-lora-zero-regex")
+
+
+@pytest.fixture(scope="session")
+def ilama_lora_files():
+    return snapshot_download(repo_id="jeeejeee/ilama-text2sql-spider")
+
+
+@pytest.fixture(scope="session")
+def minicpmv_lora_files():
+    return snapshot_download(repo_id="jeeejeee/minicpmv25-lora-pokemon")
+
+
+@pytest.fixture(scope="session")
+def qwen2vl_lora_files():
+    return snapshot_download(repo_id="jeeejeee/qwen2-vl-lora-pokemon")
+
+
+@pytest.fixture(scope="session")
+def qwen25vl_base_huggingface_id():
+    # used as a base model for testing with qwen25vl lora adapter
+    return "Qwen/Qwen2.5-VL-3B-Instruct"
+
+
+@pytest.fixture(scope="session")
+def qwen25vl_lora_files():
+    return snapshot_download(repo_id="jeeejeee/qwen25-vl-lora-pokemon")
+
+
+@pytest.fixture(scope="session")
+def qwen2vl_language_lora_files():
+    return snapshot_download(repo_id="prashanth058/qwen2vl-flickr-lora-language")
+
+
+@pytest.fixture(scope="session")
+def qwen2vl_vision_tower_connector_lora_files():
+    return snapshot_download(repo_id="prashanth058/qwen2vl-flickr-lora-tower-connector")
+
+
+@pytest.fixture(scope="session")
+def qwen2vl_vision_tower_lora_files():
+    return snapshot_download(repo_id="prashanth058/qwen2vl-flickr-lora-tower")
+
+
+@pytest.fixture(scope="session")
+def qwen25vl_vision_lora_files():
+    return snapshot_download(repo_id="EpochEcho/qwen2.5-3b-vl-lora-vision-connector")
+
+
+@pytest.fixture(scope="session")
+def qwen3vl_vision_lora_files():
+    return snapshot_download(repo_id="EpochEcho/qwen3-4b-vl-lora-vision-connector")
+
+
+@pytest.fixture(scope="session")
+def qwen3_meowing_lora_files():
+    """Download Qwen3 Meow LoRA files once per test session."""
+    return snapshot_download(repo_id="Jackmin108/Qwen3-0.6B-Meow-LoRA")
+
+
+@pytest.fixture(scope="session")
+def qwen3_woofing_lora_files():
+    """Download Qwen3 Woof LoRA files once per test session."""
+    return snapshot_download(repo_id="Jackmin108/Qwen3-0.6B-Woof-LoRA")
+
+
+@pytest.fixture(scope="session")
+def tinyllama_lora_files():
+    return snapshot_download(repo_id="jashing/tinyllama-colorist-lora")
+
+
+@pytest.fixture(scope="session")
+def deepseekv2_lora_files():
+    return snapshot_download(repo_id="wuchen01/DeepSeek-V2-Lite-Chat-All-LoRA")
+
+
+@pytest.fixture(scope="session")
+def gptoss20b_lora_files():
+    return snapshot_download(repo_id="jeeejeee/gpt-oss-20b-lora-adapter-text2sql")
+
+
+@pytest.fixture(scope="session")
+def qwen3moe_lora_files():
+    return snapshot_download(repo_id="jeeejeee/qwen3-moe-text2sql-spider")
+
+
+@pytest.fixture(scope="session")
+def olmoe_lora_files():
+    return snapshot_download(repo_id="jeeejeee/olmoe-instruct-text2sql-spider")
+
+
+@pytest.fixture(scope="session")
+def qwen3_lora_files():
+    return snapshot_download(repo_id="charent/self_cognition_Alice")
+
+
+@pytest.fixture(scope="session")
+def llama32_lora_huggingface_id():
+    # huggingface repo id is used to test lora runtime downloading.
+    return "jeeejeee/llama32-3b-text2sql-spider"
+
+
+@pytest.fixture(scope="session")
+def llama32_lora_files(llama32_lora_huggingface_id):
+    return snapshot_download(repo_id=llama32_lora_huggingface_id)
+
+
+@pytest.fixture(scope="session")
+def whisper_lora_files():
+    return snapshot_download(repo_id="chengyili2005/whisper-small-mandarin-lora")
+
+
+@pytest.fixture
+def reset_default_device():
+    """
+    Some tests, such as `test_punica_ops.py`, explicitly set the
+    default device, which can affect subsequent tests. Adding this fixture
+    helps avoid this problem.
+    """
+    original_device = torch.get_default_device()
+    yield
+    torch.set_default_device(original_device)

third_party/vllm/tests/lora/test_add_lora.py

@@ -0,0 +1,113 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import asyncio
+import time
+
+import pytest
+
+from vllm.engine.arg_utils import AsyncEngineArgs
+from vllm.entrypoints.openai.api_server import (
+    build_async_engine_client_from_engine_args,
+)
+from vllm.inputs import TextPrompt
+from vllm.lora.request import LoRARequest
+from vllm.sampling_params import SamplingParams
+from vllm.utils.async_utils import merge_async_iterators
+
+MODEL_PATH = "zai-org/chatglm3-6b"
+LORA_RANK = 64
+DEFAULT_MAX_LORAS = 4 * 3
+
+
+def get_lora_requests(lora_path) -> list[LoRARequest]:
+    lora_requests: list[LoRARequest] = [
+        LoRARequest(lora_name=f"{i}", lora_int_id=i, lora_path=lora_path)
+        for i in range(1, DEFAULT_MAX_LORAS + 1)
+    ]
+    return lora_requests
+
+
+async def requests_processing_time(llm, lora_requests: list[LoRARequest]) -> float:
+    sampling_params = SamplingParams(
+        n=1, temperature=0.0, top_p=1.0, ignore_eos=True, max_tokens=1
+    )
+
+    generators = []
+    start = time.perf_counter()
+
+    for lora_request in lora_requests:
+        lora_int_id = lora_request.lora_int_id
+        generator = llm.generate(
+            prompt=TextPrompt(prompt=f"hello {lora_int_id}", multi_modal_data=None),  # type: ignore
+            sampling_params=sampling_params,
+            lora_request=lora_request,
+            request_id=f"test{lora_int_id}",
+        )
+        generators.append(generator)
+
+    all_gens = merge_async_iterators(*generators)
+    async for i, res in all_gens:
+        pass
+
+    end = time.perf_counter()
+    return end - start
+
+
+@pytest.mark.asyncio
+async def test_add_lora(chatglm3_lora_files):
+    """
+    The add_lora function is used to preload some LoRA adapters into the
+    engine in anticipation of future requests using these adapters. To test
+    this functionality, we use the async engine to process some requests - We
+    do it twice, once with add_lora() preloading and once without.
+
+    We measure the request processing time in both cases and expect the time
+    to be lesser in the case with add_lora() calls.
+    """
+    lora_requests: list[LoRARequest] = get_lora_requests(chatglm3_lora_files)
+
+    max_loras = len(set([lr.lora_int_id for lr in lora_requests]))
+    # Create engine in eager-mode. Due to high max_loras, the CI can
+    # OOM during cuda-graph capture.
+    engine_args = AsyncEngineArgs(
+        model=MODEL_PATH,
+        enable_lora=True,
+        max_loras=max_loras,
+        max_lora_rank=LORA_RANK,
+        max_model_len=128,
+        gpu_memory_utilization=0.8,  # avoid OOM
+        trust_remote_code=True,
+        enforce_eager=True,
+    )
+
+    # split lora_requests into 3 parts
+    part_size = len(lora_requests) // 3
+    dummy_run_requests = lora_requests[:part_size]
+    warmup_run_requests = lora_requests[part_size : part_size * 2]
+    cold_run_requests = lora_requests[part_size * 2 :]
+
+    async with build_async_engine_client_from_engine_args(engine_args) as llm:
+        # Dummy run - So any 1-time functionality like triton kernel compilation
+        # is complete here.
+        await requests_processing_time(llm, dummy_run_requests)
+
+        # Run with warmup
+        add_lora_tasks = [llm.add_lora(lr) for lr in warmup_run_requests]
+        add_lora_results = await asyncio.gather(*add_lora_tasks)
+
+        # Test that all all_lora calls are successful.
+        assert all(add_lora_results)
+
+        time_with_add_lora = await requests_processing_time(llm, warmup_run_requests)
+
+        # Run without any warmup
+        time_cold_start = await requests_processing_time(llm, cold_run_requests)
+
+    print(f"time hot-start {time_with_add_lora} vs time cold-start {time_cold_start} ")
+
+    assert time_with_add_lora < time_cold_start, (
+        f"time_with_add_lora={time_with_add_lora}, "
+        f"time_cold_start={time_cold_start}"
+        "The engine request processing time with LoRA pre-loading "
+        "must be less than the version that does on-demand LoRA loading."
+    )

third_party/vllm/tests/lora/test_chatglm3_tp.py

@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import vllm
+import vllm.config
+from vllm.lora.request import LoRARequest
+
+from ..utils import create_new_process_for_each_test, multi_gpu_test
+
+MODEL_PATH = "zai-org/chatglm3-6b"
+
+PROMPT_TEMPLATE = """I want you to act as a SQL terminal in front of an example database, you need only to return the sql command to me.Below is an instruction that describes a task, Write a response that appropriately completes the request.\n"\n##Instruction:\nconcert_singer contains tables such as stadium, singer, concert, singer_in_concert. Table stadium has columns such as Stadium_ID, Location, Name, Capacity, Highest, Lowest, Average. Stadium_ID is the primary key.\nTable singer has columns such as Singer_ID, Name, Country, Song_Name, Song_release_year, Age, Is_male. Singer_ID is the primary key.\nTable concert has columns such as concert_ID, concert_Name, Theme, Stadium_ID, Year. concert_ID is the primary key.\nTable singer_in_concert has columns such as concert_ID, Singer_ID. concert_ID is the primary key.\nThe Stadium_ID of concert is the foreign key of Stadium_ID of stadium.\nThe Singer_ID of singer_in_concert is the foreign key of Singer_ID of singer.\nThe concert_ID of singer_in_concert is the foreign key of concert_ID of concert.\n\n###Input:\n{query}\n\n###Response:"""  # noqa: E501
+
+EXPECTED_LORA_OUTPUT = [
+    "SELECT count(*) FROM singer",
+    "SELECT avg(age) ,  min(age) ,  max(age) FROM singer WHERE country  =  'France'",
+    "SELECT name ,  country ,  age FROM singer ORDER BY age",
+]
+
+
+def do_sample(llm: vllm.LLM, lora_path: str, lora_id: int) -> list[str]:
+    prompts = [
+        PROMPT_TEMPLATE.format(query="How many singers do we have?"),
+        PROMPT_TEMPLATE.format(
+            query=(
+                "What is the average, minimum, and maximum "
+                "age of all singers from France?"
+            )
+        ),
+        PROMPT_TEMPLATE.format(
+            query=(
+                "Show name, country, age for all singers ordered "
+                "by age from the oldest to the youngest."
+            )
+        ),
+    ]
+    sampling_params = vllm.SamplingParams(temperature=0, max_tokens=32)
+    outputs = llm.generate(
+        prompts,
+        sampling_params,
+        lora_request=LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None,
+    )
+    # Print the outputs.
+    generated_texts: list[str] = []
+    for output in outputs:
+        prompt = output.prompt
+        generated_text = output.outputs[0].text.strip()
+        generated_texts.append(generated_text)
+        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
+    return generated_texts
+
+
+@create_new_process_for_each_test()
+def test_chatglm3_lora(chatglm3_lora_files):
+    llm = vllm.LLM(
+        MODEL_PATH,
+        max_model_len=512,
+        enable_lora=True,
+        max_loras=2,
+        max_num_seqs=16,
+        max_lora_rank=64,
+        trust_remote_code=True,
+    )
+
+    output1 = do_sample(llm, chatglm3_lora_files, lora_id=1)
+    for i in range(len(EXPECTED_LORA_OUTPUT)):
+        assert output1[i] == EXPECTED_LORA_OUTPUT[i]
+    output2 = do_sample(llm, chatglm3_lora_files, lora_id=2)
+    for i in range(len(EXPECTED_LORA_OUTPUT)):
+        assert output2[i] == EXPECTED_LORA_OUTPUT[i]
+
+
+@multi_gpu_test(num_gpus=4)
+def test_chatglm3_lora_tp4(chatglm3_lora_files):
+    llm = vllm.LLM(
+        MODEL_PATH,
+        max_model_len=512,
+        enable_lora=True,
+        max_loras=2,
+        max_lora_rank=64,
+        max_num_seqs=16,
+        tensor_parallel_size=4,
+        trust_remote_code=True,
+        fully_sharded_loras=False,
+        compilation_config=vllm.config.CompilationConfig(  # Avoid OOM
+            cudagraph_specialize_lora=False,
+        ),
+    )
+
+    output1 = do_sample(llm, chatglm3_lora_files, lora_id=1)
+    for i in range(len(EXPECTED_LORA_OUTPUT)):
+        assert output1[i] == EXPECTED_LORA_OUTPUT[i]
+    output2 = do_sample(llm, chatglm3_lora_files, lora_id=2)
+    for i in range(len(EXPECTED_LORA_OUTPUT)):
+        assert output2[i] == EXPECTED_LORA_OUTPUT[i]
+
+
+@multi_gpu_test(num_gpus=4)
+def test_chatglm3_lora_tp4_fully_sharded_loras(chatglm3_lora_files):
+    # https://github.com/NVIDIA/nccl/issues/1790, set a lower value for
+    # gpu_memory_utilization here because NCCL >= 2.26.3 seems to use
+    # more GPU memory causing vLLM to OOM
+    llm = vllm.LLM(
+        MODEL_PATH,
+        max_model_len=512,
+        enable_lora=True,
+        max_loras=2,
+        max_lora_rank=64,
+        tensor_parallel_size=4,
+        trust_remote_code=True,
+        fully_sharded_loras=True,
+        gpu_memory_utilization=0.8,
+        compilation_config=vllm.config.CompilationConfig(  # Avoid OOM
+            cudagraph_specialize_lora=False,
+        ),
+    )
+    output1 = do_sample(llm, chatglm3_lora_files, lora_id=1)
+    for i in range(len(EXPECTED_LORA_OUTPUT)):
+        assert output1[i] == EXPECTED_LORA_OUTPUT[i]
+    output2 = do_sample(llm, chatglm3_lora_files, lora_id=2)
+    for i in range(len(EXPECTED_LORA_OUTPUT)):
+        assert output2[i] == EXPECTED_LORA_OUTPUT[i]

third_party/vllm/tests/lora/test_deepseekv2_tp.py

@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# NOTE To avoid overloading the CI pipeline, this test script will
+# not be triggered on CI and is primarily intended for local testing
+# and verification.
+
+import vllm
+from vllm.lora.request import LoRARequest
+
+from ..utils import multi_gpu_test
+
+MODEL_PATH = "deepseek-ai/DeepSeek-V2-Lite-Chat"
+
+PROMPT_TEMPLATE = "<｜begin▁of▁sentence｜>You are a helpful assistant.\n\nUser: {context}\n\nAssistant:"  # noqa: E501
+
+
+def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int):
+    prompts = [
+        PROMPT_TEMPLATE.format(context="Who are you?"),
+    ]
+    sampling_params = vllm.SamplingParams(temperature=0, max_tokens=64)
+    outputs = llm.generate(
+        prompts,
+        sampling_params,
+        lora_request=LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None,
+    )
+    # Print the outputs.
+    generated_texts: list[str] = []
+    for output in outputs:
+        prompt = output.prompt
+        generated_text = output.outputs[0].text.strip()
+        generated_texts.append(generated_text)
+        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
+    # return generated_texts
+    expected_lora_output = [
+        "I am \u5f20\u5b50\u8c6a, an AI assistant developed by \u9648\u58eb\u680b.",  # noqa: E501
+    ]
+    for i in range(len(expected_lora_output)):
+        assert generated_texts[i].startswith(expected_lora_output[i])
+
+
+def test_deepseekv2_lora(deepseekv2_lora_files):
+    # We enable enforce_eager=True here to reduce VRAM usage for lora-test CI,
+    # Otherwise, the lora-test will fail due to CUDA OOM.
+    llm = vllm.LLM(
+        MODEL_PATH,
+        max_model_len=1024,
+        enable_lora=True,
+        max_loras=4,
+        enforce_eager=True,
+        trust_remote_code=True,
+        enable_chunked_prefill=True,
+    )
+    generate_and_test(llm, deepseekv2_lora_files, 1)
+
+
+def test_deepseekv2(deepseekv2_lora_files):
+    # We enable enforce_eager=True here to reduce VRAM usage for lora-test CI,
+    # Otherwise, the lora-test will fail due to CUDA OOM.
+    llm = vllm.LLM(
+        MODEL_PATH,
+        max_model_len=1024,
+        enable_lora=True,
+        max_loras=4,
+        enforce_eager=True,
+        trust_remote_code=True,
+    )
+    generate_and_test(llm, deepseekv2_lora_files, 1)
+
+
+@multi_gpu_test(num_gpus=2)
+def test_deepseekv2_tp2(deepseekv2_lora_files):
+    # We enable enforce_eager=True here to reduce VRAM usage for lora-test CI,
+    # Otherwise, the lora-test will fail due to CUDA OOM.
+    llm = vllm.LLM(
+        MODEL_PATH,
+        max_model_len=1024,
+        enable_lora=True,
+        max_loras=4,
+        enforce_eager=True,
+        trust_remote_code=True,
+        tensor_parallel_size=2,
+    )
+    generate_and_test(llm, deepseekv2_lora_files, 2)
+
+
+@multi_gpu_test(num_gpus=4)
+def test_deepseekv2_tp4(deepseekv2_lora_files):
+    # We enable enforce_eager=True here to reduce VRAM usage for lora-test CI,
+    # Otherwise, the lora-test will fail due to CUDA OOM.
+    llm = vllm.LLM(
+        MODEL_PATH,
+        max_model_len=1024,
+        enable_lora=True,
+        max_loras=4,
+        enforce_eager=True,
+        trust_remote_code=True,
+        tensor_parallel_size=4,
+    )
+    generate_and_test(llm, deepseekv2_lora_files, 2)

third_party/vllm/tests/lora/test_default_mm_loras.py

@@ -0,0 +1,157 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Tests for applying default registered multimodal loras.
+"""
+
+import os
+import unittest.mock as mock
+
+import pytest
+from huggingface_hub import snapshot_download
+
+from vllm.lora.request import LoRARequest
+
+from ..conftest import AudioTestAssets, VllmRunner
+from ..utils import create_new_process_for_each_test
+
+MODEL_PATH = snapshot_download("microsoft/Phi-4-multimodal-instruct")
+AUDIO_LORA_PATH = os.path.join(MODEL_PATH, "speech-lora")
+IMAGE_LORA_PATH = os.path.join(MODEL_PATH, "vision-lora")
+
+AUDIO_PROMPT = "<|user|><|audio_1|>Can you transcribe this audio?<|end|><|assistant|>"  # noqa: E501
+
+# Responses are greedy decoded; we just check the end of
+# the generated text. If the lora is inactive, this model
+# generates commentary on the transcription.
+RESPONSE_SUFFIX_WITH_LORA = "Spoken text: The first words I spoke in the original chronograph, a little piece of practical poetry. Mary had a little lamb, it slept with quite a snow, and everywhere that Mary went, the lamb was sure to go."  # noqa: E501
+RESPONSE_SUFFIX_WITHOUT_LORA = "Certainly! Here is the transcription of the audio you provided:\n\nThe first words I spoke in the original phonograph record: A little piece of practical poetry. Mary had a little lamb; its fleece was white as snow, and everywhere that Mary went, the lamb was sure to go."  # noqa: E501
+
+VLLM_RUNNER_BASE_KWARGS = {
+    "model_name": MODEL_PATH,
+    "dtype": "half",
+    "enable_lora": "True",
+    "max_num_seqs": 2,
+    "max_lora_rank": 320,
+    # Keep these LoRA tests on short-RoPE for determinism post-LongRoPE change.
+    "max_model_len": 4096,
+    "gpu_memory_utilization": 0.8,
+    "limit_mm_per_prompt": {"audio": 1},
+    "enforce_eager": True,
+}
+
+
+def run_test(vllm_runner, audio_assets, lora_request, expected_suffix, **kwargs):
+    inputs = [([AUDIO_PROMPT], [audio_assets[0].audio_and_sample_rate[0]])]
+
+    # Apply any additional kwargs as overrides to the base kwargs
+    vllm_runner_kwargs = {**VLLM_RUNNER_BASE_KWARGS, **kwargs}
+
+    with vllm_runner(**vllm_runner_kwargs) as vllm_model:
+        vllm_outputs_with_default_lora = [
+            vllm_model.generate_greedy(
+                prompts,
+                max_tokens=128,
+                audios=audios,
+                lora_request=lora_request,
+            )
+            for prompts, audios in inputs
+        ]
+
+        assert vllm_outputs_with_default_lora[-1][-1][-1].endswith(expected_suffix)
+
+
+@create_new_process_for_each_test()
+def test_active_default_mm_lora(
+    vllm_runner: type[VllmRunner],
+    audio_assets: AudioTestAssets,
+):
+    """Ensure that we can use the default audio lora."""
+    run_test(
+        vllm_runner,
+        audio_assets,
+        lora_request=None,
+        default_mm_loras={"audio": AUDIO_LORA_PATH},
+        expected_suffix=RESPONSE_SUFFIX_WITH_LORA,
+    )
+
+
+@create_new_process_for_each_test()
+def test_inactive_default_mm_lora(
+    vllm_runner: type[VllmRunner],
+    audio_assets: AudioTestAssets,
+):
+    """Ensure that modalities are filtered properly."""
+    # Default image lora won't be active since we only pass audio
+    run_test(
+        vllm_runner,
+        audio_assets,
+        lora_request=None,
+        default_mm_loras={"image": IMAGE_LORA_PATH},
+        expected_suffix=RESPONSE_SUFFIX_WITHOUT_LORA,
+    )
+
+
+@create_new_process_for_each_test()
+def test_default_mm_lora_succeeds_with_redundant_lora_request(
+    vllm_runner: type[VllmRunner],
+    audio_assets: AudioTestAssets,
+):
+    """Ensure that redundantly providing the lora works."""
+    run_test(
+        vllm_runner,
+        audio_assets,
+        lora_request=LoRARequest("audio", 1, AUDIO_LORA_PATH),
+        default_mm_loras={"audio": AUDIO_LORA_PATH},
+        expected_suffix=RESPONSE_SUFFIX_WITH_LORA,
+    )
+
+
+@create_new_process_for_each_test()
+def test_default_mm_lora_fails_with_overridden_lora_request(
+    vllm_runner: type[VllmRunner],
+    audio_assets: AudioTestAssets,
+):
+    """Ensure that if the lora_request conflicts with default_mm_loras,
+    we use the lora_request."""
+    run_test(
+        vllm_runner,
+        audio_assets,
+        lora_request=LoRARequest("speech", 2, AUDIO_LORA_PATH),
+        default_mm_loras={"audio": IMAGE_LORA_PATH},
+        expected_suffix=RESPONSE_SUFFIX_WITH_LORA,
+    )
+
+
+@create_new_process_for_each_test()
+def test_default_mm_lora_does_not_expand_string_reqs(vllm_runner):
+    class MockEngineException(Exception):
+        pass
+
+    # Regression test for ensuring default multimodal lora resolution
+    # does not expand the lora req if the prompt type is a string.
+    vllm_runner_kwargs = {
+        **VLLM_RUNNER_BASE_KWARGS,
+        **{"default_mm_loras": {"audio": AUDIO_LORA_PATH}},
+    }
+
+    # Avoid the full generation call since these tests are expensive;
+    # just check what lora request is actually submitted to the engine
+    mock_err = "Engine is mocked for this test"
+
+    with (
+        mock.patch(
+            "vllm.v1.engine.llm_engine.LLMEngine.add_request",
+            side_effect=MockEngineException(mock_err),
+        ) as mock_add_request,
+        vllm_runner(**vllm_runner_kwargs) as vllm_model,
+    ):
+        # Die once we actually submit the request to the engine
+        with pytest.raises(MockEngineException):
+            vllm_model.llm.generate(prompts=AUDIO_PROMPT)
+
+        # Then check to make sure the submitted lora request
+        # and text prompt were zipped together correctly
+        engine_args, engine_kwargs = mock_add_request.call_args
+        assert engine_args[1]["prompt"] == AUDIO_PROMPT
+        assert engine_kwargs["lora_request"] is None

third_party/vllm/tests/lora/test_fused_moe_lora_kernel.py

@@ -0,0 +1,746 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import os
+import random
+
+import pytest
+import torch
+
+from tests.utils import ensure_current_vllm_config, multi_gpu_test
+from vllm import _custom_ops as ops
+from vllm.distributed import (
+    init_distributed_environment,
+    initialize_model_parallel,
+    tensor_model_parallel_all_gather,
+    tensor_model_parallel_all_reduce,
+)
+from vllm.distributed.parallel_state import (
+    get_tensor_model_parallel_world_size,
+)
+from vllm.lora.ops.triton_ops import fused_moe_lora
+from vllm.platforms import current_platform
+from vllm.utils.network_utils import get_open_port
+from vllm.utils.torch_utils import set_random_seed
+
+
+@pytest.fixture(autouse=True)
+def reset_device(reset_default_device):
+    pass
+
+
+def round_up(x, base):
+    return ((x + base - 1) // base) * base
+
+
+def CEILDIV(x, y):
+    return (x + y - 1) // y
+
+
+def assign_loras_to_tokens(num_tokens: int, num_sequences: int, max_loras: int):
+    """
+    Split `num_tokens` into `num_sequences` sequences.
+    Each sequence randomly selects 1 LoRA index from [0, max_loras),
+    and all tokens in that sequence are assigned this LoRA index.
+
+    Args:
+        num_tokens (int): Total number of tokens.
+        num_sequences (int): Number of sequences to split the tokens into.
+        max_loras (int): Total number of available LoRA modules.
+
+    Returns:
+        torch.Tensor: 1D tensor of shape [num_tokens], where each value
+                      is the LoRA index assigned to that token.
+    """
+    assert num_sequences > 0 and max_loras > 0
+    assert num_tokens >= num_sequences, "num_tokens must be >= num_sequences"
+
+    # Compute token distribution per sequence (distribute remainder evenly)
+    tokens_per_seq = num_tokens // num_sequences
+    remainder = num_tokens % num_sequences
+
+    token_lora_mapping = torch.empty(num_tokens, dtype=torch.int32)
+
+    start = 0
+    for seq_idx in range(num_sequences):
+        # Determine the token range for this sequence
+        end = start + tokens_per_seq + (1 if seq_idx < remainder else 0)
+
+        # Randomly select one LoRA ID for this sequence
+        lora_id = random.randint(0, max_loras - 1)
+
+        # Assign the same LoRA ID to all tokens in this sequence
+        token_lora_mapping[start:end] = lora_id
+
+        start = end
+
+    return token_lora_mapping
+
+
+def assign_experts_to_tokens(num_tokens: int, num_experts: int, top_k_num: int):
+    """
+    For each token, randomly select `top_k_num` distinct experts out of `num_experts`,
+    and assign normalized random weights that sum to 1.
+
+    Args:
+        num_tokens (int): Total number of tokens.
+        num_experts (int): Total number of available experts.
+        top_k_num (int): Number of experts to select per token.
+
+    Returns:
+        expert_indices (torch.Tensor): shape [num_tokens, top_k_num],
+                                       expert index for each token.
+        expert_weights (torch.Tensor): shape [num_tokens, top_k_num],
+                                       normalized weights (sum = 1 per row).
+    """
+    assert top_k_num <= num_experts, "top_k_num must be <= num_experts"
+
+    # Randomly select top_k_num distinct experts for each token
+    expert_indices = torch.empty((num_tokens, top_k_num), dtype=torch.int32)
+    for i in range(num_tokens):
+        # Randomly choose unique expert indices
+        selected = torch.randperm(num_experts)[:top_k_num]
+        expert_indices[i] = selected
+
+    # Generate random weights and normalize along dim=1
+    expert_weights = torch.rand((num_tokens, top_k_num), dtype=torch.float32)
+    expert_weights = expert_weights / expert_weights.sum(dim=1, keepdim=True)
+
+    return expert_indices, expert_weights
+
+
+def sample_data(
+    num_tokens: int,
+    num_sequences: int,
+    max_loras: int,
+    num_experts: int,
+    top_k_num: int,
+):
+    topk_ids, topk_weights = assign_experts_to_tokens(
+        num_tokens, num_experts, top_k_num
+    )
+    token_lora_mapping = assign_loras_to_tokens(num_tokens, num_sequences, max_loras)
+    active_lora_ids = torch.full((max_loras + 1,), -1, dtype=torch.int32)
+    lora_ids = torch.unique(token_lora_mapping, sorted=True)
+    active_lora_ids[: lora_ids.size(0)].copy_(lora_ids, non_blocking=True)
+    return topk_ids, topk_weights, token_lora_mapping, active_lora_ids
+
+
+def use_fused_moe_lora_kernel(
+    topk_ids,
+    topk_weights,
+    token_lora_mapping,
+    max_lora_rank,
+    top_k_num,
+    lora_ids,
+    lora_a_stacked,
+    lora_b_stacked,
+    hidden_states,
+    output,
+    max_loras,
+    num_experts,
+    block_size,
+    fully_sharded=False,
+    offset=0,
+):
+    max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
+    max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
+    max_num_m_blocks = CEILDIV(max_num_tokens_padded, block_size)
+
+    # init output tensors
+    sorted_token_ids = torch.empty(
+        (max_loras * max_num_tokens_padded,),
+        dtype=torch.int32,
+    )
+    expert_ids = torch.empty((max_loras * max_num_m_blocks,), dtype=torch.int32)
+    num_tokens_post_padded = torch.empty((max_loras,), dtype=torch.int32)
+    adapter_enabled = torch.ones(max_loras + 1, dtype=torch.int32)
+
+    # call kernel
+    ops.moe_lora_align_block_size(
+        topk_ids,
+        token_lora_mapping,
+        num_experts,
+        block_size,
+        max_loras,
+        max_num_tokens_padded,
+        max_num_m_blocks,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        adapter_enabled,
+        lora_ids,
+    )
+
+    config = {
+        "BLOCK_SIZE_M": block_size,
+        "BLOCK_SIZE_N": 32,
+        "BLOCK_SIZE_K": 64,
+        "GROUP_SIZE_M": 1,
+        "NUM_WARPS": 4,
+        "NUM_STAGES": 3,
+        "SPLIT_K": 1,
+    }
+
+    mul_routed_weight = False
+    expert_ids = expert_ids.view(max_loras, -1)
+    sorted_token_ids = sorted_token_ids.view(max_loras, -1)
+
+    # num_active_loras is the number of active LoRAs
+    # (max_loras + 1 to include no-lora case)
+    # Stored as CPU tensor to match the kernel API (torch.compile compatibility)
+    num_active_loras = torch.tensor([max_loras + 1], dtype=torch.int32, device="cpu")
+
+    fused_moe_lora(
+        output,
+        hidden_states,
+        lora_a_stacked,
+        lora_b_stacked,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        token_lora_mapping,
+        max_lora_rank,
+        top_k_num,
+        lora_ids,
+        num_active_loras,
+        adapter_enabled,
+        config["BLOCK_SIZE_M"],
+        config["BLOCK_SIZE_N"],
+        config["BLOCK_SIZE_K"],
+        config["GROUP_SIZE_M"],
+        config["NUM_WARPS"],
+        config["NUM_STAGES"],
+        config["SPLIT_K"],
+        config["BLOCK_SIZE_M"],
+        config["BLOCK_SIZE_N"],
+        config["BLOCK_SIZE_K"],
+        config["GROUP_SIZE_M"],
+        config["NUM_WARPS"],
+        config["NUM_STAGES"],
+        config["SPLIT_K"],
+        mul_routed_weight,
+        fully_sharded=fully_sharded,
+        offset=offset,
+    )
+
+
+def use_torch(
+    hidden_states,
+    token_lora_mapping,
+    topk_ids,
+    lora_a_stacked,
+    lora_b_stacked,
+    top_k_num,
+    num_slices=1,
+):
+    outputs = []
+    for i in range(hidden_states.shape[0]):
+        slice_tensors = []
+        for slice_id in range(num_slices):
+            lora_idx = token_lora_mapping[i]
+            expert_ids = topk_ids[i]
+            lora_a = lora_a_stacked[slice_id][lora_idx][expert_ids]
+            lora_b = lora_b_stacked[slice_id][lora_idx][expert_ids]
+            tensors = [
+                hidden_states[i] @ lora_a[x].T @ lora_b[x].T for x in range(top_k_num)
+            ]
+            slice_tensors.append(torch.stack(tensors, dim=0))
+
+        outputs.append(torch.concat(slice_tensors, dim=-1))
+    return torch.stack(outputs, dim=0)
+
+
+DEVICE_TYPE = current_platform.device_type
+DTYPES = [torch.float16, torch.bfloat16]
+DEVICES = [f"{DEVICE_TYPE}:{0}"]
+SEED = [42]
+
+
+@pytest.mark.parametrize("num_tokens", [100])
+@pytest.mark.parametrize("top_k_num", [6, 12])
+@pytest.mark.parametrize("num_experts", [64])
+@pytest.mark.parametrize("max_loras", [4, 6, 16])
+@pytest.mark.parametrize("N", [1408])
+@pytest.mark.parametrize("K", [2048])
+@pytest.mark.parametrize("max_lora_rank", [16, 32, 64])
+@pytest.mark.parametrize("block_size", [16])
+@pytest.mark.parametrize("num_slices", [1, 2])
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("seed", SEED)
+def test_fused_moe_lora_kernel(
+    num_tokens,
+    top_k_num,
+    num_experts,
+    max_loras,
+    N,
+    K,
+    max_lora_rank,
+    block_size,
+    num_slices,
+    dtype,
+    device,
+    seed,
+):
+    torch.set_default_device(device)
+    set_random_seed(seed)
+    # the number of randomly generated sentences.
+    num_sequences = 10
+    # generate data
+    topk_ids, topk_weights, token_lora_mapping, lora_ids = sample_data(
+        num_tokens, num_sequences, max_loras, num_experts, top_k_num
+    )
+
+    # init lora weights
+    lora_a_stacked = [
+        torch.rand(
+            (
+                max_loras,
+                num_experts,
+                max_lora_rank,
+                K,
+            ),
+            dtype=dtype,
+        )
+        for _ in range(num_slices)
+    ]
+    lora_b_stacked = [
+        torch.rand(
+            (
+                max_loras,
+                num_experts,
+                N // num_slices,
+                max_lora_rank,
+            ),
+            dtype=dtype,
+        )
+        for _ in range(num_slices)
+    ]
+    hidden_states = torch.rand(
+        (
+            num_tokens,
+            K,
+        ),
+        dtype=dtype,
+    )
+
+    # fused_moe_lora_kernel output
+    output = torch.zeros((num_tokens, top_k_num, N), dtype=dtype)
+    use_fused_moe_lora_kernel(
+        topk_ids,
+        topk_weights,
+        token_lora_mapping,
+        max_lora_rank,
+        top_k_num,
+        lora_ids,
+        lora_a_stacked,
+        lora_b_stacked,
+        hidden_states,
+        output,
+        max_loras,
+        num_experts,
+        block_size,
+    )
+    # pytorch output
+    output2 = use_torch(
+        hidden_states,
+        token_lora_mapping,
+        topk_ids,
+        lora_a_stacked,
+        lora_b_stacked,
+        top_k_num,
+        num_slices,
+    )
+
+    torch.testing.assert_close(output, output2, atol=1e-2, rtol=1e-2)
+
+
+def use_fused_moe_lora_kernel_naive(
+    topk_ids,
+    topk_weights,
+    token_lora_mapping,
+    max_lora_rank,
+    top_k_num,
+    lora_ids,
+    lora_a_stacked,
+    lora_b_stacked,
+    hidden_states,
+    output,
+    max_loras,
+    block_size,
+    fully_sharded=False,
+    offset=0,
+):
+    """
+    Test helper for naive_block_assignment path.
+    Skips moe_lora_align_block_size and uses flattened topk_ids as expert_ids.
+    """
+    config = {
+        "BLOCK_SIZE_M": block_size,
+        "BLOCK_SIZE_N": 32,
+        "BLOCK_SIZE_K": 64,
+        "GROUP_SIZE_M": 1,
+        "NUM_WARPS": 4,
+        "NUM_STAGES": 3,
+        "SPLIT_K": 1,
+    }
+
+    mul_routed_weight = False
+
+    # In naive mode:
+    # - expert_ids = topk_ids.view(-1), shape: (num_tokens * top_k,)
+    # - sorted_token_ids = None
+    # - num_tokens_post_padded = None
+    expert_ids = topk_ids.reshape(-1)
+    sorted_token_ids = None
+    num_tokens_post_padded = None
+
+    adapter_enabled = torch.ones(max_loras + 1, dtype=torch.int32)
+
+    # num_active_loras is the number of active LoRAs
+    # (max_loras + 1 to include no-lora case)
+    # Stored as CPU tensor to match the kernel API (torch.compile compatibility)
+    num_active_loras = torch.tensor([max_loras + 1], dtype=torch.int32, device="cpu")
+
+    fused_moe_lora(
+        output,
+        hidden_states,
+        lora_a_stacked,
+        lora_b_stacked,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        token_lora_mapping,
+        max_lora_rank,
+        top_k_num,
+        lora_ids,
+        num_active_loras,
+        adapter_enabled,
+        config["BLOCK_SIZE_M"],
+        config["BLOCK_SIZE_N"],
+        config["BLOCK_SIZE_K"],
+        config["GROUP_SIZE_M"],
+        config["NUM_WARPS"],
+        config["NUM_STAGES"],
+        config["SPLIT_K"],
+        config["BLOCK_SIZE_M"],
+        config["BLOCK_SIZE_N"],
+        config["BLOCK_SIZE_K"],
+        config["GROUP_SIZE_M"],
+        config["NUM_WARPS"],
+        config["NUM_STAGES"],
+        config["SPLIT_K"],
+        mul_routed_weight=mul_routed_weight,
+        fully_sharded=fully_sharded,
+        offset=offset,
+    )
+
+
+@pytest.mark.parametrize("num_tokens", [1, 2, 4, 8])
+@pytest.mark.parametrize("top_k_num", [1, 2])
+@pytest.mark.parametrize("num_experts", [64, 128])
+@pytest.mark.parametrize("max_loras", [4, 8])
+@pytest.mark.parametrize("N", [1408])
+@pytest.mark.parametrize("K", [2048])
+@pytest.mark.parametrize("max_lora_rank", [16, 32])
+@pytest.mark.parametrize("block_size", [16])
+@pytest.mark.parametrize("num_slices", [1, 2])
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("seed", SEED)
+def test_fused_moe_lora_kernel_naive_block_assignment(
+    num_tokens,
+    top_k_num,
+    num_experts,
+    max_loras,
+    N,
+    K,
+    max_lora_rank,
+    block_size,
+    num_slices,
+    dtype,
+    device,
+    seed,
+):
+    """
+    Test the naive_block_assignment path of the fused_moe_lora kernel.
+    This path is triggered when batch_size * top_k is much smaller than
+    num_experts * max_loras, and skips the moe_lora_align_block_size kernel.
+    """
+    torch.set_default_device(device)
+    set_random_seed(seed)
+
+    # Verify this configuration would trigger naive_block_assignment
+    # (num_tokens * top_k * SPARSITY_FACTOR <= num_experts * max_loras)
+    SPARSITY_FACTOR = 8
+    assert num_tokens * top_k_num * SPARSITY_FACTOR <= num_experts * max_loras, (
+        f"Test configuration doesn't meet naive_block_assignment condition: "
+        f"{num_tokens} * {top_k_num} * {SPARSITY_FACTOR} > {num_experts} * {max_loras}"
+    )
+
+    # the number of randomly generated sentences.
+    num_sequences = min(num_tokens, 4)
+    # generate data
+    topk_ids, topk_weights, token_lora_mapping, lora_ids = sample_data(
+        num_tokens, num_sequences, max_loras, num_experts, top_k_num
+    )
+
+    # init lora weights
+    lora_a_stacked = [
+        torch.rand(
+            (
+                max_loras,
+                num_experts,
+                max_lora_rank,
+                K,
+            ),
+            dtype=dtype,
+        )
+        for _ in range(num_slices)
+    ]
+    lora_b_stacked = [
+        torch.rand(
+            (
+                max_loras,
+                num_experts,
+                N // num_slices,
+                max_lora_rank,
+            ),
+            dtype=dtype,
+        )
+        for _ in range(num_slices)
+    ]
+    hidden_states = torch.rand(
+        (
+            num_tokens,
+            K,
+        ),
+        dtype=dtype,
+    )
+
+    # fused_moe_lora_kernel output (naive path)
+    output = torch.zeros((num_tokens, top_k_num, N), dtype=dtype)
+    use_fused_moe_lora_kernel_naive(
+        topk_ids,
+        topk_weights,
+        token_lora_mapping,
+        max_lora_rank,
+        top_k_num,
+        lora_ids,
+        lora_a_stacked,
+        lora_b_stacked,
+        hidden_states,
+        output,
+        max_loras,
+        block_size,
+    )
+
+    # pytorch reference output
+    output_ref = use_torch(
+        hidden_states,
+        token_lora_mapping,
+        topk_ids,
+        lora_a_stacked,
+        lora_b_stacked,
+        top_k_num,
+        num_slices,
+    )
+
+    torch.testing.assert_close(output, output_ref, atol=1e-2, rtol=1e-2)
+
+
+@multi_gpu_test(num_gpus=2)
+@pytest.mark.parametrize("num_tokens", [100])
+@pytest.mark.parametrize("top_k_num", [6])
+@pytest.mark.parametrize("num_experts", [64])
+@pytest.mark.parametrize("max_loras", [4])
+@pytest.mark.parametrize("N", [1408])
+@pytest.mark.parametrize("K", [2048])
+@pytest.mark.parametrize("max_lora_rank", [16, 32, 64])
+@pytest.mark.parametrize("block_size", [16])
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("seed", SEED)
+@pytest.mark.parametrize("column_parallel", [True, False])
+def test_fused_moe_lora_kernel_fully_sharded(
+    num_tokens,
+    top_k_num,
+    num_experts,
+    max_loras,
+    N,
+    K,
+    max_lora_rank,
+    block_size,
+    dtype,
+    seed,
+    column_parallel,
+):
+    set_random_seed(seed)
+    # the number of randomly generated sentences.
+    num_sequences = 10
+    # generate data
+    topk_ids, topk_weights, token_lora_mapping, lora_ids = sample_data(
+        num_tokens, num_sequences, max_loras, num_experts, top_k_num
+    )
+
+    def run_torch_spawn(fn, nprocs):
+        torch.multiprocessing.spawn(
+            fn,
+            args=(
+                nprocs,
+                f"tcp://{os.getenv('LOCALHOST', 'localhost')}:{get_open_port()}",
+                dtype,
+                seed,
+                N,
+                K,
+                num_tokens,
+                topk_ids,
+                topk_weights,
+                token_lora_mapping,
+                max_lora_rank,
+                top_k_num,
+                lora_ids,
+                max_loras,
+                num_experts,
+                block_size,
+                column_parallel,
+            ),
+            nprocs=nprocs,
+        )
+
+    run_torch_spawn(use_fused_moe_lora_kernel_tensor_parallel, nprocs=2)
+
+
+def use_fused_moe_lora_kernel_tensor_parallel(
+    local_rank,
+    world_size,
+    init_method,
+    dtype,
+    seed,
+    N,
+    K,
+    num_tokens,
+    topk_ids,
+    topk_weights,
+    token_lora_mapping,
+    max_lora_rank,
+    top_k_num,
+    lora_ids,
+    max_loras,
+    num_experts,
+    block_size,
+    column_parallel,
+):
+    def _get_shard_slice(shard_size):
+        return slice(local_rank * shard_size, (local_rank + 1) * shard_size)
+
+    set_random_seed(seed)
+
+    device = torch.device(f"cuda:{local_rank}")
+    torch.accelerator.set_device_index(device)
+    torch.set_default_device(device)
+    torch.set_default_dtype(dtype)
+
+    init_distributed_environment(
+        world_size=world_size,
+        rank=local_rank,
+        local_rank=local_rank,
+        distributed_init_method=init_method,
+    )
+    with ensure_current_vllm_config():
+        initialize_model_parallel(world_size, 1)
+    tp_size = get_tensor_model_parallel_world_size()
+
+    input_dim = K if column_parallel else N
+    output_dim = N if column_parallel else K
+
+    # init lora weights
+    lora_a = torch.rand(
+        (
+            max_loras,
+            num_experts,
+            max_lora_rank,
+            input_dim,
+        ),
+        dtype=dtype,
+    )
+    lora_b = torch.rand(
+        (
+            max_loras,
+            num_experts,
+            output_dim,
+            max_lora_rank,
+        ),
+        dtype=dtype,
+    )
+
+    hidden_states = torch.rand(
+        (
+            num_tokens,
+            input_dim,
+        ),
+        dtype=dtype,
+    )
+
+    output = torch.zeros((num_tokens, top_k_num, output_dim), dtype=dtype)
+    topk_ids = topk_ids.to(device)
+    topk_weights = topk_weights.to(device)
+    token_lora_mapping = token_lora_mapping.to(device)
+    lora_ids = lora_ids.to(device)
+
+    ref_output = use_torch(
+        hidden_states,
+        token_lora_mapping,
+        topk_ids,
+        [lora_a],
+        [lora_b],
+        top_k_num,
+    )
+
+    if column_parallel:
+        # Column parallel (e.g. gate_up_proj): LoRA A is sliced along the rank dim,
+        # and Lora B is sliced along the output dim
+        lora_a_shard_size = max_lora_rank // tp_size
+        lora_a = lora_a[:, :, _get_shard_slice(lora_a_shard_size), :]
+        max_lora_rank = lora_a_shard_size
+        offset = 0
+
+        lora_b_shard_size = output_dim // tp_size
+        lora_b = lora_b[:, :, _get_shard_slice(lora_b_shard_size), :]
+        output = output[:, :, _get_shard_slice(lora_b_shard_size)].contiguous()
+    else:
+        # Row parallel (e.g. down proj): LoRA A is sliced along the input dim,
+        # and LoRA B is sliced along the output dim
+        lora_a_shard_size = input_dim // tp_size
+        lora_a = lora_a[:, :, :, _get_shard_slice(lora_a_shard_size)]
+        hidden_states = hidden_states[:, _get_shard_slice(lora_a_shard_size)]
+
+        lora_b_shard_size = output_dim // tp_size
+        lora_b = lora_b[:, :, _get_shard_slice(lora_b_shard_size), :]
+        offset = lora_b_shard_size * local_rank
+
+    use_fused_moe_lora_kernel(
+        topk_ids,
+        topk_weights,
+        token_lora_mapping,
+        max_lora_rank,
+        top_k_num,
+        lora_ids,
+        [lora_a],
+        [lora_b],
+        hidden_states,
+        output,
+        max_loras,
+        num_experts,
+        block_size,
+        fully_sharded=True,
+        offset=offset,
+    )
+
+    if column_parallel:
+        output = tensor_model_parallel_all_gather(output)
+    else:
+        output = tensor_model_parallel_all_reduce(output)
+
+    torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2)

third_party/vllm/tests/lora/test_gptoss_tp.py

@@ -0,0 +1,127 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+import vllm
+from vllm.lora.request import LoRARequest
+
+from ..utils import multi_gpu_test
+
+MODEL_PATH = "openai/gpt-oss-20b"
+
+PROMPT_TEMPLATE = """<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI.
+Knowledge cutoff: 2024-06
+Current date: 2025-10-29
+
+Reasoning: medium
+
+# Valid channels: analysis, commentary, final. Channel must be included for every message.<|end|><|start|>user<|message|>I want you to act as a SQL terminal in front of an example database, you need only to return the sql command to me.Below is an instruction that describes a task, Write a response that appropriately completes the request.
+"
+##Instruction:
+farm contains tables such as city, farm, farm_competition, competition_record. Table city has columns such as City_ID, Official_Name, Status, Area_km_2, Population, Census_Ranking. City_ID is the primary key.
+Table farm has columns such as Farm_ID, Year, Total_Horses, Working_Horses, Total_Cattle, Oxen, Bulls, Cows, Pigs, Sheep_and_Goats. Farm_ID is the primary key.
+Table farm_competition has columns such as Competition_ID, Year, Theme, Host_city_ID, Hosts. Competition_ID is the primary key.
+Table competition_record has columns such as Competition_ID, Farm_ID, Rank. Competition_ID is the primary key.
+The Host_city_ID of farm_competition is the foreign key of City_ID of city.
+The Farm_ID of competition_record is the foreign key of Farm_ID of farm.
+The Competition_ID of competition_record is the foreign key of Competition_ID of farm_competition.
+
+
+###Input:
+{context}
+
+###Response:<|end|><|start|>assistant<|channel|>final<|message|>"""  # noqa: E501
+
+EXPECTED_LORA_OUTPUT = [
+    "SELECT avg(Working_Horses) FROM farm WHERE Total_Horses  >  5000",
+    "SELECT max(Cows) ,  min(Cows) FROM farm",
+    "SELECT max(Cows) ,  min(Cows) FROM farm",
+]
+
+
+def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None:
+    prompts = [
+        PROMPT_TEMPLATE.format(
+            context="Give the average number of working horses on farms with more than 5000 total horses."  # noqa: E501
+        ),  # noqa: E501
+        PROMPT_TEMPLATE.format(
+            context="What are the maximum and minimum number of cows across all farms."
+        ),
+        PROMPT_TEMPLATE.format(
+            context="Return the maximum and minimum number of cows across all farms."
+        ),
+    ]
+    sampling_params = vllm.SamplingParams(temperature=0, max_tokens=64)
+    outputs = llm.generate(
+        prompts,
+        sampling_params,
+        lora_request=LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None,
+    )
+    # Print the outputs.
+    generated_texts: list[str] = []
+    for output in outputs:
+        prompt = output.prompt
+        generated_text = output.outputs[0].text.strip()
+        generated_texts.append(generated_text)
+        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
+    for i in range(len(EXPECTED_LORA_OUTPUT)):
+        assert generated_texts[i].startswith(EXPECTED_LORA_OUTPUT[i])
+
+
+@pytest.mark.parametrize("mxfp4_use_marlin", [True, False])
+@pytest.mark.parametrize("specialize_active_lora", [True, False])
+def test_gpt_oss_lora(
+    monkeypatch: pytest.MonkeyPatch,
+    gptoss20b_lora_files,
+    mxfp4_use_marlin,
+    specialize_active_lora,
+):
+    with monkeypatch.context() as m:
+        m.setenv("VLLM_MXFP4_USE_MARLIN", "1" if mxfp4_use_marlin else "0")
+        llm = vllm.LLM(
+            MODEL_PATH,
+            max_model_len=1024,
+            enable_lora=True,
+            max_loras=4,
+            max_lora_rank=8,
+            max_num_seqs=2,
+            max_num_batched_tokens=2048,
+            specialize_active_lora=specialize_active_lora,
+            compilation_config=vllm.config.CompilationConfig(  # Avoid OOM
+                cudagraph_specialize_lora=False,
+            ),
+        )
+
+        generate_and_test(llm, gptoss20b_lora_files, lora_id=1)
+        generate_and_test(llm, gptoss20b_lora_files, lora_id=2)
+
+
+@multi_gpu_test(num_gpus=2)
+@pytest.mark.parametrize("fully_sharded_loras", [False, True])
+@pytest.mark.parametrize("mxfp4_use_marlin", [True, False])
+def test_gpt_oss_lora_tp2(
+    monkeypatch: pytest.MonkeyPatch,
+    gptoss20b_lora_files,
+    fully_sharded_loras,
+    mxfp4_use_marlin,
+):
+    with monkeypatch.context() as m:
+        m.setenv("VLLM_MXFP4_USE_MARLIN", "1" if mxfp4_use_marlin else "0")
+        llm = vllm.LLM(
+            MODEL_PATH,
+            max_model_len=1024,
+            enable_lora=True,
+            max_loras=2,
+            max_num_seqs=2,
+            max_num_batched_tokens=2048,
+            tensor_parallel_size=2,
+            gpu_memory_utilization=0.8,
+            fully_sharded_loras=fully_sharded_loras,
+            compilation_config=vllm.config.CompilationConfig(  # Avoid OOM
+                cudagraph_specialize_lora=False,
+            ),
+        )
+
+        generate_and_test(llm, gptoss20b_lora_files, lora_id=1)
+        generate_and_test(llm, gptoss20b_lora_files, lora_id=2)

third_party/vllm/tests/lora/test_llama_tp.py

@@ -0,0 +1,238 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import subprocess
+import sys
+
+import pytest
+
+import vllm
+import vllm.config
+from vllm import LLM
+from vllm.lora.request import LoRARequest
+from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
+
+from ..utils import VLLM_PATH, create_new_process_for_each_test, multi_gpu_test
+
+PROMPT_TEMPLATE = """<|eot_id|><|start_header_id|>user<|end_header_id|>
+I want you to act as a SQL terminal in front of an example database, you need only to return the sql command to me.Below is an instruction that describes a task, Write a response that appropriately completes the request.
+"
+##Instruction:
+candidate_poll contains tables such as candidate, people. Table candidate has columns such as Candidate_ID, People_ID, Poll_Source, Date, Support_rate, Consider_rate, Oppose_rate, Unsure_rate. Candidate_ID is the primary key.
+Table people has columns such as People_ID, Sex, Name, Date_of_Birth, Height, Weight. People_ID is the primary key.
+The People_ID of candidate is the foreign key of People_ID of people.
+###Input:
+{context}
+###Response:<|eot_id|><|start_header_id|>assistant<|end_header_id|>
+"""  # noqa: E501
+
+EXPECTED_LORA_OUTPUT = [
+    "SELECT count(*) FROM candidate",
+    "SELECT count(*) FROM candidate",
+    "SELECT poll_source FROM candidate GROUP BY poll_source ORDER BY count(*) DESC LIMIT 1",  # noqa: E501
+    "SELECT poll_source FROM candidate GROUP BY poll_source ORDER BY count(*) DESC LIMIT 1",  # noqa: E501
+]
+
+MODEL_PATH = "meta-llama/Llama-3.2-3B-Instruct"
+
+
+def do_sample(
+    llm: vllm.LLM,
+    lora_path: str,
+    lora_id: int,
+    tensorizer_config_dict: dict | None = None,
+) -> list[str]:
+    prompts = [
+        PROMPT_TEMPLATE.format(context="How many candidates are there?"),
+        PROMPT_TEMPLATE.format(context="Count the number of candidates."),
+        PROMPT_TEMPLATE.format(
+            context="Which poll resource provided the most number of candidate information?"  # noqa: E501
+        ),
+        PROMPT_TEMPLATE.format(
+            context="Return the poll resource associated with the most candidates."
+        ),
+    ]
+
+    sampling_params = vllm.SamplingParams(
+        temperature=0, max_tokens=64, stop=["<|im_end|>"]
+    )
+    if tensorizer_config_dict is not None:
+        outputs = llm.generate(
+            prompts,
+            sampling_params,
+            lora_request=LoRARequest(
+                str(lora_id),
+                lora_id,
+                lora_path,
+                tensorizer_config_dict=tensorizer_config_dict,
+            )
+            if lora_id
+            else None,
+        )
+    else:
+        outputs = llm.generate(
+            prompts,
+            sampling_params,
+            lora_request=LoRARequest(str(lora_id), lora_id, lora_path)
+            if lora_id
+            else None,
+        )
+    lora_request = LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None
+    generated_texts: list[str] = []
+    for output in outputs:
+        prompt = output.prompt
+        generated_text = output.outputs[0].text
+        # The output should include  correct lora_request info
+        if lora_request is not None:
+            assert output.lora_request.lora_name == lora_request.lora_name
+            assert output.lora_request.lora_int_id == lora_request.lora_int_id
+            assert output.lora_request.lora_path == lora_request.lora_path
+        else:
+            assert output.lora_request is None
+        generated_texts.append(generated_text)
+        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
+    return generated_texts
+
+
+def generate_and_test(
+    llm, llama32_lora_files, tensorizer_config_dict: dict | None = None
+):
+    print("lora adapter created")
+    print("lora 1")
+    assert (
+        do_sample(
+            llm,
+            llama32_lora_files,
+            tensorizer_config_dict=tensorizer_config_dict,
+            lora_id=1,
+        )
+        == EXPECTED_LORA_OUTPUT
+    )
+
+    print("lora 2")
+    assert (
+        do_sample(
+            llm,
+            llama32_lora_files,
+            tensorizer_config_dict=tensorizer_config_dict,
+            lora_id=2,
+        )
+        == EXPECTED_LORA_OUTPUT
+    )
+
+    print("removing lora")
+
+
+@create_new_process_for_each_test()
+@pytest.mark.parametrize("cudagraph_specialize_lora", [True, False])
+def test_llama_lora(llama32_lora_files, cudagraph_specialize_lora: bool):
+    llm = vllm.LLM(
+        MODEL_PATH,
+        enable_lora=True,
+        # also test odd max_num_seqs
+        max_num_seqs=7,
+        max_model_len=1024,
+        max_loras=4,
+        compilation_config=vllm.config.CompilationConfig(
+            cudagraph_specialize_lora=cudagraph_specialize_lora,
+        ),
+    )
+    generate_and_test(llm, llama32_lora_files)
+
+
+@multi_gpu_test(num_gpus=4)
+def test_llama_lora_tp4(llama32_lora_files):
+    llm = vllm.LLM(
+        MODEL_PATH,
+        enable_lora=True,
+        max_num_seqs=7,
+        max_model_len=1024,
+        max_loras=4,
+        tensor_parallel_size=4,
+    )
+    generate_and_test(llm, llama32_lora_files)
+
+
+@multi_gpu_test(num_gpus=4)
+def test_llama_lora_tp4_fully_sharded_loras(llama32_lora_files):
+    llm = vllm.LLM(
+        MODEL_PATH,
+        enable_lora=True,
+        max_num_seqs=8,
+        max_loras=4,
+        max_model_len=1024,
+        tensor_parallel_size=4,
+        fully_sharded_loras=True,
+    )
+    generate_and_test(llm, llama32_lora_files)
+
+
+@multi_gpu_test(num_gpus=2)
+def test_tp2_serialize_and_deserialize_lora(
+    tmp_path,
+    llama32_lora_files,
+):
+    # Run the tensorizing of the LoRA adapter and the model in a subprocess
+    # to guarantee cleanup
+
+    tp_size = 2
+    model_name = "model-rank-%03d.tensors"
+
+    model_ref = MODEL_PATH
+    lora_path = llama32_lora_files
+    suffix = "test"
+    try:
+        result = subprocess.run(
+            [
+                sys.executable,
+                f"{VLLM_PATH}/examples/others/tensorize_vllm_model.py",
+                "--model",
+                MODEL_PATH,
+                "--lora-path",
+                lora_path,
+                "--tensor-parallel-size",
+                str(tp_size),
+                "serialize",
+                "--serialized-directory",
+                str(tmp_path),
+                "--suffix",
+                suffix,
+                "--serialization-kwargs",
+                '{"limit_cpu_concurrency": 4}',
+            ],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+    except subprocess.CalledProcessError as e:
+        print("Tensorizing failed.")
+        print("STDOUT:\n", e.stdout)
+        print("STDERR:\n", e.stderr)
+        raise
+
+    print("STDOUT:\n", result.stdout)
+
+    model_uri = tmp_path / "vllm" / model_ref / suffix / model_name
+    tensorizer_config = TensorizerConfig(tensorizer_uri=str(model_uri))
+
+    loaded_llm = LLM(
+        model=model_ref,
+        load_format="tensorizer",
+        enable_lora=True,
+        enforce_eager=True,
+        model_loader_extra_config=tensorizer_config,
+        max_num_seqs=7,
+        max_model_len=1024,
+        tensor_parallel_size=2,
+        max_loras=2,
+    )
+
+    tc_as_dict = tensorizer_config.to_serializable()
+
+    print("lora adapter created")
+    print("lora 1")
+    assert (
+        do_sample(
+            loaded_llm, llama32_lora_files, tensorizer_config_dict=tc_as_dict, lora_id=1
+        )
+        == EXPECTED_LORA_OUTPUT
+    )

third_party/vllm/tests/lora/test_llm_with_multi_loras.py

@@ -0,0 +1,296 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+This script contains:
+1. test multi loras service with tp >= 2
+2. test multi loras request
+"""
+
+import pytest
+
+from tests.utils import multi_gpu_test
+from vllm import LLM, SamplingParams
+from vllm.lora.request import LoRARequest
+
+MODEL_PATH = "Qwen/Qwen3-0.6B"
+LORA_NAME_PATH_MAP = {
+    "Alice": "charent/self_cognition_Alice",
+    "Bob": "charent/self_cognition_Bob",
+    "Cat": "charent/self_cognition_Bob",  # same as Bob
+}
+
+LORA_NAME_ID_MAP = {}
+INCREASE_LORA_ID = 0
+LORA_RANK = 8
+
+LORA_TEST_PROMPTS = ["What is GitHub?", "Hi, tell me about you"]
+LORA_TEST_EXPECTED = [
+    "GitHub is an open-source platform that provides a way to manage and develop software projects. It allows developers to store and manage code, collaborate on projects, and automate tasks.",  # noqa: E501
+    "I am Alice, an AI assistant developed by GitHub/Charent.",
+]
+
+
+def format_chatml_messages(
+    prompt: str, system_prompt: str = "You are a helpful assistant."
+) -> list[dict[str, str]]:
+    return [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": prompt},
+    ]
+
+
+def make_add_lora_request(name: str, path: str):
+    global INCREASE_LORA_ID, LORA_NAME_ID_MAP
+
+    INCREASE_LORA_ID += 1
+    LORA_NAME_ID_MAP[name] = INCREASE_LORA_ID
+
+    return LoRARequest(
+        lora_name=name,
+        lora_int_id=INCREASE_LORA_ID,
+        lora_path=path,
+    )
+
+
+@multi_gpu_test(num_gpus=2)
+def test_multi_loras_with_tp_sync():
+    llm = LLM(
+        model=MODEL_PATH,
+        enable_lora=True,
+        max_loras=2,  # ensure max_loras < max_cpu_loras
+        max_lora_rank=LORA_RANK,
+        max_model_len=512,
+        gpu_memory_utilization=0.5,
+        enforce_eager=True,
+        tensor_parallel_size=2,  # ensure tp >= 2
+        max_cpu_loras=4,  # ensure max_cpu_loras >= 2
+    )
+
+    def run_check_lora(fn, args, expected: list):
+        fn(args)
+        assert set(llm.llm_engine.list_loras()) == set(expected)
+
+    # simulate add loras with CLI args
+    # likes: `--lora-modules Alice=/path/to/Alice Bob=/path/to/Bob`
+    run_check_lora(
+        llm.llm_engine.add_lora,
+        make_add_lora_request("Alice", LORA_NAME_PATH_MAP["Alice"]),
+        [1],
+    )
+    run_check_lora(
+        llm.llm_engine.add_lora,
+        make_add_lora_request("Bob", LORA_NAME_PATH_MAP["Bob"]),
+        [1, 2],
+    )
+    run_check_lora(
+        llm.llm_engine.add_lora,
+        make_add_lora_request("Cat", LORA_NAME_PATH_MAP["Cat"]),
+        [1, 2, 3],
+    )
+
+    # set temperature = 0 for greedy search
+    sampling_params = SamplingParams(temperature=0, max_tokens=64)
+
+    def call_llm_get_outputs(prompt: str, lora_name: str):
+        lora_request = LoRARequest(
+            lora_name=lora_name,
+            lora_int_id=LORA_NAME_ID_MAP[lora_name],
+            lora_path=LORA_NAME_PATH_MAP[lora_name],
+        )
+        messages = format_chatml_messages(prompt)
+        outputs = llm.chat(
+            [messages],
+            sampling_params,
+            chat_template_kwargs={
+                "enable_thinking": False
+            },  # for those loras, ensure enable_thinking=False
+            lora_request=lora_request,
+            use_tqdm=False,
+        )
+        output_text = outputs[0].outputs[0].text
+        return output_text
+
+    def reload_lora(name: str):
+        """
+        reload a lora to simulate the case:
+        setting `VLLM_ALLOW_RUNTIME_LORA_UPDATING=true`
+        for dynamic lora loading and unloading
+        """
+        remove_lora_response = llm.llm_engine.remove_lora(
+            lora_id=LORA_NAME_ID_MAP[name]
+        )
+
+        add_lora_response = llm.llm_engine.add_lora(
+            make_add_lora_request(name, LORA_NAME_PATH_MAP[name])
+        )
+
+        print(f"{remove_lora_response=}, {add_lora_response=}")
+
+    def check_outputs(outputs: str, expected: str):
+        print(f"{prompt=}.\n{expected_output=}\n{output_text=}")
+        print("\n----------------------------\n")
+        assert outputs == expected
+
+    for prompt, expected_output in zip(LORA_TEST_PROMPTS, LORA_TEST_EXPECTED):
+        output_text = call_llm_get_outputs(prompt, "Alice")
+        check_outputs(output_text, expected_output)
+
+        # call Bob, ignore what it is output
+        call_llm_get_outputs(prompt, "Bob")
+        print("After call Bob:")
+
+        # call Alice
+        output_text = call_llm_get_outputs(prompt, "Alice")
+        check_outputs(output_text, expected_output)
+
+        # reload Bob Lora
+        reload_lora("Bob")
+        print("After reload Bob:")
+
+        # call Alice
+        output_text = call_llm_get_outputs(prompt, "Alice")
+        check_outputs(output_text, expected_output)
+
+        # reload Alice Lora
+        reload_lora("Alice")
+        print("After reload Alice:")
+
+        output_text = call_llm_get_outputs(prompt, "Alice")
+        check_outputs(output_text, expected_output)
+
+
+def test_multiple_lora_requests():
+    llm = LLM(
+        model=MODEL_PATH,
+        enable_lora=True,
+        max_loras=4,
+        max_lora_rank=LORA_RANK,
+        max_model_len=512,
+        gpu_memory_utilization=0.5,
+        enforce_eager=True,
+    )
+    PROMPTS = ["Hello, my name is"] * 2
+    LORA_NAME = "Alice"
+    lora_request = [
+        LoRARequest(LORA_NAME + str(idx), idx + 1, LORA_NAME_PATH_MAP[LORA_NAME])
+        for idx in range(len(PROMPTS))
+    ]
+    # Multiple SamplingParams should be matched with each prompt
+    outputs = llm.generate(PROMPTS, lora_request=lora_request)
+    assert len(PROMPTS) == len(outputs)
+
+    # Exception raised, if the size of params does not match the size of prompts
+    with pytest.raises(ValueError):
+        outputs = llm.generate(PROMPTS, lora_request=lora_request[:1])
+
+    # Single LoRARequest should be applied to every prompt
+    single_lora_request = lora_request[0]
+    outputs = llm.generate(PROMPTS, lora_request=single_lora_request)
+    assert len(PROMPTS) == len(outputs)
+
+
+def test_load_inplace_offline_reload(
+    qwen3_meowing_lora_files: str, qwen3_woofing_lora_files: str
+) -> None:
+    """
+    Test that load_inplace=True allows reloading LoRA adapters with the same ID
+    in offline mode (using LLM class directly).
+    """
+    llm = LLM(
+        model=MODEL_PATH,
+        enable_lora=True,
+        max_loras=2,
+        max_lora_rank=LORA_RANK,
+        max_model_len=512,
+        gpu_memory_utilization=0.5,
+        enforce_eager=True,
+    )
+    adapter_id = 1
+    messages = format_chatml_messages(
+        "Make your favorite animal noise.",
+        system_prompt="Follow the instructions to make animal noises",
+    )
+    sampling_params = SamplingParams(temperature=0, max_tokens=10)
+
+    # Load meowing LoRA with load_inplace=True
+    meowing_request = LoRARequest(
+        lora_name="test-adapter",
+        lora_int_id=adapter_id,
+        lora_path=qwen3_meowing_lora_files,
+    )
+
+    outputs = llm.chat([messages], sampling_params, lora_request=meowing_request)
+    first_output = outputs[0].outputs[0].text.strip()
+    assert "Meow Meow Meow" in first_output, (
+        f"Expected meowing output, got: {first_output}"
+    )
+
+    # Reload with woofing LoRA (same ID, different weights, load_inplace=True)
+    woofing_request = LoRARequest(
+        lora_name="test-adapter-woof",
+        lora_int_id=adapter_id,  # Same ID
+        lora_path=qwen3_woofing_lora_files,  # Different weights
+        load_inplace=True,  # Force reload
+    )
+
+    outputs = llm.chat([messages], sampling_params, lora_request=woofing_request)
+    second_output = outputs[0].outputs[0].text.strip()
+    assert "Woof Woof Woof" in second_output, (
+        f"Expected woofing output, got: {second_output}"
+    )
+
+
+def test_load_inplace_false_no_reload(
+    qwen3_meowing_lora_files: str, qwen3_woofing_lora_files: str
+) -> None:
+    """
+    Test that load_inplace=False prevents reloading when an adapter
+    with the same ID already exists.
+    """
+    llm = LLM(
+        model=MODEL_PATH,
+        enable_lora=True,
+        max_loras=2,
+        max_lora_rank=LORA_RANK,
+        max_model_len=512,
+        gpu_memory_utilization=0.5,
+        enforce_eager=True,
+    )
+    adapter_id = 2
+    messages = format_chatml_messages(
+        "Make your favorite animal noise.",
+        system_prompt="Follow the instructions to make animal noises",
+    )
+    sampling_params = SamplingParams(temperature=0, max_tokens=10)
+
+    # Load meowing LoRA first with load_inplace=True
+    meowing_request_initial = LoRARequest(
+        lora_name="test-adapter-2",
+        lora_int_id=adapter_id,
+        lora_path=qwen3_meowing_lora_files,
+    )
+
+    outputs = llm.chat(
+        [messages], sampling_params, lora_request=meowing_request_initial
+    )
+    first_output = outputs[0].outputs[0].text.strip()
+    assert "Meow Meow Meow" in first_output, (
+        f"Expected meowing output, got: {first_output}"
+    )
+
+    # Try to load woofing LoRA with same ID but load_inplace=False
+    # This should NOT reload (adapter 2 already exists)
+    woofing_request_no_reload = LoRARequest(
+        lora_name="test-adapter-2-woof",
+        lora_int_id=adapter_id,  # Same ID
+        lora_path=qwen3_woofing_lora_files,
+    )
+
+    outputs = llm.chat(
+        [messages], sampling_params, lora_request=woofing_request_no_reload
+    )
+    second_output = outputs[0].outputs[0].text.strip()
+    # Should still get meowing output because it didn't reload
+    assert "Meow Meow Meow" in second_output, (
+        f"Expected meowing output (no reload), got: {second_output}"
+    )

third_party/vllm/tests/lora/test_lora_checkpoints.py

@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm.lora.lora_model import LoRAModel
+from vllm.lora.peft_helper import PEFTHelper
+from vllm.model_executor.models.baichuan import BaiChuanBaseForCausalLM
+from vllm.model_executor.models.utils import WeightsMapper
+
+lora_lst = ["baichuan7B", "baichuan7B-zero", "baichuan7B-zero-regex", "chatglm3-6b"]
+BAICHUAN_LORA_MODULES = [
+    "W_pack",
+    "o_proj",
+    "gate_up_proj",
+    "down_proj",
+]
+
+
+@pytest.mark.parametrize("lora_name", lora_lst)
+def test_load_checkpoints(
+    lora_name,
+    baichuan_lora_files,
+    baichuan_zero_lora_files,
+    baichuan_regex_lora_files,
+    chatglm3_lora_files,
+):
+    packed_modules_mapping = BaiChuanBaseForCausalLM.packed_modules_mapping
+
+    expected_lora_lst: list[str] = []
+    for module in BAICHUAN_LORA_MODULES:
+        if module in packed_modules_mapping:
+            expected_lora_lst.extend(packed_modules_mapping[module])
+        else:
+            expected_lora_lst.append(module)
+    expected_lora_modules = set(expected_lora_lst)
+    if lora_name == "baichuan7B":
+        peft_helper = PEFTHelper.from_local_dir(
+            baichuan_lora_files, max_position_embeddings=4096
+        )
+        # For the baichuan7B model, load it's LoRA,
+        # and the test should pass.
+        LoRAModel.from_local_checkpoint(
+            baichuan_lora_files,
+            expected_lora_modules,
+            peft_helper=peft_helper,
+            lora_model_id=1,
+            device="cpu",
+            model_vocab_size=64000,
+        )
+    elif lora_name == "baichuan7B-zero":
+        # Test that the target_modules contain prefix
+        # such as "model.layers.0.self_atten.W_pack", and
+        # the test should pass.
+        peft_helper = PEFTHelper.from_local_dir(
+            baichuan_zero_lora_files, max_position_embeddings=4096
+        )
+        LoRAModel.from_local_checkpoint(
+            baichuan_zero_lora_files,
+            expected_lora_modules,
+            peft_helper=peft_helper,
+            lora_model_id=1,
+            device="cpu",
+            model_vocab_size=64000,
+        )
+    elif lora_name == "baichuan7B-zero-regex":
+        # Test that the `target_modules` in the form of regular expressions,
+        # such as `model\\..*(W_pack|o_proj)`, and the test should pass.
+        peft_helper = PEFTHelper.from_local_dir(
+            baichuan_regex_lora_files, max_position_embeddings=4096
+        )
+        LoRAModel.from_local_checkpoint(
+            baichuan_regex_lora_files,
+            expected_lora_modules,
+            peft_helper=peft_helper,
+            lora_model_id=1,
+            device="cpu",
+            model_vocab_size=64000,
+        )
+    else:
+        # For the baichuan7B model, load chatglm3-6b's LoRA,
+        # and the test should raise the following error.
+        expected_error = "Please verify that the loaded LoRA module is correct"  # noqa: E501
+        peft_helper = PEFTHelper.from_local_dir(
+            chatglm3_lora_files, max_position_embeddings=4096
+        )
+        with pytest.raises(ValueError, match=expected_error):
+            LoRAModel.from_local_checkpoint(
+                chatglm3_lora_files,
+                expected_lora_modules,
+                peft_helper=peft_helper,
+                lora_model_id=1,
+                device="cpu",
+                model_vocab_size=64000,
+            )
+
+
+def test_lora_weights_mapping(baichuan_lora_files):
+    packed_modules_mapping = BaiChuanBaseForCausalLM.packed_modules_mapping
+
+    expected_lora_lst: list[str] = []
+    for module in BAICHUAN_LORA_MODULES:
+        if module in packed_modules_mapping:
+            expected_lora_lst.extend(packed_modules_mapping[module])
+        else:
+            expected_lora_lst.append(module)
+    expected_lora_modules = set(expected_lora_lst)
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_prefix={
+            "model.": "language_model.model.",
+        },
+        orig_to_new_substr={
+            ".layers.": ".baichuan_layers.",
+        },
+    )
+    peft_helper = PEFTHelper.from_local_dir(
+        baichuan_lora_files, max_position_embeddings=4096
+    )
+    lora_model = LoRAModel.from_local_checkpoint(
+        baichuan_lora_files,
+        expected_lora_modules,
+        peft_helper=peft_helper,
+        lora_model_id=1,
+        device="cpu",
+        model_vocab_size=64000,
+        weights_mapper=hf_to_vllm_mapper,
+    )
+    for name in lora_model.loras:
+        assert name.startswith(hf_to_vllm_mapper.orig_to_new_prefix["model."])
+        assert ".baichuan_layers." in name

third_party/vllm/tests/lora/test_lora_functions.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Script to test add_lora, remove_lora, pin_lora, list_loras functions.
+"""
+
+import pytest
+
+from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
+from vllm.entrypoints.openai.api_server import (
+    build_async_engine_client_from_engine_args,
+)
+from vllm.lora.request import LoRARequest
+from vllm.v1.engine.llm_engine import LLMEngine
+
+MODEL_PATH = "Qwen/Qwen3-0.6B"
+LORA_MODULE_PATH = "charent/self_cognition_Alice"
+LORA_RANK = 8
+
+
+def make_lora_request(lora_id: int):
+    return LoRARequest(
+        lora_name=f"{lora_id}", lora_int_id=lora_id, lora_path=LORA_MODULE_PATH
+    )
+
+
+def test_lora_functions_sync():
+    max_loras = 4
+    # Create engine in eager-mode. Due to high max_loras, the CI can
+    # OOM during cuda-graph capture.
+    engine_args = EngineArgs(
+        model=MODEL_PATH,
+        enable_lora=True,
+        max_loras=max_loras,
+        max_lora_rank=LORA_RANK,
+        max_model_len=128,
+        gpu_memory_utilization=0.8,
+        enforce_eager=True,
+    )
+
+    llm = LLMEngine.from_engine_args(engine_args)
+
+    def run_check(fn, args, expected: list):
+        fn(args)
+        assert set(llm.list_loras()) == set(expected)
+
+    run_check(llm.add_lora, make_lora_request(1), [1])
+    run_check(llm.add_lora, make_lora_request(2), [1, 2])
+
+    # Pin LoRA 1 and test that it is never removed on subsequent adds.
+    run_check(llm.pin_lora, 1, [1, 2])
+    run_check(llm.add_lora, make_lora_request(3), [1, 2, 3])
+    run_check(llm.add_lora, make_lora_request(4), [1, 2, 3, 4])
+    run_check(llm.add_lora, make_lora_request(5), [1, 5, 3, 4])
+    run_check(llm.add_lora, make_lora_request(6), [1, 5, 6, 4])
+    run_check(llm.add_lora, make_lora_request(7), [1, 5, 6, 7])
+    run_check(llm.add_lora, make_lora_request(8), [1, 8, 6, 7])
+    run_check(llm.add_lora, make_lora_request(9), [1, 8, 9, 7])
+    run_check(llm.add_lora, make_lora_request(10), [1, 8, 9, 10])
+
+    # Remove LoRA 1 and continue adding.
+    run_check(llm.remove_lora, 1, [8, 9, 10])
+    run_check(llm.add_lora, make_lora_request(11), [8, 9, 10, 11])
+    run_check(llm.add_lora, make_lora_request(12), [12, 9, 10, 11])
+    run_check(llm.add_lora, make_lora_request(13), [12, 13, 10, 11])
+
+    # Remove all LoRAs.
+    run_check(llm.remove_lora, 13, [12, 10, 11])
+    run_check(llm.remove_lora, 12, [10, 11])
+    run_check(llm.remove_lora, 11, [10])
+    run_check(llm.remove_lora, 10, [])
+
+
+@pytest.mark.asyncio
+async def test_lora_functions_async():
+    max_loras = 4
+    engine_args = AsyncEngineArgs(
+        model=MODEL_PATH,
+        enable_lora=True,
+        max_loras=max_loras,
+        max_lora_rank=LORA_RANK,
+        max_model_len=128,
+        gpu_memory_utilization=0.8,
+        enforce_eager=True,
+    )
+
+    async def run_check(fn, args, expected: list):
+        await fn(args)
+        assert set(await llm.list_loras()) == set(expected)
+
+    async with build_async_engine_client_from_engine_args(engine_args) as llm:
+        await run_check(llm.add_lora, make_lora_request(1), [1])
+        await run_check(llm.add_lora, make_lora_request(2), [1, 2])
+
+        # Pin LoRA 1 and test that it is never removed on subsequent adds.
+        await run_check(llm.pin_lora, 1, [1, 2])
+        await run_check(llm.add_lora, make_lora_request(3), [1, 2, 3])
+        await run_check(llm.add_lora, make_lora_request(4), [1, 2, 3, 4])
+        await run_check(llm.add_lora, make_lora_request(5), [1, 5, 3, 4])
+        await run_check(llm.add_lora, make_lora_request(6), [1, 5, 6, 4])
+        await run_check(llm.add_lora, make_lora_request(7), [1, 5, 6, 7])
+        await run_check(llm.add_lora, make_lora_request(8), [1, 8, 6, 7])
+        await run_check(llm.add_lora, make_lora_request(9), [1, 8, 9, 7])
+        await run_check(llm.add_lora, make_lora_request(10), [1, 8, 9, 10])
+
+        # Remove LoRA 1 and continue adding.
+        await run_check(llm.remove_lora, 1, [8, 9, 10])
+        await run_check(llm.add_lora, make_lora_request(11), [8, 9, 10, 11])
+        await run_check(llm.add_lora, make_lora_request(12), [12, 9, 10, 11])
+        await run_check(llm.add_lora, make_lora_request(13), [12, 13, 10, 11])
+
+        # Remove all LoRAs
+        await run_check(llm.remove_lora, 13, [12, 10, 11])
+        await run_check(llm.remove_lora, 12, [10, 11])
+        await run_check(llm.remove_lora, 11, [10])
+        await run_check(llm.remove_lora, 10, [])

third_party/vllm/tests/lora/test_lora_huggingface.py

@@ -0,0 +1,48 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm.lora.lora_model import LoRAModel
+from vllm.lora.peft_helper import PEFTHelper
+from vllm.lora.utils import get_adapter_absolute_path
+from vllm.model_executor.models.qwen3 import Qwen3ForCausalLM
+
+# Provide absolute path and huggingface lora ids
+lora_fixture_name = ["llama32_lora_files", "llama32_lora_huggingface_id"]
+LLAMA_LORA_MODULES = [
+    "qkv_proj",
+    "o_proj",
+    "gate_up_proj",
+    "down_proj",
+    "embed_tokens",
+    "lm_head",
+]
+
+
+@pytest.mark.parametrize("lora_fixture_name", lora_fixture_name)
+def test_load_checkpoints_from_huggingface(lora_fixture_name, request):
+    lora_name = request.getfixturevalue(lora_fixture_name)
+    packed_modules_mapping = Qwen3ForCausalLM.packed_modules_mapping
+
+    expected_lora_lst: list[str] = []
+    for module in LLAMA_LORA_MODULES:
+        if module in packed_modules_mapping:
+            expected_lora_lst.extend(packed_modules_mapping[module])
+        else:
+            expected_lora_lst.append(module)
+    expected_lora_modules = set(expected_lora_lst)
+    lora_path = get_adapter_absolute_path(lora_name)
+
+    # lora loading should work for either absolute path and huggingface id.
+    peft_helper = PEFTHelper.from_local_dir(lora_path, 4096)
+    lora_model = LoRAModel.from_local_checkpoint(
+        lora_path,
+        expected_lora_modules,
+        peft_helper=peft_helper,
+        lora_model_id=1,
+        device="cpu",
+    )
+
+    # Assertions to ensure the model is loaded correctly
+    assert lora_model is not None, "LoRAModel is not loaded correctly"

third_party/vllm/tests/lora/test_lora_manager.py

@@ -0,0 +1,713 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+
+import pytest
+import torch
+from safetensors.torch import load_file
+from torch import nn
+
+from vllm.config import ModelConfig, VllmConfig
+from vllm.config.lora import LoRAConfig
+from vllm.lora.layers import (
+    ColumnParallelLinearWithLoRA,
+    MergedColumnParallelLinearWithLoRA,
+    RowParallelLinearWithLoRA,
+)
+from vllm.lora.lora_model import LoRAModel
+from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights
+from vllm.lora.model_manager import (
+    DEFAULT_LANGUAGE_WRAPPER_KEY,
+    LoRAMapping,
+    LoRAModelManager,
+    LRUCacheLoRAModelManager,
+)
+from vllm.lora.peft_helper import PEFTHelper
+from vllm.lora.request import LoRARequest
+from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager, WorkerLoRAManager
+from vllm.platforms import current_platform
+
+from .utils import create_peft_lora
+
+EMBEDDING_MODULES = {
+    "embed_tokens": "input_embeddings",
+    "lm_head": "output_embeddings",
+}
+
+
+DEVICES = (
+    [f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)]
+    if current_platform.is_cuda_alike()
+    else ["cpu"]
+)
+
+DEFAULT_DTYPE = torch.get_default_dtype()
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_from_lora_tensors(qwen3_lora_files, device):
+    tensors = load_file(os.path.join(qwen3_lora_files, "adapter_model.safetensors"))
+
+    peft_helper = PEFTHelper.from_local_dir(
+        qwen3_lora_files, max_position_embeddings=4096
+    )
+    lora_model = LoRAModel.from_lora_tensors(
+        1,
+        tensors,
+        peft_helper=peft_helper,
+        device=device,
+    )
+    for module_name, lora in lora_model.loras.items():
+        assert lora.module_name == module_name
+        assert lora.rank == 8
+        assert lora.lora_alpha == 32
+        assert lora.lora_a is not None
+        assert lora.lora_b is not None
+        assert lora.lora_a.device == torch.device(device)
+        assert lora.lora_b.device == torch.device(device)
+        assert lora.lora_a.shape[0] == lora.lora_b.shape[1], (
+            f"{lora.lora_a.shape=}, {lora.lora_b.shape=}"
+        )
+        assert lora.lora_a.shape[0] == 8
+
+
+def create_lora(
+    lora_id: int, model: nn.Module, sub_modules: list[str], device: torch.device
+) -> LoRAModel:
+    loras: dict[str, LoRALayerWeights] = {}
+    for name in sub_modules:
+        w = model.get_submodule(name).weight
+        loras[name] = LoRALayerWeights(
+            name,
+            8,
+            16,
+            torch.rand([8, w.shape[1]], device=device),
+            torch.rand([w.shape[0], 8], device=device),
+        )
+    return LoRAModel(lora_id, 8, loras)
+
+
+def create_packed_lora(
+    lora_id: int,
+    model: nn.Module,
+    module_name,
+    replaced_module_names,
+    device: torch.device,
+    empty_replaced_module_name=None,
+) -> LoRAModel:
+    w = model.get_submodule(module_name).weight
+    loras: dict[str, LoRALayerWeights] = {}
+    for replaced_module_name in replaced_module_names:
+        if replaced_module_name == empty_replaced_module_name:
+            continue
+        loras[replaced_module_name] = LoRALayerWeights(
+            replaced_module_name,
+            8,
+            16,
+            torch.rand([8, w.shape[1]], device=device),
+            torch.rand([w.shape[0] // len(replaced_module_names), 8], device=device),
+        )
+    return LoRAModel(lora_id, 8, loras)
+
+
+def test_replace_submodules(default_vllm_config, dist_init, dummy_model):
+    model = dummy_model
+    manager = LoRAModelManager(
+        model,
+        1,
+        1,
+        1,
+        LoRAConfig(
+            max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
+        ),
+        torch.device(DEVICES[0]),
+    )
+    model = manager.model
+    assert isinstance(model.get_submodule("dense1"), ColumnParallelLinearWithLoRA)
+    assert isinstance(
+        model.get_submodule("layer1.dense1"), ColumnParallelLinearWithLoRA
+    )
+    assert isinstance(model.get_submodule("dense2"), RowParallelLinearWithLoRA)
+    assert isinstance(model.get_submodule("layer1.dense2"), RowParallelLinearWithLoRA)
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_lora_model_manager(default_vllm_config, dist_init, dummy_model, device):
+    model = dummy_model
+    model_lora1 = create_lora(
+        1, model, ["layer1.dense1", "dense2", "lm_head"], device=device
+    )
+    model_lora2 = create_lora(2, model, ["dense1", "dense2", "lm_head"], device=device)
+    model_lora3 = create_lora(3, model, ["dense1", "dense2", "lm_head"], device=device)
+    manager = LoRAModelManager(
+        model,
+        2,
+        2,
+        2,
+        LoRAConfig(
+            max_lora_rank=8, max_cpu_loras=3, max_loras=2, lora_dtype=DEFAULT_DTYPE
+        ),
+        device=device,
+    )
+    assert all(x is None for x in manager.lora_index_to_id)
+    assert manager.add_adapter(model_lora1)
+    assert manager.activate_adapter(1)
+    assert manager.lora_index_to_id[0] == 1
+    assert not manager.add_adapter(model_lora1)
+    assert not manager.activate_adapter(1)
+    assert manager.add_adapter(model_lora2)
+    assert manager.activate_adapter(2)
+    assert manager.lora_index_to_id[0] == 1
+    assert manager.lora_index_to_id[1] == 2
+    assert not manager.add_adapter(model_lora2)
+    assert not manager.activate_adapter(2)
+    assert manager.add_adapter(model_lora3)
+    assert manager.lora_index_to_id[0] == 1
+    assert manager.lora_index_to_id[1] == 2
+    with pytest.raises(ValueError):
+        assert manager.activate_adapter(3)
+    assert manager.lora_index_to_id[0] == 1
+    assert manager.lora_index_to_id[1] == 2
+    assert manager.remove_adapter(model_lora2.id)
+    assert manager.lora_index_to_id[1] is None
+    assert not manager.remove_adapter(model_lora2.id)
+    assert manager.remove_adapter(model_lora1.id)
+    assert not manager.remove_adapter(model_lora1.id)
+    assert manager.add_adapter(model_lora1)
+    assert manager.lora_index_to_id[0] is None
+    assert manager.lora_index_to_id[1] is None
+    assert manager.add_adapter(model_lora2)
+    assert manager.activate_adapter(3)
+    assert manager.lora_index_to_id[0] == 3
+    assert manager.lora_index_to_id[1] is None
+    assert manager.activate_adapter(2)
+    assert manager.lora_index_to_id[0] == 3
+    assert manager.lora_index_to_id[1] == 2
+    assert manager.device == device
+    assert (
+        manager.punica_wrapper_mapping.get(DEFAULT_LANGUAGE_WRAPPER_KEY).device
+        == device
+    )
+    assert hasattr(manager, "supported_lora_modules")
+    assert sorted(manager.supported_lora_modules) == [
+        "dense1",
+        "dense2",
+        "lm_head",
+        "output",
+    ]
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_lora_lru_cache_model_manager(
+    default_vllm_config, dist_init, dummy_model, device
+):
+    model = dummy_model
+    model_lora1 = create_lora(
+        1, model, ["layer1.dense1", "dense2", "lm_head"], device=device
+    )
+    model_lora2 = create_lora(2, model, ["dense1", "dense2", "lm_head"], device=device)
+    model_lora3 = create_lora(3, model, ["dense1", "dense2", "lm_head"], device=device)
+    manager = LRUCacheLoRAModelManager(
+        model,
+        2,
+        2,
+        2,
+        LoRAConfig(
+            max_lora_rank=8, max_cpu_loras=3, max_loras=2, lora_dtype=DEFAULT_DTYPE
+        ),
+        device=device,
+    )
+    assert all(x is None for x in manager.lora_index_to_id)
+    assert manager.add_adapter(model_lora1)
+    assert manager.activate_adapter(1)
+    assert manager.lora_index_to_id[0] == 1
+    assert not manager.add_adapter(model_lora1)
+    assert not manager.activate_adapter(1)
+    assert manager.add_adapter(model_lora2)
+    assert manager.activate_adapter(2)
+    assert manager.lora_index_to_id[0] == 1
+    assert manager.lora_index_to_id[1] == 2
+    assert not manager.add_adapter(model_lora2)
+    assert not manager.activate_adapter(2)
+    assert manager.add_adapter(model_lora3)
+    assert manager.lora_index_to_id[0] == 1
+    assert manager.lora_index_to_id[1] == 2
+    assert manager.activate_adapter(3)
+    assert manager.lora_index_to_id[0] == 3
+    assert manager.lora_index_to_id[1] == 2
+    assert manager.remove_adapter(model_lora2.id)
+    assert manager.lora_index_to_id[1] is None
+    assert not manager.remove_adapter(model_lora2.id)
+    assert manager.remove_adapter(model_lora1.id)
+    assert not manager.remove_adapter(model_lora1.id)
+    assert manager.add_adapter(model_lora1)
+    assert manager.activate_adapter(1)
+    assert manager.lora_index_to_id[0] == 3
+    assert manager.lora_index_to_id[1] == 1
+    assert manager.add_adapter(model_lora2)
+    assert manager.deactivate_adapter(3)
+    assert manager.lora_index_to_id[0] is None
+    assert manager.lora_index_to_id[1] == 1
+    assert manager.activate_adapter(2)
+    assert manager.lora_index_to_id[0] == 2
+    assert manager.lora_index_to_id[1] == 1
+    assert manager.activate_adapter(3)
+    assert manager.lora_index_to_id[0] == 2
+    assert manager.lora_index_to_id[1] == 3
+    assert manager.pin_adapter(2)
+    assert manager.lora_index_to_id[0] == 2
+    assert manager.lora_index_to_id[1] == 3
+    assert manager.activate_adapter(1)
+    assert manager.lora_index_to_id[0] == 2
+    assert manager.lora_index_to_id[1] == 1
+    assert manager.deactivate_adapter(2)
+    assert manager.lora_index_to_id[0] is None
+    assert manager.lora_index_to_id[1] == 1
+    assert manager.activate_adapter(3)
+    assert manager.lora_index_to_id[0] == 3
+    assert manager.lora_index_to_id[1] == 1
+    assert manager.pin_adapter(3)
+    assert manager.pin_adapter(1)
+    with pytest.raises(RuntimeError):
+        assert manager.pin_adapter(2)
+    assert manager.lora_index_to_id[0] == 3
+    assert manager.lora_index_to_id[1] == 1
+    with pytest.raises(RuntimeError):
+        assert manager.activate_adapter(2)
+
+    assert manager.deactivate_adapter(3)
+    assert manager.pin_adapter(2)
+    assert manager.lora_index_to_id[0] == 2
+    assert manager.lora_index_to_id[1] == 1
+    assert manager.remove_adapter(3)
+    with pytest.raises(ValueError):
+        assert manager.pin_adapter(3)
+    assert (
+        manager.punica_wrapper_mapping.get(DEFAULT_LANGUAGE_WRAPPER_KEY).device
+        == device
+    )
+    assert manager.device == device
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_lru_lora_model_manager(default_vllm_config, dist_init, dummy_model, device):
+    # This tests just the LRU cache functionality, everything else is
+    # tested in test_lora_model_manager
+    model = dummy_model
+    model_lora1 = create_lora(
+        1, model, ["layer1.dense1", "dense2", "lm_head"], device=device
+    )
+    model_lora2 = create_lora(2, model, ["dense1", "dense2", "lm_head"], device=device)
+    model_lora3 = create_lora(3, model, ["dense1", "dense2", "lm_head"], device=device)
+    model_lora4 = create_lora(4, model, ["dense1", "dense2", "lm_head"], device=device)
+    manager = LRUCacheLoRAModelManager(
+        model,
+        2,
+        2,
+        2,
+        LoRAConfig(
+            max_lora_rank=8, max_cpu_loras=2, max_loras=2, lora_dtype=DEFAULT_DTYPE
+        ),
+        device=device,
+    )
+    assert all(x is None for x in manager.lora_index_to_id)
+
+    # Add up to capacity
+    assert manager.add_adapter(model_lora1)
+    assert manager.add_adapter(model_lora2)
+    assert manager.activate_adapter(1)
+    assert manager.activate_adapter(2)
+
+    assert set(manager.list_adapters()) == {1, 2}
+    assert manager.lora_index_to_id[0] == 1
+    assert manager.lora_index_to_id[1] == 2
+
+    # Add over capacity
+    assert manager.add_adapter(model_lora3)
+    assert manager.add_adapter(model_lora4)
+    assert manager.activate_adapter(3)
+    assert manager.activate_adapter(4)
+
+    assert set(manager.list_adapters()) == {3, 4}
+    assert manager.lora_index_to_id[0] == 3
+    assert manager.lora_index_to_id[1] == 4
+
+    # Add 3 again to move it to the top and then add 2
+    # should return false since it's in already
+    assert not manager.add_adapter(model_lora3)
+    assert not manager.activate_adapter(3)
+    assert manager.add_adapter(model_lora2)
+    assert manager.activate_adapter(2)
+
+    assert set(manager.list_adapters()) == {3, 2}
+    assert manager.lora_index_to_id[0] == 3
+    assert manager.lora_index_to_id[1] == 2
+
+    # Remove manually
+    assert manager.remove_adapter(3)
+    assert not manager.remove_adapter(3)
+
+    assert set(manager.list_adapters()) == {2}
+    assert manager.lora_index_to_id[0] is None
+    assert manager.lora_index_to_id[1] == 2
+
+    assert manager.add_adapter(model_lora3)
+    assert manager.activate_adapter(3)
+    assert manager.add_adapter(model_lora4)
+    assert manager.activate_adapter(4)
+
+    assert set(manager.list_adapters()) == {3, 4}
+    assert manager.lora_index_to_id[0] == 3
+    assert manager.lora_index_to_id[1] == 4
+
+    assert manager.remove_oldest_adapter()
+    assert set(manager.list_adapters()) == {4}
+    assert manager.lora_index_to_id[0] is None
+    assert manager.lora_index_to_id[1] == 4
+
+    assert manager.remove_oldest_adapter()
+    assert set(manager.list_adapters()) == set()
+    assert all(x is None for x in manager.lora_index_to_id)
+
+    assert not manager.remove_oldest_adapter()
+    assert set(manager.list_adapters()) == set()
+    assert all(x is None for x in manager.lora_index_to_id)
+
+    # pinning
+    assert manager.add_adapter(model_lora3)
+    assert manager.activate_adapter(3)
+    assert manager.add_adapter(model_lora4)
+    assert manager.activate_adapter(4)
+    assert set(manager.list_adapters()) == {3, 4}
+    with pytest.raises(ValueError):
+        assert manager.pin_adapter(1)
+    assert manager.pin_adapter(3)
+    # Remove manually
+    assert manager.remove_adapter(3)
+    assert not manager.remove_adapter(3)
+
+    assert set(manager.list_adapters()) == {4}
+    assert manager.lora_index_to_id[0] is None
+    assert manager.lora_index_to_id[1] == 4
+
+    assert manager.add_adapter(model_lora1)
+    assert manager.pin_adapter(1)
+    assert manager.add_adapter(model_lora2)
+    assert manager.activate_adapter(2)
+
+    assert set(manager.list_adapters()) == {1, 2}
+    assert manager.lora_index_to_id[0] == 1
+    assert manager.lora_index_to_id[1] == 2
+
+    assert manager.remove_oldest_adapter()
+    assert set(manager.list_adapters()) == {1}
+    assert manager.lora_index_to_id[0] == 1
+    assert manager.lora_index_to_id[1] is None
+
+    with pytest.raises(RuntimeError):
+        assert manager.remove_oldest_adapter()
+
+    assert set(manager.list_adapters()) == {1}
+    assert (
+        manager.punica_wrapper_mapping.get(DEFAULT_LANGUAGE_WRAPPER_KEY).device
+        == device
+    )
+    assert manager.device == device
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_lru_cache_worker_adapter_manager(
+    default_vllm_config, dist_init, dummy_model, device, tmp_path
+):
+    lora_config = LoRAConfig(
+        max_lora_rank=8, max_cpu_loras=4, max_loras=4, lora_dtype=DEFAULT_DTYPE
+    )
+
+    dummy_lora_files = f"{tmp_path}/lora_adapter"
+    os.makedirs(dummy_lora_files, exist_ok=True)
+    create_peft_lora(
+        dummy_model,
+        save_dir=dummy_lora_files,
+        target_modules=["layer1.dense1", "dense2"],
+        lora_dtype=DEFAULT_DTYPE,
+    )
+
+    model_config = ModelConfig(max_model_len=16)
+    vllm_config = VllmConfig(model_config=model_config, lora_config=lora_config)
+
+    vllm_config.scheduler_config.max_num_seqs = 4
+    vllm_config.scheduler_config.max_num_batched_tokens = 2
+    worker_adapter_manager = LRUCacheWorkerLoRAManager(
+        vllm_config, device, EMBEDDING_MODULES
+    )
+
+    worker_adapter_manager.max_num_seqs = 4
+    worker_adapter_manager.max_num_batched_tokens = 2
+
+    worker_adapter_manager.create_lora_manager(dummy_model)
+
+    mapping = LoRAMapping([], [])
+    worker_adapter_manager.set_active_adapters(
+        [LoRARequest("1", 1, dummy_lora_files), LoRARequest("2", 2, dummy_lora_files)],
+        mapping,
+    )
+    assert worker_adapter_manager.list_adapters() == {1, 2}
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2
+
+    worker_adapter_manager.set_active_adapters(
+        [
+            LoRARequest("1", 1, dummy_lora_files),
+            LoRARequest("3", 3, dummy_lora_files),
+            LoRARequest("4", 4, dummy_lora_files),
+        ],
+        mapping,
+    )
+    assert worker_adapter_manager.list_adapters() == {1, 2, 3, 4}
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 3
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[3] == 4
+
+    worker_adapter_manager.set_active_adapters(
+        [
+            LoRARequest("1", 1, dummy_lora_files),
+            LoRARequest("2", 2, dummy_lora_files),
+            LoRARequest("5", 5, dummy_lora_files),
+        ],
+        mapping,
+    )
+    assert worker_adapter_manager.list_adapters() == {1, 2, 4, 5}
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 5
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[3] == 4
+
+    worker_adapter_manager.set_active_adapters(
+        [
+            LoRARequest("1", 1, dummy_lora_files),
+            LoRARequest("1", 1, dummy_lora_files),
+            LoRARequest("1", 1, dummy_lora_files),
+        ],
+        mapping,
+    )
+    assert worker_adapter_manager.list_adapters() == {1, 2, 4, 5}
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 5
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[3] == 4
+
+    worker_adapter_manager.set_active_adapters(
+        [
+            LoRARequest("6", 6, dummy_lora_files),
+            LoRARequest("7", 7, dummy_lora_files),
+            LoRARequest("8", 8, dummy_lora_files),
+        ],
+        mapping,
+    )
+    assert worker_adapter_manager.list_adapters() == {1, 6, 7, 8}
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 7
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 8
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[3] == 6
+
+    # Over capacity
+    with pytest.raises(RuntimeError):
+        worker_adapter_manager.set_active_adapters(
+            [
+                LoRARequest("10", 10, dummy_lora_files),
+                LoRARequest("11", 11, dummy_lora_files),
+                LoRARequest("12", 12, dummy_lora_files),
+                LoRARequest("13", 13, dummy_lora_files),
+                LoRARequest("14", 14, dummy_lora_files),
+            ],
+            mapping,
+        )
+
+    assert worker_adapter_manager.device == device
+    punica_wrapper = worker_adapter_manager._adapter_manager.punica_wrapper_mapping.get(
+        DEFAULT_LANGUAGE_WRAPPER_KEY
+    )
+    assert punica_wrapper.device == device
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_worker_adapter_manager(
+    default_vllm_config, dist_init, dummy_model_gate_up, device, tmp_path
+):
+    # Should remove every LoRA not specified in the request.
+    lora_config = LoRAConfig(
+        max_lora_rank=8, max_cpu_loras=4, max_loras=4, lora_dtype=DEFAULT_DTYPE
+    )
+
+    model_config = ModelConfig(max_model_len=16)
+    vllm_config = VllmConfig(model_config=model_config, lora_config=lora_config)
+
+    vllm_config.scheduler_config.max_num_seqs = 4
+    vllm_config.scheduler_config.max_num_batched_tokens = 2
+
+    worker_adapter_manager = WorkerLoRAManager(vllm_config, device, EMBEDDING_MODULES)
+    worker_adapter_manager.vocab_size = dummy_model_gate_up.unpadded_vocab_size
+    worker_adapter_manager.create_lora_manager(dummy_model_gate_up)
+
+    dummy_lora_files = f"{tmp_path}/lora_adapter"
+    os.makedirs(dummy_lora_files, exist_ok=True)
+    create_peft_lora(
+        dummy_model_gate_up,
+        save_dir=dummy_lora_files,
+        target_modules=["layer1.dense1", "dense2"],
+        lora_dtype=DEFAULT_DTYPE,
+    )
+
+    mapping = LoRAMapping([], [])
+    worker_adapter_manager.set_active_adapters(
+        [LoRARequest("1", 1, dummy_lora_files), LoRARequest("2", 2, dummy_lora_files)],
+        mapping,
+    )
+    assert worker_adapter_manager.list_adapters() == {1, 2}
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2
+
+    worker_adapter_manager.set_active_adapters(
+        [
+            LoRARequest("1", 1, dummy_lora_files),
+            LoRARequest("3", 3, dummy_lora_files),
+            LoRARequest("4", 4, dummy_lora_files),
+        ],
+        mapping,
+    )
+    assert worker_adapter_manager.list_adapters() == {1, 3, 4}
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 3
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 4
+
+    worker_adapter_manager.set_active_adapters(
+        [
+            LoRARequest("1", 1, dummy_lora_files),
+            LoRARequest("2", 2, dummy_lora_files),
+            LoRARequest("5", 5, dummy_lora_files),
+        ],
+        mapping,
+    )
+    assert worker_adapter_manager.list_adapters() == {1, 2, 5}
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 5
+
+    worker_adapter_manager.set_active_adapters(
+        [
+            LoRARequest("1", 1, dummy_lora_files),
+            LoRARequest("1", 1, dummy_lora_files),
+            LoRARequest("1", 1, dummy_lora_files),
+        ],
+        mapping,
+    )
+    assert worker_adapter_manager.list_adapters() == {1}
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] is None
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] is None
+
+    worker_adapter_manager.set_active_adapters(
+        [
+            LoRARequest("6", 6, dummy_lora_files),
+            LoRARequest("7", 7, dummy_lora_files),
+            LoRARequest("8", 8, dummy_lora_files),
+        ],
+        mapping,
+    )
+    assert worker_adapter_manager.list_adapters() == {6, 7, 8}
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 8
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 6
+    assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 7
+
+    # Over capacity
+    with pytest.raises(RuntimeError):
+        worker_adapter_manager.set_active_adapters(
+            [
+                LoRARequest("10", 10, dummy_lora_files),
+                LoRARequest("11", 11, dummy_lora_files),
+                LoRARequest("12", 12, dummy_lora_files),
+                LoRARequest("13", 13, dummy_lora_files),
+                LoRARequest("14", 14, dummy_lora_files),
+            ],
+            mapping,
+        )
+
+    assert worker_adapter_manager.device == device
+    punica_wrapper = worker_adapter_manager._adapter_manager.punica_wrapper_mapping.get(
+        DEFAULT_LANGUAGE_WRAPPER_KEY
+    )
+    assert punica_wrapper.device == device
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_packed_loras(default_vllm_config, dist_init, dummy_model_gate_up, device):
+    model = dummy_model_gate_up
+    model_lora = create_packed_lora(
+        1,
+        model,
+        module_name="gate_up_proj",
+        replaced_module_names=["gate_proj", "up_proj"],
+        device=device,
+    )
+    model_lora1 = create_packed_lora(
+        2,
+        model,
+        module_name="gate_up_proj",
+        replaced_module_names=["gate_proj", "up_proj"],
+        device=device,
+        empty_replaced_module_name="gate_proj",
+    )
+
+    manager = LoRAModelManager(
+        model,
+        2,
+        2,
+        2,
+        LoRAConfig(
+            max_lora_rank=8, max_cpu_loras=2, max_loras=2, lora_dtype=DEFAULT_DTYPE
+        ),
+        device=device,
+    )
+    model = manager.model
+
+    assert isinstance(
+        model.get_submodule("gate_up_proj"), MergedColumnParallelLinearWithLoRA
+    )
+    # Verify packed lora is correct
+    model_lora_clone = model_lora.clone(1)
+    model_lora_clone1 = model_lora1.clone(1)
+    assert manager.add_adapter(model_lora)
+    assert manager.add_adapter(model_lora1)
+
+    assert model_lora.get_lora("gate_proj") is None
+    assert model_lora.get_lora("up_proj") is None
+    assert model_lora1.get_lora("up_proj") is None
+    packed_lora = model_lora.get_lora("gate_up_proj")
+    assert packed_lora and isinstance(packed_lora, PackedLoRALayerWeights)
+
+    torch.testing.assert_close(
+        packed_lora.lora_a[0], model_lora_clone.get_lora("gate_proj").lora_a
+    )
+    torch.testing.assert_close(
+        packed_lora.lora_b[0], model_lora_clone.get_lora("gate_proj").lora_b
+    )
+    torch.testing.assert_close(
+        packed_lora.lora_a[1], model_lora_clone.get_lora("up_proj").lora_a
+    )
+    torch.testing.assert_close(
+        packed_lora.lora_b[1], model_lora_clone.get_lora("up_proj").lora_b
+    )
+
+    packed_lora1 = model_lora1.get_lora("gate_up_proj")
+    assert packed_lora1 and isinstance(packed_lora1, PackedLoRALayerWeights)
+
+    assert packed_lora1.lora_a[0] is None
+    assert packed_lora1.lora_b[0] is None
+    torch.testing.assert_close(
+        packed_lora1.lora_a[1], model_lora_clone1.get_lora("up_proj").lora_a
+    )
+    torch.testing.assert_close(
+        packed_lora1.lora_b[1], model_lora_clone1.get_lora("up_proj").lora_b
+    )

third_party/vllm/tests/lora/test_minicpmv_tp.py

@@ -0,0 +1,121 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+import vllm
+from vllm.assets.image import ImageAsset
+from vllm.lora.request import LoRARequest
+from vllm.platforms import current_platform
+
+from ..utils import multi_gpu_test
+
+MODEL_PATH = "openbmb/MiniCPM-Llama3-V-2_5"
+
+PROMPT_TEMPLATE = (
+    "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n"
+    "(<image>./</image>)\nWhat is in the image?<|eot_id|>"
+    "<|start_header_id|>assistant<|end_header_id|>\n\n"
+)
+
+IMAGE_ASSETS = [
+    ImageAsset("stop_sign"),
+]
+
+# After fine-tuning with LoRA, all generated content should start begin `A`.
+EXPECTED_OUTPUT = [
+    "A red and white stop sign with a Chinese archway in the background featuring red lanterns and gold accents.",  # noqa: E501
+]
+
+
+def do_sample(llm: vllm.LLM, lora_path: str, lora_id: int) -> list[str]:
+    sampling_params = vllm.SamplingParams(
+        temperature=0,
+        max_tokens=5,
+        stop_token_ids=[128001, 128009],  # eos_id, eot_id
+    )
+
+    inputs = [
+        {
+            "prompt": PROMPT_TEMPLATE,
+            "multi_modal_data": {"image": asset.pil_image},
+        }
+        for asset in IMAGE_ASSETS
+    ]
+
+    outputs = llm.generate(
+        inputs,
+        sampling_params,
+        lora_request=LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None,
+    )
+    # Print the outputs.
+    generated_texts: list[str] = []
+    for output in outputs:
+        generated_text = output.outputs[0].text.strip()
+        generated_texts.append(generated_text)
+        print(f"Generated text: {generated_text!r}")
+    return generated_texts
+
+
+def test_minicpmv_lora(minicpmv_lora_files):
+    llm = vllm.LLM(
+        MODEL_PATH,
+        max_num_seqs=2,
+        enable_lora=True,
+        max_loras=2,
+        max_lora_rank=8,
+        enforce_eager=True,
+        max_model_len=2048,
+        limit_mm_per_prompt={"image": 2, "video": 0},
+        trust_remote_code=True,
+    )
+    output1 = do_sample(llm, minicpmv_lora_files, lora_id=1)
+    for i in range(len(EXPECTED_OUTPUT)):
+        assert EXPECTED_OUTPUT[i].startswith(output1[i])
+    output2 = do_sample(llm, minicpmv_lora_files, lora_id=2)
+    for i in range(len(EXPECTED_OUTPUT)):
+        assert EXPECTED_OUTPUT[i].startswith(output2[i])
+
+
+@pytest.mark.skipif(
+    current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
+)
+@multi_gpu_test(num_gpus=4)
+def test_minicpmv_tp4_wo_fully_sharded_loras(minicpmv_lora_files):
+    llm = vllm.LLM(
+        MODEL_PATH,
+        enable_lora=True,
+        max_num_seqs=2,
+        max_loras=4,
+        max_lora_rank=64,
+        tensor_parallel_size=4,
+        limit_mm_per_prompt={"image": 2, "video": 0},
+        trust_remote_code=True,
+    )
+    output_tp = do_sample(llm, minicpmv_lora_files, lora_id=1)
+    for i in range(len(EXPECTED_OUTPUT)):
+        assert EXPECTED_OUTPUT[i].startswith(output_tp[i])
+
+
+@pytest.mark.skipif(
+    current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
+)
+@multi_gpu_test(num_gpus=4)
+def test_minicpmv_tp4_fully_sharded_loras(minicpmv_lora_files):
+    llm = vllm.LLM(
+        MODEL_PATH,
+        enable_lora=True,
+        max_num_seqs=2,
+        max_loras=2,
+        max_lora_rank=8,
+        tensor_parallel_size=4,
+        trust_remote_code=True,
+        limit_mm_per_prompt={"image": 1, "video": 0},
+        fully_sharded_loras=True,
+    )
+    output_tp = do_sample(llm, minicpmv_lora_files, lora_id=1)
+    for i in range(len(EXPECTED_OUTPUT)):
+        assert EXPECTED_OUTPUT[i].startswith(output_tp[i])
+    output_tp = do_sample(llm, minicpmv_lora_files, lora_id=2)
+    for i in range(len(EXPECTED_OUTPUT)):
+        assert EXPECTED_OUTPUT[i].startswith(output_tp[i])

third_party/vllm/tests/lora/test_mixtral.py

@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+import vllm
+from vllm.lora.request import LoRARequest
+from vllm.platforms import current_platform
+
+MODEL_PATH = "mistralai/Mixtral-8x7B-Instruct-v0.1"
+
+
+def do_sample(
+    llm: vllm.LLM, lora_path: str, lora_id: int, prompts: list[str]
+) -> list[str]:
+    sampling_params = vllm.SamplingParams(temperature=0, max_tokens=256)
+    outputs = llm.generate(
+        prompts,
+        sampling_params,
+        lora_request=LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None,
+    )
+    # Print the outputs.
+    generated_texts: list[str] = []
+    for output in outputs:
+        prompt = output.prompt
+        generated_text = output.outputs[0].text.strip()
+        generated_texts.append(generated_text)
+        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
+    return generated_texts
+
+
+@pytest.mark.parametrize("tp_size", [4])
+def test_mixtral_lora(mixtral_lora_files, tp_size):
+    """Original test, the LoRA model has the common target modules, not all"""
+    if (
+        torch.accelerator.device_count() < tp_size
+        and tp_size > 1
+        and current_platform.is_cuda_alike()
+    ):
+        pytest.skip(f"Not enough GPUs for tensor parallelism {tp_size}")
+
+    prompts = [
+        "[system] Given a target sentence construct the underlying meaning representation\nof the input sentence as a single function with attributes and attribute\nvalues. This function should describe the target string accurately and the\nfunction must be one of the following ['inform', 'request', 'give_opinion',\n'confirm', 'verify_attribute', 'suggest', 'request_explanation',\n'recommend', 'request_attribute'].\n\nThe attributes must be one of the following:\n['name', 'exp_release_date', 'release_year', 'developer', 'esrb', 'rating',\n'genres', 'player_perspective', 'has_multiplayer', 'platforms',\n'available_on_steam', 'has_linux_release', 'has_mac_release', 'specifier'] [/system] [user] Here is the target sentence:\nSpellForce 3 is a pretty bad game. The developer Grimlore Games is clearly a bunch of no-talent hacks, and 2017 was a terrible year for games anyway. [/user] [assistant]",  # noqa: E501
+        "[system] Given a target sentence construct the underlying meaning representation\nof the input sentence as a single function with attributes and attribute\nvalues. This function should describe the target string accurately and the\nfunction must be one of the following ['inform', 'request', 'give_opinion',\n'confirm', 'verify_attribute', 'suggest', 'request_explanation',\n'recommend', 'request_attribute'].\n\nThe attributes must be one of the following:\n['name', 'exp_release_date', 'release_year', 'developer', 'esrb', 'rating',\n'genres', 'player_perspective', 'has_multiplayer', 'platforms',\n'available_on_steam', 'has_linux_release', 'has_mac_release', 'specifier'] [/system] [user] Here is the target sentence:\nI wanted to like Grimlore Games' 2017 entry, but in SpellForce 3 they just didn't get anything right. [/user] [assistant]",  # noqa: E501
+        "[system] Given a target sentence construct the underlying meaning representation\nof the input sentence as a single function with attributes and attribute\nvalues. This function should describe the target string accurately and the\nfunction must be one of the following ['inform', 'request', 'give_opinion',\n'confirm', 'verify_attribute', 'suggest', 'request_explanation',\n'recommend', 'request_attribute'].\n\nThe attributes must be one of the following:\n['name', 'exp_release_date', 'release_year', 'developer', 'esrb', 'rating',\n'genres', 'player_perspective', 'has_multiplayer', 'platforms',\n'available_on_steam', 'has_linux_release', 'has_mac_release', 'specifier'] [/system] [user] Here is the target sentence:\nBioShock is a good role-playing, action-adventure, shooter that released for PlayStation, Xbox, and PC in 2007. It is available on Steam, and it has a Mac release but not a Linux release. [/user] [assistant]",  # noqa: E501
+    ]
+
+    llm = vllm.LLM(
+        MODEL_PATH,
+        enable_lora=True,
+        max_num_seqs=16,
+        max_loras=4,
+        distributed_executor_backend="ray",
+        tensor_parallel_size=tp_size,
+    )
+
+    expected_lora_output = [
+        [
+            "give_opinion(name[SpellForce 3], release_year[2017], developer[Grimlore Games], rating[poor])"  # noqa: E501
+        ],
+        [
+            "give_opinion(name[SpellForce 3], developer[Grimlore Games], release_year[2017], rating[poor])",  # noqa: E501
+            "give_opinion(name[SpellForce 3], release_year[2017], developer[Grimlore Games], rating[poor])",  # noqa: E501
+        ],
+        [
+            "inform(name[BioShock], release_year[2007], rating[good], genres[action-adventure, role-playing, shooter], platforms[PlayStation, Xbox, PC], available_on_steam[yes], has_linux_release[no], has_mac_release[yes])"  # noqa: E501
+        ],
+    ]
+
+    def check_outputs(generated: list[str]):
+        assert len(generated) == len(expected_lora_output)
+        for gen, gt_choices in zip(generated, expected_lora_output):
+            assert gen in gt_choices
+
+    check_outputs(do_sample(llm, mixtral_lora_files, lora_id=1, prompts=prompts))
+    check_outputs(do_sample(llm, mixtral_lora_files, lora_id=2, prompts=prompts))

third_party/vllm/tests/lora/test_moe_lora_align_sum.py

@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import random
+
+import pytest
+import torch
+
+from vllm import _custom_ops as ops
+
+
+def round_up(x, base):
+    return ((x + base - 1) // base) * base
+
+
+def CEILDIV(x, y):
+    return (x + y - 1) // y
+
+
+def sample_data(num_experts, max_loras, num_tokens, topk_num):
+    topk_ids = torch.zeros((num_tokens, topk_num), dtype=torch.int32)
+    token_lora_mapping = torch.zeros((num_tokens,), dtype=torch.int32)
+
+    for i in range(num_tokens):
+        pool = list(range(num_experts))
+        random.shuffle(pool)
+        for j in range(topk_num):
+            topk_ids[i, j] = pool[j]
+        token_lora_mapping[i] = random.randint(0, max_loras - 1)
+
+    return topk_ids.to("cuda"), token_lora_mapping.to("cuda")
+
+
+@pytest.mark.parametrize("num_tokens", [100, 200, 1024, 4096])  # 81920
+@pytest.mark.parametrize("topk_num", [6])
+@pytest.mark.parametrize("num_experts", [64, 128, 256, 512])
+@pytest.mark.parametrize("max_loras", [2, 32])
+@pytest.mark.parametrize("block_size", [16])
+def test_moe_lora_align_block_size(
+    num_tokens, topk_num, num_experts, max_loras, block_size
+):
+    # sample data
+    random.seed(1)
+    topk_ids, token_lora_mapping = sample_data(
+        num_experts, max_loras, num_tokens, topk_num
+    )
+
+    # compute paddings
+    max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
+    max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
+    if topk_ids.numel() < num_experts:
+        max_num_tokens_padded = topk_ids.numel() * block_size
+    max_num_m_blocks = CEILDIV(max_num_tokens_padded, block_size)
+
+    # init output tensors
+    sorted_token_ids = torch.full(
+        (max_loras * max_num_tokens_padded,),
+        topk_ids.numel(),
+        dtype=torch.int32,
+        device="cuda",
+    )
+    expert_ids = torch.full(
+        (max_loras * max_num_m_blocks,), num_experts, dtype=torch.int32, device="cuda"
+    )
+    num_tokens_post_pad = torch.zeros((max_loras,), dtype=torch.int32, device="cuda")
+    adapter_enabled = torch.ones((max_loras + 1,), dtype=torch.int32, device="cuda")
+    lora_ids = torch.arange(max_loras + 2, dtype=torch.int32, device="cuda")
+
+    # call kernel
+    ops.moe_lora_align_block_size(
+        topk_ids,
+        token_lora_mapping,
+        num_experts,
+        block_size,
+        max_loras,
+        max_num_tokens_padded,
+        max_num_m_blocks,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_pad,
+        adapter_enabled,
+        lora_ids,
+    )
+
+    # verify values
+    expert_ids = expert_ids.view(max_loras, -1)
+    sorted_token_ids = sorted_token_ids.view(max_loras, -1, block_size)
+
+    for lora_idx in range(max_loras):
+        for token_idx in range(sorted_token_ids.size(1)):
+            block = sorted_token_ids[lora_idx][token_idx]
+            indices = block[block != topk_ids.numel()]
+            if indices.numel() > 0:
+                expert_id = expert_ids[lora_idx][token_idx]
+                assert torch.all(topk_ids.view(-1)[indices] == expert_id)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])

third_party/vllm/tests/lora/test_olmoe_tp.py

@@ -0,0 +1,217 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import shutil
+from collections.abc import Sequence
+
+import pytest
+import torch
+from safetensors.torch import load_file, save_file
+
+import vllm
+from vllm.lora.request import LoRARequest
+
+from ..utils import multi_gpu_test
+
+MODEL_PATH = "allenai/OLMoE-1B-7B-0125-Instruct"
+
+PROMPT_TEMPLATE = """I want you to act as a SQL terminal in front of an example database, you need only to return the sql command to me. Do not return any additional explanation. Below is an instruction that describes a task, Write a response that appropriately completes the request.
+"
+##Instruction:
+candidate_poll contains tables such as candidate, people. Table candidate has columns such as Candidate_ID, People_ID, Poll_Source, Date, Support_rate, Consider_rate, Oppose_rate, Unsure_rate. Candidate_ID is the primary key.
+Table people has columns such as People_ID, Sex, Name, Date_of_Birth, Height, Weight. People_ID is the primary key.
+The People_ID of candidate is the foreign key of People_ID of people.
+
+
+###Input:
+{context}
+
+###Response:"""  # noqa: E501
+
+EXPECTED_LORA_OUTPUT = [
+    "SELECT count(*) FROM candidate",
+    "SELECT count(*) FROM candidate",
+    "SELECT poll_source FROM candidate GROUP BY poll_source ORDER BY count(*) DESC LIMIT 1",  # noqa: E501
+    "SELECT poll_source FROM candidate GROUP BY poll_source ORDER BY count(*) DESC LIMIT 1",  # noqa: E501
+]
+
+EXPECTED_BASE_MODEL_OUTPUT = [
+    "SELECT COUNT(Candidate_ID) FROM candidate",
+    "SELECT COUNT(Candidate_ID) FROM candidate",
+    "SELECT Candidate_ID, COUNT(*) as Total_Candidates\nFROM candidate\nINNER JOIN people ON candidate.People_ID = people.People_ID",  # noqa: E501
+    # There are multiple acceptable responses
+    (
+        "SELECT Candidate_ID, Poll_Source FROM candidate WHERE People_ID IN (SELECT People_ID FROM people) ORDER BY COUNT(*) DESC LIMIT 1",  # noqa: E501
+        "SELECT Candidate_ID, Poll_Source FROM candidate WHERE COUNT(People_ID) = (SELECT COUNT(People_ID) FROM people) ORDER BY Candidate_ID DESC LIMIT 1",  # noqa: E501
+    ),
+]
+
+
+def _output_matches(generated: str, accepted: str | Sequence[str]) -> bool:
+    if isinstance(accepted, str):
+        accepted = (accepted,)
+    return any(generated.startswith(s) for s in accepted)
+
+
+def generate_and_test(
+    llm: vllm.LLM,
+    lora_path: str,
+    lora_id: list[int | None] | int | None,
+    compare_lower: bool = False,
+) -> None:
+    prompts = [
+        PROMPT_TEMPLATE.format(context="How many candidates are there?"),
+        PROMPT_TEMPLATE.format(context="Count the number of candidates."),
+        PROMPT_TEMPLATE.format(
+            context="Which poll resource provided the most number of candidate information?"  # noqa: E501
+        ),
+        PROMPT_TEMPLATE.format(
+            context="Return the poll resource associated with the most candidates."
+        ),
+    ]
+
+    lora_request = None
+    if isinstance(lora_id, int):
+        lora_request = LoRARequest(str(lora_id), lora_id, lora_path)
+    elif isinstance(lora_id, list):
+        lora_request = [
+            LoRARequest(str(i), i, lora_path) if i is not None else None
+            for i in lora_id
+        ]
+
+    sampling_params = vllm.SamplingParams(temperature=0, max_tokens=64)
+    outputs = llm.generate(prompts, sampling_params, lora_request=lora_request)
+    # Print the outputs.
+    generated_texts: list[str] = []
+    for output in outputs:
+        prompt = output.prompt
+        generated_text = output.outputs[0].text.strip()
+        generated_texts.append(generated_text)
+        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
+
+    for i in range(len(EXPECTED_LORA_OUTPUT)):
+        req_lora_id = lora_id[i] if isinstance(lora_id, list) else lora_id
+        generated_text = generated_texts[i]
+        expected_output = (
+            EXPECTED_LORA_OUTPUT[i]
+            if req_lora_id is not None
+            else EXPECTED_BASE_MODEL_OUTPUT[i]
+        )
+
+        if compare_lower:
+            generated_text = generated_text.lower()
+            if isinstance(expected_output, str):
+                expected_output = (expected_output.lower(),)
+            else:
+                expected_output = tuple(s.lower() for s in expected_output)
+        assert _output_matches(generated_text, expected_output), (
+            f"Output {i}: {generated_text!r} does not match any of {expected_output!r}"
+        )
+
+
+def test_olmoe_lora(olmoe_lora_files):
+    # We enable enforce_eager=True here to reduce VRAM usage for lora-test CI,
+    # Otherwise, the lora-test will fail due to CUDA OOM.
+    llm = vllm.LLM(
+        MODEL_PATH,
+        max_model_len=1024,
+        enable_lora=True,
+        max_loras=4,
+        enforce_eager=True,
+        trust_remote_code=True,
+        enable_chunked_prefill=True,
+    )
+
+    generate_and_test(llm, olmoe_lora_files, lora_id=1)
+    generate_and_test(llm, olmoe_lora_files, lora_id=2)
+
+
+def test_olmoe_lora_mixed(olmoe_lora_files):
+    llm = vllm.LLM(
+        MODEL_PATH,
+        max_model_len=1024,
+        enable_lora=True,
+        max_loras=4,
+        enforce_eager=True,
+        trust_remote_code=True,
+        enable_chunked_prefill=True,
+    )
+
+    generate_and_test(llm, olmoe_lora_files, lora_id=[1, None, 3, None])
+
+
+def test_olmoe_lora_mixed_random(olmoe_lora_files, tmp_path):
+    # Create a dummy LoRA with random weights based on the real one
+    random_lora_path = tmp_path / "random_lora"
+    shutil.copytree(olmoe_lora_files, random_lora_path)
+
+    weights_path = random_lora_path / "adapter_model.safetensors"
+    weights = load_file(str(weights_path))
+    random_weights = {k: torch.randn_like(v) for k, v in weights.items()}
+    save_file(random_weights, str(weights_path))
+
+    llm = vllm.LLM(
+        MODEL_PATH,
+        max_model_len=1024,
+        enable_lora=True,
+        max_loras=4,
+        enforce_eager=True,
+        trust_remote_code=True,
+        enable_chunked_prefill=True,
+    )
+
+    prompts = [
+        PROMPT_TEMPLATE.format(context="How many candidates are there?"),
+        PROMPT_TEMPLATE.format(context="Count the number of candidates."),
+    ]
+
+    lora_requests = [
+        LoRARequest("real", 1, olmoe_lora_files),
+        LoRARequest("random", 2, str(random_lora_path)),
+    ]
+
+    sampling_params = vllm.SamplingParams(temperature=0, max_tokens=64)
+    outputs = llm.generate(prompts, sampling_params, lora_request=lora_requests)
+    assert outputs[0].outputs[0].text.strip().startswith(EXPECTED_LORA_OUTPUT[0])
+
+
+@pytest.mark.parametrize("fully_sharded_loras", [False, True])
+@multi_gpu_test(num_gpus=2)
+def test_olmoe_lora_tp2(olmoe_lora_files, fully_sharded_loras):
+    llm = vllm.LLM(
+        MODEL_PATH,
+        max_model_len=1024,
+        enable_lora=True,
+        max_loras=4,
+        enforce_eager=True,
+        trust_remote_code=True,
+        enable_chunked_prefill=True,
+        tensor_parallel_size=2,
+        fully_sharded_loras=fully_sharded_loras,
+    )
+
+    generate_and_test(llm, olmoe_lora_files, lora_id=1)
+    generate_and_test(llm, olmoe_lora_files, lora_id=2)
+
+
+@pytest.mark.parametrize("fully_sharded_loras", [False, True])
+@multi_gpu_test(num_gpus=4)
+def test_olmoe_lora_tp4(olmoe_lora_files, fully_sharded_loras):
+    llm = vllm.LLM(
+        MODEL_PATH,
+        max_model_len=1024,
+        enable_lora=True,
+        max_loras=4,
+        enforce_eager=True,
+        trust_remote_code=True,
+        enable_chunked_prefill=True,
+        tensor_parallel_size=4,
+        fully_sharded_loras=fully_sharded_loras,
+    )
+    generate_and_test(
+        llm, olmoe_lora_files, lora_id=1, compare_lower=fully_sharded_loras
+    )
+    generate_and_test(
+        llm, olmoe_lora_files, lora_id=2, compare_lower=fully_sharded_loras
+    )

third_party/vllm/tests/lora/test_peft_helper.py

@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+import math
+import shutil
+
+import pytest
+
+from vllm.config.lora import LoRAConfig
+from vllm.lora.peft_helper import PEFTHelper
+
+ERROR_CASES = [
+    (
+        "test_rank",
+        {"r": 1024},
+        "is greater than max_lora_rank",
+    ),
+    ("test_dora", {"use_dora": True}, "does not yet support DoRA"),
+    (
+        "test_modules_to_save",
+        {"modules_to_save": ["lm_head"]},
+        "only supports modules_to_save being None",
+    ),
+]
+
+
+def test_peft_helper_pass(llama32_lora_files, tmp_path):
+    peft_helper = PEFTHelper.from_local_dir(
+        llama32_lora_files, max_position_embeddings=4096
+    )
+    lora_config = LoRAConfig(max_lora_rank=16, max_cpu_loras=3, max_loras=2)
+    peft_helper.validate_legal(lora_config)
+    assert peft_helper.r == 8
+    assert peft_helper.lora_alpha == 32
+    target_modules = sorted(peft_helper.target_modules)
+
+    assert target_modules == [
+        "down_proj",
+        "embed_tokens",
+        "gate_proj",
+        "k_proj",
+        "lm_head",
+        "o_proj",
+        "q_proj",
+        "up_proj",
+        "v_proj",
+    ]
+    assert peft_helper.vllm_max_position_embeddings == 4096
+
+    # test RSLoRA
+    rslora_config = dict(use_rslora=True)
+    test_dir = tmp_path / "test_rslora"
+    shutil.copytree(llama32_lora_files, test_dir)
+
+    # Load and modify configuration
+    config_path = test_dir / "adapter_config.json"
+    with open(config_path) as f:
+        adapter_config = json.load(f)
+    # Apply configuration changes
+    adapter_config.update(rslora_config)
+
+    # Save modified configuration
+    with open(config_path, "w") as f:
+        json.dump(adapter_config, f)
+
+    peft_helper = PEFTHelper.from_local_dir(test_dir, max_position_embeddings=4096)
+    peft_helper.validate_legal(lora_config)
+    scaling = peft_helper.lora_alpha / math.sqrt(peft_helper.r)
+    assert abs(peft_helper.vllm_lora_scaling_factor - scaling) < 1e-3
+
+
+@pytest.mark.parametrize("test_name,config_change,expected_error", ERROR_CASES)
+def test_peft_helper_error(
+    llama32_lora_files,
+    tmp_path,
+    test_name: str,
+    config_change: dict,
+    expected_error: str,
+):
+    test_dir = tmp_path / test_name
+    shutil.copytree(llama32_lora_files, test_dir)
+
+    # Load and modify configuration
+    config_path = test_dir / "adapter_config.json"
+    with open(config_path) as f:
+        adapter_config = json.load(f)
+    # Apply configuration changes
+    adapter_config.update(config_change)
+
+    # Save modified configuration
+    with open(config_path, "w") as f:
+        json.dump(adapter_config, f)
+    lora_config = LoRAConfig(max_lora_rank=16, max_cpu_loras=3, max_loras=2)
+    # Test loading the adapter
+    with pytest.raises(ValueError, match=expected_error):
+        PEFTHelper.from_local_dir(
+            test_dir, max_position_embeddings=4096
+        ).validate_legal(lora_config)

third_party/vllm/tests/lora/test_punica_ops.py

@@ -0,0 +1,477 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from threading import Lock
+
+import pytest
+import torch
+
+import vllm.lora.ops.torch_ops as torch_ops
+import vllm.lora.ops.triton_ops as triton_ops
+from vllm.lora.ops.triton_ops import LoRAKernelMeta
+from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT, _LORA_B_PTR_DICT
+from vllm.utils.torch_utils import set_random_seed
+
+from .utils import PunicaTensors, assert_close, generate_data_for_nslices
+
+
+@pytest.fixture(autouse=True)
+def reset_device(reset_default_device):
+    pass
+
+
+# Utility shrink and expand operations used as reference implementations.
+def sgmv_shrink_for_nslices(
+    nslices: int,
+    inputs_tensor: torch.Tensor,
+    lora_weights_lst: list[torch.Tensor],
+    out_tensor: torch.Tensor,
+    b_seq_start_loc: torch.Tensor,
+    seq_len_tensor: torch.Tensor,
+    prompt_lora_mapping: torch.Tensor,
+    batches: int,
+    max_seq_length: int,
+    num_tokens: int,
+    scaling: float,
+):
+    """
+    Wrapper around torch_ops.sgmv_shrink that handles any nslices.
+    """
+    for index in range(nslices):
+        torch_ops.sgmv_shrink(
+            inputs_tensor,
+            lora_weights_lst[index],
+            out_tensor[index],
+            b_seq_start_loc,
+            seq_len_tensor,
+            prompt_lora_mapping,
+            batches,
+            max_seq_length,
+            num_tokens,
+            scaling,
+        )
+
+
+def sgmv_expand_for_nslices(
+    nslices: int,
+    hidden_size: int,
+    inputs_tensor: torch.Tensor,
+    lora_weights_lst: list[torch.Tensor],
+    out_tensor: torch.Tensor,
+    b_seq_start_loc: torch.Tensor,
+    seq_len_tensor: torch.Tensor,
+    prompt_lora_mapping: torch.Tensor,
+    batches: int,
+    max_seq_length: int,
+    num_tokens: int,
+    add_inputs: bool,
+) -> None:
+    """
+    Wrapper around torch_ops.sgmv_expand that handles any nslices.
+    """
+    if nslices == 1:
+        # Verify the torch's sgmv_expand op
+        torch_ops.sgmv_expand(
+            inputs_tensor[0],
+            lora_weights_lst[0],
+            out_tensor,
+            b_seq_start_loc,
+            seq_len_tensor,
+            prompt_lora_mapping,
+            batches,
+            max_seq_length,
+            num_tokens,
+            add_inputs=add_inputs,
+        )
+    else:
+        slice_offset = 0
+        for index in range(nslices):
+            lora_weights = lora_weights_lst[index]
+            torch_ops.sgmv_expand_slice(
+                inputs_tensor[index],
+                lora_weights,
+                out_tensor,
+                b_seq_start_loc,
+                seq_len_tensor,
+                prompt_lora_mapping,
+                batches,
+                max_seq_length,
+                num_tokens,
+                slice_offset,
+                hidden_size,
+                add_inputs=add_inputs,
+            )
+            slice_offset += hidden_size
+
+
+_dict_lock = Lock()
+
+
+def check_lora_shrink_kernel(
+    batches: int,
+    num_loras: int,
+    rank: int,
+    hidden_size: int,
+    nslices: int,
+    dtype: torch.dtype,
+    device: str,
+    seq_length: int,
+    scaling: float,
+):
+    """
+    Compare outputs of torch_ops.sgmv_shrink and triton_ops.lora_shrink
+    kernels.
+    """
+    data: PunicaTensors = generate_data_for_nslices(
+        batches,
+        hidden_size,
+        num_loras,
+        rank,
+        seq_length,
+        nslices,
+        dtype,
+        "shrink",
+        device,
+    )
+    max_seq_length, token_nums = data.meta()
+
+    # Setup metadata information for SGMV and reference kernels
+    sgmv_meta_args = (
+        data.b_seq_start_loc,
+        data.seq_len_tensor,
+        data.prompt_lora_mapping,
+        batches,
+        max_seq_length,
+        token_nums,
+    )
+
+    # Setup metadata information for the LoRA kernel.
+    lora_meta = LoRAKernelMeta.make(
+        max_loras=num_loras, max_num_tokens=token_nums, device="cuda"
+    )
+    lora_meta.prepare_tensors(data.token_lora_mapping)
+
+    ref_out_tensor = data.ref_out_tensor
+    out_tensor = data.our_out_tensor.clone()
+
+    # Preventing cache error pointer.
+    with _dict_lock:
+        # lora_shrink kernel
+        _LORA_A_PTR_DICT.clear()
+        triton_ops.lora_shrink(
+            data.inputs_tensor,
+            data.lora_weights,
+            out_tensor,
+            *lora_meta.meta_args(token_nums=token_nums, specialize_active_lora=False),
+            scaling,
+        )
+
+    # Reference
+    sgmv_shrink_for_nslices(
+        nslices,
+        data.inputs_tensor,
+        data.lora_weights,
+        ref_out_tensor,
+        *sgmv_meta_args,
+        scaling,
+    )
+
+    assert_close(out_tensor, ref_out_tensor)
+
+
+def check_lora_expand_kernel(
+    batches: int,
+    num_loras: int,
+    rank: int,
+    hidden_size: int,
+    nslices: int,
+    dtype: torch.dtype,
+    device: str,
+    seq_length: int,
+    add_inputs: bool,
+):
+    """
+    Compare outputs of torch_ops.sgmv_expand and triton_ops.lora_expand
+    kernels.
+    """
+    data: PunicaTensors = generate_data_for_nslices(
+        batches,
+        hidden_size,
+        num_loras,
+        rank,
+        seq_length,
+        nslices,
+        dtype,
+        "expand",
+        device,
+    )
+
+    max_seq_length, token_nums = data.meta()
+
+    # Setup metadata information for SGMV and reference kernels
+    sgmv_meta_args = (
+        data.b_seq_start_loc,
+        data.seq_len_tensor,
+        data.prompt_lora_mapping,
+        batches,
+        max_seq_length,
+        token_nums,
+    )
+
+    # Setup metadata information for the LoRA kernel.
+    lora_meta = LoRAKernelMeta.make(
+        max_loras=num_loras, max_num_tokens=token_nums, device="cuda"
+    )
+    lora_meta.prepare_tensors(data.token_lora_mapping)
+
+    # Setup output tensors
+    ref_out_tensor = data.ref_out_tensor
+    out_tensor = data.our_out_tensor.clone()
+
+    with _dict_lock:
+        # lora_expand kernel
+        _LORA_B_PTR_DICT.clear()
+        triton_ops.lora_expand(
+            data.inputs_tensor,
+            data.lora_weights,
+            out_tensor,
+            *lora_meta.meta_args(token_nums=token_nums, specialize_active_lora=False),
+            offset_start=0,
+            add_inputs=add_inputs,
+        )
+
+    # Reference
+    sgmv_expand_for_nslices(
+        nslices,
+        hidden_size,
+        data.inputs_tensor,
+        data.lora_weights,
+        ref_out_tensor,
+        *sgmv_meta_args,
+        add_inputs=add_inputs,
+    )
+
+    assert_close(out_tensor, ref_out_tensor)
+
+
+# Tests
+# We test the punica kernels along 2 verticals mainly.
+# 1. Variations in hidden_dim size
+# 2. Variations in all other parameters like (batch_size, max_rank, num_loras
+#  etc.)
+
+# We have collected the hidden_sizes included in the LoRA models
+# currently supported by vLLM. It tests whether the corresponding Triton
+# kernel can run normally when tensor parallelism is set to
+# [1, 2, 4, 8, 16, 32, 64].
+HIDDEN_SIZES = [
+    128,
+    256,
+    512,
+    896,
+    1024,
+    1152,
+    1216,
+    1280,
+    1536,
+    1664,
+    2048,
+    2240,
+    2304,
+    2368,
+    2432,
+    2560,
+    2752,
+    3072,
+    3328,
+    3456,
+    3584,
+    3712,
+    4096,
+    4480,
+    4608,
+    4736,
+    4864,
+    5120,
+    5504,
+    5632,
+    5888,
+    6144,
+    6400,
+    6848,
+    6912,
+    7168,
+    7424,
+    8192,
+    8960,
+    9216,
+    9472,
+    10240,
+    11008,
+    11264,
+    13824,
+    14336,
+    14784,
+    14848,
+    15360,
+    18944,
+    22016,
+    22528,
+    24576,
+    27392,
+    27648,
+    29568,
+    29696,
+    32000,
+    32256,
+    32512,
+    32768,
+    33024,
+    36864,
+    43264,
+    49152,
+    49408,
+    60544,
+    60672,
+    64000,
+    64256,
+    102400,
+    102656,
+    128000,
+    128256,
+]
+# The size of TP
+divisibility = [1, 2, 8, 16, 64]
+
+all_hidden_size = []
+for div in divisibility:
+    for hidden_size in HIDDEN_SIZES:
+        all_hidden_size.append(hidden_size // div)
+
+HIDDEN_SIZES = list(set(all_hidden_size))
+
+# Test params that focuses on hidden_size variation.
+hs_test_params = {
+    "hidden_sizes": HIDDEN_SIZES,
+    "batches": [4],
+    "num_loras": [4],
+    "max_ranks": [32],
+}
+
+# General tests params that tests for variations in all dimensions
+# except hidden_size.
+test_params = {
+    "hidden_sizes": [2049],
+    "batches": [1, 4, 16, 32],
+    "num_loras": [1, 8, 32, 128],
+    "max_ranks": [1, 4, 8, 16, 32, 64, 128, 256],
+}
+
+DTYPES = [torch.float16, torch.bfloat16]
+DEVICES = [f"cuda:{0}"]
+SEED = [0]
+
+
+@pytest.mark.parametrize("batches", test_params["batches"])
+@pytest.mark.parametrize("num_loras", test_params["num_loras"])
+@pytest.mark.parametrize("rank", test_params["max_ranks"])
+@pytest.mark.parametrize("hidden_size", test_params["hidden_sizes"])
+@pytest.mark.parametrize("nslices", [1, 2, 3])
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("seed", SEED)
+@pytest.mark.parametrize("op_type", ["shrink", "expand"])
+def test_kernels(
+    batches: int,
+    num_loras: int,
+    rank: int,
+    hidden_size: int,
+    nslices: int,
+    dtype: torch.dtype,
+    device: str,
+    seed: int,
+    op_type: str,
+):
+    """
+    Tests LoRA kernels.
+    """
+    torch.set_default_device(device)
+    torch.accelerator.set_device_index(device)
+    set_random_seed(seed)
+
+    if op_type == "shrink":
+        check_lora_shrink_kernel(
+            batches=batches,
+            num_loras=num_loras,
+            rank=rank,
+            hidden_size=hidden_size,
+            nslices=nslices,
+            dtype=dtype,
+            device=device,
+            seq_length=128,
+            scaling=0.5,
+        )
+    else:
+        check_lora_expand_kernel(
+            batches=batches,
+            num_loras=num_loras,
+            rank=rank,
+            hidden_size=hidden_size,
+            nslices=nslices,
+            dtype=dtype,
+            device=device,
+            seq_length=128,
+            add_inputs=True,
+        )
+
+
+@pytest.mark.parametrize("batches", hs_test_params["batches"])
+@pytest.mark.parametrize("num_loras", hs_test_params["num_loras"])
+@pytest.mark.parametrize("rank", hs_test_params["max_ranks"])
+@pytest.mark.parametrize("hidden_size", hs_test_params["hidden_sizes"])
+@pytest.mark.parametrize("nslices", [1, 2, 3])
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("seed", SEED)
+@pytest.mark.parametrize("op_type", ["shrink", "expand"])
+def test_kernels_hidden_size(
+    batches: int,
+    num_loras: int,
+    rank: int,
+    hidden_size: int,
+    nslices: int,
+    dtype: torch.dtype,
+    device: str,
+    seed: int,
+    op_type: str,
+):
+    """
+    Tests SGMV and LoRA kernels.
+    """
+    torch.set_default_device(device)
+    torch.accelerator.set_device_index(device)
+    set_random_seed(seed)
+
+    if op_type == "shrink":
+        check_lora_shrink_kernel(
+            batches=batches,
+            num_loras=num_loras,
+            rank=rank,
+            hidden_size=hidden_size,
+            nslices=nslices,
+            dtype=dtype,
+            device=device,
+            seq_length=128,
+            scaling=0.5,
+        )
+    else:
+        check_lora_expand_kernel(
+            batches=batches,
+            num_loras=num_loras,
+            rank=rank,
+            hidden_size=hidden_size,
+            nslices=nslices,
+            dtype=dtype,
+            device=device,
+            seq_length=128,
+            add_inputs=True,
+        )

third_party/vllm/tests/lora/test_punica_ops_fp8.py

@@ -0,0 +1,999 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""FP8 accuracy tests for LoRA shrink and expand kernels.
+
+Tests the FP8 kernels by:
+1. Quantizing bf16 inputs/weights to FP8
+2. Dequantizing them back to bf16
+3. Running the bf16 reference (sgmv_shrink/sgmv_expand) with dequantized values
+4. Comparing FP8 kernel output against this dequantized reference
+
+This isolates kernel correctness from quantization precision loss,
+allowing much tighter tolerances than comparing against the original bf16.
+"""
+
+import math
+from threading import Lock
+
+import pytest
+import torch
+
+import vllm.lora.ops.torch_ops as torch_ops
+import vllm.lora.ops.triton_ops as triton_ops
+from vllm.lora.ops.triton_ops import LoRAKernelMeta
+from vllm.lora.ops.triton_ops.lora_expand_fp8_op import (
+    _EXPAND_LORA_SCALE_PTR_DICT,
+)
+from vllm.lora.ops.triton_ops.lora_shrink_fp8_op import (
+    _SHRINK_LORA_SCALE_PTR_DICT,
+)
+from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT, _LORA_B_PTR_DICT
+from vllm.utils.torch_utils import set_random_seed
+
+DEVICES = [f"cuda:{0}"]
+SEED = [0]
+
+_dict_lock = Lock()
+
+
+@pytest.fixture(autouse=True)
+def reset_device(reset_default_device):
+    pass
+
+
+# ============================================================================
+# Reference implementations (bf16 baseline)
+# ============================================================================
+
+
+def sgmv_shrink_for_nslices(
+    nslices,
+    inputs_tensor,
+    lora_weights_lst,
+    out_tensor,
+    b_seq_start_loc,
+    seq_len_tensor,
+    prompt_lora_mapping,
+    batches,
+    max_seq_length,
+    num_tokens,
+    scaling,
+):
+    """Wrapper around torch_ops.sgmv_shrink that handles any nslices."""
+    for index in range(nslices):
+        torch_ops.sgmv_shrink(
+            inputs_tensor,
+            lora_weights_lst[index],
+            out_tensor[index],
+            b_seq_start_loc,
+            seq_len_tensor,
+            prompt_lora_mapping,
+            batches,
+            max_seq_length,
+            num_tokens,
+            scaling,
+        )
+
+
+def sgmv_expand_for_nslices(
+    nslices,
+    hidden_size,
+    inputs_tensor,
+    lora_weights_lst,
+    out_tensor,
+    b_seq_start_loc,
+    seq_len_tensor,
+    prompt_lora_mapping,
+    batches,
+    max_seq_length,
+    num_tokens,
+    add_inputs,
+):
+    """Wrapper around torch_ops.sgmv_expand that handles any nslices."""
+    if nslices == 1:
+        torch_ops.sgmv_expand(
+            inputs_tensor[0],
+            lora_weights_lst[0],
+            out_tensor,
+            b_seq_start_loc,
+            seq_len_tensor,
+            prompt_lora_mapping,
+            batches,
+            max_seq_length,
+            num_tokens,
+            add_inputs=add_inputs,
+        )
+    else:
+        slice_offset = 0
+        for index in range(nslices):
+            torch_ops.sgmv_expand_slice(
+                inputs_tensor[index],
+                lora_weights_lst[index],
+                out_tensor,
+                b_seq_start_loc,
+                seq_len_tensor,
+                prompt_lora_mapping,
+                batches,
+                max_seq_length,
+                num_tokens,
+                slice_offset,
+                hidden_size,
+                add_inputs=add_inputs,
+            )
+            slice_offset += hidden_size
+
+
+# ============================================================================
+# FP8 Quantization Helpers
+# ============================================================================
+
+FP8_DTYPE = torch.float8_e4m3fn
+FP8_MAX = torch.finfo(FP8_DTYPE).max
+FP8_MIN = torch.finfo(FP8_DTYPE).min
+
+
+def quantize_to_fp8_per_tensor(
+    tensor: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize a tensor to FP8 with per-tensor scaling."""
+    amax = tensor.abs().float().max().clamp(min=1e-12)
+    scale = (amax / FP8_MAX).to(torch.float32)
+    fp8_tensor = (tensor.float() / scale).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE)
+    return fp8_tensor, scale.reshape(1)
+
+
+def quantize_to_fp8_per_channel(
+    tensor: torch.Tensor,
+    channel_dim: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize a tensor to FP8 with per-channel scaling.
+
+    For shrink lora_a weights of shape (num_loras, rank, hidden_size):
+        channel_dim=1 gives per-rank scaling -> scale shape (num_loras, rank)
+    For expand lora_b weights of shape (num_loras, hidden_size, rank):
+        channel_dim=1 gives per-hidden scaling -> scale shape (num_loras, hidden_size)
+    """
+    # Compute amax along all dims except the leading dims up to channel_dim+1
+    reduce_dims = list(range(channel_dim + 1, tensor.ndim))
+    if reduce_dims:
+        amax = tensor.abs().float().amax(dim=reduce_dims).clamp(min=1e-12)
+    else:
+        amax = tensor.abs().float().clamp(min=1e-12)
+    scale = (amax / FP8_MAX).to(torch.float32)
+
+    # Expand scale for broadcasting
+    for _ in reduce_dims:
+        scale = scale.unsqueeze(-1)
+    fp8_tensor = (tensor.float() / scale).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE)
+    scale = scale.squeeze()
+    if scale.ndim == 0:
+        scale = scale.unsqueeze(0)
+    return fp8_tensor, scale
+
+
+def quantize_to_fp8_per_token(
+    tensor: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize a 2D tensor to FP8 with per-token (per-row) scaling.
+
+    Input shape: (num_tokens, hidden_size)
+    Returns: (fp8_tensor, scale) where scale shape is (num_tokens, 1)
+    """
+    assert tensor.ndim == 2
+    amax = tensor.abs().float().amax(dim=1, keepdim=True).clamp(min=1e-12)
+    scale = (amax / FP8_MAX).to(torch.float32)
+    fp8_tensor = (tensor.float() / scale).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE)
+    return fp8_tensor, scale
+
+
+def quantize_to_fp8_blockwise(
+    tensor: torch.Tensor,
+    group_n: int,
+    group_k: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize a 2D or 3D tensor to FP8 with block-wise scaling.
+
+    For a 2D tensor (num_tokens, hidden_size):
+        Blocks of size (1, group_k) ->
+            scale shape (num_tokens, ceil(hidden_size/group_k))
+
+    For a 3D tensor (num_loras, N, K):
+        Blocks of size (group_n, group_k) ->
+            scale shape (num_loras, ceil(N/group_n), ceil(K/group_k))
+    """
+    if tensor.ndim == 2:
+        M, K = tensor.shape
+        n_blocks_k = math.ceil(K / group_k)
+        scale = torch.zeros(M, n_blocks_k, dtype=torch.float32, device=tensor.device)
+        fp8_tensor = torch.zeros_like(tensor, dtype=FP8_DTYPE)
+        for m in range(M):
+            for bk in range(n_blocks_k):
+                k_start = bk * group_k
+                k_end = min(k_start + group_k, K)
+                block = tensor[m, k_start:k_end].float()
+                amax = block.abs().max().clamp(min=1e-12)
+                s = (amax / FP8_MAX).to(torch.float32)
+                scale[m, bk] = s
+                fp8_tensor[m, k_start:k_end] = (
+                    (block / s).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE)
+                )
+        return fp8_tensor, scale
+    elif tensor.ndim == 3:
+        L, N, K = tensor.shape
+        n_blocks_n = math.ceil(N / group_n)
+        n_blocks_k = math.ceil(K / group_k)
+        scale = torch.zeros(
+            L, n_blocks_n, n_blocks_k, dtype=torch.float32, device=tensor.device
+        )
+        fp8_tensor = torch.zeros_like(tensor, dtype=FP8_DTYPE)
+        for li in range(L):
+            for bn in range(n_blocks_n):
+                for bk in range(n_blocks_k):
+                    n_start = bn * group_n
+                    n_end = min(n_start + group_n, N)
+                    k_start = bk * group_k
+                    k_end = min(k_start + group_k, K)
+                    block = tensor[li, n_start:n_end, k_start:k_end].float()
+                    amax = block.abs().max().clamp(min=1e-12)
+                    s = (amax / FP8_MAX).to(torch.float32)
+                    scale[li, bn, bk] = s
+                    fp8_tensor[li, n_start:n_end, k_start:k_end] = (
+                        (block / s).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE)
+                    )
+        return fp8_tensor, scale
+    else:
+        raise ValueError(f"Unsupported tensor ndim: {tensor.ndim}")
+
+
+# ============================================================================
+# FP8 Dequantization Helpers
+# ============================================================================
+
+
+def dequantize_fp8_per_tensor(
+    fp8_tensor: torch.Tensor,
+    scale: torch.Tensor,
+    output_dtype: torch.dtype = torch.bfloat16,
+) -> torch.Tensor:
+    """Dequantize FP8 tensor with per-tensor scale back to output_dtype."""
+    return (fp8_tensor.float() * scale.float()).to(output_dtype)
+
+
+def dequantize_fp8_per_channel(
+    fp8_tensor: torch.Tensor,
+    scale: torch.Tensor,
+    channel_dim: int,
+    output_dtype: torch.dtype = torch.bfloat16,
+) -> torch.Tensor:
+    """Dequantize FP8 tensor with per-channel scale back to output_dtype.
+
+    For 3D tensor (num_loras, N, K) with channel_dim=1:
+        scale shape is (num_loras, N), broadcast over K.
+    """
+    expand_scale = scale.float()
+    # Add trailing dims for broadcasting
+    for _ in range(channel_dim + 1, fp8_tensor.ndim):
+        expand_scale = expand_scale.unsqueeze(-1)
+    return (fp8_tensor.float() * expand_scale).to(output_dtype)
+
+
+def dequantize_fp8_per_token(
+    fp8_tensor: torch.Tensor,
+    scale: torch.Tensor,
+    output_dtype: torch.dtype = torch.bfloat16,
+) -> torch.Tensor:
+    """Dequantize FP8 2D tensor with per-token scale back to output_dtype.
+
+    fp8_tensor: (num_tokens, hidden_size), scale: (num_tokens, 1)
+    """
+    return (fp8_tensor.float() * scale.float()).to(output_dtype)
+
+
+def dequantize_fp8_blockwise(
+    fp8_tensor: torch.Tensor,
+    scale: torch.Tensor,
+    group_n: int,
+    group_k: int,
+    output_dtype: torch.dtype = torch.bfloat16,
+) -> torch.Tensor:
+    """Dequantize FP8 tensor with block-wise scale back to output_dtype."""
+    if fp8_tensor.ndim == 2:
+        M, K = fp8_tensor.shape
+        out = torch.zeros(M, K, dtype=output_dtype, device=fp8_tensor.device)
+        n_blocks_k = math.ceil(K / group_k)
+        for m in range(M):
+            for bk in range(n_blocks_k):
+                k_start = bk * group_k
+                k_end = min(k_start + group_k, K)
+                out[m, k_start:k_end] = (
+                    fp8_tensor[m, k_start:k_end].float() * scale[m, bk].float()
+                ).to(output_dtype)
+        return out
+    elif fp8_tensor.ndim == 3:
+        L, N, K = fp8_tensor.shape
+        out = torch.zeros(L, N, K, dtype=output_dtype, device=fp8_tensor.device)
+        n_blocks_n = math.ceil(N / group_n)
+        n_blocks_k = math.ceil(K / group_k)
+        for l_idx in range(L):
+            for bn in range(n_blocks_n):
+                for bk in range(n_blocks_k):
+                    n_start = bn * group_n
+                    n_end = min(n_start + group_n, N)
+                    k_start = bk * group_k
+                    k_end = min(k_start + group_k, K)
+                    out[l_idx, n_start:n_end, k_start:k_end] = (
+                        fp8_tensor[l_idx, n_start:n_end, k_start:k_end].float()
+                        * scale[l_idx, bn, bk].float()
+                    ).to(output_dtype)
+        return out
+    else:
+        raise ValueError(f"Unsupported tensor ndim: {fp8_tensor.ndim}")
+
+
+# ============================================================================
+# FP8 Data Generation
+# ============================================================================
+
+
+def generate_fp8_shrink_data(
+    batches: int,
+    hidden_size: int,
+    num_loras: int,
+    rank: int,
+    seq_length: int,
+    nslices: int,
+    dtype: torch.dtype,
+    device: str,
+    quant_mode: str,  # "per_tensor", "per_channel", "blockwise"
+    group_k: int = 128,
+    group_n: int = 128,
+):
+    """Generate test data for FP8 shrink kernel.
+
+    Shrink: output = input @ lora_a^T * scaling
+    input: (num_tokens, hidden_size) -> quantized to FP8
+    lora_a: (num_loras, rank, hidden_size) -> quantized to FP8
+
+    Returns bf16 reference tensors, FP8 quantized tensors with scales,
+    and dequantized bf16 tensors for accurate reference computation.
+    """
+    seq_len_tensor = torch.randint(seq_length, seq_length + 1, (batches,)).to(device)
+    b_seq_start_loc = torch.cumsum(
+        torch.tensor([0] + seq_len_tensor[:-1].tolist(), dtype=torch.long),
+        dim=0,
+    ).to(device)
+    total_tokens = seq_len_tensor.sum().item()
+
+    # Generate bf16 reference data
+    inputs_bf16 = torch.randn(total_tokens, hidden_size, dtype=dtype, device=device)
+
+    lora_a_weights_bf16 = []
+    for _ in range(nslices):
+        lora_a_weights_bf16.append(
+            torch.randn(num_loras, rank, hidden_size, dtype=dtype, device=device)
+        )
+
+    # Quantize inputs to FP8 and dequantize back for reference
+    if quant_mode == "blockwise":
+        inputs_fp8, a_scale = quantize_to_fp8_blockwise(
+            inputs_bf16, group_n=1, group_k=group_k
+        )
+        inputs_dequant = dequantize_fp8_blockwise(
+            inputs_fp8,
+            a_scale,
+            group_n=1,
+            group_k=group_k,
+            output_dtype=dtype,
+        )
+    elif quant_mode == "per_tensor":
+        # Per-tensor: kernel loads a single scalar from a_scale_ptr
+        inputs_fp8, a_scale = quantize_to_fp8_per_tensor(inputs_bf16)
+        inputs_dequant = dequantize_fp8_per_tensor(
+            inputs_fp8,
+            a_scale,
+            output_dtype=dtype,
+        )
+    else:
+        # per_channel: kernel loads per-token a_scale via ram indexing
+        inputs_fp8, a_scale = quantize_to_fp8_per_token(inputs_bf16)
+        inputs_dequant = dequantize_fp8_per_token(
+            inputs_fp8,
+            a_scale,
+            output_dtype=dtype,
+        )
+
+    # Quantize lora_a weights to FP8 and dequantize back for reference
+    b_scales = []
+    lora_a_weights_fp8 = []
+    lora_a_weights_dequant = []
+    for w in lora_a_weights_bf16:
+        if quant_mode == "per_tensor":
+            w_fp8, w_scale = quantize_to_fp8_per_tensor(w)
+            w_dequant = dequantize_fp8_per_tensor(w_fp8, w_scale, output_dtype=dtype)
+            # Scale shape: (1,) -> need (num_loras,) for the kernel
+            w_scale = w_scale.expand(num_loras).contiguous()
+            lora_a_weights_fp8.append(w_fp8)
+            b_scales.append(w_scale)
+            lora_a_weights_dequant.append(w_dequant)
+        elif quant_mode == "per_channel":
+            # Per-channel along rank dim: scale shape (num_loras, rank)
+            w_fp8, w_scale = quantize_to_fp8_per_channel(w, channel_dim=1)
+            w_dequant = dequantize_fp8_per_channel(
+                w_fp8,
+                w_scale,
+                channel_dim=1,
+                output_dtype=dtype,
+            )
+            lora_a_weights_fp8.append(w_fp8)
+            b_scales.append(w_scale)
+            lora_a_weights_dequant.append(w_dequant)
+        elif quant_mode == "blockwise":
+            w_fp8, w_scale = quantize_to_fp8_blockwise(
+                w, group_n=group_n, group_k=group_k
+            )
+            w_dequant = dequantize_fp8_blockwise(
+                w_fp8,
+                w_scale,
+                group_n=group_n,
+                group_k=group_k,
+                output_dtype=dtype,
+            )
+            lora_a_weights_fp8.append(w_fp8)
+            b_scales.append(w_scale)
+            lora_a_weights_dequant.append(w_dequant)
+
+    # Output tensor (float32 for shrink)
+    out_tensor = torch.zeros(
+        nslices, total_tokens, rank, dtype=torch.float32, device=device
+    )
+    ref_out_tensor = out_tensor.clone()
+
+    # Token-to-lora mapping
+    lora_indices_tensor = torch.randint(0, max(num_loras - 1, 1), (batches,)).to(device)
+    token_lora_mapping = torch.zeros(total_tokens, dtype=torch.long, device=device)
+    current_offset = 0
+    for b_id in range(batches):
+        lora_index = lora_indices_tensor[b_id]
+        sl = seq_len_tensor[b_id].item()
+        token_lora_mapping[current_offset : current_offset + sl] = lora_index
+        current_offset += sl
+
+    return {
+        "inputs_bf16": inputs_bf16,
+        "inputs_fp8": inputs_fp8,
+        "inputs_dequant": inputs_dequant,
+        "lora_a_bf16": lora_a_weights_bf16,
+        "lora_a_fp8": lora_a_weights_fp8,
+        "lora_a_dequant": lora_a_weights_dequant,
+        "a_scale": a_scale,
+        "b_scales": b_scales,
+        "out_tensor": out_tensor,
+        "ref_out_tensor": ref_out_tensor,
+        "token_lora_mapping": token_lora_mapping,
+        "seq_len_tensor": seq_len_tensor,
+        "b_seq_start_loc": b_seq_start_loc,
+        "lora_indices_tensor": lora_indices_tensor,
+        "total_tokens": total_tokens,
+    }
+
+
+def generate_fp8_expand_data(
+    batches: int,
+    hidden_size: int,
+    num_loras: int,
+    rank: int,
+    seq_length: int,
+    nslices: int,
+    dtype: torch.dtype,
+    device: str,
+    quant_mode: str,  # "per_tensor", "per_channel", "blockwise"
+    group_k: int = 128,
+    group_n: int = 128,
+):
+    """Generate test data for FP8 expand kernel (w8a8).
+
+    Expand: output += input @ lora_b^T
+    input: (nslices, num_tokens, rank) -> quantized to FP8 (activations)
+    lora_b: (num_loras, hidden_size, rank) -> quantized to FP8 (weights)
+
+    In w8a8 mode, both activations and weights are FP8.
+    Returns bf16 reference tensors, FP8 quantized tensors with scales,
+    and dequantized bf16 tensors for accurate reference computation.
+    """
+    seq_len_tensor = torch.randint(seq_length, seq_length + 1, (batches,)).to(device)
+    b_seq_start_loc = torch.cumsum(
+        torch.tensor([0] + seq_len_tensor[:-1].tolist(), dtype=torch.long),
+        dim=0,
+    ).to(device)
+    total_tokens = seq_len_tensor.sum().item()
+
+    # Generate bf16 input (shrink output) and quantize to FP8
+    inputs_bf16 = torch.randn(nslices, total_tokens, rank, dtype=dtype, device=device)
+
+    # Quantize input to FP8 and dequantize back for reference
+    inputs_2d_all = inputs_bf16.reshape(-1, rank)
+    if quant_mode == "blockwise":
+        # For blockwise, the kernel indexes a_scale by token id (0..total_tokens-1)
+        # shared across slices. Compute shared scale across slices, then quantize.
+        # First compute per-token-per-block scale across all slices
+        n_blocks_k = math.ceil(rank / group_k)
+        a_scale = torch.zeros(
+            total_tokens, n_blocks_k, dtype=torch.float32, device=device
+        )
+        for m in range(total_tokens):
+            for bk in range(n_blocks_k):
+                k_start = bk * group_k
+                k_end = min(k_start + group_k, rank)
+                # Max across all slices for this token and block
+                block_amax = torch.tensor(0.0, device=device)
+                for s in range(nslices):
+                    block = inputs_bf16[s, m, k_start:k_end].float()
+                    block_amax = torch.max(
+                        block_amax, block.abs().max().clamp(min=1e-12)
+                    )
+                a_scale[m, bk] = (block_amax / FP8_MAX).to(torch.float32)
+
+        # Quantize all slices with the shared scale
+        inputs_fp8_list = []
+        inputs_dequant_list = []
+        for s in range(nslices):
+            slice_2d = inputs_bf16[s]  # (total_tokens, rank)
+            fp8_slice = torch.zeros_like(slice_2d, dtype=FP8_DTYPE)
+            dequant_slice = torch.zeros_like(slice_2d)
+            for m in range(total_tokens):
+                for bk in range(n_blocks_k):
+                    k_start = bk * group_k
+                    k_end = min(k_start + group_k, rank)
+                    block = slice_2d[m, k_start:k_end].float()
+                    s_val = a_scale[m, bk]
+                    fp8_slice[m, k_start:k_end] = (
+                        (block / s_val).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE)
+                    )
+                    dequant_slice[m, k_start:k_end] = (
+                        fp8_slice[m, k_start:k_end].float() * s_val.float()
+                    ).to(dtype)
+            inputs_fp8_list.append(fp8_slice)
+            inputs_dequant_list.append(dequant_slice)
+        inputs_fp8 = torch.stack(inputs_fp8_list, dim=0)
+        inputs_dequant = torch.stack(inputs_dequant_list, dim=0)
+    elif quant_mode == "per_tensor":
+        # Per-tensor: kernel loads a single scalar from a_scale_ptr
+        inputs_fp8_2d, a_scale = quantize_to_fp8_per_tensor(inputs_2d_all)
+        inputs_dequant_2d = dequantize_fp8_per_tensor(
+            inputs_fp8_2d,
+            a_scale,
+            output_dtype=dtype,
+        )
+        inputs_fp8 = inputs_fp8_2d.reshape(nslices, total_tokens, rank)
+        inputs_dequant = inputs_dequant_2d.reshape(nslices, total_tokens, rank)
+    else:
+        # per_channel: kernel loads per-token a_scale via ram indexing.
+        # The kernel uses the same a_scale for all slices (indexed by token
+        # id 0..total_tokens-1), so we compute a shared per-token scale
+        # across all slices, then quantize each slice with that shared scale.
+        per_slice_views = [inputs_bf16[s] for s in range(nslices)]
+        # (nslices, total_tokens, rank) -> max across slices per token
+        stacked = torch.stack(per_slice_views, dim=0)  # (nslices, tokens, rank)
+        amax = stacked.abs().float().amax(dim=(0, 2), keepdim=False).clamp(min=1e-12)
+        # amax shape: (total_tokens,)
+        a_scale = (amax / FP8_MAX).to(torch.float32).unsqueeze(1)  # (tokens, 1)
+        # Quantize all slices with the shared scale
+        inputs_fp8_2d = (
+            (inputs_2d_all.float() / a_scale.repeat(nslices, 1))
+            .clamp(FP8_MIN, FP8_MAX)
+            .to(FP8_DTYPE)
+        )
+        inputs_dequant_2d = (
+            inputs_fp8_2d.float() * a_scale.repeat(nslices, 1).float()
+        ).to(dtype)
+        inputs_fp8 = inputs_fp8_2d.reshape(nslices, total_tokens, rank)
+        inputs_dequant = inputs_dequant_2d.reshape(nslices, total_tokens, rank)
+
+    # Generate bf16 LoRA B weights
+    lora_b_weights_bf16 = []
+    for _ in range(nslices):
+        lora_b_weights_bf16.append(
+            torch.randn(num_loras, hidden_size, rank, dtype=dtype, device=device)
+        )
+
+    # Quantize LoRA B weights to FP8 and dequantize back for reference
+    b_scales = []
+    lora_b_weights_fp8 = []
+    lora_b_weights_dequant = []
+    for w in lora_b_weights_bf16:
+        if quant_mode == "per_tensor":
+            w_fp8, w_scale = quantize_to_fp8_per_tensor(w)
+            w_dequant = dequantize_fp8_per_tensor(w_fp8, w_scale, output_dtype=dtype)
+            w_scale = w_scale.expand(num_loras).contiguous()
+            lora_b_weights_fp8.append(w_fp8)
+            b_scales.append(w_scale)
+            lora_b_weights_dequant.append(w_dequant)
+        elif quant_mode == "per_channel":
+            # Per-channel along hidden_size dim: scale (num_loras, hidden_size)
+            w_fp8, w_scale = quantize_to_fp8_per_channel(w, channel_dim=1)
+            w_dequant = dequantize_fp8_per_channel(
+                w_fp8,
+                w_scale,
+                channel_dim=1,
+                output_dtype=dtype,
+            )
+            lora_b_weights_fp8.append(w_fp8)
+            b_scales.append(w_scale)
+            lora_b_weights_dequant.append(w_dequant)
+        elif quant_mode == "blockwise":
+            w_fp8, w_scale = quantize_to_fp8_blockwise(
+                w, group_n=group_n, group_k=group_k
+            )
+            w_dequant = dequantize_fp8_blockwise(
+                w_fp8,
+                w_scale,
+                group_n=group_n,
+                group_k=group_k,
+                output_dtype=dtype,
+            )
+            lora_b_weights_fp8.append(w_fp8)
+            b_scales.append(w_scale)
+            lora_b_weights_dequant.append(w_dequant)
+
+    # Output tensor (initialized randomly for add_inputs)
+    out_tensor = torch.randn(
+        total_tokens, hidden_size * nslices, dtype=dtype, device=device
+    )
+    ref_out_tensor = out_tensor.clone()
+
+    # Token-to-lora mapping
+    lora_indices_tensor = torch.randint(0, max(num_loras - 1, 1), (batches,)).to(device)
+    token_lora_mapping = torch.zeros(total_tokens, dtype=torch.long, device=device)
+    current_offset = 0
+    for b_id in range(batches):
+        lora_index = lora_indices_tensor[b_id]
+        sl = seq_len_tensor[b_id].item()
+        token_lora_mapping[current_offset : current_offset + sl] = lora_index
+        current_offset += sl
+
+    return {
+        "inputs_bf16": inputs_bf16,
+        "inputs_fp8": inputs_fp8,
+        "inputs_dequant": inputs_dequant,
+        "a_scale": a_scale,
+        "lora_b_bf16": lora_b_weights_bf16,
+        "lora_b_fp8": lora_b_weights_fp8,
+        "lora_b_dequant": lora_b_weights_dequant,
+        "b_scales": b_scales,
+        "out_tensor": out_tensor,
+        "ref_out_tensor": ref_out_tensor,
+        "token_lora_mapping": token_lora_mapping,
+        "seq_len_tensor": seq_len_tensor,
+        "b_seq_start_loc": b_seq_start_loc,
+        "lora_indices_tensor": lora_indices_tensor,
+        "total_tokens": total_tokens,
+    }
+
+
+# ============================================================================
+# FP8 Shrink Kernel Check
+# ============================================================================
+
+
+def check_lora_shrink_fp8_kernel(
+    batches: int,
+    num_loras: int,
+    rank: int,
+    hidden_size: int,
+    nslices: int,
+    dtype: torch.dtype,
+    device: str,
+    seq_length: int,
+    scaling: float,
+    quant_mode: str,
+    group_k: int = 128,
+    group_n: int = 128,
+):
+    """Test FP8 shrink kernel against dequantized bf16 reference.
+
+    Instead of comparing FP8 kernel output against the original bf16 reference
+    (which conflates quantization error with kernel error), we:
+    1. Quantize bf16 inputs/weights to FP8
+    2. Dequantize them back to bf16
+    3. Run the bf16 reference (sgmv_shrink) with the dequantized values
+    4. Compare FP8 kernel output against this dequantized reference
+
+    This isolates kernel correctness from quantization precision loss,
+    allowing much tighter tolerances.
+    """
+    data = generate_fp8_shrink_data(
+        batches,
+        hidden_size,
+        num_loras,
+        rank,
+        seq_length,
+        nslices,
+        dtype,
+        device,
+        quant_mode,
+        group_k,
+        group_n,
+    )
+
+    total_tokens = data["total_tokens"]
+
+    # Setup LoRA kernel metadata
+    lora_meta = LoRAKernelMeta.make(
+        max_loras=num_loras, max_num_tokens=total_tokens, device=device
+    )
+    lora_meta.prepare_tensors(data["token_lora_mapping"])
+
+    out_tensor = data["out_tensor"]
+
+    # Determine quantization params for the kernel
+    per_channel = quant_mode == "per_channel"
+    gk = group_k if quant_mode == "blockwise" else 0
+    gn = group_n if quant_mode == "blockwise" else 0
+
+    with _dict_lock:
+        _LORA_A_PTR_DICT.clear()
+        _SHRINK_LORA_SCALE_PTR_DICT.clear()
+        triton_ops.lora_shrink_fp8(
+            data["inputs_fp8"],
+            data["lora_a_fp8"],
+            out_tensor,
+            *lora_meta.meta_args(token_nums=total_tokens, specialize_active_lora=False),
+            scaling,
+            data["b_scales"],
+            a_scale=data["a_scale"],
+            group_k=gk,
+            group_n=gn,
+            use_fp8_w8a8=True,
+            per_channel_quant=per_channel,
+        )
+
+    # Compute reference using dequantized (round-tripped) tensors.
+    # This means the reference sees the same quantization error as the kernel,
+    # so any difference is purely kernel error.
+    ref_out_tensor = data["ref_out_tensor"]
+    max_seq_length = data["seq_len_tensor"].max().item()
+    sgmv_shrink_for_nslices(
+        nslices,
+        data["inputs_dequant"],
+        data["lora_a_dequant"],
+        ref_out_tensor,
+        data["b_seq_start_loc"],
+        data["seq_len_tensor"],
+        data["lora_indices_tensor"],
+        batches,
+        max_seq_length,
+        total_tokens,
+        scaling,
+    )
+
+    # With dequantized reference, we can use much tighter tolerances
+    # since we're only measuring kernel error, not quantization error.
+    # Blockwise accumulation order differs from the bf16 reference, so
+    # allow a slightly larger margin for sporadic rounding outliers.
+    rtol, atol = 0.1, 0.25
+    torch.testing.assert_close(
+        out_tensor.to(dtype), ref_out_tensor.to(dtype), rtol=rtol, atol=atol
+    )
+
+
+# ============================================================================
+# FP8 Expand Kernel Check
+# ============================================================================
+
+
+def check_lora_expand_fp8_kernel(
+    batches: int,
+    num_loras: int,
+    rank: int,
+    hidden_size: int,
+    nslices: int,
+    dtype: torch.dtype,
+    device: str,
+    seq_length: int,
+    add_inputs: bool,
+    quant_mode: str,
+    group_k: int = 128,
+    group_n: int = 128,
+):
+    """Test FP8 expand kernel (w8a8) against dequantized bf16 reference.
+
+    Instead of comparing FP8 kernel output against the original bf16 reference
+    (which conflates quantization error with kernel error), we:
+    1. Quantize bf16 inputs/weights to FP8
+    2. Dequantize them back to bf16
+    3. Run the bf16 reference (sgmv_expand) with the dequantized values
+    4. Compare FP8 kernel output against this dequantized reference
+
+    This isolates kernel correctness from quantization precision loss,
+    allowing much tighter tolerances.
+    """
+    data = generate_fp8_expand_data(
+        batches,
+        hidden_size,
+        num_loras,
+        rank,
+        seq_length,
+        nslices,
+        dtype,
+        device,
+        quant_mode,
+        group_k,
+        group_n,
+    )
+
+    total_tokens = data["total_tokens"]
+
+    # Setup LoRA kernel metadata
+    lora_meta = LoRAKernelMeta.make(
+        max_loras=num_loras, max_num_tokens=total_tokens, device=device
+    )
+    lora_meta.prepare_tensors(data["token_lora_mapping"])
+
+    out_tensor = data["out_tensor"]
+
+    # Determine quantization params for the kernel
+    per_channel = quant_mode == "per_channel"
+    gk = group_k if quant_mode == "blockwise" else 0
+    gn = group_n if quant_mode == "blockwise" else 0
+
+    with _dict_lock:
+        _LORA_B_PTR_DICT.clear()
+        _EXPAND_LORA_SCALE_PTR_DICT.clear()
+        triton_ops.lora_expand_fp8(
+            data["inputs_fp8"],
+            data["lora_b_fp8"],
+            out_tensor,
+            *lora_meta.meta_args(token_nums=total_tokens, specialize_active_lora=False),
+            data["b_scales"],
+            a_scale=data["a_scale"],
+            offset_start=0,
+            add_inputs=add_inputs,
+            group_k=gk,
+            group_n=gn,
+            use_fp8_w8a8=True,
+            per_channel_quant=per_channel,
+        )
+
+    # Compute reference using dequantized (round-tripped) tensors.
+    ref_out_tensor = data["ref_out_tensor"]
+    max_seq_length = data["seq_len_tensor"].max().item()
+    sgmv_expand_for_nslices(
+        nslices,
+        hidden_size,
+        data["inputs_dequant"],
+        data["lora_b_dequant"],
+        ref_out_tensor,
+        data["b_seq_start_loc"],
+        data["seq_len_tensor"],
+        data["lora_indices_tensor"],
+        batches,
+        max_seq_length,
+        total_tokens,
+        add_inputs=add_inputs,
+    )
+
+    # With dequantized reference, we can use much tighter tolerances
+    # since we're only measuring kernel error, not quantization error.
+    rtol, atol = 0.1, 0.15
+    torch.testing.assert_close(out_tensor, ref_out_tensor, rtol=rtol, atol=atol)
+
+
+# ============================================================================
+# FP8 Test Parameters
+# ============================================================================
+
+fp8_test_params = {
+    "hidden_sizes": [512, 1024, 2048],
+    "batches": [1, 4, 16],
+    "num_loras": [1, 4, 8],
+    "max_ranks": [8, 16, 32, 64],
+}
+
+
+# ============================================================================
+# FP8 Shrink Tests
+# ============================================================================
+
+
+@pytest.mark.parametrize("batches", fp8_test_params["batches"])
+@pytest.mark.parametrize("num_loras", fp8_test_params["num_loras"])
+@pytest.mark.parametrize("rank", fp8_test_params["max_ranks"])
+@pytest.mark.parametrize("hidden_size", fp8_test_params["hidden_sizes"])
+@pytest.mark.parametrize("nslices", [1, 2, 3])
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("seed", SEED)
+@pytest.mark.parametrize("quant_mode", ["per_tensor", "per_channel", "blockwise"])
+def test_lora_shrink_fp8(
+    batches: int,
+    num_loras: int,
+    rank: int,
+    hidden_size: int,
+    nslices: int,
+    dtype: torch.dtype,
+    device: str,
+    seed: int,
+    quant_mode: str,
+):
+    """Test FP8 shrink kernel with per-tensor, per-channel, and block-wise
+    quantization, comparing against the bf16 baseline."""
+    torch.set_default_device(device)
+    set_random_seed(seed)
+
+    # For blockwise, group sizes must divide evenly or be handled by the kernel
+    group_k = 128
+    group_n = 128
+
+    # Adjust group sizes if they're larger than the dimensions
+    if quant_mode == "blockwise":
+        group_k = min(group_k, hidden_size)
+        group_n = min(group_n, rank)
+
+    check_lora_shrink_fp8_kernel(
+        batches=batches,
+        num_loras=num_loras,
+        rank=rank,
+        hidden_size=hidden_size,
+        nslices=nslices,
+        dtype=dtype,
+        device=device,
+        seq_length=128,
+        scaling=0.5,
+        quant_mode=quant_mode,
+        group_k=group_k,
+        group_n=group_n,
+    )
+
+
+# ============================================================================
+# FP8 Expand Tests
+# ============================================================================
+
+
+@pytest.mark.parametrize("batches", fp8_test_params["batches"])
+@pytest.mark.parametrize("num_loras", fp8_test_params["num_loras"])
+@pytest.mark.parametrize("rank", fp8_test_params["max_ranks"])
+@pytest.mark.parametrize("hidden_size", fp8_test_params["hidden_sizes"])
+@pytest.mark.parametrize("nslices", [1, 2, 3])
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("seed", SEED)
+@pytest.mark.parametrize("quant_mode", ["per_tensor", "per_channel", "blockwise"])
+def test_lora_expand_fp8(
+    batches: int,
+    num_loras: int,
+    rank: int,
+    hidden_size: int,
+    nslices: int,
+    dtype: torch.dtype,
+    device: str,
+    seed: int,
+    quant_mode: str,
+):
+    """Test FP8 expand kernel with per-tensor, per-channel, and block-wise
+    quantization, comparing against the bf16 baseline."""
+    torch.set_default_device(device)
+    set_random_seed(seed)
+
+    group_k = 128
+    group_n = 128
+
+    # Adjust group sizes if they're larger than the dimensions
+    if quant_mode == "blockwise":
+        group_k = min(group_k, rank)
+        group_n = min(group_n, hidden_size)
+
+    check_lora_expand_fp8_kernel(
+        batches=batches,
+        num_loras=num_loras,
+        rank=rank,
+        hidden_size=hidden_size,
+        nslices=nslices,
+        dtype=dtype,
+        device=device,
+        seq_length=128,
+        add_inputs=True,
+        quant_mode=quant_mode,
+        group_k=group_k,
+        group_n=group_n,
+    )

third_party/vllm/tests/lora/test_punica_xpu_ops.py

@@ -0,0 +1,298 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from tests.lora.utils import (
+    PunicaTensors,
+    assert_close,
+    generate_data,
+    generate_data_for_expand_nslices,
+)
+from vllm.lora.ops.xpu_ops import bgmv_expand, bgmv_expand_slice, bgmv_shrink
+from vllm.platforms import current_platform
+
+
+def torch_bgmv_expand(
+    inputs: torch.Tensor,
+    lora_b_weights: torch.Tensor,
+    output_tensor: torch.Tensor,
+    lora_indices_tensor: torch.Tensor,
+    add_inputs: bool = True,
+):
+    selected_loras = lora_b_weights[lora_indices_tensor].to(dtype=output_tensor.dtype)
+    if len(selected_loras.shape) == 4:
+        selected_loras = selected_loras.squeeze(dim=1)
+    inputs = inputs.to(dtype=output_tensor.dtype)
+    outputs = torch.einsum("bi, boi -> bo", inputs, selected_loras)
+
+    limit = output_tensor.shape[0]
+    if outputs.shape[0] == 1 and output_tensor.shape[0] != 1:
+        limit = 1
+
+    # LoRA adapter and model may add different amounts of padding to output
+    common_len = min(outputs.shape[1], output_tensor.shape[1])
+
+    if add_inputs:
+        output_tensor[:, :common_len] += outputs[:limit, :common_len]
+    else:
+        output_tensor[:, :common_len] = outputs[:limit, :common_len]
+
+
+def torch_bgmv_shrink(
+    inputs: torch.Tensor,
+    lora_b_weights: torch.Tensor,
+    output_tensor: torch.Tensor,
+    lora_indices_tensor: torch.Tensor,
+    scaling: float = 1.0,
+):
+    selected_loras = lora_b_weights[lora_indices_tensor].to(dtype=output_tensor.dtype)
+    if len(selected_loras.shape) == 4:
+        selected_loras = selected_loras.squeeze(dim=1)
+    inputs = inputs.to(dtype=output_tensor.dtype)
+    outputs = torch.einsum("bi, boi -> bo", inputs, selected_loras)
+
+    output_tensor[:, : outputs.shape[1]] = scaling * outputs[:]
+
+
+def torch_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,
+):
+    selected_loras = lora_b_weights[lora_indices_tensor].to(dtype=output_tensor.dtype)
+    inputs = inputs.to(dtype=output_tensor.dtype)
+    if len(selected_loras.shape) == 4:
+        selected_loras = selected_loras.squeeze(dim=1)
+    outputs = torch.einsum("bi, boi -> bo", inputs, selected_loras)
+
+    if add_inputs:
+        output_tensor[:, slice_offset : slice_offset + slice_size] += outputs[:]
+    else:
+        output_tensor[:, slice_offset : slice_offset + slice_size] = outputs[:]
+
+
+def check_bgmv_shrink(
+    batches: int,
+    num_loras: int,
+    rank: int,
+    hidden_size: int,
+    dtype: torch.dtype,
+    device: str,
+    scaling: float,
+):
+    """
+    Compare vllm.bgmv_shrink against a reference implementation.
+    """
+    seq_length = 1
+    data: PunicaTensors = generate_data(
+        batches,
+        hidden_size,
+        num_loras,
+        rank,
+        seq_length,
+        dtype,
+        "shrink",
+        device,
+    )
+
+    bgmv_shrink(
+        data.inputs_tensor,
+        data.lora_weights,
+        data.our_out_tensor,
+        data.token_lora_mapping,
+        scaling,
+    )
+
+    torch_bgmv_shrink(
+        data.inputs_tensor,
+        data.lora_weights,
+        data.ref_out_tensor,
+        data.token_lora_mapping,
+        scaling,
+    )
+
+    data.ref_out_tensor = data.ref_out_tensor.to(torch.float32)
+    assert_close(data.our_out_tensor, data.ref_out_tensor)
+
+
+def check_bgmv_expand(
+    batches: int,
+    num_loras: int,
+    rank: int,
+    hidden_size: int,
+    dtype: torch.dtype,
+    device: str,
+    add_inputs: bool,
+):
+    """
+    Compare vllm.bgmv_expand against a reference implementation.
+    """
+    seq_length = 1
+    data: PunicaTensors = generate_data(
+        batches,
+        hidden_size,
+        num_loras,
+        rank,
+        seq_length,
+        dtype,
+        "expand",
+        device,
+    )
+
+    bgmv_expand(
+        data.inputs_tensor,
+        data.lora_weights,
+        data.our_out_tensor,
+        data.token_lora_mapping,
+        add_inputs=add_inputs,
+    )
+    torch_bgmv_expand(
+        data.inputs_tensor,
+        data.lora_weights,
+        data.ref_out_tensor,
+        data.token_lora_mapping,
+        add_inputs=add_inputs,
+    )
+    assert_close(data.ref_out_tensor, data.our_out_tensor)
+
+
+def check_bgmv_expand_slice(
+    batches: int,
+    num_loras: int,
+    rank: int,
+    hidden_size: int,
+    nslices: int,
+    dtype: torch.dtype,
+    device: str,
+    add_inputs: bool,
+):
+    """
+    Compare vllm.bgmv_expand_slice against a reference implementation.
+    """
+    seq_length = 1
+    data: PunicaTensors = generate_data_for_expand_nslices(
+        batches,
+        hidden_size,
+        num_loras,
+        rank,
+        seq_length,
+        dtype,
+        nslices,
+        device,
+    )
+
+    slice_offset = 0
+    for index in range(nslices):
+        bgmv_expand_slice(
+            data.inputs_tensor,
+            data.lora_weights[index],
+            data.our_out_tensor,
+            data.token_lora_mapping,
+            slice_offset,
+            slice_size=hidden_size,
+            add_inputs=add_inputs,
+        )
+        torch_bgmv_expand_slice(
+            data.inputs_tensor,
+            data.lora_weights[index],
+            data.ref_out_tensor,
+            data.token_lora_mapping,
+            slice_offset,
+            slice_size=hidden_size,
+            add_inputs=add_inputs,
+        )
+
+        slice_offset += hidden_size
+    assert_close(data.ref_out_tensor, data.our_out_tensor)
+
+
+# General tests params that tests for variations in all dimensions
+# except hidden_size.
+test_params = {
+    "hidden_sizes": [2049],
+    "batches": [4],
+    "num_loras": [4],
+    "max_ranks": [32],
+}
+
+DTYPES = [torch.float16, torch.bfloat16]
+DEVICES = [f"xpu:{0}"]
+SEED = [0]
+
+
+@pytest.mark.parametrize("batches", test_params["batches"])
+@pytest.mark.parametrize("num_loras", test_params["num_loras"])
+@pytest.mark.parametrize("rank", test_params["max_ranks"])
+@pytest.mark.parametrize("hidden_size", test_params["hidden_sizes"])
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("seed", SEED)
+@pytest.mark.parametrize("op_type", ["shrink", "expand"])
+@pytest.mark.skipif(not current_platform.is_xpu(), reason="skip for non xpu platform")
+def test_bgmv(
+    batches: int,
+    num_loras: int,
+    rank: int,
+    hidden_size: int,
+    dtype: torch.dtype,
+    device: str,
+    seed: int,
+    op_type: str,
+):
+    if op_type == "shrink":
+        check_bgmv_shrink(
+            batches=batches,
+            num_loras=num_loras,
+            rank=rank,
+            hidden_size=hidden_size,
+            dtype=dtype,
+            device=device,
+            scaling=0.5,
+        )
+    else:
+        check_bgmv_expand(
+            batches=batches,
+            num_loras=num_loras,
+            rank=rank,
+            hidden_size=hidden_size,
+            dtype=dtype,
+            device=device,
+            add_inputs=True,
+        )
+
+
+@pytest.mark.parametrize("batches", test_params["batches"])
+@pytest.mark.parametrize("num_loras", test_params["num_loras"])
+@pytest.mark.parametrize("rank", test_params["max_ranks"])
+@pytest.mark.parametrize("hidden_size", test_params["hidden_sizes"])
+@pytest.mark.parametrize("nslices", [2, 3])
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("seed", SEED)
+@pytest.mark.skipif(not current_platform.is_xpu(), reason="skip for non xpu platform")
+def test_bgmv_expand_nslices(
+    batches: int,
+    num_loras: int,
+    rank: int,
+    hidden_size: int,
+    nslices: int,
+    dtype: torch.dtype,
+    device: str,
+    seed: int,
+):
+    check_bgmv_expand_slice(
+        batches=batches,
+        num_loras=num_loras,
+        rank=rank,
+        hidden_size=hidden_size,
+        nslices=nslices,
+        dtype=dtype,
+        device=device,
+        add_inputs=True,
+    )

third_party/vllm/tests/lora/test_quant_model.py

@@ -0,0 +1,167 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Adapted from
+# https://github.com/fmmoret/vllm/blob/fm-support-lora-on-quantized-models/tests/lora/test_llama.py
+from dataclasses import dataclass
+
+import pytest
+
+import vllm
+from vllm.distributed import cleanup_dist_env_and_memory
+from vllm.lora.request import LoRARequest
+from vllm.platforms import current_platform
+
+
+@dataclass
+class ModelWithQuantization:
+    model_path: str
+    quantization: str
+
+
+MODELS: list[ModelWithQuantization]
+# AWQ quantization is currently not supported in ROCm.
+if current_platform.is_rocm():
+    MODELS = [
+        ModelWithQuantization(
+            model_path="TheBloke/TinyLlama-1.1B-Chat-v0.3-GPTQ", quantization="gptq"
+        ),
+    ]
+else:
+    MODELS = [
+        ModelWithQuantization(
+            model_path="TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ", quantization="awq"
+        ),
+        ModelWithQuantization(
+            model_path="TheBloke/TinyLlama-1.1B-Chat-v0.3-GPTQ", quantization="gptq"
+        ),
+    ]
+
+
+def do_sample(
+    llm: vllm.LLM, lora_path: str, lora_id: int, max_tokens: int = 256
+) -> list[str]:
+    raw_prompts = [
+        "Give me an orange-ish brown color",
+        "Give me a neon pink color",
+    ]
+
+    def format_prompt_tuples(prompt):
+        return f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
+
+    prompts = [format_prompt_tuples(p) for p in raw_prompts]
+
+    sampling_params = vllm.SamplingParams(
+        temperature=0, max_tokens=max_tokens, stop=["<|im_end|>"]
+    )
+    outputs = llm.generate(
+        prompts,
+        sampling_params,
+        lora_request=LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None,
+    )
+    # Print the outputs.
+    generated_texts: list[str] = []
+    for output in outputs:
+        prompt = output.prompt
+        generated_text = output.outputs[0].text
+        generated_texts.append(generated_text)
+        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
+    return generated_texts
+
+
+@pytest.mark.parametrize("model", MODELS)
+def test_quant_model_lora(tinyllama_lora_files, model):
+    llm = vllm.LLM(
+        model=model.model_path,
+        enable_lora=True,
+        max_num_seqs=16,
+        max_loras=4,
+        max_model_len=400,
+        gpu_memory_utilization=0.2,  # avoid OOM
+        quantization=model.quantization,
+        trust_remote_code=True,
+        enable_chunked_prefill=True,
+        tokenizer=tinyllama_lora_files,
+    )
+
+    if model.quantization is None:
+        expected_lora_output = [
+            "#ff8050",
+            "#ff8080",
+        ]
+    elif model.quantization == "awq":
+        expected_lora_output = [
+            "#f07700: A v",
+            "#f00000: A v",
+        ]
+    elif model.quantization == "gptq":
+        expected_lora_output = [
+            "#f08800: This is",
+            "#f07788 \n#",
+        ]
+
+    def expect_match(output, expected_output):
+        # HACK: GPTQ lora outputs are just incredibly unstable.
+        # Assert that the outputs changed.
+        if model.quantization == "gptq" and expected_output is expected_lora_output:
+            for i, o in enumerate(output):
+                assert o.startswith("#"), (
+                    f"Expected example {i} to start with # but got {o}"
+                )
+            return
+        assert output == expected_output
+
+    max_tokens = 10
+
+    print("lora adapter created")
+    print("lora 1")
+    output = do_sample(llm, tinyllama_lora_files, lora_id=1, max_tokens=max_tokens)
+    expect_match(output, expected_lora_output)
+
+    print("lora 2")
+    output = do_sample(llm, tinyllama_lora_files, lora_id=2, max_tokens=max_tokens)
+    expect_match(output, expected_lora_output)
+
+    print("removing lora")
+
+    del llm
+    cleanup_dist_env_and_memory()
+
+
+@pytest.mark.parametrize("model", MODELS)
+def test_quant_model_tp_equality(tinyllama_lora_files, num_gpus_available, model):
+    if num_gpus_available < 2:
+        pytest.skip(f"Not enough GPUs for tensor parallelism {2}")
+    if model.quantization == "gptq":
+        pytest.skip("GPTQ lora outputs are just incredibly unstable")
+    llm_tp1 = vllm.LLM(
+        model=model.model_path,
+        enable_lora=True,
+        max_num_seqs=16,
+        max_loras=4,
+        gpu_memory_utilization=0.2,  # avoid OOM
+        quantization=model.quantization,
+        trust_remote_code=True,
+        enable_chunked_prefill=True,
+    )
+    output_tp1 = do_sample(llm_tp1, tinyllama_lora_files, lora_id=1)
+
+    del llm_tp1
+    cleanup_dist_env_and_memory()
+
+    llm_tp2 = vllm.LLM(
+        model=model.model_path,
+        enable_lora=True,
+        max_num_seqs=16,
+        max_loras=4,
+        tensor_parallel_size=2,
+        gpu_memory_utilization=0.2,  # avoid OOM
+        quantization=model.quantization,
+        enable_chunked_prefill=True,
+    )
+    output_tp2 = do_sample(llm_tp2, tinyllama_lora_files, lora_id=1)
+
+    del llm_tp2
+    cleanup_dist_env_and_memory()
+
+    assert output_tp1 == output_tp2

third_party/vllm/tests/lora/test_qwen3_unembed.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Tests for Qwen3 unembed LoRA support.
+
+This test creates synthetic LoRA weights that include lm_head (output embedding)
+to verify that Qwen3 properly supports LoRA on the unembed/lm_head layer.
+"""
+
+import json
+import os
+import tempfile
+
+import numpy as np
+import torch
+from safetensors.torch import save_file
+
+from vllm import LLM, SamplingParams
+from vllm.lora.request import LoRARequest
+
+MODEL_PATH = "Qwen/Qwen3-0.6B"
+HIDDEN_SIZE = 1024
+VOCAB_SIZE = 151936
+
+
+def create_qwen3_lora_with_lm_head(save_dir: str, rank: int = 8) -> None:
+    """Create synthetic Qwen3 LoRA weights with lm_head."""
+    lora_weights = {}
+    for module in ["q_proj", "v_proj"]:
+        lora_A = torch.from_numpy(
+            np.random.randn(rank, HIDDEN_SIZE).astype(np.float16) * 0.01
+        )
+        lora_B = torch.zeros(HIDDEN_SIZE, rank, dtype=torch.float16)
+        key_prefix = f"base_model.model.model.layers.0.self_attn.{module}"
+        lora_weights[f"{key_prefix}.lora_A.weight"] = lora_A
+        lora_weights[f"{key_prefix}.lora_B.weight"] = lora_B
+
+    # lm_head LoRA weights
+    lora_weights["base_model.model.lm_head.lora_A.weight"] = torch.from_numpy(
+        np.random.randn(rank, HIDDEN_SIZE).astype(np.float16) * 0.01
+    )
+    lora_weights["base_model.model.lm_head.lora_B.weight"] = torch.zeros(
+        VOCAB_SIZE, rank, dtype=torch.float16
+    )
+
+    adapter_config = {
+        "peft_type": "LORA",
+        "base_model_name_or_path": MODEL_PATH,
+        "task_type": "CAUSAL_LM",
+        "inference_mode": True,
+        "r": rank,
+        "lora_alpha": rank * 2,
+        "lora_dropout": 0.0,
+        "bias": "none",
+        "target_modules": ["q_proj", "v_proj", "lm_head"],
+    }
+
+    os.makedirs(save_dir, exist_ok=True)
+    with open(os.path.join(save_dir, "adapter_config.json"), "w") as f:
+        json.dump(adapter_config, f)
+    save_file(lora_weights, os.path.join(save_dir, "adapter_model.safetensors"))
+
+
+def test_qwen3_unembed_lora():
+    """Verify Qwen3 can load and generate with LoRA adapters with lm_head."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # Initialize engine first (before creating torch tensors)
+        llm = LLM(
+            model=MODEL_PATH,
+            enable_lora=True,
+            max_loras=4,
+            max_lora_rank=8,
+            max_model_len=128,
+            gpu_memory_utilization=0.8,
+            enforce_eager=True,
+        )
+
+        # Create LoRA weights after engine init
+        create_qwen3_lora_with_lm_head(tmpdir, rank=8)
+
+        lora_request = LoRARequest("lm_head_lora", 1, tmpdir)
+        llm.llm_engine.add_lora(lora_request)
+
+        assert 1 in llm.llm_engine.list_loras(), "lm_head LoRA should be loaded"
+
+        # Test generation
+        sampling_params = SamplingParams(temperature=0, max_tokens=32)
+        prompts = ["Hello, my name is"]
+
+        # Generate with base model (no LoRA)
+        base_outputs = llm.generate(prompts, sampling_params, use_tqdm=False)
+        assert len(base_outputs) == 1
+        assert len(base_outputs[0].outputs[0].text) > 0
+
+        # Generate with lm_head LoRA
+        lora_outputs = llm.generate(
+            prompts, sampling_params, lora_request=lora_request, use_tqdm=False
+        )
+        assert len(lora_outputs) == 1
+        assert len(lora_outputs[0].outputs[0].text) > 0

third_party/vllm/tests/lora/test_qwen3moe_tp.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+# NOTE To avoid overloading the CI pipeline, this test script will not
+# be triggered on CI and is primarily intended for local testing and verification.
+
+import vllm
+from vllm.lora.request import LoRARequest
+
+from ..utils import multi_gpu_test
+
+MODEL_PATH = "Qwen/Qwen3-30B-A3B"
+
+PROMPT_TEMPLATE = """<|im_start|>user
+I want you to act as a SQL terminal in front of an example database, you need only to return the sql command to me.Below is an instruction that describes a task, Write a response that appropriately completes the request.
+"
+##Instruction:
+candidate_poll contains tables such as candidate, people. Table candidate has columns such as Candidate_ID, People_ID, Poll_Source, Date, Support_rate, Consider_rate, Oppose_rate, Unsure_rate. Candidate_ID is the primary key.
+Table people has columns such as People_ID, Sex, Name, Date_of_Birth, Height, Weight. People_ID is the primary key.
+The People_ID of candidate is the foreign key of People_ID of people.
+
+
+###Input:
+{context}
+
+###Response:<|im_end|>
+<|im_start|>assistant"""  # noqa: E501
+
+EXPECTED_LORA_OUTPUT = [
+    "<think>\n\n</think>\n\nSELECT count(*) FROM candidate",
+    "<think>\n\n</think>\n\nSELECT count(*) FROM candidate",
+    "<think>\n\n</think>\n\nSELECT poll_source FROM candidate GROUP BY poll_source ORDER BY count(*) DESC LIMIT 1",  # noqa: E501
+    "<think>\n\n</think>\n\nSELECT poll_source FROM candidate GROUP BY poll_source ORDER BY count(*) DESC LIMIT 1",  # noqa: E501
+]
+
+
+def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None:
+    prompts = [
+        PROMPT_TEMPLATE.format(context="How many candidates are there?"),
+        PROMPT_TEMPLATE.format(context="Count the number of candidates."),
+        PROMPT_TEMPLATE.format(
+            context="Which poll resource provided the most number of candidate information?"  # noqa: E501
+        ),
+        PROMPT_TEMPLATE.format(
+            context="Return the poll resource associated with the most candidates."
+        ),
+    ]
+    sampling_params = vllm.SamplingParams(temperature=0, max_tokens=64)
+    outputs = llm.generate(
+        prompts,
+        sampling_params,
+        lora_request=LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None,
+    )
+    # Print the outputs.
+    generated_texts: list[str] = []
+    for output in outputs:
+        prompt = output.prompt
+        generated_text = output.outputs[0].text.strip()
+        generated_texts.append(generated_text)
+        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
+
+    for i in range(len(EXPECTED_LORA_OUTPUT)):
+        assert generated_texts[i].startswith(EXPECTED_LORA_OUTPUT[i])
+
+
+def test_qwen3moe_lora(qwen3moe_lora_files):
+    # We enable enforce_eager=True here to reduce VRAM usage for lora-test CI,
+    # Otherwise, the lora-test will fail due to CUDA OOM.
+    llm = vllm.LLM(
+        MODEL_PATH,
+        max_model_len=1024,
+        enable_lora=True,
+        max_loras=4,
+        enforce_eager=True,
+        trust_remote_code=True,
+        enable_chunked_prefill=True,
+    )
+
+    generate_and_test(llm, qwen3moe_lora_files, lora_id=1)
+    generate_and_test(llm, qwen3moe_lora_files, lora_id=2)
+
+
+@multi_gpu_test(num_gpus=2)
+def test_qwen3moe_lora_tp2(qwen3moe_lora_files):
+    llm = vllm.LLM(
+        MODEL_PATH,
+        max_model_len=1024,
+        enable_lora=True,
+        max_loras=4,
+        enforce_eager=True,
+        trust_remote_code=True,
+        enable_chunked_prefill=True,
+        tensor_parallel_size=2,
+    )
+
+    generate_and_test(llm, qwen3moe_lora_files, lora_id=1)
+    generate_and_test(llm, qwen3moe_lora_files, lora_id=2)
+
+
+@multi_gpu_test(num_gpus=4)
+def test_qwen3moe_lora_tp4(qwen3moe_lora_files):
+    llm = vllm.LLM(
+        MODEL_PATH,
+        max_model_len=1024,
+        enable_lora=True,
+        max_loras=4,
+        enforce_eager=True,
+        trust_remote_code=True,
+        enable_chunked_prefill=True,
+        tensor_parallel_size=4,
+    )
+
+    generate_and_test(llm, qwen3moe_lora_files, lora_id=1)
+    generate_and_test(llm, qwen3moe_lora_files, lora_id=2)

third_party/vllm/tests/lora/test_qwenvl.py

@@ -0,0 +1,312 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import dataclass
+
+from packaging.version import Version
+from transformers import __version__ as TRANSFORMERS_VERSION
+
+import vllm
+from vllm.assets.image import ImageAsset
+from vllm.lora.request import LoRARequest
+from vllm.sampling_params import BeamSearchParams
+
+
+@dataclass
+class TestConfig:
+    model_path: str
+    lora_path: str
+    max_num_seqs: int = 2
+    max_loras: int = 2
+    max_lora_rank: int = 32
+    enable_tower_connector_lora: bool = False
+    max_model_len: int = 8192
+    gpu_memory_utilization: float = 0.85
+    mm_processor_kwargs: dict[str, object] | None = None
+    mm_processor_cache_gb: float = 4
+
+    def __post_init__(self):
+        if self.mm_processor_kwargs is None:
+            # There is a bug in transformers v4 where size is ignored by
+            # `Qwen2VLProcessor.__call__`
+            if Version(TRANSFORMERS_VERSION) < Version("5.2.0"):
+                self.mm_processor_kwargs = {
+                    "min_pixels": 28 * 28,
+                    "max_pixels": 1280 * 28 * 28,
+                }
+            else:
+                self.mm_processor_kwargs = {
+                    "size": {
+                        "shortest_edge": 28 * 28,
+                        "longest_edge": 1280 * 28 * 28,
+                    }
+                }
+
+
+class Qwen2VLTester:
+    """Test helper for Qwen2 VL models with LoRA"""
+
+    PROMPT_TEMPLATE = (
+        "<|im_start|>system\nYou are a helpful assistant.<|im_end|>"
+        "\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
+        "What is in the image?<|im_end|>\n"
+        "<|im_start|>assistant\n"
+    )
+
+    def __init__(self, config: TestConfig):
+        self.config = config
+        self.llm = self._initialize_llm()
+
+    def _initialize_llm(self) -> vllm.LLM:
+        """Initialize the LLM with given configuration"""
+        return vllm.LLM(
+            model=self.config.model_path,
+            max_num_seqs=self.config.max_num_seqs,
+            enable_lora=True,
+            max_loras=self.config.max_loras,
+            max_lora_rank=self.config.max_lora_rank,
+            enable_tower_connector_lora=self.config.enable_tower_connector_lora,
+            trust_remote_code=True,
+            gpu_memory_utilization=self.config.gpu_memory_utilization,
+            mm_processor_kwargs=self.config.mm_processor_kwargs,
+            mm_processor_cache_gb=self.config.mm_processor_cache_gb,
+            max_model_len=self.config.max_model_len,
+        )
+
+    def run_test(
+        self,
+        images: list[ImageAsset],
+        expected_outputs: list[str],
+        lora_id: int | None = None,
+        lora_name: str | None = None,
+        temperature: float = 0,
+        max_tokens: int = 5,
+    ):
+        sampling_params = vllm.SamplingParams(
+            temperature=temperature,
+            max_tokens=max_tokens,
+        )
+        inputs = [
+            {
+                "prompt": self.PROMPT_TEMPLATE,
+                "multi_modal_data": {"image": asset.pil_image},
+            }
+            for asset in images
+        ]
+
+        lora_request = LoRARequest(
+            lora_name if lora_name else str(lora_id), lora_id, self.config.lora_path
+        )
+        outputs = self.llm.generate(inputs, sampling_params, lora_request=lora_request)
+        generated_texts = [output.outputs[0].text.strip() for output in outputs]
+        # Validate outputs
+        for generated, expected in zip(generated_texts, expected_outputs):
+            assert expected.startswith(generated), (
+                f"Generated text {generated} doesn't match expected pattern {expected}"
+            )
+
+    def run_beam_search_test(
+        self,
+        images: list[ImageAsset],
+        expected_outputs: list[list[str]],
+        lora_id: int | None = None,
+        temperature: float = 0,
+        beam_width: int = 2,
+        max_tokens: int = 5,
+    ):
+        beam_search_params = BeamSearchParams(
+            beam_width=beam_width, max_tokens=max_tokens, temperature=temperature
+        )
+
+        inputs = [
+            {
+                "prompt": self.PROMPT_TEMPLATE,
+                "multi_modal_data": {"image": asset.pil_image},
+            }
+            for asset in images
+        ]
+
+        lora_request = LoRARequest(str(lora_id), lora_id, self.config.lora_path)
+        outputs = self.llm.beam_search(
+            inputs, beam_search_params, lora_request=lora_request
+        )
+
+        for output_obj, expected_texts in zip(outputs, expected_outputs):
+            output_texts = [seq.text for seq in output_obj.sequences]
+
+            for output_text, expected_text in zip(output_texts, expected_texts):
+                # NOTE beam search .text contains the whole text including inputs
+                assert output_text.endswith(expected_text), (
+                    f"Generated {output_text} does not match expected {expected_text}"
+                )
+
+
+TEST_IMAGES = [
+    ImageAsset("stop_sign"),
+    ImageAsset("cherry_blossom"),
+]
+
+EXPECTED_OUTPUTS = [
+    "A red stop sign stands prominently in the foreground, with a traditional Chinese gate and a black SUV in the background, illustrating a blend of modern and cultural elements.",  # noqa: E501
+    "A majestic skyscraper stands tall, partially obscured by a vibrant canopy of cherry blossoms, against a clear blue sky.",  # noqa: E501
+]
+
+EXPECTED_OUTPUTS_LANGUAGE = [
+    "A stop sign is shown in an Asian city, with buildings and a car in the "
+    "background.",
+    "The Tokyo Skytree can be seen behind the pink blossoms of the cherry trees.",
+]
+
+EXPECTED_OUTPUTS_VISION = [
+    "A stop sign in front of oriental buildings.",
+    "A tree with pink flowers in front of it and a blue sky behind the flowers.",
+]
+
+EXPECTED_OUTPUTS_VISION_NO_CONNECTOR = [
+    "A stop sign is located on the street of a Chinese neighborhood.",
+    "A closeup shot of the Tokyo Skytree with pink flowers in the foreground.",
+]
+
+EXPECTED_BEAM_SEARCH_OUTPUTS = [
+    [
+        "A majestic skyscraper stands",
+        "A majestic tower stands tall",
+    ],
+]
+
+QWEN2VL_MODEL_PATH = "Qwen/Qwen2-VL-2B-Instruct"
+QWEN25VL_MODEL_PATH = "Qwen/Qwen2.5-VL-3B-Instruct"
+QWEN3VL_MODEL_PATH = "Qwen/Qwen3-VL-4B-Instruct"
+
+
+def test_qwen2vl_lora(qwen2vl_lora_files):
+    """Test Qwen 2.0 VL model with LoRA"""
+    config = TestConfig(model_path=QWEN2VL_MODEL_PATH, lora_path=qwen2vl_lora_files)
+    tester = Qwen2VLTester(config)
+
+    # Test with different LoRA IDs
+    for lora_id in [1, 2]:
+        tester.run_test(TEST_IMAGES, expected_outputs=EXPECTED_OUTPUTS, lora_id=lora_id)
+
+
+def test_qwen2vl_lora_beam_search(qwen2vl_lora_files):
+    """Test Qwen 2.0 VL model with LoRA through beam search."""
+    config = TestConfig(model_path=QWEN2VL_MODEL_PATH, lora_path=qwen2vl_lora_files)
+    tester = Qwen2VLTester(config)
+
+    # Test with different LoRA IDs
+    for lora_id in [1, 2]:
+        # NOTE currently, we only test cherry blossom since stop sign
+        # output is slightly different for v1; - the root cause is likely
+        # independent of the intent of this test, which is to ensure beam
+        # search passes through lora through correctly.
+        tester.run_beam_search_test(
+            [ImageAsset("cherry_blossom")],
+            expected_outputs=EXPECTED_BEAM_SEARCH_OUTPUTS,
+            lora_id=lora_id,
+        )
+
+
+def test_qwen25vl_lora(qwen25vl_lora_files):
+    """Test Qwen 2.5 VL model with LoRA"""
+    config = TestConfig(model_path=QWEN25VL_MODEL_PATH, lora_path=qwen25vl_lora_files)
+    tester = Qwen2VLTester(config)
+
+    # Test with different LoRA IDs
+    for lora_id in [1, 2]:
+        tester.run_test(TEST_IMAGES, expected_outputs=EXPECTED_OUTPUTS, lora_id=lora_id)
+
+
+def test_qwen25vl_vision_lora(qwen25vl_vision_lora_files):
+    config = TestConfig(
+        model_path=QWEN25VL_MODEL_PATH,
+        lora_path=qwen25vl_vision_lora_files,
+        # Currently, tower_connector_lora is incompatible with
+        # the multi-modal processor cache.
+        # TODO: Remove this restriction
+        mm_processor_cache_gb=0,
+        enable_tower_connector_lora=True,
+    )
+    tester = Qwen2VLTester(config)
+    for lora_id in [1, 2]:
+        tester.run_test(
+            TEST_IMAGES,
+            expected_outputs=EXPECTED_OUTPUTS,
+            lora_id=lora_id,
+        )
+
+
+def test_qwen3vl_vision_lora(qwen3vl_vision_lora_files):
+    config = TestConfig(
+        model_path=QWEN3VL_MODEL_PATH,
+        lora_path=qwen3vl_vision_lora_files,
+        # Currently, tower_connector_lora is incompatible with
+        # the multi-modal processor cache.
+        # TODO: Remove this restriction
+        mm_processor_cache_gb=0,
+        enable_tower_connector_lora=True,
+    )
+    tester = Qwen2VLTester(config)
+    for lora_id in [1, 2]:
+        tester.run_test(
+            TEST_IMAGES,
+            expected_outputs=EXPECTED_OUTPUTS,
+            lora_id=lora_id,
+        )
+
+
+def test_qwen2vl_multiple_lora_types(
+    qwen2vl_language_lora_files,
+    qwen2vl_vision_tower_connector_lora_files,
+    qwen2vl_vision_tower_lora_files,
+):
+    """
+    Test multiple LoRA adapter types (language, vision tower + connector,
+    vision tower only) using the same LLM instance to verify mm_encoder_cache
+    behavior with different LoRA requests.
+
+    By reusing the same LLM instance across different LoRA requests, we ensure that
+    the multimodal encoder cache correctly manages state transitions between
+    language-only and vision-enabled LoRA adapters.
+    """
+    config = TestConfig(
+        model_path=QWEN2VL_MODEL_PATH,
+        # We'll override the lora_path for each specific test, but need to provide
+        # an initial path for initialization
+        lora_path=qwen2vl_language_lora_files,
+        # Currently, tower_connector_lora is incompatible with
+        # the multi-modal processor cache.
+        # TODO: Remove this restriction
+        mm_processor_cache_gb=0,
+        enable_tower_connector_lora=True,
+    )
+    tester = Qwen2VLTester(config)
+
+    # Test 1: Language-only LoRA adapter
+    tester.config.lora_path = qwen2vl_language_lora_files
+    for lora_id in [1, 2]:
+        tester.run_test(
+            TEST_IMAGES,
+            expected_outputs=EXPECTED_OUTPUTS_LANGUAGE,
+            lora_id=lora_id,
+            lora_name="language_only",
+        )
+
+    # Test 2: Vision tower + connector LoRA adapter
+    tester.config.lora_path = qwen2vl_vision_tower_connector_lora_files
+    for lora_id in [3, 4]:
+        tester.run_test(
+            TEST_IMAGES,
+            expected_outputs=EXPECTED_OUTPUTS_VISION,
+            lora_id=lora_id,
+            lora_name="vision_tower_connector",
+        )
+
+    # Test 3: Vision tower only LoRA adapter (no connector)
+    tester.config.lora_path = qwen2vl_vision_tower_lora_files
+    for lora_id in [5, 6]:
+        tester.run_test(
+            TEST_IMAGES,
+            expected_outputs=EXPECTED_OUTPUTS_VISION_NO_CONNECTOR,
+            lora_id=lora_id,
+            lora_name="vision_tower",
+        )

third_party/vllm/tests/lora/test_resolver.py

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import pytest
+
+from vllm.lora.request import LoRARequest
+from vllm.lora.resolver import LoRAResolver, LoRAResolverRegistry
+
+
+class DummyLoRAResolver(LoRAResolver):
+    """A dummy LoRA resolver for testing."""
+
+    async def resolve_lora(
+        self, base_model_name: str, lora_name: str
+    ) -> LoRARequest | None:
+        if lora_name == "test_lora":
+            return LoRARequest(
+                lora_name=lora_name,
+                lora_path=f"/dummy/path/{base_model_name}/{lora_name}",
+                lora_int_id=abs(hash(lora_name)),
+            )
+        return None
+
+
+def test_resolver_registry_registration():
+    """Test basic resolver registration functionality."""
+    registry = LoRAResolverRegistry
+    resolver = DummyLoRAResolver()
+
+    # Register a new resolver
+    registry.register_resolver("dummy", resolver)
+    assert "dummy" in registry.get_supported_resolvers()
+
+    # Get registered resolver
+    retrieved_resolver = registry.get_resolver("dummy")
+    assert retrieved_resolver is resolver
+
+
+def test_resolver_registry_duplicate_registration():
+    """Test registering a resolver with an existing name."""
+    registry = LoRAResolverRegistry
+    resolver1 = DummyLoRAResolver()
+    resolver2 = DummyLoRAResolver()
+
+    registry.register_resolver("dummy", resolver1)
+    registry.register_resolver("dummy", resolver2)
+
+    assert registry.get_resolver("dummy") is resolver2
+
+
+def test_resolver_registry_unknown_resolver():
+    """Test getting a non-existent resolver."""
+    registry = LoRAResolverRegistry
+
+    with pytest.raises(KeyError, match="not found"):
+        registry.get_resolver("unknown_resolver")
+
+
+@pytest.mark.asyncio
+async def test_dummy_resolver_resolve():
+    """Test the dummy resolver's resolve functionality."""
+    dummy_resolver = DummyLoRAResolver()
+    base_model_name = "base_model_test"
+    lora_name = "test_lora"
+
+    # Test successful resolution
+    result = await dummy_resolver.resolve_lora(base_model_name, lora_name)
+    assert isinstance(result, LoRARequest)
+    assert result.lora_name == lora_name
+    assert result.lora_path == f"/dummy/path/{base_model_name}/{lora_name}"
+
+    # Test failed resolution
+    result = await dummy_resolver.resolve_lora(base_model_name, "nonexistent_lora")
+    assert result is None

third_party/vllm/tests/lora/test_transformers_model.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+import vllm
+from vllm.lora.request import LoRARequest
+from vllm.platforms import current_platform
+
+from ..utils import create_new_process_for_each_test, multi_gpu_test
+
+MODEL_PATH = "hmellor/Ilama-3.2-1B"
+
+PROMPT_TEMPLATE = """I want you to act as a SQL terminal in front of an example database, you need only to return the sql command to me.Below is an instruction that describes a task, Write a response that appropriately completes the request.\n"\n##Instruction:\nconcert_singer contains tables such as stadium, singer, concert, singer_in_concert. Table stadium has columns such as Stadium_ID, Location, Name, Capacity, Highest, Lowest, Average. Stadium_ID is the primary key.\nTable singer has columns such as Singer_ID, Name, Country, Song_Name, Song_release_year, Age, Is_male. Singer_ID is the primary key.\nTable concert has columns such as concert_ID, concert_Name, Theme, Stadium_ID, Year. concert_ID is the primary key.\nTable singer_in_concert has columns such as concert_ID, Singer_ID. concert_ID is the primary key.\nThe Stadium_ID of concert is the foreign key of Stadium_ID of stadium.\nThe Singer_ID of singer_in_concert is the foreign key of Singer_ID of singer.\nThe concert_ID of singer_in_concert is the foreign key of concert_ID of concert.\n\n###Input:\n{query}\n\n###Response:"""  # noqa: E501
+
+EXPECTED_LORA_OUTPUT = [
+    "SELECT count(*) FROM singer",
+    "SELECT avg(age) ,  min(age) ,  max(age) FROM singer WHERE country  =  'France'",  # noqa: E501
+    "SELECT DISTINCT Country FROM singer WHERE Age  >  20",
+]
+
+
+def do_sample(llm: vllm.LLM, lora_path: str, lora_id: int) -> list[str]:
+    prompts = [
+        PROMPT_TEMPLATE.format(query="How many singers do we have?"),
+        PROMPT_TEMPLATE.format(
+            query="What is the average, minimum, and maximum age of all singers from France?"  # noqa: E501
+        ),
+        PROMPT_TEMPLATE.format(
+            query="What are all distinct countries where singers above age 20 are from?"  # noqa: E501
+        ),
+    ]
+    sampling_params = vllm.SamplingParams(temperature=0, max_tokens=32)
+    outputs = llm.generate(
+        prompts,
+        sampling_params,
+        lora_request=LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None,
+    )
+    # Print the outputs.
+    generated_texts: list[str] = []
+    for output in outputs:
+        prompt = output.prompt
+        generated_text = output.outputs[0].text.strip()
+        generated_texts.append(generated_text)
+        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
+    return generated_texts
+
+
+def test_ilama_lora(ilama_lora_files):
+    llm = vllm.LLM(
+        MODEL_PATH,
+        max_model_len=1024,
+        enable_lora=True,
+        max_loras=4,
+        max_lora_rank=16,
+        trust_remote_code=True,
+        enable_chunked_prefill=True,
+    )
+
+    output1 = do_sample(llm, ilama_lora_files, lora_id=1)
+    for i in range(len(EXPECTED_LORA_OUTPUT)):
+        assert output1[i] == EXPECTED_LORA_OUTPUT[i]
+    output2 = do_sample(llm, ilama_lora_files, lora_id=2)
+    for i in range(len(EXPECTED_LORA_OUTPUT)):
+        assert output2[i] == EXPECTED_LORA_OUTPUT[i]
+
+
+@pytest.mark.skipif(
+    current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
+)
+@multi_gpu_test(num_gpus=4)
+@create_new_process_for_each_test()
+def test_ilama_lora_tp4(ilama_lora_files):
+    llm = vllm.LLM(
+        MODEL_PATH,
+        max_model_len=1024,
+        enable_lora=True,
+        max_loras=4,
+        max_lora_rank=16,
+        tensor_parallel_size=4,
+        trust_remote_code=True,
+        fully_sharded_loras=False,
+        enable_chunked_prefill=True,
+    )
+
+    output1 = do_sample(llm, ilama_lora_files, lora_id=1)
+    for i in range(len(EXPECTED_LORA_OUTPUT)):
+        assert output1[i] == EXPECTED_LORA_OUTPUT[i]
+    output2 = do_sample(llm, ilama_lora_files, lora_id=2)
+    for i in range(len(EXPECTED_LORA_OUTPUT)):
+        assert output2[i] == EXPECTED_LORA_OUTPUT[i]
+
+
+@pytest.mark.skipif(
+    current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
+)
+@multi_gpu_test(num_gpus=4)
+@create_new_process_for_each_test()
+def test_ilama_lora_tp4_fully_sharded_loras(ilama_lora_files):
+    llm = vllm.LLM(
+        MODEL_PATH,
+        max_model_len=1024,
+        enable_lora=True,
+        max_loras=4,
+        max_lora_rank=16,
+        tensor_parallel_size=4,
+        trust_remote_code=True,
+        fully_sharded_loras=True,
+        enable_chunked_prefill=True,
+    )
+    output1 = do_sample(llm, ilama_lora_files, lora_id=1)
+    for i in range(len(EXPECTED_LORA_OUTPUT)):
+        assert output1[i] == EXPECTED_LORA_OUTPUT[i]
+    output2 = do_sample(llm, ilama_lora_files, lora_id=2)
+    for i in range(len(EXPECTED_LORA_OUTPUT)):
+        assert output2[i] == EXPECTED_LORA_OUTPUT[i]

third_party/vllm/tests/lora/test_utils.py

@@ -0,0 +1,201 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections import OrderedDict
+from typing import NamedTuple
+from unittest.mock import MagicMock, patch
+
+import pytest
+from huggingface_hub.utils import HfHubHTTPError
+from torch import nn
+
+from vllm.lora.utils import (
+    get_adapter_absolute_path,
+    parse_fine_tuned_lora_name,
+    replace_submodule,
+)
+from vllm.model_executor.models.utils import WeightsMapper
+
+
+class LoRANameParserTestConfig(NamedTuple):
+    name: str
+    module_name: str
+    is_lora_a: bool
+    weights_mapper: WeightsMapper | None = None
+
+
+def test_parse_fine_tuned_lora_name_valid():
+    fixture = [
+        LoRANameParserTestConfig(
+            "base_model.model.lm_head.lora_A.weight", "lm_head", True, False
+        ),
+        LoRANameParserTestConfig(
+            "base_model.model.lm_head.lora_B.weight", "lm_head", False, False
+        ),
+        LoRANameParserTestConfig(
+            "base_model.model.model.embed_tokens.lora_embedding_A",
+            "model.embed_tokens",
+            True,
+        ),
+        LoRANameParserTestConfig(
+            "base_model.model.model.embed_tokens.lora_embedding_B",
+            "model.embed_tokens",
+            False,
+        ),
+        LoRANameParserTestConfig(
+            "base_model.model.model.layers.9.mlp.down_proj.lora_A.weight",
+            "model.layers.9.mlp.down_proj",
+            True,
+        ),
+        LoRANameParserTestConfig(
+            "base_model.model.model.layers.9.mlp.down_proj.lora_B.weight",
+            "model.layers.9.mlp.down_proj",
+            False,
+        ),
+        LoRANameParserTestConfig(
+            "language_model.layers.9.mlp.down_proj.lora_A.weight",
+            "language_model.layers.9.mlp.down_proj",
+            True,
+        ),
+        LoRANameParserTestConfig(
+            "language_model.layers.9.mlp.down_proj.lora_B.weight",
+            "language_model.layers.9.mlp.down_proj",
+            False,
+        ),
+        # Test with WeightsMapper
+        LoRANameParserTestConfig(
+            "base_model.model.model.layers.9.mlp.down_proj.lora_A.weight",
+            "language_model.model.layers.9.mlp.down_proj",
+            True,
+            weights_mapper=WeightsMapper(
+                orig_to_new_prefix={"model.": "language_model.model."}
+            ),
+        ),
+        LoRANameParserTestConfig(
+            "base_model.model.model.layers.9.mlp.down_proj.lora_B.weight",
+            "language_model.model.layers.9.mlp.down_proj",
+            False,
+            weights_mapper=WeightsMapper(
+                orig_to_new_prefix={"model.": "language_model.model."}
+            ),
+        ),
+        LoRANameParserTestConfig(
+            "model.layers.9.mlp.down_proj.lora_A.weight",
+            "language_model.model.layers.9.mlp.down_proj",
+            True,
+            weights_mapper=WeightsMapper(
+                orig_to_new_prefix={"model.": "language_model.model."}
+            ),
+        ),
+        LoRANameParserTestConfig(
+            "model.layers.9.mlp.down_proj.lora_B.weight",
+            "language_model.model.layers.9.mlp.down_proj",
+            False,
+            weights_mapper=WeightsMapper(
+                orig_to_new_prefix={"model.": "language_model.model."}
+            ),
+        ),
+    ]
+    for name, module_name, is_lora_a, weights_mapper in fixture:
+        assert (module_name, is_lora_a) == parse_fine_tuned_lora_name(
+            name, weights_mapper
+        )
+
+
+def test_parse_fine_tuned_lora_name_invalid():
+    fixture = {
+        "base_model.weight",
+        "base_model.model.weight",
+    }
+    for name in fixture:
+        with pytest.raises(ValueError, match="unsupported LoRA weight"):
+            parse_fine_tuned_lora_name(name)
+
+
+def test_replace_submodule():
+    model = nn.Sequential(
+        OrderedDict(
+            [
+                ("dense1", nn.Linear(764, 100)),
+                ("act1", nn.ReLU()),
+                ("dense2", nn.Linear(100, 50)),
+                (
+                    "seq1",
+                    nn.Sequential(
+                        OrderedDict(
+                            [
+                                ("dense1", nn.Linear(100, 10)),
+                                ("dense2", nn.Linear(10, 50)),
+                            ]
+                        )
+                    ),
+                ),
+                ("act2", nn.ReLU()),
+                ("output", nn.Linear(50, 10)),
+                ("outact", nn.Sigmoid()),
+            ]
+        )
+    )
+
+    sigmoid = nn.Sigmoid()
+
+    replace_submodule(model, "act1", sigmoid)
+    assert dict(model.named_modules())["act1"] == sigmoid
+
+    dense2 = nn.Linear(1, 5)
+    replace_submodule(model, "seq1.dense2", dense2)
+    assert dict(model.named_modules())["seq1.dense2"] == dense2
+
+
+# Unit tests for get_adapter_absolute_path
+@patch("os.path.isabs")
+def test_get_adapter_absolute_path_absolute(mock_isabs):
+    path = "/absolute/path/to/lora"
+    mock_isabs.return_value = True
+    assert get_adapter_absolute_path(path) == path
+
+
+@patch("os.path.expanduser")
+def test_get_adapter_absolute_path_expanduser(mock_expanduser):
+    # Path with ~ that needs to be expanded
+    path = "~/relative/path/to/lora"
+    absolute_path = "/home/user/relative/path/to/lora"
+    mock_expanduser.return_value = absolute_path
+    assert get_adapter_absolute_path(path) == absolute_path
+
+
+@patch("os.path.exists")
+@patch("os.path.abspath")
+def test_get_adapter_absolute_path_local_existing(mock_abspath, mock_exist):
+    # Relative path that exists locally
+    path = "relative/path/to/lora"
+    absolute_path = "/absolute/path/to/lora"
+    mock_exist.return_value = True
+    mock_abspath.return_value = absolute_path
+    assert get_adapter_absolute_path(path) == absolute_path
+
+
+@patch("huggingface_hub.snapshot_download")
+@patch("os.path.exists")
+def test_get_adapter_absolute_path_huggingface(mock_exist, mock_snapshot_download):
+    # Hugging Face model identifier
+    path = "org/repo"
+    absolute_path = "/mock/snapshot/path"
+    mock_exist.return_value = False
+    mock_snapshot_download.return_value = absolute_path
+    assert get_adapter_absolute_path(path) == absolute_path
+
+
+@patch("huggingface_hub.snapshot_download")
+@patch("os.path.exists")
+def test_get_adapter_absolute_path_huggingface_error(
+    mock_exist, mock_snapshot_download
+):
+    # Hugging Face model identifier with download error
+    path = "org/repo"
+    mock_exist.return_value = False
+    mock_snapshot_download.side_effect = HfHubHTTPError(
+        "failed to query model info",
+        response=MagicMock(),
+    )
+    assert get_adapter_absolute_path(path) == path

third_party/vllm/tests/lora/test_whisper.py

@@ -0,0 +1,153 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Integration tests for Whisper models with LoRA adapters.
+
+These tests verify that Whisper models can correctly load and use LoRA adapters
+for speech-to-text transcription tasks.
+"""
+
+import pytest
+
+import vllm
+from vllm.assets.audio import AudioAsset
+from vllm.lora.request import LoRARequest
+
+from ..utils import create_new_process_for_each_test
+
+# Model configuration
+WHISPER_MODEL = "openai/whisper-small"
+
+# Test prompts for Whisper transcription
+WHISPER_PROMPT = "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>"
+
+# Note: whisper_lora_files fixture is defined in conftest.py
+
+
+@pytest.fixture(autouse=True)
+def use_spawn_for_whisper(monkeypatch):
+    """Whisper has issues with forked workers, use spawn instead."""
+    monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+
+
+def create_whisper_llm(enable_lora: bool = True, max_loras: int = 2):
+    """Create a Whisper LLM instance with optional LoRA support."""
+    return vllm.LLM(
+        model=WHISPER_MODEL,
+        enable_lora=enable_lora,
+        max_loras=max_loras if enable_lora else 1,
+        max_lora_rank=64,
+        max_model_len=448,
+        dtype="half",
+        enforce_eager=True,  # For stability in tests
+    )
+
+
+def run_whisper_inference(
+    llm: vllm.LLM,
+    lora_path: str | None = None,
+    lora_id: int = 1,
+) -> list[str]:
+    """Run Whisper inference with optional LoRA adapter."""
+    # Load test audio
+    audio_asset = AudioAsset("mary_had_lamb")
+    audio_data = audio_asset.audio_and_sample_rate
+
+    inputs = [
+        {
+            "prompt": WHISPER_PROMPT,
+            "multi_modal_data": {"audio": audio_data},
+        }
+    ]
+
+    sampling_params = vllm.SamplingParams(
+        temperature=0,
+        max_tokens=200,
+    )
+
+    # Prepare LoRA request if adapter path is provided
+    lora_request = None
+    if lora_path:
+        lora_request = LoRARequest(
+            lora_name=f"whisper_lora_{lora_id}",
+            lora_int_id=lora_id,
+            lora_path=lora_path,
+        )
+
+    outputs = llm.generate(inputs, sampling_params, lora_request=lora_request)
+
+    return [output.outputs[0].text for output in outputs]
+
+
+@create_new_process_for_each_test()
+def test_whisper_lora_inference(whisper_lora_files):
+    """Test basic Whisper inference with a LoRA adapter.
+
+    This test verifies that:
+    1. Whisper model can be loaded with LoRA support enabled
+    2. A LoRA adapter can be applied during inference
+    3. The model produces valid transcription output
+    """
+    llm = create_whisper_llm(enable_lora=True)
+
+    # Run inference with LoRA
+    outputs = run_whisper_inference(llm, lora_path=whisper_lora_files, lora_id=1)
+
+    # Verify we got a non-empty transcription
+    assert len(outputs) == 1
+    assert len(outputs[0]) > 0, "Expected non-empty transcription output"
+
+    # The output should contain some recognizable words from the audio
+    # (Mary had a little lamb)
+    print(f"Transcription output: {outputs[0]}")
+
+
+@create_new_process_for_each_test()
+def test_whisper_multi_lora(whisper_lora_files):
+    """Test Whisper with multiple LoRA adapter IDs.
+
+    This test verifies that the same LoRA adapter can be loaded with
+    different IDs and produce consistent results.
+    """
+    llm = create_whisper_llm(enable_lora=True, max_loras=4)
+
+    # Test with different LoRA IDs using the same adapter
+    outputs_lora1 = run_whisper_inference(llm, lora_path=whisper_lora_files, lora_id=1)
+    outputs_lora2 = run_whisper_inference(llm, lora_path=whisper_lora_files, lora_id=2)
+
+    # Both should produce valid outputs
+    assert len(outputs_lora1[0]) > 0
+    assert len(outputs_lora2[0]) > 0
+
+    # Same adapter with different IDs should produce same output
+    assert outputs_lora1 == outputs_lora2, (
+        f"Expected same outputs for same adapter with different IDs. "
+        f"Got: {outputs_lora1} vs {outputs_lora2}"
+    )
+
+
+@create_new_process_for_each_test()
+def test_whisper_with_and_without_lora(whisper_lora_files):
+    """Test that Whisper produces different outputs with and without LoRA.
+
+    This test verifies that the LoRA adapter actually affects the model output.
+    """
+    llm = create_whisper_llm(enable_lora=True)
+
+    # Run with LoRA
+    outputs_with_lora = run_whisper_inference(
+        llm, lora_path=whisper_lora_files, lora_id=1
+    )
+
+    # Run without LoRA (base model only)
+    outputs_without_lora = run_whisper_inference(llm, lora_path=None)
+
+    # Both should produce valid outputs
+    assert len(outputs_with_lora[0]) > 0
+    assert len(outputs_without_lora[0]) > 0
+
+    print(f"Output with LoRA: {outputs_with_lora[0]}")
+    print(f"Output without LoRA: {outputs_without_lora[0]}")
+
+    # Note: Outputs may or may not differ depending on the adapter
+    # The main verification is that both configurations work

third_party/vllm/tests/lora/test_worker.py

@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+import random
+import tempfile
+from unittest.mock import patch
+
+from vllm.config import (
+    CacheConfig,
+    DeviceConfig,
+    ModelConfig,
+    ParallelConfig,
+    SchedulerConfig,
+    VllmConfig,
+    set_current_vllm_config,
+)
+from vllm.config.load import LoadConfig
+from vllm.config.lora import LoRAConfig
+from vllm.lora.model_manager import LoRAMapping
+from vllm.lora.request import LoRARequest
+from vllm.v1.worker.gpu_worker import Worker
+
+MODEL_PATH = "Qwen/Qwen3-0.6B"
+NUM_LORAS = 16
+
+
+@patch.dict(os.environ, {"RANK": "0"})
+def test_worker_apply_lora(qwen3_lora_files):
+    def set_active_loras(worker: Worker, lora_requests: list[LoRARequest]):
+        lora_mapping = LoRAMapping([], [])
+
+        worker.model_runner.lora_manager.set_active_adapters(
+            lora_requests, lora_mapping
+        )
+
+    model_config = ModelConfig(
+        MODEL_PATH,
+        seed=0,
+        dtype="float16",
+        max_model_len=127,
+        enforce_eager=True,
+    )
+
+    vllm_config = VllmConfig(
+        model_config=model_config,
+        load_config=LoadConfig(
+            download_dir=None,
+            load_format="dummy",
+        ),
+        parallel_config=ParallelConfig(
+            pipeline_parallel_size=1,
+            tensor_parallel_size=1,
+            data_parallel_size=1,
+        ),
+        scheduler_config=SchedulerConfig(
+            max_model_len=model_config.max_model_len,
+            is_encoder_decoder=model_config.is_encoder_decoder,
+            runner_type="generate",
+            max_num_batched_tokens=32,
+            max_num_seqs=32,
+            max_num_partial_prefills=32,
+        ),
+        device_config=DeviceConfig("cuda"),
+        cache_config=CacheConfig(
+            block_size=16,
+            cache_dtype="auto",
+        ),
+        lora_config=LoRAConfig(
+            max_lora_rank=8, max_cpu_loras=NUM_LORAS, max_loras=NUM_LORAS
+        ),
+    )
+    worker = Worker(
+        vllm_config=vllm_config,
+        local_rank=0,
+        rank=0,
+        distributed_init_method=f"file://{tempfile.mkstemp()[1]}",
+    )
+
+    with set_current_vllm_config(vllm_config):
+        worker.init_device()
+        worker.load_model()
+
+    set_active_loras(worker, [])
+    assert worker.list_loras() == set()
+
+    lora_requests = [
+        LoRARequest(str(i + 1), i + 1, qwen3_lora_files) for i in range(NUM_LORAS)
+    ]
+
+    set_active_loras(worker, lora_requests)
+    assert worker.list_loras() == {
+        lora_request.lora_int_id for lora_request in lora_requests
+    }
+
+    for i in range(NUM_LORAS):
+        random.seed(i)
+        iter_lora_requests = random.choices(
+            lora_requests, k=random.randint(1, NUM_LORAS)
+        )
+        random.shuffle(iter_lora_requests)
+        iter_lora_requests = iter_lora_requests[: -random.randint(0, NUM_LORAS)]
+        set_active_loras(worker, lora_requests)
+        assert worker.list_loras().issuperset(
+            {lora_request.lora_int_id for lora_request in iter_lora_requests}
+        )

third_party/vllm/tests/lora/utils.py

@@ -0,0 +1,407 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+import os
+from dataclasses import dataclass
+
+import torch
+from safetensors.torch import save_file
+
+from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights
+
+
+class DummyLoRAManager:
+    def __init__(self, device: torch.device = "cuda:0"):
+        super().__init__()
+        self._loras: dict[str, LoRALayerWeights] = {}
+        self._device = device
+
+    def set_module_lora(self, module_name: str, lora: LoRALayerWeights):
+        self._loras[module_name] = lora
+
+    def get_module_lora(self, module_name: str) -> LoRALayerWeights:
+        return self._loras[module_name]
+
+    def init_random_lora(
+        self,
+        module_name: str,
+        weight: torch.Tensor,
+        rank: int = 8,
+    ):
+        lora = LoRALayerWeights(
+            module_name,
+            rank=rank,
+            lora_alpha=1,
+            lora_a=torch.rand(
+                [rank, weight.shape[1]], dtype=weight.dtype, device=self._device
+            ),
+            lora_b=torch.rand(
+                [weight.shape[0], rank], dtype=weight.dtype, device=self._device
+            ),
+        )
+        self.set_module_lora(module_name, lora)
+
+        return lora
+
+    def init_lora(
+        self,
+        module_name: str,
+        input_dim: int,
+        output_dim: int,
+        rank=8,
+        noop=False,
+        embeddings_tensor=None,
+    ):
+        lora = LoRALayerWeights(
+            module_name,
+            rank=rank,
+            lora_alpha=1,
+            lora_a=torch.rand([rank, input_dim], device="cuda"),
+            lora_b=torch.rand([output_dim, input_dim], device="cuda"),
+            embeddings_tensor=embeddings_tensor,
+        )
+        self.set_module_lora(module_name, lora)
+        return lora
+
+    def reset_lora(self):
+        self._loras = {}
+
+    def init_packed_lora(
+        self,
+        module_name: str,
+        input_dim: int,
+        output_dims: list[int],
+        noop_lora_index: list[int] | None = None,
+        rank: int = 8,
+    ):
+        base_loras: list[LoRALayerWeights] = []
+        noop_lora_index_set = set(noop_lora_index or [])
+
+        for i, out_dim in enumerate(output_dims):
+            base_lora = self.init_lora(
+                module_name + "_000_" + str(i),
+                input_dim,
+                out_dim,
+                rank=rank,
+                noop=i in noop_lora_index_set,
+            )
+            base_loras.append(base_lora)
+        packed_lora = PackedLoRALayerWeights.pack(base_loras)
+        self.set_module_lora(module_name, packed_lora)
+        return packed_lora
+
+
+def assert_close(a, b):
+    rtol, atol = {
+        torch.float16: (6e-2, 6e-2),
+        torch.bfloat16: (6e-2, 6e-2),
+        torch.float32: (1e-2, 1e-2),
+    }[a.dtype]
+    torch.testing.assert_close(a, b, rtol=rtol, atol=atol)
+
+
+@dataclass
+class PunicaTensors:
+    inputs_tensor: torch.Tensor
+    lora_weights: torch.Tensor | list[torch.Tensor]
+    our_out_tensor: torch.Tensor
+    ref_out_tensor: torch.Tensor
+    b_seq_start_loc: torch.Tensor
+    prompt_lora_mapping: torch.Tensor
+    seq_len_tensor: torch.Tensor
+    token_lora_mapping: torch.Tensor
+
+    def meta(self) -> tuple[int, int]:
+        """
+        Infer max_seq_length and token_nums from the tensors
+        and return them.
+        """
+        max_seq_length = self.seq_len_tensor.max()
+        token_nums = self.seq_len_tensor.sum().item()
+        if isinstance(max_seq_length, tuple):
+            max_seq_length = max_seq_length[0].item()
+        else:
+            max_seq_length = max_seq_length.item()
+        return max_seq_length, token_nums
+
+
+def generate_data(
+    batches,
+    hidden_size,
+    lora_nums,
+    max_rank,
+    seq_length,
+    dtype,
+    op_type,
+    device,
+) -> PunicaTensors:
+    seq_len_tensor = torch.randint(seq_length, seq_length + 1, (batches,)).to(device)
+    b_seq_start_loc = torch.cumsum(
+        torch.tensor([0] + seq_len_tensor[:-1].tolist(), dtype=torch.long),
+        dim=0,
+    ).to(device)
+    total_tokens = seq_len_tensor.sum()
+    if op_type == "shrink":
+        inputs_tensor = torch.rand((total_tokens, hidden_size), dtype=dtype).to(device)
+        lora_weights = torch.rand(
+            (lora_nums, max_rank, hidden_size),  # col-major
+            dtype=dtype,
+        ).to(device)
+        # shrink op need atomic_add, so output is initinized by 0
+        ref_out_tensor = torch.zeros(
+            (total_tokens, max_rank), dtype=dtype, device=inputs_tensor.device
+        )
+        # NOTE  shrink kernel using torch.float32 as output type
+        our_out_tensor = torch.zeros((total_tokens, max_rank), dtype=torch.float32).to(
+            device
+        )
+    else:
+        inputs_tensor = torch.rand(
+            (total_tokens, max_rank),
+            dtype=dtype,
+        ).to(device)
+        lora_weights = torch.rand(
+            (lora_nums, hidden_size, max_rank),  # col-major
+            dtype=dtype,
+        ).to(device)
+        # expand op needs to complete y+=a@lora_b, so output is
+        # initinized randomly
+        ref_out_tensor = torch.rand(
+            (total_tokens, hidden_size),
+            dtype=dtype,
+        ).to(device)
+        # Ensure the same input.
+        our_out_tensor = ref_out_tensor.clone()
+    lora_indices_tensor = torch.randint(
+        0, lora_nums - 1 if lora_nums > 1 else 1, (batches,)
+    ).to(device)
+    indices = torch.zeros((total_tokens), dtype=torch.long).to(device)
+    current_offset = 0
+    for b_id in range(batches):
+        lora_index = lora_indices_tensor[b_id]
+        indices[current_offset : current_offset + seq_len_tensor[b_id]].copy_(
+            lora_index
+        )
+        current_offset += seq_len_tensor[b_id].item()
+
+    return PunicaTensors(
+        inputs_tensor,
+        lora_weights,
+        our_out_tensor,
+        ref_out_tensor,
+        b_seq_start_loc,
+        lora_indices_tensor,
+        seq_len_tensor,
+        indices,
+    )
+
+
+def generate_data_for_expand_nslices(
+    batches,
+    hidden_size,
+    lora_nums,
+    max_rank,
+    seq_length,
+    dtype,
+    nslices,
+    device,
+) -> PunicaTensors:
+    seq_len_tensor = torch.randint(seq_length, seq_length + 1, (batches,)).to(device)
+    b_seq_start_loc = torch.cumsum(
+        torch.tensor([0] + seq_len_tensor[:-1].tolist(), dtype=torch.long),
+        dim=0,
+    ).to(device)
+    total_tokens = seq_len_tensor.sum()
+    inputs_tensor = torch.rand(
+        (total_tokens, max_rank),
+        dtype=dtype,
+    ).to(device)
+    lora_weights_lst = []
+    for _ in range(nslices):
+        lora_weights_lst.append(
+            torch.rand(
+                (lora_nums, hidden_size, max_rank),  # col-major
+                dtype=dtype,
+            ).to(device)
+        )
+    # expand op needs to complete y+=a@lora_b, so output is
+    # initinized randomly
+    ref_out_tensor = torch.rand((total_tokens, hidden_size * nslices), dtype=dtype).to(
+        device
+    )
+    # Ensure the same input.
+    our_out_tensor = ref_out_tensor.clone()
+    lora_indices_tensor = torch.randint(
+        0, lora_nums - 1 if lora_nums > 1 else 1, (batches,)
+    )
+    indices = torch.zeros((total_tokens), dtype=torch.long).to(device)
+    current_offset = 0
+    for b_id in range(batches):
+        lora_index = lora_indices_tensor[b_id]
+        indices[current_offset : current_offset + seq_len_tensor[b_id]] = (
+            lora_index.item()
+        )
+        current_offset += seq_len_tensor[b_id].item()
+
+    lora_indices_tensor = lora_indices_tensor.to(device)
+    return PunicaTensors(
+        inputs_tensor,
+        lora_weights_lst,
+        our_out_tensor,
+        ref_out_tensor,
+        b_seq_start_loc,
+        lora_indices_tensor,
+        seq_len_tensor,
+        indices,
+    )
+
+
+def generate_data_for_nslices(
+    batches,
+    hidden_size,
+    lora_nums,
+    max_rank,
+    seq_length,
+    nslices,
+    dtype,
+    op_type,
+    device,
+) -> PunicaTensors:
+    seq_len_tensor = torch.randint(seq_length, seq_length + 1, (batches,)).to(device)
+    b_seq_start_loc = torch.cumsum(
+        torch.tensor([0] + seq_len_tensor[:-1].tolist(), dtype=torch.long),
+        dim=0,
+    ).to(device)
+    total_tokens = seq_len_tensor.sum()
+
+    lora_weights_lst = []
+    if op_type == "shrink":
+        inputs_tensor = torch.rand((total_tokens, hidden_size), dtype=dtype).to(device)
+
+        for _ in range(nslices):
+            if op_type == "shrink":
+                lora_weights_lst.append(
+                    torch.rand(
+                        (lora_nums, max_rank, hidden_size),  # col-major
+                        dtype=dtype,
+                    ).to(device)
+                )
+        # NOTE  shrink kernel using torch.float32 as output type
+        # shrink op need atomic_add, so output is initinized by 0
+        our_out_tensor = torch.zeros(
+            (nslices, total_tokens, max_rank),
+            dtype=torch.float32,
+        ).to(device)
+    else:
+        inputs_tensor = torch.rand(
+            (nslices, total_tokens, max_rank),
+            dtype=dtype,
+        ).to(device)
+        for _ in range(nslices):
+            lora_weights_lst.append(
+                torch.rand(
+                    (lora_nums, hidden_size, max_rank),  # col-major
+                    dtype=dtype,
+                ).to(device)
+            )
+        # expand op needs to complete y+=a@lora_b, so output is
+        # initinized randomly
+        our_out_tensor = torch.rand(
+            (total_tokens, hidden_size * nslices), dtype=dtype
+        ).to(device)
+
+    # Ensure the same input.
+    ref_out_tensor = our_out_tensor.clone()
+    lora_indices_tensor = torch.randint(
+        0, lora_nums - 1 if lora_nums > 1 else 1, (batches,)
+    )
+    indices = torch.zeros((total_tokens), dtype=torch.long).to(device)
+    current_offset = 0
+    for b_id in range(batches):
+        lora_index = lora_indices_tensor[b_id]
+        indices[current_offset : current_offset + seq_len_tensor[b_id]] = (
+            lora_index.item()
+        )
+        current_offset += seq_len_tensor[b_id].item()
+
+    lora_indices_tensor = lora_indices_tensor.to(device)
+    return PunicaTensors(
+        inputs_tensor,
+        lora_weights_lst,
+        our_out_tensor,
+        ref_out_tensor,
+        b_seq_start_loc,
+        lora_indices_tensor,
+        seq_len_tensor,
+        indices,
+    )
+
+
+def create_peft_lora(
+    model: torch.nn.Module,
+    save_dir: str,
+    target_modules: list[str],
+    rank: int = 8,
+    alpha: int = 16,
+    dropout: float = 0.1,
+    lora_dtype: torch.dtype = torch.float16,
+) -> dict[str, torch.Tensor]:
+    lora_weights = {}
+    adapter_config = {
+        "peft_type": "LORA",
+        "auto_mapping": None,
+        "base_model_name_or_path": "dummy_model",
+        "revision": None,
+        "task_type": "CAUSAL_LM",
+        "inference_mode": False,
+        "r": rank,
+        "lora_alpha": alpha,
+        "lora_dropout": dropout,
+        "fan_in_fan_out": False,
+        "bias": "none",
+        "modules_to_save": None,
+        "init_lora_weights": True,
+        "layers_to_transform": None,
+        "layers_pattern": None,
+        "target_modules": target_modules,
+        "exclude_modules": None,
+        "use_rslora": False,
+        "use_dora": False,
+        "loftq_config": None,
+    }
+
+    for module_name in target_modules:
+        module = model
+        for attr in module_name.split("."):
+            module = getattr(module, attr)
+
+        if hasattr(module, "input_size") and hasattr(module, "output_size"):
+            in_features = module.input_size
+            out_features = module.output_size
+
+        elif hasattr(module, "embedding_dim") and hasattr(module, "num_embeddings"):
+            # ParallelLMHead
+            in_features = module.embedding_dim
+            out_features = module.num_embeddings
+        else:
+            raise ValueError(f"Unable to determine dimensions for module {module_name}")
+
+        lora_A = torch.randn(rank, in_features, dtype=lora_dtype)
+
+        torch.nn.init.kaiming_uniform_(lora_A, a=5**0.5)
+
+        lora_B = torch.zeros(out_features, rank, dtype=lora_dtype)
+
+        # PEFT style
+        lora_weights[f"base_model.model.{module_name}.lora_A.weight"] = lora_A
+        lora_weights[f"base_model.model.{module_name}.lora_B.weight"] = lora_B
+
+    config_path = os.path.join(save_dir, "adapter_config.json")
+    with open(config_path, "w", encoding="utf-8") as f:
+        json.dump(adapter_config, f, indent=2, ensure_ascii=False)
+
+    weights_path = os.path.join(save_dir, "adapter_model.safetensors")
+    save_file(lora_weights, weights_path)
+
+    return lora_weights