Feat/pcp qwen36 hybrid

[Yopry]

PR Type

New Features

Description

Enables tensor-parallel + Prefill Context Parallel (PCP) serving of Qwen3.6-style hybrid models (interleaved full-attention +
GatedDeltaNet/linear-attention layers, e.g. Qwen3.6-35B-A3B) on an unmodified vLLM 0.20.2, entirely through monkey-patches installed by the plugin.

vLLM 0.20.2 hard-rejects context parallelism (pcp/dcp > 1) for hybrid/mamba KV-cache layouts (multiple block sizes). This PR lifts that restriction
for Qwen3.6 and adds the sequence-splitting + cross-rank gather logic needed to make prefill run correctly under PCP. Every patch self-gates on
prefill_context_parallel_size > 1, so it is a no-op for non-PCP runs.

Scope is prefill correctness (verified: first-token argmax matches the tp=4, pcp=1 baseline). Decode is out of scope in this PR.

Known limitations:

• Verified only up to a world size of 4 (tp*pcp ≤ 4: tp2pcp2 / tp4pcp1 / tp1pcp4). Larger configurations are untested.
• generate() teardown can stall on the max_tokens=1 abort path (use max_tokens ≥ 2); the prefill result itself is correct.
• Decode under PCP is not handled: full-attention KV is sequence-split, so multi-token generation drifts / can produce NaN (batched decode path).
• Multi-request batches with unequal lengths are not yet handled (the current gather/last-token logic assumes equal-length batches).

Changes

• Add vllm_fl/patches/qwen36_pcp.py: the full PCP patch set —
  • Guard patches: bypass the three vLLM hybrid+PCP hard-guards (resolve_kv_cache_block_sizes, HybridKVCacheCoordinator.init,
    MambaManager.find_longest_cache_hit) by presenting pcp==1 to block-size accounting; sequence splitting is done separately at the model-runner level.
  • Sequence splitting: patch the V1 GPUModelRunner._prepare_inputs to give each PCP rank a contiguous local token slice (local num_scheduled_tokens

• shifted num_computed_tokens), plus a logits gather so the final prompt token's hidden state feeds the LM head.
  • Full-attention layers: sequence-split with a K/V all-gather across PCP ranks + end-aligned causal flash_attn_varlen_func.
  • GDN/linear-attention layers: full-sequence "island" (all-gather hidden into per-request full sequences, run the original GDN scan, slice the
    local segment back out).

• vllm_fl/init.py: install the patches, including in register_model so worker processes pick them up.

Prefill Context Parallel (PCP) for Qwen3.6 hybrid models

The plugin adds Prefill Context Parallel for the Qwen3.6 hybrid architecture
(Gated DeltaNet linear-attention layers + periodic full-attention layers). The
prompt is split across the PCP ranks; full-attention layers all-gather K/V and
run causal attention over the reconstructed sequence, while the GDN /
linear-attention layers run a redundant full-sequence scan so every rank ends
with the correct recurrent state.

Enable it with --prefill-context-parallel-size N (vllm serve) or
prefill_context_parallel_size=N (LLM(...)).

# PCP currently supports the FlashAttention backend only. FL routes attention
# to FlagGems (Triton) by default, which does not declare PCP support and makes
# engine init fail with "TritonAttentionImpl does not support PCP". Disable the
# FlagGems preference so attention falls back to FlashAttention:
export VLLM_FL_PREFER_ENABLED=0

# If flag_gems fails to detect the device at plugin load ("No device were
# detected"), pin the vendor explicitly:
export GEMS_VENDOR=nvidia

VLLM_PLUGINS=fl vllm serve /path/to/Qwen3.6-35B-A3B \
    --tensor-parallel-size 2 \
    --prefill-context-parallel-size 2 \
    --enforce-eager --trust-remote-code

Set VLLM_PCP_DEBUG=1 to print per-layer [PCP-DBG] traces (sequence split,
K/V all-gather, GDN full-sequence island) for verifying that PCP actually runs.

If you run a CUDA-13 build on a driver that only supports CUDA <= 12.8 and hit
cudaErrorInsufficientDriver, put the bundled forward-compatibility libcuda on
the loader path: export LD_LIBRARY_PATH=/usr/local/cuda-13.0/compat:$LD_LIBRARY_PATH.

Scope / limitations

• Prefill correctness is verified (e.g. tp=2/pcp=2 matches the non-parallel
  baseline first token, within bf16 rounding).
• Decode under PCP is a known limitation and not yet supported.
• PCP currently supports the FlashAttention backend only; the FlagGems Triton
  attention backend does not yet declare supports_pcp.

Testing

• Hardware: 4×H200, enforce_eager, FlashAttention 2, GDN prefill backend = triton.
• Correctness: first-token argmax parity between PCP and the tp=4, pcp=1 baseline on Qwen3.6-35B-A3B, verified across context lengths (16k / 64k /
  128k / 256k) and across all three world-size-4 layouts (tp2pcp2, tp4pcp1, tp1pcp4). Matches the baseline golden token.
• Not covered by the repo test suite: the existing tests/ (incl. the qwen3_6 tp2 configs) run with pcp=1, where these patches are a no-op, so they do
  not exercise the PCP path. Correctness was verified with standalone scripts as above.

Checklist

• I have run the existing tests and they pass
• I have added tests for my changes (if applicable)
• I have updated the documentation (if applicable)

[CLAassistant]

All committers have signed the CLA.