Day0 Inference Support and Performance Optimization for Hybrid Architecture LLMs

Background: Making “Day0 Support” a Mandatory Requirement for Inference Engines

The iteration pace of Large Language Models (LLMs) is accelerating, accompanied by a rise in architectural innovations. State-recurrent linear attention is back in the mainstream, index-selection-based sparse attention is being deployed, Mixture of Experts (MoE) has become standard for massive models, and speculative decoding is a routine method for accelerating inference. For inference engines, each architectural change introduces new operators, quantization schemes, and parallelization strategies.

“Day0 support” means that on the day new model weights are released, the inference platform can complete end-to-end adaptation. Developers can deploy and experience the model immediately without waiting weeks for operator development and tuning. Achieving this requires a systematic engineering approach: proactively turning general capabilities into reusable assets while focusing engineering efforts on quickly bridging architectural differences.

New architectural models typically introduce three types of challenges for inference engines:

• New Operators: Linear recurrent operators, index scoring operators, and sparse attention kernels cannot directly utilize traditional full-attention kernels.
• New States: Beyond the KV cache, recurrent states and index data become “first-class citizens” requiring fine-grained memory and bandwidth management.
• New Parallelism: Context Parallelism (CP) for ultra-long sequences cannot simply copy the partitioning methods of full attention when dealing with the recurrent dependencies of linear attention and the uneven causal index workloads of sparse attention.

Below, we use the adaptation of a hybrid architecture model—combining linear attention, sparse attention with learnable indexing, and MoE—as an example to fully deconstruct the technical path for Day0 inference support and performance optimization.

Target Architecture Overview: Interleaved Design of Two Attention Types

Before diving into optimization, it is crucial to understand the target object. The core design of this hybrid architecture model interleaves two types of attention modules within the same network:

1. Linear Attention: Processes sequences based on state recurrence. It compresses historical information into a fixed-size state, updating sequentially token by token. Its compute and memory overhead scales near-linearly with sequence length, making it naturally suited for the “broad memory” of long contexts.
2. Sparse Attention with Learnable Indexing: Selects first, reads later. Based on latent attention (representing historical KV via low-rank compression), it introduces an index selection mechanism. A lightweight indexer scores the relevance of historical positions and filters the Top-K most relevant ones. The sparse attention kernel then reads only the KV corresponding to these selected positions to complete the computation. This transforms a “full KV scan” into “targeted indexed reading,” significantly reducing attention compute and memory access while retaining precise retrieval capabilities.
3. Cross-Layer Residual Aggregation: Aggregates attention representations from previous layers back to the current layer, improving information flow in deep networks and preventing deeper layers from “forgetting” features extracted by shallower ones.
4. MoE Feed-Forward Layer: Uses sparse activation of expert networks to handle the model’s primary parameters, activating only a few experts per token during inference.

The two attention types are explicitly complementary: linear attention uses a fixed-size state to “remember everything but without exact precision,” while sparse attention uses indexing to “precisely retrieve a few key pieces of information.” This defines two main threads for inference optimization—linear attention focuses on state management and recurrence parallelization, while sparse attention focuses on deep fusion of the indexing path and quantization.

Day0 Adaptation Strategy: Reuse Existing Assets, Focus on Increments

Day0 does not mean developing from scratch. On the contrary, Day0 speed depends on how many reusable assets were沉淀 (accumulated) during the adaptation cycle of previous similar architectures. Most capabilities can be directly inherited:

Operator Layer:

• Linear attention operators reuse two granularities of fused versions: a deep chunked fusion version for standard scenarios, and a large-granularity fusion version for high throughput, covering optimal execution paths across different batch sizes.
• Cross-layer residual aggregation continues using a two-stage fusion scheme, connecting aggregation computation with adjacent attention computation to reduce intermediate memory writes.
• The MoE side reuses large-granularity fused operators and compute-communication collaborative operators that overlap communication with expert computation.

System Layer:

• A hybrid parallel framework (free combination of Tensor Parallelism, Pipeline Parallelism, Expert Parallelism, and Context Parallelism).
• Communication optimization and multi-stream orchestration, allowing compute, communication, and memory access to run in parallel across different hardware units.

Compiler Layer:

• Ultra-large granularity operator fusion to reduce kernel launch overhead and intermediate tensors.
• Graph-mode full-graph dispatch to compress host-side scheduling overhead.
• Compilation caching to avoid redundant compilation of identical compute graphs.
• Tiled weight layouts optimized for hardware memory access units to improve weight loading efficiency.

Decoding Layer:

• A speculative decoding framework to accelerate autoregressive decoding.
• State replay mechanism: the recurrent state of linear attention can be rolled back and replayed if speculative decoding verification fails, avoiding recomputation of the entire sequence.

