merge v0.4.0

tests/worker/spec_decode/test_multi_step_worker.py → tests/spec_decode/test_multi_step_worker.py

-import torch
 import random
-import pytest
 from unittest.mock import MagicMock
 
-from vllm.worker.spec_decode.multi_step_worker import MultiStepWorker
-from vllm.worker.worker import Worker
+import pytest
+import torch
+
 from vllm.model_executor.utils import set_random_seed
+from vllm.sequence import SamplerOutput
+from vllm.spec_decode.multi_step_worker import (DraftModelTop1Proposer,
+                                                MultiStepWorker)
+from vllm.worker.worker import Worker
 
-from .utils import (create_execute_model_data, create_worker,
-                    create_seq_group_metadata_from_prompts, zero_kv_cache,
-                    patch_execute_model_with_seeds,
-                    assert_logprobs_dict_allclose)
+from .utils import (assert_logprobs_dict_allclose, create_batch,
+                    create_execute_model_data,
+                    create_seq_group_metadata_from_prompts, create_worker,
+                    patch_execute_model_with_seeds, zero_kv_cache)
 
 
 @pytest.mark.parametrize('num_steps', list(range(1, 17)))
@@ -90,8 +93,8 @@ def test_same_output_for_single_step():
         num_gpu_blocks,
         seed,
     )
-    multi_step_worker.model_runner = worker.model_runner
-    multi_step_worker.cache_engine = worker.cache_engine
+    # multi_step_worker.model_runner = worker.model_runner
+    # multi_step_worker.cache_engine = worker.cache_engine
 
     num_steps = 1
 
@@ -259,3 +262,160 @@ def test_same_output_for_multi_step():
             multi_step_output_logprobs, single_step_output_logprobs):
         assert_logprobs_dict_allclose(multi_step_logprobs,
                                       single_step_logprobs)
+
+
+@torch.inference_mode()
+def test_draft_proposals_full_speculation_len():
+    """Verify DraftModelTop1Proposer correctly handles case where all sequences
+    can speculate.
+    """
+    k = 10
+    batch_size = 32
+    vocab_size = 32_000
+    device = 'cuda:0'
+
+    draft_worker = MagicMock()
+    proposer = DraftModelTop1Proposer(
+        draft_worker=draft_worker,
+        device=device,
+        max_model_len=2048,
+        vocab_size=vocab_size,
+    )
+    draft_worker.execute_model_multi_step.return_value = [
+        SamplerOutput(
+            outputs=[],
+            sampled_token_probs=torch.rand(batch_size,
+                                           vocab_size,
+                                           device=device,
+                                           dtype=torch.float32),
+            sampled_token_ids=torch.randint(low=0,
+                                            high=vocab_size,
+                                            size=(batch_size, ),
+                                            device=device,
+                                            dtype=torch.long),
+        ) for _ in range(k)
+    ]
+
+    execute_model_data, _, _ = create_batch(batch_size, k)
+
+    proposals = proposer.get_proposals(
+        **execute_model_data.to_dict(),
+        max_proposal_len=k,
+    )
+
+    assert torch.is_tensor(proposals.proposal_token_ids)
+    assert torch.is_tensor(proposals.proposal_probs)
+
+    assert proposals.proposal_token_ids.shape == torch.Size([batch_size, k])
+    assert proposals.proposal_probs.shape[:-1] == torch.Size([batch_size, k])
+
+    assert proposals.proposal_lens.shape == torch.Size([batch_size])
+    assert proposals.proposal_lens.tolist() == [k for _ in range(batch_size)]
+
+
+@torch.inference_mode()
+def test_draft_proposals_no_speculations():
+    """Verify DraftModelTop1Proposer correctly handles case where no sequences
+    can speculate.
+    """
+    k = 10
+    batch_size = 32
+    vocab_size = 32_000
+    device = 'cuda:0'
+    prompt_len = 10
+
+    draft_worker = MagicMock()
+    proposer = DraftModelTop1Proposer(
+        draft_worker=draft_worker,
+        device=device,
+        max_model_len=prompt_len + k - 1,
+        vocab_size=vocab_size,
+    )
+
+    execute_model_data, _, _ = create_batch(batch_size,
+                                            k,
+                                            prompt_len=prompt_len)
+
+    proposals = proposer.get_proposals(
+        **execute_model_data.to_dict(),
+        max_proposal_len=k,
+    )
+
+    assert torch.is_tensor(proposals.proposal_token_ids)
+    assert torch.is_tensor(proposals.proposal_probs)
+
+    assert proposals.proposal_token_ids.shape == torch.Size([0, k])
+    assert proposals.proposal_probs.shape[:-1] == torch.Size([0, k])
+
+    assert proposals.proposal_lens.shape == torch.Size([batch_size])
+    assert proposals.proposal_lens.tolist() == [0 for _ in range(batch_size)]
+
+
+@torch.inference_mode()
+def test_draft_proposals_mixed_k():
+    """Verify DraftModelTop1Proposer correctly handles case some sequences can
+    speculate and some can't.
+    """
+    k = 10
+    batch_size = 32
+    vocab_size = 32_000
+    device = 'cuda:0'
+
+    small_prompt_len = 5
+    long_prompt_len = 10
+    prev_output_token_len = 20
+
+    expected_num_proposal_seqs = 6
+    expected_num_no_proposal_seqs = batch_size - expected_num_proposal_seqs
+
+    prompt_len = [
+        small_prompt_len for _ in range(expected_num_proposal_seqs - 1)
+    ] + [long_prompt_len
+         for _ in range(expected_num_no_proposal_seqs)] + [small_prompt_len]
+
+    draft_worker = MagicMock()
+    proposer = DraftModelTop1Proposer(
+        draft_worker=draft_worker,
+        device=device,
+        max_model_len=long_prompt_len + prev_output_token_len + k - 1,
+        vocab_size=vocab_size,
+    )
+
+    draft_worker.execute_model_multi_step.return_value = [
+        SamplerOutput(
+            outputs=[],
+            sampled_token_probs=torch.rand(expected_num_proposal_seqs,
+                                           vocab_size,
+                                           device=device,
+                                           dtype=torch.float32),
+            sampled_token_ids=torch.randint(
+                low=0,
+                high=vocab_size,
+                size=(expected_num_proposal_seqs, ),
+                device=device,
+                dtype=torch.long),
+        ) for _ in range(k)
+    ]
+
+    execute_model_data, _, _ = create_batch(
+        batch_size,
+        k,
+        prompt_len=prompt_len,
+        prev_output_token_len=prev_output_token_len,
+    )
+
+    proposals = proposer.get_proposals(
+        **execute_model_data.to_dict(),
+        max_proposal_len=k,
+    )
+
+    assert torch.is_tensor(proposals.proposal_token_ids)
+    assert torch.is_tensor(proposals.proposal_probs)
+
+    assert proposals.proposal_token_ids.shape == torch.Size([batch_size, k])
+    assert proposals.proposal_probs.shape[:-1] == torch.Size([batch_size, k])
+
+    assert proposals.proposal_lens.shape == torch.Size([batch_size])
+    assert proposals.proposal_lens.tolist() == [
+        k for _ in range(expected_num_proposal_seqs - 1)
+    ] + [0 for _ in range(expected_num_no_proposal_seqs)] + [k]

tests/spec_decode/test_spec_decode_worker.py

