Phase 7.1: Multi-request batching for virtual batch attention

[AlonKellner-RedHat]

Summary

Extend Phase 7 virtual batch attention to support multi-request batching (max_num_seqs > 1). The Phase 7 MVP enforces single-request batching due to complexity of handling heterogeneous prefix lengths across multiple requests.

Background

Current limitation (Phase 7 MVP):

• Virtual batch attention raises NotImplementedError when num_reqs > 1
• Server requires --max-num-seqs 1 configuration
• Single-request processing reduces GPU utilization in high-throughput scenarios

Technical blocker:
Virtual batch transformation (make_block_attention_virtual_batches()) currently assumes homogeneous prefix lengths. With multiple concurrent requests, each request may have different num_prefix_tokens, requiring per-request metadata transformation.

Scope

In Scope

1. Per-request virtual batch transformation

   • Remove num_reqs > 1 check in dllm_plugin/attention/virtual_batches.py:56
   • Implement heterogeneous prefix length handling
   • Transform metadata separately for each request in the batch

2. Metadata structure updates

   • Support variable-length prefix chunks across requests
   • Correct KV cache slicing for heterogeneous batches
   • Maintain separate query_start_loc offsets per request

3. Integration testing

   • Add test with 2+ concurrent requests to tests/test_virtual_batch_multi_request.py
   • Validate correct attention computation with heterogeneous prefix lengths
   • End-to-end integration test with real LLaDA2.0 model

4. Performance benchmarking

   • Measure throughput improvement with multi-request batching
   • Compare max_num_seqs=1 vs max_num_seqs=4 vs max_num_seqs=8
   • Document GPU utilization improvements

5. Documentation updates

   • Update docs/OPERATOR_LLaDA2.md to reflect multi-request support
   • Remove single-request workarounds from operator guide
   • Add multi-request configuration examples

Out of Scope

• Pipeline parallelism (PP > 1) - remains unsupported (separate issue)
• Tensor parallelism enhancements - TP already supported in Phase 7
• Attention optimizations (single-pass, CUTLASS, FlashInfer) - covered by Phase 8.x

Technical Design

Current Architecture (Phase 7 MVP)

def make_block_attention_virtual_batches(
    attn_metadata: CommonAttentionMetadata,
    num_prefix_tokens: int,  # Single value for all requests
    block_size: int,
) -> tuple[CommonAttentionMetadata | None, CommonAttentionMetadata]:
    if attn_metadata.num_reqs > 1:
        raise NotImplementedError("multi-request not supported")
    # ... transform assuming single request

Proposed Architecture (Phase 7.1)

def make_block_attention_virtual_batches(
    attn_metadata: CommonAttentionMetadata,
    num_prefix_tokens: list[int] | int,  # Per-request or single value
    block_size: int,
) -> tuple[CommonAttentionMetadata | None, CommonAttentionMetadata]:
    # Handle both single-request and multi-request cases
    if isinstance(num_prefix_tokens, int):
        # Single request path (backward compatible)
        num_prefix_tokens = [num_prefix_tokens]
    
    # Per-request transformation
    for req_idx in range(attn_metadata.num_reqs):
        req_prefix_tokens = num_prefix_tokens[req_idx]
        # Transform metadata for this request
        # Update query_start_loc, seq_lens, block_tables per request
    
    # Combine transformed metadata for all requests
    # ...

Dependencies