The actual increments that need to be newly written are concentrated in the indexer, sparse attention-related operators, and the partitioning scheme for linear attention under context parallelism—which are the focus of the three optimization sections below.

Optimization 1: Indexer Fusion and Multi-Stream Orchestration to Reduce Decoding Wait

During the decoding phase, waits primarily stem from two types of overhead: operator launches and intermediate result write-backs, and serial dependencies during input preparation. The indexing path happens to hit both.

1. “Three-in-One” Indexer Fusion

The indexer’s workflow consists of three steps: scoring historical positions for relevance, weighted summarization by head, and selecting Top-K positions. If implemented as three independent operators, it would result in two intermediate write-backs and three kernel launches. The fusion scheme merges these into a single kernel:

• Scoring results remain on-chip to directly complete weighting and Top-K selection without writing to global memory.
• It directly reads the paged K cache, eliminating the need to organize a separate contiguous K copy for the indexing path and saving a data rearrangement.
• Intermediate write-backs and kernel launch overhead are eliminated simultaneously.

Schematic pseudocode:

# Before fusion: three operators in series with intermediate write-backs
scores   = matmul(q_idx, k_idx)        # Relevance scoring
weights  = head_weighted_sum(scores)   # Head weight weighting
topk_idx = topk(weights, K)            # Top-K selection
# After fusion: single kernel completes, directly reads paged K cache
topk_idx = fused_indexer(q_idx, paged_k_cache, head_w, K)

2. Sparse Attention Kernel

After obtaining the Top-K indices, the sparse attention kernel reads the corresponding KV from the main attention’s KV cache based on the indices and performs tiled sparse attention computation. Because only a few selected positions are read, the compute and memory access overhead during the attention phase drops from “proportional to historical length” to “proportional to Top-K size.” Both Prefill and Decode paths support BF16 and quantized dual paths for the indexer: the former serves as an accuracy baseline, while the latter targets peak performance, allowing developers to choose based on their accuracy vs. throughput trade-offs.

3. Three-Stream Coordination for Indexer Pre-processing

Before entering the main indexer computation, four types of inputs must be prepared: Q projection, K normalization, Q quantization, and head weight projection. If executed serially on a single stream, this preparation time fully stacks into the decoding critical path. The optimization orchestrates these steps in parallel across three streams, allowing independent preparation steps to run simultaneously: Stream 1 executes Q projection, Stream 2 executes K normalization, and Stream 3 executes Q quantization and head weight projection in parallel. Once the serial wait in the preparation phase is absorbed by parallel computation, the decoding critical path is significantly shortened. This “pre-processing multi-streaming” technique is low-cost but yields significant benefits for decoding scenarios composed of many small sequential steps.

Optimization 2: Layered Quantization to Reduce Cache and Transfer Overhead

The bottleneck in long-context inference is often memory access rather than compute. The core principle of quantization here is “layering by precision sensitivity”: core computations sensitive to precision retain high precision, while memory-intensive paths with higher precision tolerance are aggressively quantized.

1. Full Index Path Quantization: MXFP8 Group Dynamic Quantization

The indexer entirely adopts the MXFP8 format, applying dynamic quantization to the index Q/K in a per-group manner with a group size of 32, meaning every 32 elements share a scaling factor. MXFP8’s block floating-point scaling design maintains numerical stability at a fine granularity, making it well-suited for indexing tasks: the absolute score values do not matter, what matters is that the ranking results remain stable and reliable.

2. KV Cache FP8 Tiled Quantization

The KV cache read by sparse attention uses FP8, with a quantization granularity of per-token-head-tile and a tile size of 128. Each token and each head are quantized in chunks of 128 elements. Compared to coarse-grained per-layer or per-head schemes, this granularity restricts the scaling factor’s scope to a very small local area, making quantization error more controllable while keeping the storage and management overhead of scaling factors negligible. The KV cache typically occupies the largest share of memory in long-context inference; this strategy directly reduces the storage and transfer overhead of index data and the KV cache.

3. Retaining BF16 Projection Computation

The main attention projection computation still retains the BF16 path. This layered strategy—“save where you can, keep what you must”—ensures that quantization benefits are realized on the memory access side without affecting core computation precision.

Optimization 3: Prefill Context Parallelism for Long Sequence Partitioning and Load Balancing

When input sequences reach ultra-long scales, a single GPU’s memory and compute cannot independently support prefill, necessitating Context Parallelism (CP). Two key questions arise on this path: how to partition, and how to balance the load post-partition.

1. CP×TP Combined Deployment