+import random
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.rejection_sampler import RejectionSampler
+from vllm.model_executor.utils import set_random_seed
+from vllm.spec_decode.interfaces import SpeculativeProposals
+from vllm.spec_decode.metrics import (AsyncMetricsCollector,
+                                      SpecDecodeWorkerMetrics)
+from vllm.spec_decode.multi_step_worker import MultiStepWorker
+from vllm.spec_decode.spec_decode_worker import (SpecDecodeWorker,
+                                                 split_num_cache_blocks_evenly)
+
+from .utils import (ExecuteModelData, create_batch, create_sampler_output_list,
+                    mock_worker)
+
+
+@pytest.mark.parametrize('k', [1, 2, 6])
+@pytest.mark.parametrize('batch_size', [1, 2, 32])
+@torch.inference_mode()
+def test_correctly_calls_draft_model(k: int, batch_size: int):
+    """Verify SpecDecodeWorker calls the draft worker with correct
+    inputs. Everything else is mocked out.
+    """
+    draft_worker = mock_worker(cls=MultiStepWorker)
+    target_worker = mock_worker()
+    rejection_sampler = MagicMock(spec=RejectionSampler)
+    metrics_collector = MagicMock(spec=AsyncMetricsCollector)
+    worker = SpecDecodeWorker(draft_worker, target_worker, rejection_sampler,
+                              metrics_collector)
+
+    exception_secret = 'artifical stop'
+    draft_worker.get_spec_proposals.side_effect = ValueError(exception_secret)
+
+    execute_model_data, _, _ = create_batch(batch_size, k)
+
+    with pytest.raises(ValueError, match=exception_secret):
+        worker.execute_model(**execute_model_data.to_dict(), num_spec_tokens=k)
+
+    call_args_list = draft_worker.get_spec_proposals.call_args_list
+    assert len(call_args_list) == 1
+
+    for args, _ in call_args_list:
+        (seq_group_metadata_list, blocks_to_swap_in, blocks_to_swap_out,
+         blocks_to_copy, actual_k) = args
+        actual_execute_model_data = ExecuteModelData(seq_group_metadata_list,
+                                                     blocks_to_swap_in,
+                                                     blocks_to_swap_out,
+                                                     blocks_to_copy)
+        assert actual_execute_model_data == execute_model_data
+        assert actual_k == k
+
+
+@pytest.mark.parametrize('k', [1, 2, 6])
+@pytest.mark.parametrize('batch_size', [1, 2, 32])
+@torch.inference_mode()
+def test_correctly_calls_target_model(k: int, batch_size: int):
+    """Verify SpecDecodeWorker calls the target model with correct
+    inputs. Everything else is mocked out.
+    """
+    draft_worker = mock_worker(cls=MultiStepWorker)
+    target_worker = mock_worker()
+    rejection_sampler = MagicMock(spec=RejectionSampler)
+    rejection_sampler.token_id_dtype = torch.int64
+    metrics_collector = MagicMock(spec=AsyncMetricsCollector)
+
+    draft_worker.device = 'cuda'
+    target_worker.device = 'cuda'
+
+    set_random_seed(1)
+
+    worker = SpecDecodeWorker(draft_worker, target_worker, rejection_sampler,
+                              metrics_collector)
+    worker.init_device()
+
+    vocab_size = 32_000
+
+    proposal_token_ids = torch.randint(low=0,
+                                       high=vocab_size,
+                                       size=(batch_size, k),
+                                       dtype=torch.int64,
+                                       device='cuda')
+    proposal_probs = torch.rand(batch_size,
+                                k,
+                                vocab_size,
+                                dtype=torch.float32,
+                                device='cuda')
+    proposal_lens = torch.ones(batch_size, dtype=torch.int64,
+                               device='cuda') * k
+
+    execute_model_data, prompts, prev_output_tokens = create_batch(
+        batch_size, k)
+
+    draft_worker.get_spec_proposals.return_value = SpeculativeProposals(
+        proposal_token_ids=proposal_token_ids,
+        proposal_probs=proposal_probs,
+        proposal_lens=proposal_lens)
+
+    exception_secret = 'artifical stop'
+    target_worker.execute_model.side_effect = ValueError(exception_secret)
+
+    with pytest.raises(ValueError, match=exception_secret):
+        worker.execute_model(**execute_model_data.to_dict(), num_spec_tokens=k)
+
+    seen_contexts = []
+
+    call_args_list = target_worker.execute_model.call_args_list
+    assert len(call_args_list) == 1
+    for args, kwargs in call_args_list:
+        target_execute_model_data = ExecuteModelData.from_dict(kwargs)
+
+        assert len(target_execute_model_data.seq_group_metadata_list) == (
+            k + 1) * batch_size
+        for seq_group_metadata in (
+                target_execute_model_data.seq_group_metadata_list):
+            for seq_data in seq_group_metadata.seq_data.values():
+                seen_contexts.append(seq_data.get_token_ids())
+
+    expected_seen_contexts = []
+
+    for prompt, prev_generated, draft_tokens in zip(
+            prompts, prev_output_tokens, proposal_token_ids.tolist()):
+
+        for i in range(len(draft_tokens) + 1):
+            expected_seen_contexts.append(prompt + prev_generated +
+                                          draft_tokens[:i])
+
+    seen_contexts.sort()
+    expected_seen_contexts.sort()
+    assert expected_seen_contexts == seen_contexts
+
+
+@pytest.mark.parametrize('k', [1, 2, 6])
+@pytest.mark.parametrize('batch_size', [1, 2, 32])
+@torch.inference_mode()
+def test_correctly_calls_rejection_sampler(k: int, batch_size: int):
+    """Verify SpecDecodeWorker calls the rejection sampler with
+    correct inputs. Everything else is mocked out.
+    """
+    vocab_size = 32_000
+
+    draft_worker = mock_worker(cls=MultiStepWorker, vocab_size=vocab_size)
+    target_worker = mock_worker(vocab_size=vocab_size)
+    rejection_sampler = MagicMock(spec=RejectionSampler)
+    rejection_sampler.token_id_dtype = torch.int64
+    metrics_collector = MagicMock(spec=AsyncMetricsCollector)
+    draft_worker.device = 'cuda'
+    target_worker.device = 'cuda'
+
+    set_random_seed(1)
+
+    worker = SpecDecodeWorker(draft_worker, target_worker, rejection_sampler,
+                              metrics_collector)
+    worker.init_device()
+
+    proposal_token_ids = torch.randint(low=0,
+                                       high=vocab_size,
+                                       size=(batch_size, k),
+                                       dtype=torch.int64,
+                                       device='cuda')
+    proposal_probs = torch.rand(batch_size,
+                                k,
+                                vocab_size,
+                                dtype=torch.float32,
+                                device='cuda')
+
+    proposal_lens = torch.ones(batch_size, dtype=torch.int64,
+                               device='cuda') * k
+
+    execute_model_data, _, _ = create_batch(batch_size, k)
+
+    draft_worker.get_spec_proposals.return_value = SpeculativeProposals(
+        proposal_token_ids=proposal_token_ids,
+        proposal_probs=proposal_probs,
+        proposal_lens=proposal_lens)
+
+    target_token_ids = torch.randint(low=0,
+                                     high=vocab_size,
+                                     size=(1, batch_size * (k + 1)),
+                                     dtype=torch.int64,
+                                     device='cuda')
+    target_token_probs = torch.rand(1,
+                                    batch_size * (k + 1),
+                                    vocab_size,
+                                    dtype=torch.float32,
+                                    device='cuda')
+    target_output = create_sampler_output_list(target_token_ids,
+                                               target_token_probs)
+
+    target_worker.execute_model.return_value = target_output[0]
+
+    exception_secret = 'artifical stop'
+    rejection_sampler.side_effect = ValueError(exception_secret)
+
+    with pytest.raises(ValueError, match=exception_secret):
+        worker.execute_model(**execute_model_data.to_dict(), num_spec_tokens=k)
+
+    assert len(rejection_sampler.call_args_list) == 1
+    args, _ = rejection_sampler.call_args_list[0]
+    (actual_proposal_scores, actual_bonus_token_ids, actual_proposal_probs,
+     actual_proposal_token_ids) = args
+
+    assert torch.equal(actual_bonus_token_ids,
+                       target_token_ids.reshape(batch_size, k + 1)[:, -1:])
+    assert torch.equal(
+        actual_proposal_scores,
+        target_token_probs.reshape(batch_size, k + 1, -1)[:, :-1])
+    assert torch.equal(actual_proposal_token_ids, proposal_token_ids)
+    assert torch.equal(actual_proposal_probs, proposal_probs)
+
+
+@pytest.mark.parametrize('k', [1, 2, 6])
+@pytest.mark.parametrize('batch_size', [1, 2, 32])
+@torch.inference_mode()
+def test_correctly_formats_output(k: int, batch_size: int):
+    """Verify SpecDecodeWorker formats sampler output correctly.
+    Everything else is mocked out.
+    """
+    vocab_size = 32_000
+
+    draft_worker = mock_worker(cls=MultiStepWorker, vocab_size=vocab_size)
+    target_worker = mock_worker(vocab_size=vocab_size)
+    rejection_sampler = MagicMock(spec=RejectionSampler)
+    rejection_sampler.token_id_dtype = torch.int64
+    metrics_collector = MagicMock(spec=AsyncMetricsCollector)
+    draft_worker.device = 'cuda'
+    target_worker.device = 'cuda'
+
+    set_random_seed(1)
+
+    worker = SpecDecodeWorker(draft_worker, target_worker, rejection_sampler,
+                              metrics_collector)
+    worker.init_device()
+
+    proposal_token_ids = torch.randint(low=0,
+                                       high=vocab_size,
+                                       size=(batch_size, k),
+                                       dtype=torch.int64,
+                                       device='cuda')
+    proposal_probs = torch.rand(batch_size,
+                                k,
+                                vocab_size,
+                                dtype=torch.float32,
+                                device='cuda')
+
+    proposal_lens = torch.ones(batch_size, dtype=torch.int64,
+                               device='cuda') * k
+
+    execute_model_data, _, _ = create_batch(batch_size, k)
+
+    draft_worker.get_spec_proposals.return_value = SpeculativeProposals(
+        proposal_token_ids=proposal_token_ids,
+        proposal_probs=proposal_probs,
+        proposal_lens=proposal_lens)
+
+    target_token_ids = torch.randint(low=0,
+                                     high=vocab_size,
+                                     size=(1, batch_size * (k + 1)),
+                                     dtype=torch.int64,
+                                     device='cuda')
+    target_token_probs = torch.rand(1,
+                                    batch_size * (k + 1),
+                                    vocab_size,
+                                    dtype=torch.float32,
+                                    device='cuda')
+    target_output = create_sampler_output_list(target_token_ids,
+                                               target_token_probs)
+
+    target_worker.execute_model.return_value = target_output[0]
+
+    rejection_sampler_output = torch.randint(low=0,
+                                             high=vocab_size,
+                                             size=(batch_size, k + 1),
+                                             dtype=torch.int64,
+                                             device='cuda')
+    for i in range(batch_size):
+        minimum_accepted_tokens = 1
+        rejection_sampler_output[i][
+            -random.randint(minimum_accepted_tokens, k + 1):] = -1
+
+    rejection_sampler.return_value = rejection_sampler_output
+
+    output = worker.execute_model(**execute_model_data.to_dict(),
+                                  num_spec_tokens=k)
+
+    expected_output = create_sampler_output_list(
+        rejection_sampler_output.transpose(0, 1), [None for _ in range(k + 1)])
+
+    seq_ids = [
+        next(iter(seq_group_metadata.seq_data.keys()))
+        for seq_group_metadata in execute_model_data.seq_group_metadata_list
+    ]
+    actual_output_by_seq = {seq_id: [] for seq_id in seq_ids}
+    expected_output_by_seq = {seq_id: [] for seq_id in seq_ids}
+
+    for step in output:
+        for seq_group in step:
+            for sample in seq_group.samples:
+                seq_id = sample.parent_seq_id
+                actual_output_by_seq[seq_id].append(sample)
+
+    for step in expected_output:
+        for seq_group in step:
+            for sample in seq_group.samples:
+                seq_id = sample.parent_seq_id
+                expected_output_by_seq[seq_id].append(sample)
+
+    all_seen_seq_ids = set(
+        list(actual_output_by_seq.keys()) +
+        list(expected_output_by_seq.keys()))
+    for seq_id in all_seen_seq_ids:
+        actual_by_step = actual_output_by_seq[seq_id]
+        expected_by_step = expected_output_by_seq[seq_id]
+
+        for i in range(k + 1):
+            if i >= len(actual_by_step):
+                assert expected_by_step[i].output_token == -1
+                continue
+            assert actual_by_step[i].output_token == expected_by_step[
+                i].output_token
+            assert actual_by_step[i].logprobs == expected_by_step[i].logprobs
+
+
+@pytest.mark.parametrize('k', [1, 2])
+@pytest.mark.parametrize('batch_size', [1])
+@pytest.mark.parametrize('returns_metrics', [True, False])
+@torch.inference_mode()
+def test_collects_metrics(k: int, batch_size: int, returns_metrics: bool):
+    """Verify SpecDecodeWorker collects metrics.
+    """
+    vocab_size = 32_000
+
+    draft_worker = mock_worker(cls=MultiStepWorker, vocab_size=vocab_size)
+    target_worker = mock_worker(vocab_size=vocab_size)
+    rejection_sampler = MagicMock(spec=RejectionSampler)
+    rejection_sampler.token_id_dtype = torch.int64
+    metrics_collector = MagicMock(spec=AsyncMetricsCollector)
+    draft_worker.device = 'cuda'
+    target_worker.device = 'cuda'
+
+    set_random_seed(1)
+
+    worker = SpecDecodeWorker(draft_worker, target_worker, rejection_sampler,
+                              metrics_collector)
+    worker.init_device()
+
+    proposal_token_ids = torch.randint(low=0,
+                                       high=vocab_size,
+                                       size=(batch_size, k),
+                                       dtype=torch.int64,
+                                       device='cuda')
+    proposal_probs = torch.rand(batch_size,
+                                k,
+                                vocab_size,
+                                dtype=torch.float32,
+                                device='cuda')
+
+    proposal_lens = torch.ones(batch_size, dtype=torch.int64,
+                               device='cuda') * k
+
+    execute_model_data, _, _ = create_batch(batch_size, k)
+
+    draft_worker.get_spec_proposals.return_value = SpeculativeProposals(
+        proposal_token_ids=proposal_token_ids,
+        proposal_probs=proposal_probs,
+        proposal_lens=proposal_lens)
+
+    target_token_ids = torch.randint(low=0,
+                                     high=vocab_size,
+                                     size=(1, batch_size * (k + 1)),
+                                     dtype=torch.int64,
+                                     device='cuda')
+    target_token_probs = torch.rand(1,
+                                    batch_size * (k + 1),
+                                    vocab_size,
+                                    dtype=torch.float32,
+                                    device='cuda')
+    target_output = create_sampler_output_list(target_token_ids,
+                                               target_token_probs)
+
+    target_worker.execute_model.return_value = target_output[0]
+
+    rejection_sampler_output = torch.randint(low=0,
+                                             high=vocab_size,
+                                             size=(batch_size, k + 1),
+                                             dtype=torch.int64,
+                                             device='cuda')
+    for i in range(batch_size):
+        minimum_accepted_tokens = 1
+        rejection_sampler_output[i][
+            -random.randint(minimum_accepted_tokens, k + 1):] = -1
+
+    rejection_sampler.return_value = rejection_sampler_output
+
+    mock_rejsample_metrics = MagicMock(
+        spec=SpecDecodeWorkerMetrics) if returns_metrics else None
+    metrics_collector.maybe_collect_rejsample_metrics.return_value = (
+        mock_rejsample_metrics)
+
+    output = worker.execute_model(**execute_model_data.to_dict(),
+                                  num_spec_tokens=k)
+    assert output[0].spec_decode_worker_metrics == mock_rejsample_metrics
+
+    call_args_list = (
+        metrics_collector.maybe_collect_rejsample_metrics.call_args_list)
+    assert len(call_args_list) == 1
+    args, kwargs = call_args_list[0]
+    assert args[0] == k or kwargs.get('k', -1) == k
+
+
+@pytest.mark.parametrize('k', [0])
+@pytest.mark.parametrize('batch_size', [1, 2, 32])
+@torch.inference_mode()
+def test_k_equals_zero(k: int, batch_size: int):
+    """Verify that the SpecDecodeWorker calls the draft and target workers
+    when k is zero. This happens during prefill.
+    """
+    draft_worker = mock_worker(cls=MultiStepWorker)
+    target_worker = mock_worker()
+    rejection_sampler = MagicMock(spec=RejectionSampler)
+    rejection_sampler.token_id_dtype = torch.int64
+    metrics_collector = MagicMock(spec=AsyncMetricsCollector)
+
+    draft_worker.device = 'cuda'
+    target_worker.device = 'cuda'
+
+    set_random_seed(1)
+
+    worker = SpecDecodeWorker(draft_worker, target_worker, rejection_sampler,
+                              metrics_collector)
+
+    execute_model_data, prompts, prev_output_tokens = create_batch(
+        batch_size, k, prev_output_token_len=0)
+
+    out = worker.execute_model(**execute_model_data.to_dict(),
+                               num_spec_tokens=k)
+
+    assert len(out) == 1, f"expected only one token output when {k=}"
+    assert out[0].probs is None, "expect gpu tensor references to be None"
+    assert out[
+        0].sampled_tokens is None, "expect gpu tensor references to be None"
+
+    draft_worker.execute_model.assert_called_once_with(
+        **execute_model_data.to_dict(), return_python_output=False)
+    target_worker.execute_model.assert_called_once_with(
+        **execute_model_data.to_dict())
+
+
+@pytest.mark.parametrize('k', [0, 5])
+@pytest.mark.parametrize('batch_size', [0])
+@torch.inference_mode()
+def test_empty_input_batch(k: int, batch_size: int):
+    """Verify that the SpecDecodeWorker calls the draft and target workers
+    when the input batch is empty. This can happen if the engine communicates
+    to the workers information without scheduling a batch.
+    """
+    draft_worker = mock_worker(cls=MultiStepWorker)
+    target_worker = mock_worker()
+    rejection_sampler = MagicMock(spec=RejectionSampler)
+    rejection_sampler.token_id_dtype = torch.int64
+    metrics_collector = MagicMock(spec=AsyncMetricsCollector)
+
+    draft_worker.device = 'cuda'
+    target_worker.device = 'cuda'
+
+    set_random_seed(1)
+
+    worker = SpecDecodeWorker(draft_worker, target_worker, rejection_sampler,
+                              metrics_collector)
+
+    execute_model_data, prompts, prev_output_tokens = create_batch(
+        batch_size, k, prev_output_token_len=0)
+
+    out = worker.execute_model(**execute_model_data.to_dict(),
+                               num_spec_tokens=k)
+
+    assert len(out) == 1, f"expected only one token output when {k=}"
+    assert out[0].probs is None, "expect gpu tensor references to be None"
+    assert out[
+        0].sampled_tokens is None, "expect gpu tensor references to be None"
+
+    draft_worker.execute_model.assert_called_once_with(
+        **execute_model_data.to_dict(), return_python_output=False)
+    target_worker.execute_model.assert_called_once_with(
+        **execute_model_data.to_dict())
+
+
+@torch.inference_mode()
+def test_init_device():
+    """Verify SpecDecodeWorker invokes proposer/scorer worker init_device, as
+    well as other GPU initialization.
+    """
+    draft_worker = mock_worker(cls=MultiStepWorker)
+    target_worker = mock_worker()
+    rejection_sampler = MagicMock(spec=RejectionSampler)
+    rejection_sampler.token_id_dtype = torch.int64
+    metrics_collector = MagicMock(spec=AsyncMetricsCollector)
+
+    worker = SpecDecodeWorker(draft_worker, target_worker, rejection_sampler,
+                              metrics_collector)
+
+    worker.init_device()
+
+    draft_worker.init_device.assert_called_once()
+
+    target_worker.init_device.assert_called_once()
+
+    metrics_collector.init_gpu_tensors.assert_called_once()
+    rejection_sampler.init_gpu_tensors.assert_called_once()
+
+
+@torch.inference_mode()
+def test_init_cache_engine():
+    """Verify SpecDecodeWorker invokes init_cache_engine on proposer/scorer
+    workers.
+    """
+    draft_worker = mock_worker(cls=MultiStepWorker)
+    target_worker = mock_worker()
+    rejection_sampler = MagicMock(spec=RejectionSampler)
+    rejection_sampler.token_id_dtype = torch.int64
+    metrics_collector = MagicMock(spec=AsyncMetricsCollector)
+
+    worker = SpecDecodeWorker(draft_worker, target_worker, rejection_sampler,
+                              metrics_collector)
+
+    cache_config = MagicMock()
+
+    worker.init_cache_engine(cache_config)
+
+    draft_worker.init_cache_engine.assert_called_once_with(cache_config)
+    target_worker.init_cache_engine.assert_called_once_with(cache_config)
+
+
+@pytest.mark.parametrize('available_gpu_blocks', [1, 1024])
+@pytest.mark.parametrize('available_cpu_blocks', [500])
+@pytest.mark.parametrize('target_cache_block_size_bytes', [2 * 2 * 4096])
+@pytest.mark.parametrize('draft_kv_size_bytes', [0, 2 * 2 * 768, 2 * 2 * 4096])
+@torch.inference_mode()
+def test_profile_num_available_blocks(available_gpu_blocks: int,
+                                      available_cpu_blocks: int,
+                                      target_cache_block_size_bytes: int,
+                                      draft_kv_size_bytes: int):
+    """Verify SpecDecodeWorker correctly profiles num available GPU blocks.
+    Specifically, it should run profiling in the scorer worker, and then evenly
+    split the blocks between proposer and scorer worker.
+    """
+    draft_worker = mock_worker(cls=MultiStepWorker)
+    target_worker = mock_worker()
+    rejection_sampler = MagicMock(spec=RejectionSampler)
+    rejection_sampler.token_id_dtype = torch.int64
+    metrics_collector = MagicMock(spec=AsyncMetricsCollector)
+
+    target_worker.profile_num_available_blocks.return_value = (
+        available_gpu_blocks, available_cpu_blocks)
+    target_worker.get_cache_block_size_bytes.return_value = (
+        target_cache_block_size_bytes)
+    draft_worker.get_cache_block_size_bytes.return_value = draft_kv_size_bytes
+
+    worker = SpecDecodeWorker(draft_worker, target_worker, rejection_sampler,
+                              metrics_collector)
+
+    # These values do not directly impact the adjusted block size calculation,
+    # so they can be fixed.
+    gpu_memory_utilization = 0.9
+    cpu_swap_space = 100
+    block_size = 16
+
+    num_gpu_blocks, num_cpu_blocks = worker.profile_num_available_blocks(
+        block_size, gpu_memory_utilization, cpu_swap_space, cache_dtype="auto")
+
+    target_worker.profile_num_available_blocks.assert_called_once_with(
+        block_size, gpu_memory_utilization, cpu_swap_space, "auto")
+    assert num_cpu_blocks == available_cpu_blocks
+
+    assert num_gpu_blocks == split_num_cache_blocks_evenly(
+        target_cache_block_size_bytes, draft_kv_size_bytes,
+        available_gpu_blocks)
+
+
+@pytest.mark.parametrize('available_gpu_blocks',
+                         list(range(20)) + [1024, 1024**2])
+@pytest.mark.parametrize('target_cache_block_size_bytes',
+                         [2 * 2 * 4096, 2 * 2 * 8192])
+@pytest.mark.parametrize('draft_kv_size_bytes', [0, 2 * 2 * 768, 2 * 2 * 4096])
+@torch.inference_mode()
+def test_split_num_cache_blocks_evenly(available_gpu_blocks: int,
+                                       target_cache_block_size_bytes: int,
+                                       draft_kv_size_bytes: int):
+    """Verify split_num_cache_blocks_evenly does not exceed original memory
+    allocation in bytes.
+    """
+    num_blocks = split_num_cache_blocks_evenly(target_cache_block_size_bytes,
+                                               draft_kv_size_bytes,
+                                               available_gpu_blocks)
+    assert (num_blocks * target_cache_block_size_bytes) + (
+        num_blocks * draft_kv_size_bytes) <= (available_gpu_blocks *
+                                              target_cache_block_size_bytes)

