[POC]Segmented spans

[omerpaz95]

Creating a PoC/demo for using Legolink with Spans, using https://github.com/sdavidbd/vllm/tree/feature/segmented-prefill/ for creating virtual requests for the Legolink-required 32 tokens in the beginning of each span.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors.

You ask your reviewers to trigger select CI tests on top of fastcheck CI.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

🚀

[Copilot]

Pull request overview

This pull request introduces a proof-of-concept implementation for "Segmented Spans" functionality in vLLM, which enables handling of non-contiguous KV cache through virtual gap requests and span-based attention. The implementation integrates with the Legolink system and adds support for handling gaps in external KV cache by creating virtual requests for recomputation.

Changes:

• Added segmented prefill support through virtual gap request mechanism that allows recomputation of missing KV cache blocks
• Implemented fused RoPE (Rotary Position Embedding) in Triton attention kernels to improve performance
• Extended KV connector API with get_computed_token_gaps() method to report non-contiguous cache hits
• Added environment variables and configuration for span detection and block attention features

Reviewed changes

Copilot reviewed 30 out of 30 changed files in this pull request and generated 27 comments.

Show a summary per file

File	Description
vllm/v1/core/sched/scheduler.py	Implements virtual gap request creation and scheduling logic for segmented prefill
vllm/v1/core/sched/output.py	Adds data structures for gap recomputation metadata and virtual request tracking
vllm/v1/core/sched/async_scheduler.py	Updates async scheduler to handle virtual gap requests
vllm/v1/worker/gpu_model_runner.py	Adds KV cache hash debugging, virtual request slot mapping, and cleanup logic
vllm/v1/worker/gpu/model_runner.py	Extends execute/sample workflow to track and cleanup virtual gap requests
vllm/v1/attention/ops/triton_unified_attention.py	Implements fused RoPE in attention kernels by splitting Q/K into halves
vllm/v1/attention/backends/triton_attn.py	Passes cos_sin_cache to attention ops when spans are enabled
vllm/v1/attention/backend.py	Adds cos_sin_cache field to CommonAttentionMetadata
vllm/model_executor/layers/rotary_embedding/base.py	Conditionally skips key RoPE application when spans are enabled
vllm/v1/core/kv_cache_utils.py	Implements span separator token detection for fan-in and recompute token handling
vllm/v1/core/block_pool.py	Adds deduplication logic when freeing blocks to handle shared blocks
vllm/envs.py	Adds environment variables for span configuration (VLLM_V1_SPANS_*)
vllm/distributed/kv_transfer/kv_connector/v1/base.py	Adds get_computed_token_gaps() API method to base connector
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py	Adds null checks for virtual requests in offloading connector
vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py	Fixes shape indexing bugs and adds req_id tracking to ReqMeta
examples/offline_inference/spans/*	Adds example scripts and connectors demonstrating spans functionality
examples/offline_inference/kv_segmented_prefill/*	Adds test framework for validating segmented prefill correctness

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The variable name 'firstok' appears to be a typo for 'first_tok' or 'first_token'. This should be renamed for clarity.

Suggested change

	firstok = curr_block_token_ids[0]
	if firstok == envs.VLLM_V1_SPANS_TOKEN_PLUS:
	first_tok = curr_block_token_ids[0]
	if first_tok == envs.VLLM_V1_SPANS_TOKEN_PLUS:

[Copilot]

There is inconsistent trailing whitespace after this line. Trailing whitespace should be removed for code cleanliness.

Suggested change

	extra_keys = recompute_token_handler(
	block_tokens[0], block_tokens[:start_token_idx], extra_keys)
	extra_keys = recompute_token_handler(
	block_tokens[0], block_tokens[:start_token_idx], extra_keys
	)

[Copilot]

There is inconsistent trailing whitespace after this line. Trailing whitespace should be removed for code cleanliness.

Suggested change

[Copilot]

There is inconsistent trailing whitespace after this line. Trailing whitespace should be removed for code cleanliness.

Suggested change

[Copilot]

Debug print statements should not be left in production code. All these print statements (lines 93-129, and elsewhere in this file) appear to be debugging output and should be removed or converted to proper logging.

[Copilot]

Debug print statement should not be left in production code. This should be removed or converted to proper logging.

Suggested change

	print(f"Putting positions: {positions}.")
	logger.debug("Putting positions: %s.", positions)

[Copilot]

Debug print statement should not be left in production code. This should be removed or converted to proper logging.

Suggested change

	print(f"Inserted {len(failed_kv_load_req_ids)} failed_kv_load_req_ids.")
	logger.warning(
	"Inserted %d failed_kv_load_req_ids.",
	len(failed_kv_load_req_ids),
	)

[Copilot]

The variable name 'dedup_bl' is cryptic and should be renamed to something more descriptive like 'deduplicated_blocks' or 'unique_blocks' for better code readability.

[Copilot]

There is inconsistent trailing whitespace after this line. Trailing whitespace should be removed for code cleanliness.

Suggested change

[Copilot]

There is inconsistent trailing whitespace after this line. Trailing whitespace should be removed for code cleanliness.

Suggested change

	return (num_tokens // block_size) * block_size
	return (num_tokens // block_size) * block_size

[almogtavor]

Hey, this code assumes every model layer has a self_attn attribute, but that's not universal. GPT-OSS for example uses TransformerBlock which exposes attention as .attn instead of .self_attn, so this crashes with AttributeError: 'TransformerBlock' object has no attribute 'self_attn' on the first inference request.

I reproduced it with spnl-llm-d-cuda:v0.19.0 + GPT-OSS-120B. Suggested fix: replace the direct .self_attn access with a helper that checks both naming conventions:

def _get_attn_module(layer):
    """Get attention module from a layer, handling different naming conventions."""
    for attr in ("self_attn", "attn"):
        if hasattr(layer, attr):
            return getattr(layer, attr)
    raise AttributeError(
        f"{type(layer).__name__} has no recognized attention attribute "
        f"(tried 'self_attn', 'attn')"
    )

Then the existing code becomes:

if not hasattr(self, "rotate"):
    if not isinstance(self.model.model.layers[0], PPMissingLayer):
        self.rotate = _get_attn_module(self.model.model.layers[0]).rotary_emb
    else:
        for lay in self.model.model.layers:
            if not isinstance(lay, PPMissingLayer):
                self.rotate = _get_attn_module(lay).rotary_emb
                break

[almogtavor]

this isn't used