The baseline solution is a CP and TP combination: splitting the long input along the sequence dimension and combining sequence parallelism within the TP group to reduce single-GPU activation and intermediate tensor footprint. For full attention, this partitioning is relatively mature; the difficulties lie in the causal index computation of sparse attention and the recurrent dependencies of linear attention.

2. Zigzag Layout: Resolving Causal Index Load Imbalance

In causal mode, the later the sequence partition, the longer the searchable history, and the larger the indexing compute load. If partitions are assigned sequentially to GPUs, a tail latency problem occurs where “later GPUs are overwhelmed, earlier GPUs idle.” The Zigzag layout solves this by pairing head and tail: assigning opposite head and tail sequence segments to the same GPU, so each GPU handles both a “light” and a “heavy” partition, balancing the cross-GPU indexing compute load. After computation, results are restored to their original order.

8 partitions numbered 0..7, causal load increases with the number:
Sequential: GPU0:[0,1]  GPU1:[2,3]  GPU2:[4,5]  GPU3:[6,7]   <- Severe load imbalance
Zigzag:     GPU0:[0,7]  GPU1:[1,6]  GPU2:[2,5]  GPU3:[3,4]   <- Balanced light/heavy mix

3. Fused Linear Attention CP: Pre-compute in Parallel, Then Synthesize Initial State

Linear attention’s state recurrence has strict sequential dependencies: the initial state of a later segment equals the final state of the previous segment. If only TP is used to split by attention heads, each GPU still processes the entire sequence, and Time To First Token (TTFT) remains limited in ultra-long scenarios.

The solution is to transform “recurrence” into “parallel chunking + lightweight synthesis.” Implemented via an operator development DSL, the entire process is expressed as a compilable fused kernel:

• Each GPU first pre-computes its assigned sequence chunk in parallel, producing a local state summary.
• State summaries from all chunks are exchanged via AllGather.
• Each GPU sequentially synthesizes the state contributions of all preceding chunks to obtain the initial state for its own chunk.
• Each GPU uses this initial state to complete its segment computation in parallel.

# Execution flow for each GPU i (schematic)
summary   = precompute(chunk_i)             # Parallel pre-computation of local state summary
summaries = all_gather(summary)             # AllGather to exchange state summaries
init_s    = compose(summaries[:i])          # Sequentially synthesize preceding chunks' state contributions
out_i     = chunk_forward(chunk_i, init_s)  # Parallel computation of the current chunk

The key is that synthesizing the initial state only requires processing lightweight state summaries rather than replaying the full sequence. Thus, the overall time is significantly lower than serial recurrence, reducing TTFT for ultra-long sequences while decreasing single-GPU activation and intermediate tensor footprint.

Practical Insights and Summary

Abstracting these practices yields several universally applicable lessons for adapting to new architectural models:

1. Operator asset reuse determines Day0 speed. Turning general components like linear attention, MoE, and residual aggregation into multi-granularity fused operators ahead of time means that when a new model arrives, efforts only need to focus on actual differences (like the indexer). This compresses the adaptation cycle from weeks to days.
2. Quantization must be layered, not one-size-fits-all. Memory-intensive paths with high precision tolerance, like indexing and KV cache, are suitable for fine-grained MXFP8/FP8 quantization; core projection computations should retain the BF16 baseline. Quantization benefits are primarily cashed in on memory footprint and bandwidth pressure, which is critical for long-context scenarios.
3. Parallel strategies must be combined, not isolated. CP×TP solves the “fits in memory” problem, Zigzag layout solves the “balanced compute” problem, and state pre-computation and synthesis solve the “recurrence can be parallel” problem. Stacking these three ensures fast and stable prefill for ultra-long sequences.
4. Compilation and scheduling are hidden multipliers. Graph mode, compilation caching, ultra-large granularity fusion, and multi-stream orchestration are conventional methods individually, but combined, they dictate actual throughput and decoding latency in scenarios with numerous small operators.
5. State management is a new challenge for new architectures. Recurrent states, index data, and the KV cache jointly constitute the three core states during inference. In speculative decoding scenarios, the state replay mechanism reduces the cost of “verification failure” from recomputing the entire sequence to simply rolling back and replaying the state.

In summary, Day0 inference support for new architectural models is a combination of a “reserve force” and a “blitzkrieg”: proactively building predictable general capabilities into reusable assets, and converging unpredictable architectural differences into a few modules for rapid breakthrough. Through coordinated optimization across operator fusion, layered quantization, parallelism reconstruction, and compilation scheduling, new architectures can achieve usable—even excellent—performance on their very first day of deployment. For teams planning to introduce hybrid architecture models, this methodology can be directly migrated to their own inference platform engineering practices.