tests/spec_decode/test_utils.py

+from unittest.mock import MagicMock
+
+import pytest
+
+from vllm.sequence import SequenceGroupMetadata
+from vllm.spec_decode.util import get_all_seq_ids, split_batch_by_proposal_len
+
+
+def test_get_all_seq_ids():
+    """Verify get_all_seq_ids extracts all seq ids.
+    """
+    expected_seq_ids = list(range(10)) + list(range(100, 110))
+
+    seq_group_metadata_list = [
+        SequenceGroupMetadata(
+            request_id=str(seq_id),
+            is_prompt=True,
+            seq_data={
+                seq_id: MagicMock(),
+            },
+            sampling_params=MagicMock(),
+            block_tables={
+                seq_id: MagicMock(),
+            },
+            lora_request=None,
+        ) for seq_id in expected_seq_ids
+    ]
+
+    actual_seq_ids = get_all_seq_ids(seq_group_metadata_list)
+    assert actual_seq_ids == expected_seq_ids
+
+
+@pytest.fixture
+def fake_sequence_group_metadata():
+    seq_ids = list(range(3))
+    return [
+        SequenceGroupMetadata(
+            request_id=str(i),
+            is_prompt=True,
+            seq_data={
+                i: MagicMock(),
+            },
+            sampling_params=MagicMock(),
+            block_tables={
+                i: MagicMock(),
+            },
+            lora_request=None,
+        ) for i in seq_ids
+    ]
+
+
+def test_filter_zero_length_proposals(fake_sequence_group_metadata):
+    proposal_lens = [0, 1, 0]
+    filtered_groups, indices = split_batch_by_proposal_len(
+        fake_sequence_group_metadata,
+        proposal_lens,
+        select_proposal_len_zero=True)
+
+    expected_groups = [
+        fake_sequence_group_metadata[0], fake_sequence_group_metadata[2]
+    ]
+    expected_indices = [0, 2]
+
+    assert filtered_groups == expected_groups
+    assert indices == expected_indices
+
+
+def test_filter_non_zero_length_proposals(fake_sequence_group_metadata):
+    proposal_lens = [0, 1, 2]
+    filtered_groups, indices = split_batch_by_proposal_len(
+        fake_sequence_group_metadata,
+        proposal_lens,
+        select_proposal_len_zero=False)
+
+    expected_groups = [
+        fake_sequence_group_metadata[1], fake_sequence_group_metadata[2]
+    ]
+    expected_indices = [1, 2]
+
+    assert filtered_groups == expected_groups
+    assert indices == expected_indices
+
+
+def test_empty_inputs():
+    filtered_groups, indices = split_batch_by_proposal_len(
+        [], [], select_proposal_len_zero=True)
+
+    assert filtered_groups == []
+    assert indices == []
+
+
+def test_all_zero_with_non_zero_filter(fake_sequence_group_metadata):
+    proposal_lens = [0, 0, 0]
+    filtered_groups, indices = split_batch_by_proposal_len(
+        fake_sequence_group_metadata,
+        proposal_lens,
+        select_proposal_len_zero=False)
+
+    assert filtered_groups == []
+    assert indices == []
+
+
+def test_all_non_zero_with_zero_filter(fake_sequence_group_metadata):
+    proposal_lens = [1, 1, 1]
+    filtered_groups, indices = split_batch_by_proposal_len(
+        fake_sequence_group_metadata,
+        proposal_lens,
+        select_proposal_len_zero=True)
+
+    assert filtered_groups == []
+    assert indices == []

tests/worker/spec_decode/utils.py → tests/spec_decode/utils.py

+from dataclasses import dataclass, fields
+from itertools import count
+from typing import Dict, Iterable, List, Optional, Union
+from unittest.mock import MagicMock
+
 import torch
-from typing import List, Optional, Dict
 
-from vllm.worker.worker import Worker
-from vllm.utils import get_distributed_init_method, get_ip, get_open_port
 from vllm.engine.arg_utils import EngineArgs
-from vllm.sequence import SequenceGroupMetadata, SequenceData
+from vllm.model_executor.utils import set_random_seed
 from vllm.sampling_params import SamplingParams
+from vllm.sequence import (Logprob, SamplerOutput, SequenceData,
+                           SequenceGroupMetadata, SequenceGroupOutput,
+                           SequenceOutput)
+from vllm.utils import get_distributed_init_method, get_ip, get_open_port
 from vllm.worker.cache_engine import CacheEngine
-from vllm.model_executor.utils import set_random_seed
-from dataclasses import dataclass, fields
+from vllm.worker.worker import Worker
 
 
 @dataclass
@@ -24,6 +29,11 @@ class ExecuteModelData:
         return dict(
             (field.name, getattr(self, field.name)) for field in fields(self))
 
+    @classmethod
+    def from_dict(cls, d):
+        cleaned = dict((field.name, d[field.name]) for field in fields(cls))
+        return cls(**cleaned)
+
 
 def round_up_to_next_block(seq_len: int, block_size: int) -> int:
     return (seq_len + block_size - 1) // block_size
@@ -50,6 +60,21 @@ def create_execute_model_data(
     )
 
 
+def mock_worker(cls=None,
+                vocab_size: int = 30_000,
+                max_model_len: int = 2048,
+                rank: int = 0) -> MagicMock:
+    if cls is None:
+        cls = Worker
+
+    worker = MagicMock(spec=cls)
+    worker.vocab_size = vocab_size
+    worker.max_model_len = max_model_len
+    worker.rank = rank
+    worker.device = 'cuda:0'
+    return worker
+
+
 def patch_execute_model_with_seeds(worker: Worker, rand_seeds: List[int]):
     seed_iter = iter(rand_seeds)
     original_execute_model = worker.execute_model
@@ -84,7 +109,7 @@ def create_worker(cls: type,
     )
 
     (model_config, cache_config, parallel_config, scheduler_config,
-     device_config, _) = engine_args.create_engine_configs()
+     device_config, _, _) = engine_args.create_engine_configs()
 
     distributed_init_method = get_distributed_init_method(
         get_ip(), get_open_port())
@@ -100,7 +125,7 @@ def create_worker(cls: type,
         is_driver_worker=is_driver_worker,
     )
 
-    worker.init_model()
+    worker.init_device()
     worker.load_model()
 
     cache_config.num_gpu_blocks = num_gpu_blocks
@@ -117,25 +142,12 @@ def create_seq_group_metadata_from_prompts(
     block_size: int,
     final_seq_lens: List[int],
     continuations: Optional[List[List[int]]] = None,
-    num_tokens_processed: Optional[List[int]] = None,
     seq_ids: Optional[List[int]] = None,
 ) -> List[SequenceGroupMetadata]:
 
     if continuations is None:
         continuations = [[] for _ in prompts]
 
-    if num_tokens_processed is None:
-        # Default to 1 token missing from kv cache for generation sequences.
-        num_tokens_processed = []
-        for continuation, prompt in zip(continuations, prompts):
-            # If prefill, then default to zero tokens processed.
-            if not continuation:
-                num_tokens_processed.append(0)
-            else:
-                # If generation, then default to all but one tokens processed.
-                num_tokens_processed.append(
-                    len(continuation) + len(prompt) - 1)
-
     if seq_ids is None:
         seq_ids = list(i for i, _ in enumerate(prompts))
 
@@ -155,24 +167,93 @@ def create_seq_group_metadata_from_prompts(
             is_prompt=len(cont_token_ids) == 0,
             seq_data={
                 i:
-                SequenceData(prompt_token_ids=prompt_token_ids[:] +
-                             cont_token_ids[:])
+                SequenceData(
+                    prompt_token_ids=prompt_token_ids[:],
+                    output_token_ids=cont_token_ids[:],
+                ),
             },
             sampling_params=SamplingParams(temperature=0.0, ),
             block_tables={i: block_allocations[i][:]},
-        ) for i, (prompt_token_ids, cont_token_ids, num_tokens_saved) in
-        enumerate(zip(prompts, continuations, num_tokens_processed))
+        ) for i, (prompt_token_ids,
+                  cont_token_ids) in enumerate(zip(prompts, continuations))
     ]
 
 
 def assert_logprobs_dict_allclose(
-        actual_logprobs: List[Dict[int, float]],
-        expected_logprobs: List[Dict[int, float]]) -> None:
+        actual_logprobs: List[Dict[int, Logprob]],
+        expected_logprobs: List[Dict[int, Logprob]]) -> None:
     for single_step_actual_logprobs, single_step_expected_logprobs in zip(
             actual_logprobs, expected_logprobs):
         assert set(single_step_actual_logprobs.keys()) == set(
             single_step_expected_logprobs.keys())
         for token_id in single_step_actual_logprobs:
