Hybrid linear-attention (Qwen3.5/Qwen3-Next) KV cache spec + prefix-caching guard

[Carlos-Marques]

Summary

Companion to exa-labs/neuronx-distributed-inference#1 (NxDI Qwen3.5/Qwen3-Next port). Makes the plugin's cache accounting correct for hybrid linear-attention models, where ~3/4 of layers (Gated DeltaNet) keep a constant-size per-sequence state instead of a growing KV cache.

Two changes:

1. NeuronxDistributedModelRunner.get_kv_cache_spec() no longer declares FullAttentionSpec for every layer. For models whose hf_config.layer_types contains "linear_attention", those layers now get a MambaSpec:

kv_cache_spec[f"layers.{i}.linear_attn"] = MambaSpec(
    block_size=self.block_size,
    shapes=_linear_attention_state_shapes(hf_config),  # (conv_state, recurrent_state)
    dtypes=(torch.float32, torch.float32),
    mamba_type="linear_attention",
)

so vLLM's KV-cache manager budgets the real (constant) HBM footprint instead of 4x-overcounting and refusing admissible requests. State shapes come from linear_num_{key,value}_heads / linear_{key,value}_head_dim / linear_conv_kernel_dim: conv state (2*Hk*Dk + Hv*Dv, kernel-1), recurrent state (Hv, Dk, Dv).

2. NeuronPlatform.check_and_update_config() force-disables --enable-prefix-caching for hybrid models, since DeltaNet's recurrent state is a running summary at one position — KV blocks can't be reused at branch points and prefix caching would silently produce garbage. Disabling also routes into the existing block_size = max_model_len contiguous path.

Non-hybrid models are unaffected (layer_types absent → all layers FullAttentionSpec, guard no-ops).

Testing

• New unit tests: hybrid get_kv_cache_spec (per-layer spec types + state shapes) in test/unit/worker/test_model_runner.py; prefix-caching guard on/off + non-hybrid passthrough in test/unit/test_platform.py.
• Full suite: pytest test/unit/test_platform.py test/unit/worker/test_model_runner.py → 295 passed (vLLM 0.16.0, CPU).
• End-to-end on Trainium still pending (no hardware available); the NxDI side is CPU logit-verified against HF in the companion PR.

Link to Devin session: https://app.devin.ai/sessions/853f8dae8e27497a9a8ed7a1a6d347b0
Requested by: @Carlos-Marques

[devin-ai-integration]

🤖 Devin AI Engineer

I'll be helping with this pull request! Here's what you should know:

✅ I will automatically:

• Address comments on this PR. Add '(aside)' to your comment to have me ignore it.
• Look at CI failures and help fix them

Note: I can only respond to comments from users who have write access to this repository.

⚙️ Control Options:

• Disable automatic comment, CI, and merge conflict monitoring