Add Qwen3.8 packed hybrid PagedAttention export for ORT GenAI Engine

[titaiwangms]

Summary

Add an opt-in packed hybrid PagedAttention export for Qwen/Qwen3.8-27B
(model_type=qwen3_5_text) that can run through ONNX Runtime GenAI's
continuous-batching Engine/Request APIs.

Mobius's existing Qwen3.5/Qwen3.8 implementation is correct for ordinary
dense/dynamic inference. This issue adds a separate serving ABI; it is not a
fix for the existing model math.

The reference behavior is ONNX Runtime GenAI main at
d5b40851ba80ffa8e95b6b01f921dbb9008fac80 (September 14, 2026), especially:

• src/python/py/models/builders/base.py
• src/python/py/models/builders/qwen.py
• src/python/py/models/builder.py
• src/python/py/models/README.md

This is also the missing foundation for the Qwen3.8 recipe in
microsoft/olive-recipes#615.

Background

Qwen3.8-27B is a 64-layer hybrid decoder with full_attention_interval=4:

• 16 full-attention layers at IDs 3, 7, ..., 63
• 48 GatedDeltaNet linear-attention layers

Mobius currently exports these through HybridCausalLMTask using:

• dense [batch, sequence] token tensors;
• contiguous/dynamic KV cache for full-attention layers;
• fixed convolution and recurrent states for GatedDeltaNet layers.

Mobius also supports com.microsoft::PagedAttention, but only its LATENT
absorbed-MLA mode for eligible DeepSeek/GLM models. Qwen3.8 needs the
SEPARATE K/V layout and a packed hybrid decoder contract, so the existing
_paged_mla.py path cannot be reused as-is.

Phase 1 scope

Implement an executable CUDA FP16/BF16 Qwen3.8 continuous-batching path:

1. Packed token graph ABI.
2. SEPARATE K/V PagedAttention for full-attention layers.
3. Varlen convolution and GatedDeltaNet execution for linear-attention layers.
4. Correct ORT GenAI Engine configuration and hybrid state grouping.
5. Structural, numerical, state-carry, and Engine smoke tests.

Non-goals

• INT4/INT8 paged KV cache.
• DFlash2 and compact speculative state-update capture.
• Full parity with Add verified Qwen3.8-27B INT4 KV and DFlash2 recipe microsoft/olive-recipes#615 in the first PR.
• Multimodal paged serving.
• A generic redesign of all Mobius cache/task abstractions.
• Changes to the existing dense Qwen3.8 path.

Proposed architecture

Preserve the existing Qwen modules, weight names, and dense forward() paths.
Add an explicit paged-hybrid execution path rather than inserting flags
throughout generic components.

Task

Extend HybridCausalLMTask or add a dedicated PagedHybridCausalLMTask.
The dedicated class is preferable if extending the existing task would make
its dense contract ambiguous.

Suggested interface:

PagedHybridCausalLMTask(
    *,
    paged_block_size: int = 256,
    prune_prefill_prefix: bool = False,
)

The CLI/transformers builder should resolve --features paged-attention
according to the detected architecture:

• eligible MLA model -> existing LATENT paged task;
• Qwen3.5/Qwen3.8 hybrid model -> new SEPARATE paged-hybrid task;
• unsupported model -> typed error.

Paged attention remains mutually exclusive with static cache and an
incompatible explicit task.

Components and model methods

Add a model-agnostic SEPARATE PagedAttention adapter, separate from
components/_paged_mla.py.

Possible interfaces:

class PagedAttentionState(NamedTuple):
    key_cache: ir.Value
    value_cache: ir.Value
    block_table: ir.Value
    cumulative_sequence_lengths: ir.Value
    past_sequence_lengths: ir.Value
    attention_metadata: ir.Value

class Qwen35Attention:
    def forward_paged(...): ...

class GatedDeltaNet:
    def forward_packed(...): ...

Add paged variants to the Qwen decoder layer/text model/wrapper while keeping
the current dense signatures unchanged.

All graph construction must continue to use onnxscript.nn and onnx_ir;
do not introduce explicit protobuf operations.

Packed graph contract

Let:

• N be the total scheduled query tokens.
• B be the number of active sequences.
• P be the number of cache pages.
• K be the configured page size.

Shared inputs