-            actual = torch.tensor(single_step_actual_logprobs[token_id])
-            expected = torch.tensor(single_step_expected_logprobs[token_id])
+            actual = torch.tensor(
+                single_step_actual_logprobs[token_id].logprob)
+            expected = torch.tensor(
+                single_step_expected_logprobs[token_id].logprob)
             assert torch.allclose(actual, expected)
+
+
+def create_sampler_output_list(
+        token_ids: torch.Tensor,
+        probs: Iterable[Optional[torch.Tensor]],
+        seq_ids: Optional[List[int]] = None) -> List[SamplerOutput]:
+    num_steps, batch_size = token_ids.shape
+    token_ids_by_step = token_ids.tolist()
+
+    if seq_ids is None:
+        seq_ids = list(range(batch_size))
+
+    return [
+        SamplerOutput(outputs=[
+            SequenceGroupOutput(
+                samples=[
+                    SequenceOutput(
+                        output_token=token_id,
+                        parent_seq_id=seq_ids[seq_index],
+                        logprobs={token_id: 0},
+                    )
+                ],
+                prompt_logprobs=None,
+            ) for seq_index, token_id in enumerate(token_ids_by_step[step])
+        ],
+                      sampled_token_probs=probs[step],
+                      sampled_token_ids=token_ids[step])
+        for step in range(num_steps)
+    ]
+
+
+def create_batch(batch_size,
+                 k,
+                 prompt_len: Union[int, List[int]] = 10,
+                 prev_output_token_len: int = 10,
+                 seq_ids: Optional[List[int]] = None,
+                 num_gpu_blocks: Optional[int] = None,
+                 block_size: Optional[int] = None):
+    if block_size is None:
+        block_size = 8
+
+    if num_gpu_blocks is None:
+        num_gpu_blocks = 2048 // block_size
+
+    iterator = count()
+
+    if isinstance(prompt_len, int):
+        prompt_lens = [prompt_len for _ in range(batch_size)]
+    else:
+        prompt_lens = prompt_len
+
+    prompts = [[next(iterator) for _ in range(p_len)] for p_len in prompt_lens]
+    prev_output_tokens = [[
+        next(iterator) for _ in range(prev_output_token_len)
+    ] for _ in range(batch_size)]
+    final_seq_lens = [
+        len(prompt) + len(prev_output_token) + k + 1
+        for prompt, prev_output_token in zip(prompts, prev_output_tokens)
+    ]
+
+    execute_model_data = create_execute_model_data(
+        create_seq_group_metadata_from_prompts(prompts, num_gpu_blocks,
+                                               block_size, final_seq_lens,
+                                               prev_output_tokens, seq_ids), )
+    return execute_model_data, prompts, prev_output_tokens

tests/test_cache_block_hashing.py

+"""Test hashing of cache blocks.
+
+Run `pytest tests/test_cache_block_hashing.py`.
+"""
+from typing import List, Optional
+
+import pytest
+
+from vllm.lora.request import LoRARequest
+from vllm.sequence import Sequence
+from vllm.transformers_utils.tokenizer_group import TokenizerGroup
+
+# Make two prefixes with different first blocks.
+prefix_start = [("You are an expert"), ("You are a")]
+prefix_common = (
+    " school principal, skilled in effectively managing "
+    "faculty and staff. Draft 10-15 questions for a potential first grade "
+    "Head Teacher for my K-12, all-girls', independent school that emphasizes "
+    "community, joyful discovery, and life-long learning. The candidate is "
+    "coming in for a first-round panel interview for a 8th grade Math "
+    "teaching role. They have 5 years of previous teaching experience "
+    "as an assistant teacher at a co-ed, public school with experience "
+    "in middle school math teaching. Based on this, fulfill "
+    "the following: ")
+prefixes = [start + prefix_common for start in prefix_start]
+
+# Sample prompts.
+sample_prompts = [
+    "Hello, my name is", "The president of the United States is",
+    "The capital of France is", "The future of AI is"
+]
+
+
+# Helper function.
+def flatten_2d(li):
+    return [lss for ls in li for lss in ls]
+
+
+@pytest.mark.parametrize("model", ["facebook/opt-125m"])
+@pytest.mark.parametrize("block_size", [16])
+@pytest.mark.parametrize("max_num_seqs", [256])
+@pytest.mark.parametrize("concurrent_lora_int_ids",
+                         [[None], [1], [None, 1], [None, 1, 2], [1, 2]])
+def test_auto_prefix_caching(model: str, block_size: int, max_num_seqs: int,
+                             concurrent_lora_int_ids: List[Optional[int]]):
+
+    tokenizer = TokenizerGroup(
+        tokenizer_id="facebook/opt-125m",
+        enable_lora=False,
+        max_num_seqs=max_num_seqs,
+        max_input_length=None,
+    )
+
+    hashes = []
+
+    for prefix in prefixes:
+        for lora_int_id in concurrent_lora_int_ids:
+            lora_request = None
+
+            if lora_int_id is not None:
+                lora_request = LoRARequest(
+                    f"example_lora_{lora_int_id}",
+                    lora_int_id,
+                    f"example/path/to/lora_{lora_int_id}",
+                )
+
+            hashes.append([])
+            prompts = [prefix + prompt for prompt in sample_prompts]
+            seq_id = 0
+            for prompt in prompts:
+                hashes[-1].append([])
+                prompt_token_ids = tokenizer.encode(prompt)
+                seq = Sequence(seq_id, prompt, prompt_token_ids, block_size,
+                               tokenizer.tokenizer.eos_token_id, lora_request)
+
+                num_blocks = len(prompt_token_ids) // block_size
+                for idx in range(num_blocks):
+                    hashes[-1][-1].append(seq.hash_of_block(idx))
+
+                seq_id += 1
+
+    # Check that hashes made with two prefixes with different first blocks are
+    # different everywhere.
+    for hash0, hash1 in zip(flatten_2d(hashes[0]), flatten_2d(hashes[1])):
+        assert (hash0 != hash1)
+
+    # Check that hashes of different prompts made with the same prefix are the
+    # same until the hashes that contain the prompt.
+    for hash_pref in hashes:
+        same_hashes = [tuple(h[:-1]) for h in hash_pref]
+        different_hashes = [h[-1] for h in hash_pref]
+        assert (len(set(same_hashes)) == 1)
+        assert (len(set(different_hashes)) == len(different_hashes))

tests/test_config.py

+from vllm.config import ModelConfig
+
+
+def test_get_sliding_window():
+    TEST_SLIDING_WINDOW = 4096
+    # Test that the sliding window is correctly computed.
+    # For Qwen1.5/Qwen2, get_sliding_window() should be None
+    # when use_sliding_window is False.
+    qwen2_model_config = ModelConfig(
+        "Qwen/Qwen1.5-7B",
+        "Qwen/Qwen1.5-7B",
+        tokenizer_mode="auto",
+        trust_remote_code=False,
+        download_dir=None,
+        load_format="dummy",
+        seed=0,
+        dtype="float16",
+        revision=None,
+    )
+
+    qwen2_model_config.hf_config.use_sliding_window = False
+    qwen2_model_config.hf_config.sliding_window = TEST_SLIDING_WINDOW
+    assert qwen2_model_config.get_sliding_window() is None
+
+    qwen2_model_config.hf_config.use_sliding_window = True
+    assert qwen2_model_config.get_sliding_window() == TEST_SLIDING_WINDOW
+
+    mistral_model_config = ModelConfig(
+        "mistralai/Mistral-7B-v0.1",
+        "mistralai/Mistral-7B-v0.1",
+        tokenizer_mode="auto",
+        trust_remote_code=False,
+        download_dir=None,
+        load_format="dummy",
+        seed=0,
+        dtype="float16",
+        revision=None,
+    )
+    mistral_model_config.hf_config.sliding_window = None
+    assert mistral_model_config.get_sliding_window() is None
+
+    mistral_model_config.hf_config.sliding_window = TEST_SLIDING_WINDOW
+    assert mistral_model_config.get_sliding_window() == TEST_SLIDING_WINDOW
\ No newline at end of file

tests/test_logits_processor.py

+import random
+from typing import Tuple
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.utils import set_random_seed
+from vllm.sequence import SamplingParams, SequenceData, SequenceGroupMetadata
+from vllm.worker.model_runner import ModelRunner
+
+
+class MockLogitsProcessor(LogitsProcessor):
+
+    def __init__(self, vocab_size: int, scale: float,
+                 fake_logits: torch.Tensor):
+        super().__init__(vocab_size=vocab_size, scale=scale)
+        self.fake_logits = fake_logits.clone()
+
+    def forward(self, *args, **kwargs):
+        with patch(
+                "vllm.model_executor.layers.logits_processor._prune_hidden_states",
+                lambda x, y: x
+        ), patch(
+                "vllm.model_executor.layers.logits_processor.LogitsProcessor._get_logits",
+                lambda *args, **kwargs: self.fake_logits):
+            return super().forward(*args, **kwargs)
+
+
+def _prepare_test(
+    batch_size: int
+) -> Tuple[torch.Tensor, torch.Tensor, MockLogitsProcessor, ModelRunner]:
+    vocab_size = 32000
+    input_tensor = torch.rand((batch_size, 1024), dtype=torch.float16)
+    fake_logits = torch.full((batch_size, vocab_size),
+                             1e-2,
+                             dtype=input_tensor.dtype)
+    logits_processor = MockLogitsProcessor(32000, 0.5, fake_logits)
+    model_runner = ModelRunner(None, None, None, None, None)
+    return input_tensor, fake_logits, logits_processor, model_runner
+
+
+RANDOM_SEEDS = list(range(128))
+CUDA_DEVICES = [
+    f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)
+]
+
+
+@pytest.mark.parametrize("seed", RANDOM_SEEDS)
+@pytest.mark.parametrize("device", CUDA_DEVICES)
+def test_logits_processors(seed: int, device: str):
+    set_random_seed(seed)
+    torch.set_default_device(device)
+    batch_size = random.randint(1, 256)
+    input_tensor, fake_logits, logits_processor, model_runner = _prepare_test(
+        batch_size)
+
+    # This sample logits processor gives infinite score to the i-th token,
+    # where i is the length of the input sequence.
+    # We therefore expect the output token sequence to be [0, 1, 2, ...]
+    def pick_ith(token_ids, logits):
+        logits[len(token_ids)] = float("inf")
+        return logits
+
+    seq_group_metadata_list = []
+    prompt_lens = []
+    for i in range(batch_size):
+        seq_group_metadata_list.append(
+            SequenceGroupMetadata(
+                request_id=f"test_{i}",
+                is_prompt=True,
+                seq_data={0: SequenceData([1, 2, 3])},
+                sampling_params=SamplingParams(temperature=0,
+                                               logits_processors=[pick_ith]),
+                block_tables={0: [1]},
+            ))
+        prompt_lens.append(seq_group_metadata_list[-1].seq_data[0].get_len())
+
+    sampling_metadata = model_runner._prepare_sample(seq_group_metadata_list,
+                                                     prompt_lens,
+                                                     subquery_lens=prompt_lens)
+    logits_processor_output = logits_processor(
+        embedding=None,
+        hidden_states=input_tensor,
+        sampling_metadata=sampling_metadata)
+
+    assert torch.isinf(logits_processor_output[:, 0]).all()
+
+    fake_logits *= logits_processor.scale
+    assert torch.allclose(logits_processor_output[:, 1], fake_logits[:, 1],
+                          1e-4)
+
+    del model_runner

tests/test_sequence.py

+import pytest
+
+from vllm.sequence import (SamplerOutput, SequenceData, SequenceGroupOutput,
+                           SequenceOutput)
+
+
+@pytest.fixture
+def sample_outputs():
+    return [
+        SequenceGroupOutput(samples=[
+            SequenceOutput(parent_seq_id=0, output_token=i, logprobs={})
+        ],
+                            prompt_logprobs=None) for i in range(5)
+    ]
+
+
+@pytest.fixture
+def sampler_output(sample_outputs):
+    return SamplerOutput(outputs=sample_outputs)
+
+
+def test_sampler_output_initialization(sampler_output, sample_outputs):
+    assert len(sampler_output) == len(sample_outputs)
+    assert sampler_output.sampled_token_probs is None
+    assert sampler_output.sampled_token_ids is None
+    assert sampler_output.spec_decode_worker_metrics is None
+
+
+def test_sampler_output_getitem(sampler_output, sample_outputs):
+    assert sampler_output[2] == sample_outputs[2]
+
+
+def test_sampler_output_setitem(sampler_output):
+    new_output = SequenceGroupOutput(samples=[
+        SequenceOutput(parent_seq_id=0, output_token=99, logprobs={})
+    ],
+                                     prompt_logprobs=None)
+    sampler_output[2] = new_output
+    assert sampler_output[2] == new_output
+
+
+def test_sampler_output_len(sampler_output, sample_outputs):
+    assert len(sampler_output) == len(sample_outputs)
+
+
+def test_sampler_output_eq(sample_outputs):
+    sampler_output1 = SamplerOutput(outputs=sample_outputs)
+    sampler_output2 = SamplerOutput(outputs=sample_outputs.copy())
+    sampler_output3 = SamplerOutput(outputs=sample_outputs[:-1])
+    assert sampler_output1 == sampler_output2
+    assert sampler_output1 != sampler_output3
+
+
+def test_sequence_data_prefill():
+    seq_data = SequenceData(prompt_token_ids=[1, 2, 3, 4])
+    assert seq_data.get_num_uncomputed_tokens() == 4
+    assert seq_data.get_num_computed_tokens() == 0
+    # advance by 2
+    seq_data.update_num_computed_tokens(2)
+    assert seq_data.get_num_uncomputed_tokens() == 2
+    assert seq_data.get_num_computed_tokens() == 2
+
+    # advance by 1
+    seq_data.update_num_computed_tokens(1)
+    assert seq_data.get_num_uncomputed_tokens() == 1
+    assert seq_data.get_num_computed_tokens() == 3
+
+    # append tokens and reset, simulating recompute
+    seq_data.append_token_id(1, logprob=0.0)
+    seq_data.reset_num_computed_tokens()
+    assert seq_data.get_num_uncomputed_tokens() == 5
+    assert seq_data.get_num_computed_tokens() == 0