• HARD: Phase 7 complete (Issue [MVP LLaDA2.0]: Phase 7 — Real-model integration (re-validate stack on real weights) #25, PR feat(phase7+8): Production LLaDA2.0 model + vLLM-native torch.compile optimization #38)
• SOFT: Phase 6 mock stack validation (Issue [MVP LLaDA2.0]: Integration test or manual checklist: serve LLaDA2.0 with plugin stack #17)

Acceptance Criteria

• Remove num_reqs > 1 check from virtual_batches.py
• Implement per-request metadata transformation
• Handle heterogeneous prefix lengths correctly
• Add integration test with 2+ concurrent requests
• Tests pass: pytest -v tests/test_virtual_batch_multi_request.py
• Benchmark shows throughput improvement with max_num_seqs > 1
• Documentation updated to reflect multi-request support
• Server can run with --max-num-seqs 4 (or higher) without errors

Testing Strategy

Unit Tests

def test_virtual_batch_multi_request_heterogeneous_prefix():
    """Test with 2 requests having different prefix lengths."""
    attn_metadata = CommonAttentionMetadata(
        num_reqs=2,
        # Request 0: 16 prefix tokens
        # Request 1: 24 prefix tokens
        ...
    )
    
    prefix_metadata, block_metadata = make_block_attention_virtual_batches(
        attn_metadata=attn_metadata,
        num_prefix_tokens=[16, 24],  # Heterogeneous
        block_size=32,
    )
    
    assert prefix_metadata.num_reqs == 2
    assert block_metadata.num_reqs == 2
    # Validate correct per-request slicing

Integration Tests

• End-to-end test with real LLaDA2.0-mini model
• 2 concurrent requests with different prompt lengths
• Validate output correctness for both requests
• Measure GPU utilization improvement

Benchmarks

# Baseline (Phase 7 MVP)
vllm serve inclusionAI/LLaDA2.0-mini --max-num-seqs 1
# Benchmark: X tokens/sec

# Phase 7.1
vllm serve inclusionAI/LLaDA2.0-mini --max-num-seqs 4
# Benchmark: Y tokens/sec (expect Y > X)

Performance Impact

Expected improvements:

• Throughput: +50-100% with max_num_seqs=4 (GPU-dependent)
• GPU utilization: Increase from ~30-50% to ~70-90% under load
• Latency: Per-request latency may increase slightly, but overall throughput improves

Measurement:

• Use GuideLLM with constant rate profile to measure concurrent requests
• Monitor GPU utilization with nvidia-smi dmon
• Compare TTFT and ITL distributions

Risks and Mitigations

Risk 1: Complexity of per-request transformation

Mitigation:

• Start with simple case: 2 requests with same block size
• Incremental implementation: homogeneous block sizes first, then heterogeneous
• Extensive testing at each step

Risk 2: Performance regression for single-request case

Mitigation:

• Keep fast path for single-request case (isinstance(num_prefix_tokens, int))
• Benchmark both before/after to ensure no regression
• Use profiling to identify bottlenecks

Risk 3: KV cache slicing bugs

Mitigation:

• Thorough unit tests with manual verification of offsets
• Integration tests with known-good attention outputs
• Reference implementation comparison (if available)

Migration Notes

Backward compatibility: Phase 7.1 will remain backward compatible with single-request configurations. Operators can continue using --max-num-seqs 1 if desired.

Recommended upgrade path:

1. Deploy Phase 7.1 with --max-num-seqs 1 initially (validate no regression)
2. Gradually increase to --max-num-seqs 2, then 4, monitoring performance
3. Tune based on GPU memory and throughput requirements

Related Issues

• Blocks: None (Phase 7 complete)
• Blocked by: None
• Related: Issue [MVP LLaDA2.0] Milestone orchestration: timeline, dependency graph, phased plan #19 (Phase 7 orchestration), PR feat(phase7+8): Production LLaDA2.0 model + vLLM-native torch.compile optimization #38 (Phase 7 implementation)
• Supersedes: Phase 7 single-request limitation

Timeline Estimate

• Research/Design: 1-2 days (understand heterogeneous prefix transformations)
• Implementation: 3-5 days (per-request transformation + integration)
• Testing: 2-3 days (unit tests + integration tests + benchmarks)
• Documentation: 1 day

Total: ~1.5-2 weeks for single developer

References

• Phase 7 PR: feat(phase7+8): Production LLaDA2.0 model + vLLM-native torch.compile optimization #38
• Phase 7 orchestration: [MVP LLaDA2.0] Milestone orchestration: timeline, dependency graph, phased plan #19
• Virtual batch implementation: dllm_plugin/attention/virtual_batches.py
• Test coverage: tests/test_virtual_batch_multi_request.py
• Operator guide: docs/OPERATOR_LLaDA2.md
• vLLM chunked_local_attention: https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/attention/chunked_local_attention.py

[AlonKellner-RedHat]

Completed in PR #38 (delivered alongside Phase 7).

Implementation:
Multi-request batching with heterogeneous prefix lengths was fully integrated into Phase 7 instead of being deferred to a separate Phase 7.1.

Delivered Features:

• ✅ Removed num_reqs > 1 check from virtual_batches.py
• ✅ Implemented per-request metadata transformation
• ✅ Support for heterogeneous prefix lengths across concurrent requests
• ✅ Per-request block table slicing with correct offset handling

Test Coverage:

• tests/test_virtual_batch_multi_request.py: Multi-request integration tests
• tests/test_virtual_batch_edge_cases.py: Heterogeneous prefix length validation (4 concurrent requests with prefix lengths [0, 16, 32, 48])

Benchmark Validation:
Comprehensive multi-request benchmarks documented in docs/BENCHMARK_RESULTS_MULTI_REQUEST.md:

• Tested concurrency levels: 1 → 100 concurrent requests
• Configuration: --max-num-seqs 32 (production-ready)
• Results: 10,000+ requests with zero corruption
• Performance: ~9-10 req/s sustained throughput (saturation point)
• ITL scaling: 4.0ms (sync) → 5.8ms (5 RPS) → 6.7ms (10 RPS)

Documentation:

• docs/PHASE7.1_MULTI_REQUEST_VALIDATION.md: Complete validation evidence
• docs/BENCHMARK_RESULTS_MULTI_REQUEST.md: Performance benchmarks
• docs/OPERATOR_LLaDA2.md: Updated with multi-request configuration examples

Status: Originally tracked as Phase 7.1, integrated and delivered in Phase 7 PR #38.

See PR #38 for complete implementation details.