Input	Type and shape	Meaning
input_ids	INT64 [N]	Sequence-contiguous packed tokens
position_ids	INT64 [3, N]	Qwen interleaved partial MRoPE positions
block_table	INT32 [B, max_blocks_per_sequence]	Logical-to-physical page mapping
cumulative_sequence_lengths	INT32 [B + 1]	Packed query row boundaries
past_sequence_lengths	INT32 [B]	Committed lengths before the invocation
attention_metadata	INT32 [3]	CPU-resident launch/backend bounds; exact semantics must be pinned to the target runtime

There is no attention_mask.

Following the pinned ONNX Runtime GenAI builder, do not expose graph
slot_mapping in Phase 1. Leave PagedAttention input 10 empty and let the op
derive write locations from sequence lengths and block_table.

Full-attention layer state

For each full-attention layer:

past_key_values.{i}.key   T [P, K, num_kv_heads, head_dim]
past_key_values.{i}.value T [P, K, num_kv_heads, head_dim]
present.{i}.key           same shape and type
present.{i}.value         same shape and type

The cache uses kv_cache_layout="SEPARATE" semantics. Inputs and outputs must
follow the aliasing requirements of the selected ORT PagedAttention kernel.

Linear-attention layer state

For each GatedDeltaNet layer:

past_key_values.{i}.conv_state
    T [B, conv_dim, conv_kernel_size - 1]

past_key_values.{i}.recurrent_state
    FLOAT [B, num_value_heads, value_head_dim, key_head_dim]

The exact recurrent-state axis order must be verified against the pinned
native GatedDeltaNet schema. Mobius's existing dense function path may use a
different logical ordering; do not change the dense helper globally.

Outputs use corresponding present.{i}.conv_state and
present.{i}.recurrent_state ports with fixed trailing dimensions.

Outputs

• Default: logits FLOAT [N, vocab_size].
• With prefix pruning: gather rows
  cumulative_sequence_lengths[1:] - 1 before the LM head and emit
  logits FLOAT [B, vocab_size].

Per-layer behavior

Full attention

1. Project packed hidden states.
2. Split Qwen's doubled Q projection into query and output gate per head.
3. Apply the required Q/K OffsetRMSNorm and interleaved partial MRoPE.
4. Emit com.microsoft::PagedAttention with SEPARATE K/V caches.
5. Apply sigmoid(gate) to the attention output.
6. Reuse the existing output projection, residual, and MLP paths.

Verify whether the paired runtime supports Q/K norm weights fused into
PagedAttention with external Qwen MRoPE while preserving norm-before-RoPE
ordering. If not, perform normalization externally and leave the optional
norm inputs empty. Do not normalize twice.

GatedDeltaNet

The packed path must use sequence boundaries for both recurrent operations:

1. Packed projections.
2. Varlen causal convolution using cumulative_sequence_lengths.
3. Per-head Q/K/V preparation.
4. Packed/native com.microsoft::GatedDeltaNet.
5. Existing gated output norm and output projection.

Use the kernel's Qwen gate arithmetic and expected FP32 A_log, dt_bias,
and recurrent state. Do not also apply the dense function's transformed gates.

Runtime and metadata contract

The generated package must be directly loadable by the ORT GenAI Engine
without manually editing configuration.

Publish distinct decoder state groups:

paged_kv       -> full-attention layer IDs
fixed_conv     -> linear-attention layer IDs
fixed_recurrent -> linear-attention layer IDs

The config must map:

• packed input and position IDs;
• block table and sequence-length inputs;
• CPU attention_metadata;
• per-layer paged K/V names;
• per-layer convolution/recurrent state names;
• corresponding outputs.

Add engine.dynamic_batching with at least:

• block_size;
• max_batch_size;
• either gpu_utilization_factor or an explicit page budget.

Workflow/inference metadata must describe the same three state disciplines.
Detecting a block_table input alone is insufficient to claim hybrid Engine
compatibility.

Likely files

• src/mobius/__main__.py
• src/mobius/_configs/_base.py
• src/mobius/integrations/transformers/_builder.py
• src/mobius/tasks/_causal_lm.py
• src/mobius/tasks/_cache_utils.py
• src/mobius/components/_paged_attention.py (new)
• src/mobius/components/_attention.py
• src/mobius/components/_gated_deltanet.py
• src/mobius/models/qwen35.py
• src/mobius/integrations/ort_genai/auto_export.py
• src/mobius/integrations/ort_genai/genai_config.py
• src/mobius/integrations/onnx_genai/workflow_metadata.py
• related co-located, graph-build, parity, and Engine tests