tests/tokenization/test_cached_tokenizer.py

+from copy import deepcopy
+
+from transformers import AutoTokenizer
+
+from vllm.transformers_utils.tokenizer import get_cached_tokenizer
+
+
+def test_cached_tokenizer():
+    reference_tokenizer = AutoTokenizer.from_pretrained("gpt2")
+    reference_tokenizer.add_special_tokens({"cls_token": "<CLS>"})
+    reference_tokenizer.add_special_tokens(
+        {"additional_special_tokens": ["<SEP>"]})
+    cached_tokenizer = get_cached_tokenizer(deepcopy(reference_tokenizer))
+
+    assert reference_tokenizer.encode("prompt") == cached_tokenizer.encode(
+        "prompt")
+    assert set(reference_tokenizer.all_special_ids) == set(
+        cached_tokenizer.all_special_ids)
+    assert set(reference_tokenizer.all_special_tokens) == set(
+        cached_tokenizer.all_special_tokens)
+    assert set(reference_tokenizer.all_special_tokens_extended) == set(
+        cached_tokenizer.all_special_tokens_extended)

tests/tokenization/test_detokenize.py

+from typing import Dict, List
+
+import pytest
+from transformers import AutoTokenizer
+
+from vllm.sequence import Logprob, SamplingParams, Sequence, SequenceGroup
+from vllm.transformers_utils.detokenizer import Detokenizer
+from vllm.transformers_utils.tokenizer import detokenize_incrementally
+from vllm.transformers_utils.tokenizer_group import get_tokenizer_group
+
+TRUTH = [
+    "Hello here, this is a simple test",
+    "vLLM is a high-throughput and memory-efficient inference and serving engine for LLMs. It is designed to be used in production environments, where inference and serving",  # noqa
+    "我很感谢你的热情"
+]
+TOKENIZERS = [
+    "facebook/opt-125m",
+    "gpt2",
+    "bigcode/tiny_starcoder_py",
+    "EleutherAI/gpt-j-6b",
+    "EleutherAI/pythia-70m",
+    "bigscience/bloom-560m",
+    "mosaicml/mpt-7b",
+    "tiiuae/falcon-7b",
+    "meta-llama/Llama-2-7b-hf",
+    "codellama/CodeLlama-7b-hf",
+]
+
+
+def _run_incremental_decode(tokenizer, all_input_ids,
+                            skip_special_tokens: bool, starting_index: int):
+    decoded_text = ""
+    offset = 0
+    token_offset = 0
+    prev_tokens = None
+    for i in range(starting_index, len(all_input_ids)):
+        new_tokens, text, offset, token_offset = detokenize_incrementally(
+            tokenizer,
+            all_input_ids[:i + 1],
+            prev_tokens,
+            offset,
+            token_offset,
+            skip_special_tokens=skip_special_tokens)
+        decoded_text += text
+        if prev_tokens is None:
+            prev_tokens = new_tokens
+        else:
+            prev_tokens += new_tokens
+    return decoded_text
+
+
+@pytest.mark.parametrize("truth", TRUTH)
+@pytest.mark.parametrize("with_prompt", [True, False])
+@pytest.mark.parametrize("tokenizer_id", TOKENIZERS)
+@pytest.mark.parametrize("skip_special_tokens", (True, False))
+def test_decode_streaming(tokenizer_id, truth, with_prompt,
+                          skip_special_tokens):
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
+    if with_prompt:
+        truth_tokens = tokenizer(truth, add_special_tokens=False)["input_ids"]
+        prompt_input_ids = truth_tokens[:len(truth) // 2]
+        generated_input_ids = truth_tokens[len(truth) // 2:]
+        all_input_ids = prompt_input_ids + generated_input_ids
+        starting_index = len(prompt_input_ids)
+        prompt = tokenizer.decode(prompt_input_ids,
+                                  skip_special_tokens=skip_special_tokens)
+        generated = truth[len(prompt):]
+    else:
+        generated = truth
+        starting_index = 0
+        all_input_ids = tokenizer(truth, add_special_tokens=False)["input_ids"]
+    if skip_special_tokens:
+        if tokenizer.bos_token_id is not None:
+            all_input_ids = [tokenizer.bos_token_id] + all_input_ids
+            starting_index += 1
+        all_input_ids = all_input_ids + [tokenizer.eos_token_id]
+
+    decoded_text = _run_incremental_decode(
+        tokenizer,
+        all_input_ids,
+        skip_special_tokens=skip_special_tokens,
+        starting_index=starting_index)
+
+    assert decoded_text == generated
+
+    decoded_text = _run_incremental_decode(
+        tokenizer, [len(tokenizer)],
+        skip_special_tokens=skip_special_tokens,
+        starting_index=starting_index)
+
+    assert decoded_text == ''
+
+
+@pytest.fixture
+def detokenizer(tokenizer_name: str) -> Detokenizer:
+    init_kwargs = dict(
+        tokenizer_id=tokenizer_name,
+        enable_lora=False,
+        max_num_seqs=100,
+        max_input_length=None,
+        tokenizer_mode="auto",
+        trust_remote_code=False,
+        revision=None,
+    )
+
+    tokenizer_group = get_tokenizer_group(
+        None,
+        **init_kwargs,
+    )
+
+    return Detokenizer(tokenizer_group)
+
+
+@pytest.fixture(name="complete_sequence_token_ids")
+def create_complete_sequence_token_ids(complete_sequence: str,
+                                       tokenizer_name: str) -> List[int]:
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+    complete_sequence_token_ids = tokenizer(complete_sequence)["input_ids"]
+    return complete_sequence_token_ids
+
+
+def create_sequence(prompt_token_ids=None):
+    prompt_token_ids = prompt_token_ids or [1]
+    return Sequence(
+        seq_id=0,
+        prompt="<s>",
+        prompt_token_ids=prompt_token_ids,
+        block_size=16,
+    )
+
+
+def create_dummy_logprobs(
+        complete_sequence_token_ids: List[int]) -> List[Dict[int, Logprob]]:
+    return [{
+        token_id: Logprob(logprob=0.0),
+        token_id + 1: Logprob(logprob=0.1)
+    } for token_id in complete_sequence_token_ids]
+
+
+@pytest.mark.parametrize("complete_sequence", TRUTH)
+@pytest.mark.parametrize("tokenizer_name", TOKENIZERS)
+@pytest.mark.parametrize("skip_special_tokens", [True, False])
+def test_decode_sequence_logprobs(complete_sequence: str,
+                                  complete_sequence_token_ids: List[int],
+                                  detokenizer: Detokenizer,
+                                  skip_special_tokens: bool):
+    """Verify Detokenizer decodes logprobs correctly."""
+    sampling_params = SamplingParams(skip_special_tokens=skip_special_tokens,
+                                     logprobs=2)
+
+    # Run sequentially.
+    seq = create_sequence()
+    dummy_logprobs = create_dummy_logprobs(complete_sequence_token_ids)
+    sequential_logprobs_text_chosen_token = []
+    sequential_logprobs_text_other_token = []
+    for new_token, logprobs in zip(complete_sequence_token_ids,
+                                   dummy_logprobs):
+        seq.append_token_id(new_token, logprobs)
+        detokenizer.decode_sequence_inplace(seq, sampling_params)
+        sequential_logprobs_text_chosen_token.append(
+            seq.output_logprobs[-1][new_token].decoded_token)
+        sequential_logprobs_text_other_token.append(
+            seq.output_logprobs[-1][new_token + 1].decoded_token)
+    sequential_result = seq.output_text
+
+    assert sequential_result == "".join(sequential_logprobs_text_chosen_token)
+    assert sequential_result != "".join(sequential_logprobs_text_other_token)
+
+    if skip_special_tokens:
+        # Text for logprobs for the chosen token should be the same as the
+        # generated text. Note that this will only be true if we skip
+        # special tokens.
+        assert sequential_result == complete_sequence
+
+
+@pytest.mark.parametrize("complete_sequence", TRUTH)
+@pytest.mark.parametrize("tokenizer_name", TOKENIZERS)
+@pytest.mark.parametrize("skip_special_tokens", [True])
+def test_decode_prompt_logprobs(complete_sequence: str,
+                                complete_sequence_token_ids: List[int],
+                                detokenizer: Detokenizer,
+                                skip_special_tokens: bool):
+    """Verify Detokenizer decodes prompt logprobs correctly."""
+    sampling_params = SamplingParams(skip_special_tokens=skip_special_tokens,
+                                     prompt_logprobs=1)
+
+    # Run sequentially.
+    seq = create_sequence(complete_sequence_token_ids)
+    seq_group = SequenceGroup(request_id="1",
+                              seqs=[seq],
+                              sampling_params=sampling_params,
+                              arrival_time=0.0)
+    dummy_logprobs = create_dummy_logprobs(complete_sequence_token_ids)
+    detokenizer.decode_prompt_logprobs_inplace(seq_group, dummy_logprobs)
+    decoded_prompt_logprobs = dummy_logprobs
+
+    if skip_special_tokens:
+        # Text for logprobs for the chosen token should be the same as the
+        # prompt text. Note that this will only be true if we skip
+        # special tokens.
+        assert complete_sequence == "".join([
+            logprobs[token_id].decoded_token for token_id, logprobs in zip(
+                complete_sequence_token_ids, decoded_prompt_logprobs)
+        ])
+        assert complete_sequence != "".join([
+            logprobs[token_id + 1].decoded_token for token_id, logprobs in zip(
+                complete_sequence_token_ids, decoded_prompt_logprobs)
+        ])

tests/tokenization/test_tokenizer_group.py

+import asyncio
+import os
+from unittest.mock import patch
+
+import pytest
+from transformers import AutoTokenizer, PreTrainedTokenizerBase
+
+from vllm.transformers_utils.tokenizer_group import get_tokenizer_group
+from vllm.transformers_utils.tokenizer_group.ray_tokenizer_group import (
+    RayTokenizerGroupPool)
+from vllm.transformers_utils.tokenizer_group.tokenizer_group import (
+    TokenizerGroup)
+
+from ..conftest import get_tokenizer_pool_config
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("tokenizer_group_type", [None, "ray"])
+async def test_tokenizer_group(tokenizer_group_type):
+    reference_tokenizer = AutoTokenizer.from_pretrained("gpt2")
+    tokenizer_group = get_tokenizer_group(
+        get_tokenizer_pool_config(tokenizer_group_type),
+        tokenizer_id="gpt2",
+        enable_lora=False,
+        max_num_seqs=1,
+        max_input_length=None,
+    )
+    assert reference_tokenizer.encode("prompt") == tokenizer_group.encode(
+        request_id="request_id", prompt="prompt", lora_request=None)
+    assert reference_tokenizer.encode(
+        "prompt") == await tokenizer_group.encode_async(
+            request_id="request_id", prompt="prompt", lora_request=None)
+    assert isinstance(tokenizer_group.get_lora_tokenizer(None),
+                      PreTrainedTokenizerBase)
+    assert tokenizer_group.get_lora_tokenizer(
+        None) == await tokenizer_group.get_lora_tokenizer_async(None)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("tokenizer_group_type", ["ray"])
+async def test_tokenizer_group_pool(tokenizer_group_type):
+    reference_tokenizer = AutoTokenizer.from_pretrained("gpt2")
+    tokenizer_group_pool = get_tokenizer_group(
+        get_tokenizer_pool_config(tokenizer_group_type),
+        tokenizer_id="gpt2",
+        enable_lora=False,
+        max_num_seqs=1,
+        max_input_length=None,
+    )
+    # Send multiple requests to the tokenizer group pool
+    # (more than the pool size)
+    # and check that all requests are processed correctly.
+    num_requests = tokenizer_group_pool.pool_size * 5
+    requests = [
+        tokenizer_group_pool.encode_async(request_id=str(i),
+                                          prompt=f"prompt {i}",
+                                          lora_request=None)
+        for i in range(num_requests)
+    ]
+    results = await asyncio.gather(*requests)
+    expected_results = [
+        reference_tokenizer.encode(f"prompt {i}") for i in range(num_requests)
+    ]
+    assert results == expected_results
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("tokenizer_group_type", ["ray"])
+async def test_tokenizer_group_ray_pool_env_var_propagation(
+        tokenizer_group_type):
+    """Test that env vars from caller process are propagated to
+    tokenizer Ray actors."""
+    env_var = "MY_ENV_VAR"
+
+    class EnvVarCheckerTokenizerGroup(TokenizerGroup):
+
+        def ping(self):
+            assert os.environ.get(env_var) == "1"
+            return super().ping()
+
+    class EnvVarCheckerRayTokenizerGroupPool(RayTokenizerGroupPool):
+        _worker_cls = EnvVarCheckerTokenizerGroup
+
+    tokenizer_pool_config = get_tokenizer_pool_config(tokenizer_group_type)
+    tokenizer_pool = EnvVarCheckerRayTokenizerGroupPool.from_config(
+        tokenizer_pool_config,
+        tokenizer_id="gpt2",
+        enable_lora=False,
+        max_num_seqs=1,
+        max_input_length=None)
+    with pytest.raises(AssertionError):
+        tokenizer_pool.ping()
+
+    with patch.dict(os.environ, {env_var: "1"}):
+        tokenizer_pool_config = get_tokenizer_pool_config(tokenizer_group_type)
+        tokenizer_pool = EnvVarCheckerRayTokenizerGroupPool.from_config(
+            tokenizer_pool_config,
+            tokenizer_id="gpt2",
+            enable_lora=False,
+            max_num_seqs=1,
+            max_input_length=None)
+        tokenizer_pool.ping()

tests/worker/test_model_runner.py

-import random
+import pytest
 import torch
 
+from vllm.config import ModelConfig
 from vllm.sequence import SamplingParams, SequenceData, SequenceGroupMetadata
-from vllm.worker.model_runner import ModelRunner
+from vllm.worker.model_runner import ModelRunner, _get_graph_batch_size
 
 
-def test_prepare_prompt():
+@pytest.mark.parametrize("batch_size", list(range(1, 257)))
+def test_prepare_prompt(batch_size):
     model_runner = ModelRunner(None, None, None, None, None)
     model_runner.set_block_size(16)
 
-    batch_size = random.randint(1, 256)
     prompt_lens = []
     seq_group_metadata_list = []
+    block_tables = {0: [1]}
     for i in range(batch_size):
         # make sure all tokens fit into one block
         prompt_len = i % (model_runner.block_size - 1) + 1
         prompt_lens.append(prompt_len)
-        seq_data = list(range(prompt_len))
-        seq_group_metadata_list.append(
-            SequenceGroupMetadata(
-                request_id=f"test_{i}",
-                is_prompt=True,
-                seq_data={0: SequenceData(seq_data)},
-                sampling_params=SamplingParams(temperature=0),
-                block_tables={0: [1]},
-            ))
+        seq_data = SequenceData(list(range(prompt_len)))
+        seq_group_metadata = SequenceGroupMetadata(
+            request_id=f"test_{i}",
+            is_prompt=True,
+            seq_data={0: seq_data},
+            sampling_params=SamplingParams(temperature=0),
+            block_tables=block_tables,
+        )
+        assert seq_group_metadata.token_chunk_size == seq_data.get_len()
+        seq_group_metadata_list.append(seq_group_metadata)
 
     expected_selected_token_indices = []
     selected_token_start_idx = 0
-    max_seq_len = max(prompt_lens)
     for prompt_len in prompt_lens:
         expected_selected_token_indices.append(selected_token_start_idx +
                                                prompt_len - 1)
-        selected_token_start_idx += max_seq_len
-    input_tokens, input_positions, _, return_prompt_lens, _, _, _, _ = (
-        model_runner._prepare_prompt(seq_group_metadata_list))
+        selected_token_start_idx += prompt_len
+    (input_tokens, input_positions, attn_metadata, return_prompt_lens, _, _, _,
+     _, _) = (model_runner._prepare_prompt(seq_group_metadata_list))
     assert return_prompt_lens == prompt_lens
+
+    # Verify input metadata is correct for prompts.
+    device = model_runner.device
+    assert attn_metadata.is_prompt is True
+    assert torch.allclose(attn_metadata.prompt_lens_tensor,
+                          torch.tensor(prompt_lens, device=device))
+    assert attn_metadata.prompt_lens == prompt_lens
+    assert attn_metadata.num_prompt_tokens == sum(prompt_lens)
+    assert attn_metadata.num_generation_tokens == 0
+    assert attn_metadata.max_prompt_len == max(prompt_lens)
+
+    # Test subquery start locs.
+    start_idx = 0
+    start_loc = [start_idx]
+    for prompt_len in prompt_lens:
+        start_idx += prompt_len
+        start_loc.append(start_idx)
+    assert torch.allclose(
+        attn_metadata.subquery_start_loc,
+        torch.tensor(start_loc, dtype=torch.int32, device=device))
+
+    # Test seq start locs. Note that for normal prefill it is
+    # equivalent to subquery_start_loc.
+    start_idx = 0
+    seq_start_loc = [start_idx]
+    for prompt_len in prompt_lens:
+        start_idx += prompt_len
+        seq_start_loc.append(start_idx)
+
+    assert torch.allclose(
+        attn_metadata.seq_start_loc,
+        torch.tensor(start_loc, dtype=torch.int32, device=device))
+    assert attn_metadata.max_context_len is None
+    assert torch.allclose(
+        attn_metadata.context_lens,
+        torch.zeros(attn_metadata.context_lens.shape[0],
+                    dtype=torch.int,
+                    device=device))
+
+    expected = torch.tensor([[] for _ in range(len(seq_group_metadata_list))],
+                            dtype=torch.int32,
+                            device=model_runner.device)
+    assert torch.allclose(attn_metadata.block_tables, expected)
+    # Cuda graph should not be used for prerill.
+    assert attn_metadata.use_cuda_graph is False
+    assert attn_metadata.kv_cache_dtype == "auto"
+
+    assert input_tokens.shape == (sum(prompt_lens), )
+    assert input_positions.shape == (sum(prompt_lens), )
+    torch.testing.assert_close(input_tokens, input_positions)
+
     sampling_metadata = model_runner._prepare_sample(seq_group_metadata_list,
                                                      prompt_lens,
                                                      subquery_lens=prompt_lens)
-    assert input_tokens.shape == (batch_size, max_seq_len)
-    assert input_positions.shape == (batch_size, max_seq_len)
+    assert input_tokens.shape == (sum(prompt_lens), )
+    assert input_positions.shape == (sum(prompt_lens), )
+    actual = sampling_metadata.selected_token_indices
+    expected = torch.tensor(expected_selected_token_indices,
+                            device=actual.device,
+                            dtype=actual.dtype)
+    torch.testing.assert_close(actual, expected)
+    torch.testing.assert_close(input_tokens, input_positions)
+
+    actual = sampling_metadata.selected_token_indices
+    expected = torch.tensor(expected_selected_token_indices,
+                            device=actual.device,
+                            dtype=actual.dtype)
+    torch.testing.assert_close(actual, expected)
+
+
+@pytest.mark.parametrize("batch_size", list(range(1, 257)))
+def test_prepare_decode_cuda_graph(batch_size):
+    model_config = ModelConfig(
+        "facebook/opt-125m",
+        "facebook/opt-125m",
+        tokenizer_mode="auto",
+        trust_remote_code=False,
+        download_dir=None,
+        load_format="dummy",
+        seed=0,
+        dtype="float16",
+        revision=None,
+        enforce_eager=False,
+    )
+    model_runner = ModelRunner(model_config, None, None, None, None)
+    model_runner.set_block_size(16)
+
+    prompt_lens = []
+    seq_group_metadata_list = []
+    for i in range(batch_size):
+        # make sure all tokens fit into one block
+        prompt_len = i % (model_runner.block_size - 1) + 1
+        prompt_lens.append(prompt_len)
+        seq_data = list(range(prompt_len))
+        seq_data = SequenceData(seq_data)
+        seq_group_metadata = SequenceGroupMetadata(
+            request_id=f"test_{i}",
+            is_prompt=False,
+            seq_data={0: seq_data},
+            sampling_params=SamplingParams(temperature=0),
+            block_tables={0: [1]},
+        )
+        assert seq_group_metadata.token_chunk_size == 1
+        seq_group_metadata_list.append(seq_group_metadata)
+
+    input_tokens, input_positions, attn_metadata, _, _, _ = (
+        model_runner._prepare_decode(seq_group_metadata_list))
+
+    expected_bs = _get_graph_batch_size(len(seq_group_metadata_list))
+    # Verify input metadata is correct for prompts.
+    device = model_runner.device
+    assert attn_metadata.is_prompt is False
+    assert attn_metadata.prompt_lens is None
+    assert attn_metadata.num_prompt_tokens == 0
+    assert attn_metadata.num_generation_tokens == expected_bs
+    assert attn_metadata.max_prompt_len is None
+    assert attn_metadata.subquery_start_loc is None
+    assert attn_metadata.seq_start_loc is None
+    assert attn_metadata.max_context_len == max(prompt_lens)
+    assert torch.allclose(
+        attn_metadata.context_lens[:len(prompt_lens)],
+        torch.tensor(prompt_lens, dtype=torch.int, device=device))
+
+    # block table's first index corresponds to each batch, meaning in
+    # decoding it is each token.
+    assert attn_metadata.block_tables.shape[0] == len(input_tokens)
+    # Block table's second dim correspondsd to each token's block number.
+    # It is padded up to
+    assert attn_metadata.block_tables.shape[1] == (
+        model_runner.get_max_block_per_batch())
+    # Cuda graph should not be used for prerill.
+    assert attn_metadata.use_cuda_graph is True
+    assert attn_metadata.kv_cache_dtype == "auto"
+
+    assert input_tokens.shape == (expected_bs, )
+    assert input_positions.shape == (expected_bs, )
     torch.testing.assert_close(input_tokens, input_positions)
 
+    # Verify Sampling
+    expected_selected_token_indices = []
+    selected_token_start_idx = 0
+    for prompt_len in prompt_lens:
+        expected_selected_token_indices.append(selected_token_start_idx)
+        selected_token_start_idx += 1
+    sampling_metadata = model_runner._prepare_sample(seq_group_metadata_list,
+                                                     prompt_lens,
+                                                     subquery_lens=prompt_lens)
     actual = sampling_metadata.selected_token_indices
     expected = torch.tensor(expected_selected_token_indices,
                             device=actual.device,

tests/worker/test_swap.py

+import torch
+
+from vllm.engine.arg_utils import EngineArgs
+from vllm.utils import get_distributed_init_method, get_ip, get_open_port
+from vllm.worker.worker import Worker
+
+
+def test_swap() -> None:
+    # Configure the engine.
+    engine_args = EngineArgs(model="facebook/opt-125m",
+                             dtype="half",
+                             load_format="dummy")
+    (model_config, cache_config, parallel_config, scheduler_config,
+     device_config, _, _) = engine_args.create_engine_configs()
+    cache_config.num_gpu_blocks = 100
+    cache_config.num_cpu_blocks = 100
+
+    # Create the worker.
+    distributed_init_method = get_distributed_init_method(
+        get_ip(), get_open_port())
+    worker = Worker(
+        model_config=model_config,
+        parallel_config=parallel_config,
+        scheduler_config=scheduler_config,
+        device_config=device_config,
+        local_rank=0,
+        rank=0,
+        distributed_init_method=distributed_init_method,
+        is_driver_worker=True,
+    )
+
+    # Initialize the worker.
+    worker.init_device()
+    worker.load_model()
+    worker.init_cache_engine(cache_config)
+    worker.warm_up_model()
+
+    # Randomly initialize the cache.
+    gpu_cache = worker.cache_engine.gpu_cache
+    cpu_cache = worker.cache_engine.cpu_cache
+    num_layers = len(gpu_cache)
+    for i in range(num_layers):
+        gpu_key_cache, gpu_value_cache = gpu_cache[i]
+        gpu_key_cache.random_()
+        gpu_value_cache.random_()
+        cpu_key_cache, cpu_value_cache = cpu_cache[i]
+        cpu_key_cache.random_()
+        cpu_value_cache.random_()
+
+    allclose = lambda a, b: torch.allclose(
+        a.cuda(), b.cuda(), rtol=0.0, atol=0.0)
+
+    # Test swap out.
+    blocks_to_swap_out = {3: 72, 56: 35, 84: 34}
+    worker.execute_model(seq_group_metadata_list=[],
+                         blocks_to_swap_in={},
+                         blocks_to_swap_out=blocks_to_swap_out,
+                         blocks_to_copy={})
+    for i in range(num_layers):
+        gpu_key_cache, gpu_value_cache = gpu_cache[i]
+        cpu_key_cache, cpu_value_cache = cpu_cache[i]
+        for src, dst in blocks_to_swap_out.items():
+            assert allclose(gpu_key_cache[src], cpu_key_cache[dst])
+            assert allclose(gpu_value_cache[src], cpu_value_cache[dst])
+
+    # Test swap in.
+    blocks_to_swap_in = {19: 45, 67: 23, 12: 78, 40: 99, 1: 71}
+    worker.execute_model(seq_group_metadata_list=[],
+                         blocks_to_swap_in=blocks_to_swap_in,
+                         blocks_to_swap_out={},
+                         blocks_to_copy={})
+    for i in range(num_layers):
+        gpu_key_cache, gpu_value_cache = gpu_cache[i]
+        cpu_key_cache, cpu_value_cache = cpu_cache[i]
+        for src, dst in blocks_to_swap_in.items():
+            assert allclose(gpu_key_cache[dst], cpu_key_cache[src])
+            assert allclose(gpu_value_cache[dst], cpu_value_cache[src])

vllm/__init__.py

@@ -3,13 +3,13 @@
 from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
 from vllm.engine.async_llm_engine import AsyncLLMEngine
 from vllm.engine.llm_engine import LLMEngine
-from vllm.engine.ray_utils import initialize_cluster
+from vllm.engine.ray_utils import initialize_ray_cluster
 from vllm.entrypoints.llm import LLM
 from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.sampling_params import SamplingParams
 from vllm.version import __dcu_version__
 
-__version__ = "0.3.3"
+__version__ = "0.4.0"
 
 __all__ = [
     "LLM",
@@ -20,5 +20,5 @@ __all__ = [
     "EngineArgs",
     "AsyncLLMEngine",
     "AsyncEngineArgs",
-    "initialize_cluster",
+    "initialize_ray_cluster",
 ]

vllm/attention/__init__.py

+from vllm.attention.backends.abstract import (AttentionBackend,
+                                              AttentionMetadata)
+from vllm.attention.layer import Attention
+from vllm.attention.selector import get_attn_backend
+
+__all__ = [
+    "AttentionBackend",
+    "AttentionMetadata",
+    "Attention",
+    "get_attn_backend",
+]

vllm/attention/backends/abstract.py

+from abc import ABC, abstractmethod
+from dataclasses import dataclass, fields
+from typing import Any, Dict, List, Optional, Tuple, Type
+
+import torch
+
+
+class AttentionBackend(ABC):
+    """Abstract class for attention backends."""
+
+    @staticmethod
+    @abstractmethod
+    def get_impl_cls() -> Type["AttentionImpl"]:
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def make_metadata(*args, **kwargs) -> "AttentionMetadata":
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+    ) -> Tuple[int, ...]:
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def swap_blocks(
+        src_kv_cache: torch.Tensor,
+        dst_kv_cache: torch.Tensor,
+        src_to_dst: Dict[int, int],
+    ) -> None:
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def copy_blocks(
+        kv_caches: List[torch.Tensor],
+        src_to_dists: Dict[int, List[int]],
+    ) -> None:
+        raise NotImplementedError
+
+
+@dataclass
+class AttentionMetadata:
+
+    def asdict_zerocopy(self) -> Dict[str, Any]:
+        """Similar to dataclasses.asdict, but avoids deepcopying."""
+        # Note that if we add dataclasses as fields, they will need
+        # similar handling.
+        return {
+            field.name: getattr(self, field.name)
+            for field in fields(self)
+        }
+
+
+class AttentionImpl(ABC):
+
+    @abstractmethod
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: Optional[int] = None,
+        alibi_slopes: Optional[List[float]] = None,
+        sliding_window: Optional[int] = None,
+    ) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        raise NotImplementedError

vllm/attention/backends/flash_attn.py

+"""Attention layer with Flash and PagedAttention.
+
+NOTE(woosuk): At the moment, this file includes a lot of duplicated code from
+XFormers backend. The duplicated code will be removed once we use flash-attn or
+flashinfer for all the attention operations.
+"""
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple, Type
+
+import torch
+from flash_attn import flash_attn_varlen_func
+
+from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
+                                              AttentionMetadata)
+from vllm.attention.ops.paged_attn import (PagedAttention,
+                                           PagedAttentionMetadata)
+
+
+class FlashAttentionBackend(AttentionBackend):
+
+    @staticmethod
+    def get_impl_cls() -> Type["FlashAttentionImpl"]:
+        return FlashAttentionImpl
+
+    @staticmethod
+    def make_metadata(*args, **kwargs) -> "FlashAttentionMetadata":
+        return FlashAttentionMetadata(*args, **kwargs)
+
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+    ) -> Tuple[int, ...]:
+        return PagedAttention.get_kv_cache_shape(num_blocks, block_size,
+                                                 num_kv_heads, head_size)
+
+    @staticmethod
+    def swap_blocks(
+        src_kv_cache: torch.Tensor,
+        dst_kv_cache: torch.Tensor,
+        src_to_dst: Dict[int, int],
+    ) -> None:
+        PagedAttention.swap_blocks(src_kv_cache, dst_kv_cache, src_to_dst)
+
+    @staticmethod
+    def copy_blocks(
+        kv_caches: List[torch.Tensor],
+        src_to_dists: Dict[int, List[int]],
+    ) -> None:
+        PagedAttention.copy_blocks(kv_caches, src_to_dists)
+
+
+@dataclass
+class FlashAttentionMetadata(AttentionMetadata, PagedAttentionMetadata):
+    """Metadata for FlashAttentionBackend.
+
+    NOTE: Any python object stored here is not updated when it is
+    cuda-graph replayed. If you have values that need to be changed
+    dynamically, it should be stored in tensor. The tensor has to be
+    updated from `CUDAGraphRunner.forward` API.
+    """
+    # Currently, input sequences can only contain all prompts
+    # or all decoding. True if all sequences are prompts.
+    is_prompt: bool
+    # (batch_size,). The prompt length per sequence. None if it is a decoding.
+    prompt_lens: Optional[List[int]]
+    # prompt_lens stored as a tensor.
+    prompt_lens_tensor: Optional[torch.Tensor]
+    # The number of prompt tokens. Doesn't include padding.
+    num_prompt_tokens: int
+    # The number of generation tokens. Doesn't include padding.
+    num_generation_tokens: int
+
+    # NOTE(sang): Definition of context_len, subquery_len, and seqlen.
+    # |---------- N-1 iteration --------|
+    # |---------------- N iteration ---------------------|
+    # |- tokenA -|......................|-- newTokens ---|
+    # |---------- context_len ----------|
+    # |-------------------- seqlen ----------------------|
+    #                                   |- subquery_len -|
+
+    # WARNING(sang): context_len has different definition depending on if it is
+    # prefill vs decoding. When it is prefill, it doesn't include new tokens.
+    # When it is for decoding, it includes a new token.
+
+    # Maximum subquery length in the batch.
+    max_subquery_len: Optional[int]
+    # Maximum prompt length in the batch.
+    max_prompt_len: Optional[int]
+    # (batch_size + 1,). The cumulative subquery lengths of the sequences in
+    # the batch, used to index into subquery. E.g., if the subquery length
+    # is [4, 6], it is [0, 4, 10].
+    subquery_start_loc: Optional[torch.Tensor]
+    # (batch_size + 1,). The cumulative sequence lengths of the sequences in
+    # the batch, used to index into sequence. E.g., if the sequence length is
+    # [4, 6], it is [0, 4, 10].
+    seq_start_loc: Optional[torch.Tensor]
+
+    # Whether or not if cuda graph is enabled.
+    # Cuda-graph is currently enabled for decoding only.
+    # TODO(woosuk): Move `use_cuda_graph` out since it's unrelated to attention.
+    use_cuda_graph: bool
+
+
+class FlashAttentionImpl(AttentionImpl):
+    """
+    If the input tensors contain prompt tokens, the layout is as follows:
+    |<--------------- num_prompt_tokens -------------->|	
+    |<--prompt_0-->|<--prompt_1-->|...|<--prompt_N-1-->|
+
+    Otherwise, the layout is as follows:	
+    |<------------------ num_generation_tokens (M) ----------------->|	
+    |<--generation_0-->|..........|<--generation_M-1-->|<--padding-->|
+
+    Generation tokens can contain padding when cuda-graph is used.
+    Currently, prompt tokens don't contain any padding.
+
+    The prompts might have different lengths, while the generation tokens
+    always have length 1.
+    """
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: Optional[int] = None,
+        alibi_slopes: Optional[List[float]] = None,
+        sliding_window: Optional[int] = None,
+    ) -> None:
+        self.num_heads = num_heads
+        self.head_size = head_size
+        self.scale = float(scale)
+        self.num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads
+        self.sliding_window = ((sliding_window, sliding_window)
+                               if sliding_window is not None else (-1, -1))
+        if alibi_slopes is not None:
+            alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)
+        self.alibi_slopes = alibi_slopes
+
+        assert self.num_heads % self.num_kv_heads == 0
+        self.num_queries_per_kv = self.num_heads // self.num_kv_heads
+
+        suppored_head_sizes = PagedAttention.get_supported_head_sizes()
+        if head_size not in suppored_head_sizes:
+            raise ValueError(
+                f"Head size {head_size} is not supported by PagedAttention. "
+                f"Supported head sizes are: {suppored_head_sizes}.")
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache: torch.Tensor,
+        attn_metadata: FlashAttentionMetadata,
+    ) -> torch.Tensor:
+        """Forward pass with FlashAttention and PagedAttention.
+
+        Args:
+            query: shape = [num_tokens, num_heads * head_size]
+            key: shape = [num_tokens, num_kv_heads * head_size]
+            value: shape = [num_tokens, num_kv_heads * head_size]
+            kv_cache = [2, num_blocks, block_size * num_kv_heads * head_size]
+            attn_metadata: Metadata for attention.
+        Returns:
+            shape = [num_tokens, num_heads * head_size]
+        """
+        num_tokens, hidden_size = query.shape
+        # Reshape the query, key, and value tensors.
+        query = query.view(-1, self.num_heads, self.head_size)
+        key = key.view(-1, self.num_kv_heads, self.head_size)
+        value = value.view(-1, self.num_kv_heads, self.head_size)
+
+        if kv_cache is not None:
+            key_cache, value_cache = PagedAttention.split_kv_cache(
+                kv_cache, self.num_kv_heads, self.head_size)
+
+            # Reshape the input keys and values and store them in the cache.
+            # If kv_cache is not provided, the new key and value tensors are
+            # not cached. This happens during the initial memory profiling run.
+            PagedAttention.write_to_paged_cache(key, value, key_cache,
+                                                value_cache,
+                                                attn_metadata.slot_mapping,
+                                                attn_metadata.kv_cache_dtype)
+
+        if attn_metadata.is_prompt:
+            # Prompt run.
+            if kv_cache is None or attn_metadata.block_tables.numel() == 0:
+                # normal attention
+                # When block_tables are not filled, it means q and k are the
+                # prompt, and they have the same length.
+                output = flash_attn_varlen_func(
+                    q=query,
+                    k=key,
+                    v=value,
+                    cu_seqlens_q=attn_metadata.seq_start_loc,
+                    cu_seqlens_k=attn_metadata.seq_start_loc,
+                    max_seqlen_q=attn_metadata.max_prompt_len,
+                    max_seqlen_k=attn_metadata.max_prompt_len,
+                    softmax_scale=self.scale,
+                    causal=True,
+                    window_size=self.sliding_window,
+                    alibi_slopes=self.alibi_slopes,
+                )
+            else:
+                # prefix-enabled attention
+                output = PagedAttention.forward_prefix(
+                    query,
+                    key,
+                    value,
+                    key_cache,
+                    value_cache,
+                    attn_metadata.block_tables,
+                    attn_metadata.subquery_start_loc,
+                    attn_metadata.prompt_lens_tensor,
+                    attn_metadata.context_lens,
+                    attn_metadata.max_subquery_len,
+                    self.alibi_slopes,
+                )
+        else:
+            # Decoding run.
+            output = PagedAttention.forward_decode(
+                query,
+                key_cache,
+                value_cache,
+                attn_metadata.block_tables,
+                attn_metadata.context_lens,
+                attn_metadata.max_context_len,
+                attn_metadata.kv_cache_dtype,
+                self.num_kv_heads,
+                self.scale,
+                self.alibi_slopes,
+            )
+
+        # Reshape the output tensor.
+        return output.view(num_tokens, hidden_size)

vllm/attention/backends/xformers.py

+"""Attention layer with xFormers and PagedAttention."""
+import importlib
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple, Type
+
+import torch
+from xformers import ops as xops
+from xformers.ops.fmha.attn_bias import (AttentionBias,
+                                         BlockDiagonalCausalMask,
+                                         LowerTriangularMaskWithTensorBias)
+
+from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
+                                              AttentionMetadata)
+from vllm.attention.ops.paged_attn import (PagedAttention,
+                                           PagedAttentionMetadata)
+from vllm.logger import init_logger
+from vllm.utils import is_hip
+
+logger = init_logger(__name__)
+
+
+class XFormersBackend(AttentionBackend):
+
+    @staticmethod
+    def get_impl_cls() -> Type["XFormersImpl"]:
+        return XFormersImpl
+
+    @staticmethod
+    def make_metadata(*args, **kwargs) -> "XFormersMetadata":
+        return XFormersMetadata(*args, **kwargs)
+
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+    ) -> Tuple[int, ...]:
+        return PagedAttention.get_kv_cache_shape(num_blocks, block_size,
+                                                 num_kv_heads, head_size)
+
+    @staticmethod
+    def swap_blocks(
+        src_kv_cache: torch.Tensor,
+        dst_kv_cache: torch.Tensor,
+        src_to_dst: Dict[int, int],
+    ) -> None:
+        PagedAttention.swap_blocks(src_kv_cache, dst_kv_cache, src_to_dst)
+
+    @staticmethod
+    def copy_blocks(
+        kv_caches: List[torch.Tensor],
+        src_to_dists: Dict[int, List[int]],
+    ) -> None:
+        PagedAttention.copy_blocks(kv_caches, src_to_dists)
+
+
+@dataclass
+class XFormersMetadata(AttentionMetadata, PagedAttentionMetadata):
+    """Metadata for XFormersbackend.
+
+    NOTE: Any python object stored here is not updated when it is
+    cuda-graph replayed. If you have values that need to be changed
+    dynamically, it should be stored in tensor. The tensor has to be
+    updated from `CUDAGraphRunner.forward` API.
+    """
+    # Currently, input sequences can only contain all prompts
+    # or all decoding. True if all sequences are prompts.
+    is_prompt: bool
+    # (num_tokens,). The indices of the token slots that input tokens will be
+    # stored into. E.g., if `slot_mapping` is [35, 2, 17] and the block size
+    # is 16, the three tokens are stored in the 3rd slot in block 2, 2nd slot
+    # in block 0, and 1st slot in block 1, respectively.
+    slot_mapping: torch.Tensor
+    # (batch_size,). The prompt length per sequence. None if it is a decoding.
+    prompt_lens: Optional[List[int]]
+    # prompt_lens stored as a tensor.
+    prompt_lens_tensor: Optional[torch.Tensor]
+    # The number of prompt tokens. Doesn't include padding.
+    num_prompt_tokens: int
+    # The number of generation tokens. Doesn't include padding.
+    num_generation_tokens: int
+
+    # NOTE(sang): Definition of context_len, subquery_len, and seqlen.
+    # |---------- N-1 iteration --------|
+    # |---------------- N iteration ---------------------|
+    # |- tokenA -|......................|-- newTokens ---|
+    # |---------- context_len ----------|
+    # |-------------------- seqlen ----------------------|
+    #                                   |- subquery_len -|
+
+    # WARNING(sang): context_len has different definition depending on if it is
+    # prefill vs decoding. When it is prefill, it doesn't include new tokens.
+    # When it is for decoding, it includes a new token.
+
+    # Maximum subquery length in the batch.
+    max_subquery_len: Optional[int]
+    # FIXME: It is for flash attn.
+    # Maximum prompt length in the batch.
+    max_prompt_len: Optional[int]
+    # (batch_size + 1,). The cumulative subquery lengths of the sequences in
+    # the batch, used to index into subquery. E.g., if the subquery length
+    # is [4, 6], it is [0, 4, 10].
+    subquery_start_loc: Optional[torch.Tensor]
+    # FIXME: It is for flash attn.
+    # (batch_size + 1,). The cumulative sequence lengths of the sequences in
+    # the batch, used to index into sequence. E.g., if the sequence length is
+    # [4, 6], it is [0, 4, 10].
+    seq_start_loc: Optional[torch.Tensor]
+
+    # Whether or not if cuda graph is enabled.
+    # Cuda-graph is currently enabled for decoding only.
+    # TODO(woosuk): Move `use_cuda_graph` out since it's unrelated to attention.
+    use_cuda_graph: bool
+
+    def __post_init__(self):
+        # Set during the execution of the first attention op.
+        # It is a list because it is needed to set per prompt
+        # when alibi slopes is used. It is because of the limitation
+        # from xformer API.
+        # will not appear in the __repr__ and __init__
+        self.attn_bias: Optional[List[AttentionBias]] = None
+
+
+class XFormersImpl(AttentionImpl):
+    """
+    If the input tensors contain prompt tokens, the layout is as follows:
+    |<--------------- num_prompt_tokens --------------->|	
+    |<--prompt_0-->|<--prompt_1-->|...|<--prompt_N-1--->|
+
+    Otherwise, the layout is as follows:	
+    |<------------------ num_generation_tokens (M) ----------------->|	
+    |<--generation_0-->|..........|<--generation_M-1-->|<--padding-->|
+
+    Generation tokens can contain padding when cuda-graph is used.
+    Currently, prompt tokens don't contain any padding.
+
+    The prompts might have different lengths, while the generation tokens
+    always have length 1.
+    """
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: Optional[int] = None,
+        alibi_slopes: Optional[List[float]] = None,
+        sliding_window: Optional[int] = None,
+    ) -> None:
+        self.num_heads = num_heads
+        self.head_size = head_size
+        self.scale = float(scale)
+        self.num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads
+        self.sliding_window = sliding_window
+        if alibi_slopes is not None:
+            alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)
+        self.alibi_slopes = alibi_slopes
+
+        assert self.num_heads % self.num_kv_heads == 0
+        self.num_queries_per_kv = self.num_heads // self.num_kv_heads
+
+        suppored_head_sizes = PagedAttention.get_supported_head_sizes()
+        if head_size not in suppored_head_sizes:
+            raise ValueError(
+                f"Head size {head_size} is not supported by PagedAttention. "
+                f"Supported head sizes are: {suppored_head_sizes}.")
+
+        # AMD Radeon 7900 series (gfx1100) currently does not support xFormers
+        # nor FlashAttention. As a temporary workaround, we use naive PyTorch
+        # implementation of attention.
+        self.use_naive_attention = _check_use_naive_attention()
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache: Optional[torch.Tensor],
+        attn_metadata: XFormersMetadata,
+    ) -> torch.Tensor:
+        """Forward pass with xFormers and PagedAttention.
+
+        Args:
+            query: shape = [num_tokens, num_heads * head_size]
+            key: shape = [num_tokens, num_kv_heads * head_size]
+            value: shape = [num_tokens, num_kv_heads * head_size]
+            kv_cache = [2, num_blocks, block_size * num_kv_heads * head_size]
+            attn_metadata: Metadata for attention.
+        Returns:
+            shape = [num_tokens, num_heads * head_size]
+        """
+        num_tokens, hidden_size = query.shape
+        query = query.view(-1, self.num_heads, self.head_size)
+        key = key.view(-1, self.num_kv_heads, self.head_size)
+        value = value.view(-1, self.num_kv_heads, self.head_size)
+
+        if kv_cache is not None:
+            key_cache, value_cache = PagedAttention.split_kv_cache(
+                kv_cache, self.num_kv_heads, self.head_size)
+
+            # Reshape the input keys and values and store them in the cache.
+            # If kv_cache is not provided, the new key and value tensors are
+            # not cached. This happens during the initial memory profiling run.
+            PagedAttention.write_to_paged_cache(key, value, key_cache,
+                                                value_cache,
+                                                attn_metadata.slot_mapping,
+                                                attn_metadata.kv_cache_dtype)
+
+        if attn_metadata.is_prompt:
+            # Prompt run.
+            if kv_cache is None or attn_metadata.block_tables.numel() == 0:
+                # normal attention.
+                # block tables are empty if the prompt does not have a cached
+                # prefix.
+                if self.num_kv_heads != self.num_heads:
+                    # As of Nov 2023, xformers only supports MHA. For MQA/GQA,
+                    # project the key and value tensors to the desired number of
+                    # heads.
+                    # TODO(woosuk): Use MQA/GQA kernels for higher performance.
+                    query = query.view(query.shape[0], self.num_kv_heads,
+                                       self.num_queries_per_kv,
+                                       query.shape[-1])
+                    key = key[:, :,
+                              None, :].expand(key.shape[0], self.num_kv_heads,
+                                              self.num_queries_per_kv,
+                                              key.shape[-1])
+                    value = value[:, :,
+                                  None, :].expand(value.shape[0],
+                                                  self.num_kv_heads,
+                                                  self.num_queries_per_kv,
+                                                  value.shape[-1])
+
+                if self.use_naive_attention:
+                    output = torch.empty_like(query)
+                    start = 0
+                    for _, prompt_len in enumerate(attn_metadata.prompt_lens):
+                        end = start + prompt_len
+                        out = _naive_masked_attention(
+                            query[None, start:end],
+                            key[None, start:end],
+                            value[None, start:end],
+                            self.num_heads,
+                            self.num_kv_heads,
+                            self.head_size,
+                            self.scale,
+                        )
+                        # TODO(woosuk): Unnecessary copy. Optimize.
+                        output[start:end].copy_(out)
+                        start += prompt_len
+
+                    # Using view got RuntimeError: view size is not compatible
+                    # with input tensor's size and stride (at least one
+                    # dimension spans across two contiguous subspaces).
+                    # Use reshape instead.
+                    return output.reshape(num_tokens, hidden_size)
+
+                output = self._run_memory_efficient_xformers_forward(
+                    query, key, value, attn_metadata)
+            else:
+                # prefix-enabled attention
+                output = PagedAttention.forward_prefix(
+                    query,
+                    key,
+                    value,
+                    key_cache,
+                    value_cache,
+                    attn_metadata.block_tables,
+                    attn_metadata.subquery_start_loc,
+                    attn_metadata.prompt_lens_tensor,
+                    attn_metadata.context_lens,
+                    attn_metadata.max_subquery_len,
+                    self.alibi_slopes,
+                )
+        else:
+            # Decoding run.
+            output = PagedAttention.forward_decode(
+                query,
+                key_cache,
+                value_cache,
+                attn_metadata.block_tables,
+                attn_metadata.context_lens,
+                attn_metadata.max_context_len,
+                attn_metadata.kv_cache_dtype,
+                self.num_kv_heads,
+                self.scale,
+                self.alibi_slopes,
+            )
+
+        # Reshape the output tensor.
+        return output.view(-1, self.num_heads * self.head_size)
+
+    def _run_memory_efficient_xformers_forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: XFormersMetadata,
+    ) -> torch.Tensor:
+        """Attention for 1D query of multiple prompts. Multiple prompt
+        tokens are flattened in to `query` input.
+
+        Args:
+            output: shape = [num_prompt_tokens, num_heads, head_size]
+            query: shape = [num_prompt_tokens, num_heads, head_size]
+            key: shape = [num_prompt_tokens, num_kv_heads, head_size]
+            value: shape = [num_prompt_tokens, num_kv_heads, head_size]
+            attn_metadata: Metadata for attention.
+        """
+        # Set attention bias if not provided. This typically happens at
+        # the very attention layer of every iteration.
+        # FIXME(woosuk): This is a hack.
+        if attn_metadata.attn_bias is None:
+            if self.alibi_slopes is None:
+                attn_bias = BlockDiagonalCausalMask.from_seqlens(
+                    attn_metadata.prompt_lens)
+                if self.sliding_window is not None:
+                    attn_bias = attn_bias.make_local_attention(
+                        self.sliding_window)
+                attn_metadata.attn_bias = [attn_bias]
+            else:
+                attn_metadata.attn_bias = _make_alibi_bias(
+                    self.alibi_slopes, self.num_kv_heads, query.dtype,
+                    attn_metadata.prompt_lens)
+
+        op = xops.fmha.MemoryEfficientAttentionFlashAttentionOp[0] if (
+            is_hip()) else None
+        # No alibi slopes.
+        # TODO(woosuk): Too many view operations. Let's try to reduce
+        # them in the future for code readability.
+        if self.alibi_slopes is None:
+            query = query.unsqueeze(0)
+            key = key.unsqueeze(0)
+            value = value.unsqueeze(0)
+            out = xops.memory_efficient_attention_forward(
+                query,
+                key,
+                value,
+                attn_bias=attn_metadata.attn_bias[0],
+                p=0.0,
+                scale=self.scale,
+                op=op)
+
+            return out.view_as(query)
+
+        # Attention with alibi slopes.
+        # FIXME(woosuk): Because xformers does not support dynamic sequence
+        # lengths with custom attention bias, we process each prompt one by
+        # one. This is inefficient, especially when we have many short prompts.
+        output = torch.empty_like(query)
+        start = 0
+        for i, prompt_len in enumerate(attn_metadata.prompt_lens):
+            end = start + prompt_len
+            out = xops.memory_efficient_attention_forward(
+                query[None, start:end],
+                key[None, start:end],
+                value[None, start:end],
+                attn_bias=attn_metadata.attn_bias[i],
+                p=0.0,
+                scale=self.scale,
+                op=op)
+            # TODO(woosuk): Unnecessary copy. Optimize.
+            output[start:end].copy_(out.squeeze(0))
+            start += prompt_len
+        return output
+
+
+def _make_alibi_bias(
+    alibi_slopes: torch.Tensor,
+    num_kv_heads: int,
+    dtype: torch.dtype,
+    prompt_lens: List[int],
+) -> LowerTriangularMaskWithTensorBias:
+    attn_biases = []
+    for prompt_len in prompt_lens:
+        bias = torch.arange(prompt_len, dtype=dtype)
+        # NOTE(zhuohan): HF uses
+        #     `bias = bias[None, :].repeat(prompt_len, 1)`
+        # here. We find that both biases give the same results, but
+        # the bias below more accurately follows the original ALiBi
+        # paper.
+        # Calculate a matrix where each element represents ith element- jth
+        # element.
+        bias = bias[None, :] - bias[:, None]
+
+        padded_len = (prompt_len + 7) // 8 * 8
+        num_heads = alibi_slopes.shape[0]
+        bias = torch.empty(
+            1,  # batch size
+            num_heads,
+            prompt_len,
+            padded_len,
+            device=alibi_slopes.device,
+            dtype=dtype,
+        )[:, :, :, :prompt_len].copy_(bias)
+        bias.mul_(alibi_slopes[:, None, None])
+        if num_heads != num_kv_heads:
+            bias = bias.unflatten(1, (num_kv_heads, num_heads // num_kv_heads))
+        attn_biases.append(LowerTriangularMaskWithTensorBias(bias))
+
+    return attn_biases
+
+
+def _check_use_naive_attention() -> bool:
+    if not is_hip():
+        return False
+    # For ROCm, check whether flash attention is installed or not.
+    use_naive_attention = importlib.util.find_spec("flash_attn") is None
+    if use_naive_attention:
+        logger.warning("flash_attn is not installed. Using naive attention. "
+                       "This will take significantly more GPU memory.")
+        return True
+    return False
+
+
+def _naive_masked_attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    num_heads: int,
+    num_kv_heads: int,
+    head_size: int,
+    scale: float,
+) -> torch.Tensor:
+    query = query.view(-1, num_heads, head_size)
+    key = key.view(-1, num_kv_heads, head_size)
+    value = value.view(-1, num_kv_heads, head_size)
+    seq_len, _, _ = query.shape
+    attn_mask = torch.triu(torch.ones(seq_len,
+                                      seq_len,
+                                      dtype=query.dtype,
+                                      device=query.device),
+                           diagonal=1)
+    attn_mask = attn_mask * torch.finfo(query.dtype).min
+
+    attn_weights = scale * torch.einsum("qhd,khd->hqk", query, key).float()
+    attn_weights = attn_weights + attn_mask.float()
+    attn_weights = torch.softmax(attn_weights, dim=-1).to(value.dtype)
+    out = torch.einsum("hqk,khd->qhd", attn_weights, value)
+    return out