Delivery plan

Slice 0: pin schemas and fixtures

• Pin the paired ORT and ORT GenAI revisions.
• Record PagedAttention, varlen convolution, and GatedDeltaNet operand order,
  attributes, dtypes, state layout, and aliasing requirements.
• Add tiny Qwen3.8-shaped config fixtures.

Slice 1: packed Qwen graph

• Add the packed-hybrid task and graph inputs.
• Add SEPARATE PagedAttention emission.
• Add packed Qwen attention shapes, MRoPE, output gating, and pruned LM head.

Slice 2: packed recurrent layers

• Add varlen convolution and packed GatedDeltaNet.
• Preserve FP32 recurrent state and verify the state layout.
• Add mixed-layer state input/output registration.

Slice 3: runtime packaging

• Emit ORT GenAI decoder mappings, state groups, and dynamic-batching config.
• Emit matching Mobius workflow/inference metadata.
• Run an actual Engine/Request smoke test.

Follow-up issues

• INT8/INT4 per-channel paged KV cache and calibrated scale loading.
• CUDA graph qualification.
• DFlash2 integration and compact state-update capture.
• Full Add verified Qwen3.8-27B INT4 KV and DFlash2 recipe microsoft/olive-recipes#615 parity.

Acceptance criteria

• Default dense Qwen3.5/Qwen3.8 and existing LATENT MLA exports are unchanged.
• Unsupported EP, dtype, model, task, and cache combinations fail with typed errors.
• A Qwen3.8-27B structural export contains 16 SEPARATE PagedAttention nodes and 48 packed/varlen GatedDeltaNet paths, derived from config rather than hard-coded.
• The paged graph has packed inputs, no dense attention mask, and only full-attention layers own K/V page buffers.
• Qwen output gating, Q/K norm, partial interleaved MRoPE, and residual/MLP order match the existing dense implementation.
• Tiny packed batches with unequal sequence lengths match independent dense runs for prefill and continuation within dtype-specific tolerances.
• Tests compare logits, logical paged K/V contents, convolution states, and recurrent states.
• Tests cover nonzero past state, multi-token continuation, page boundaries around the configured block size, request reordering/removal, new-request admission, and page reuse without cross-request contamination.
• Pruned logits equal the last-token rows selected from unpruned logits.
• Generated ORT GenAI configuration declares the exact paged-KV/fixed-conv/fixed-recurrent layer groups.
• A pinned ORT GenAI Engine/Request test loads the generated artifact without manual repair, overlaps unequal prompts, admits a new request during decoding, and matches isolated execution.
• Qwen3.8-27B is smoke-tested on suitable CUDA hardware before support is documented as production-ready.

Risks and open questions

1. Runtime schema compatibility: verify exact PagedAttention optional input
   positions, SEPARATE cache types, block-size constraints, BF16 support,
   aliasing, and attention_metadata semantics against the runtime used by
   Mobius CI/release artifacts.
2. Varlen operator availability: verify exact names and schemas for
   varlen causal convolution and GatedDeltaNet. Mobius currently uses local
   function bodies for dense execution; that is not proof of native packed
   kernel compatibility.
3. Recurrent-state layout: verify whether the native op uses
   [B, H_v, D_v, D_k] or another ordering and add explicit conversion tests.
4. Q/K norm and MRoPE ordering: confirm whether Qwen3.8 can use the fused
   norm inputs while retaining external [3, N] MRoPE.
5. Engine row compaction: validate that page state, conv state, recurrent
   state, RNG, and request identity receive the same row permutation.
6. Resource requirements: tiny tests establish ABI correctness but do not
   replace a real Qwen3.8-27B CUDA run.

Later recipe-parity work

microsoft/olive-recipes#615 additionally requires INT4/INT8 per-channel paged
KV, CUDA graphs, an INT4 DFlash2 drafter, cache-pool reservation, and compact
state updates. Those should build on the Phase 1 execution contract rather
than be coupled into the initial implementation.