Tackling the Dynamicity in a Production LLM Serving System with SOTA Optimizations via Hybrid Prefill/Decode/Verify Scheduling on Efficient Meta-kernels

The introduction of new technologies, such as APC, SD, and SplitFuse, increases the diversity of Attention shapes and mask structures, leading to MFU and MBU issues of the current NPU attention kernel (torch-npu 2.1 FusedInferAttentionScore [10]), as illustrated in Fig 3.

Firstly, for prefill Attention, without any optimizations, the query length equals the key-value length, resulting in a square-shaped Attention score matrix and a lower triangular mask. This shape is ideal for optimization[36], as the sparsity in the mask can be exploited to achieve high performance. The current state-of-the-art (SOTA) torch-npu kernel efficiently handles such square-shaped inputs with triangular masks, achieving an MFU of $53\%$.

However, in real-world scenarios involving prefixes, parts of the key and value are reused from the K/V cache[8, 47], transforming the Attention score shape from a square to a rectangle with arbitrary dimensions. For these rectangular shapes, the performance of the torch-npu kernel degrades significantly, with MFU dropping to $47\%$ and $30\%$, respectively. While the reuse of the K/V cache theoretically reduces computation, the decline in kernel efficiency offsets this benefit, resulting in no meaningful end-to-end performance improvement.

Secondly, unlike Linear, where tokens from different batches can share weights, decode Attention requires each token to have its own independent K/V cache, making reuse impossible. This leads to inherently memory-bounded computations. PagedAttention (PA)[31] further fragments the K/V cache access patterns, limiting access to small block-sized portions instead of large contiguous blocks. As illustrated in Fig. 3(b), the MBU decreases as block size becomes small, exacerbating the memory bottleneck of the decode stage.

Lastly, Speculative Decoding (SD)[20, 33, 38] reshapes decode Attention to resemble prefill Attention but with a shorter query length. Different SD algorithms generate tokens with varying causal structures, resulting in diverse masks during the verify stage. However, due to the absence of a verify-specific kernel for Ascend, the prefill Attention kernel—equipped with full masks—is used as a fallback. This fallback prevents verify-specific optimizations, leading to MBU lower than $30\%$.

GEMM is a core operation in LLMs. Beyond the four GEMM operations in the Linear section, the query-key ($QK$) and score-value ($SV$) computations in Attention also rely on GEMM. We summarize all of GEMM operations in Fig. 4(a). For Linear operations, the $M$ dimension is tied to the number of tokens. For Attention $QK$ and $SV$ operations, the $N$ and $K$ dimensions are dependent on the token K/V length. As previously discussed, the number of tokens in the prefill and verify stages is highly dynamic, and this variability is further amplified by the introduction of Automatic Prefix Caching (APC)[8, 47], making token lengths even more unpredictable.

In addition, for Linear GEMM, the dimensions $N$ and $K$ are influenced by $hiddenSize$, which varies with the architecture of the LLM. Consequently, the $M$, $N$, and $K$ dimensions in GEMM operations for LLM are all dynamically changing.

Designing a matrix multiplication on AI accelerators to support arbitrary shapes while ensuring high performance is inherently challenging. Most accelerators are typically optimized for specific tile sizes, such as $16\times 16$, making it difficult to fully utilize their capabilities with arbitrary shapes. Irregular shapes complicate parallelism and load balancing, causing inefficiencies in workload distribution across processing units. The requirement for flexible algorithms to adapt to diverse shapes increases algorithmic complexity. Additionally, handling boundary conditions for non-uniform matrix sizes introduces overhead, further affecting performance. As shown in Fig. 4(b), employing a general-purpose GEMM kernel (torch-npu 2.1 linear operator[9]) to accommodate all possible results in unstable performance and fails to achieve optimal utilization across all GEMM operations.

In practical systems, P/D/V stages may exist independently or coexist simultaneously, leading to arbitrary combinations of interleaved P/D/V stages within a given scheduling budget. For Linear operations, tokens from different stages can be grouped together and treated as the left matrix in a GEMM operation, sharing the same weight matrix on the right. This approach is straightforward, as tokens from different stages can reuse the same large model weights.

In contrast, handling Attention operations is significantly more complex. Stages are independent, and even within the same stage, batches are also independent. Enumerating and tailoring optimizations for each possible combination of stages and batches is highly labor-intensive and impractical.

A common alternative is a batch-by-batch execution, such as selective batching[44], which processes each batch independently by invoking the corresponding Attention kernel. However, this method introduces additional memory overhead from splitting, rearranging, and merging data, which degrades overall system performance. As shown in Fig. 3(d), the memory overhead may account for more than $50\%$. Furthermore, batch-by-batch processing leads to inefficient utilization of computational resources, further limiting system efficiency.

Built on Huawei’s DaVinci architecture [35, 34], the Huawei Ascend NPU is a high-performance AI processor. Fig. 5 illustrates the micro-architecture of the Ascend, which consists primarily of AIC and AIV components. The AIC, similar to Nvidia’s Tensor Core, handles matrix computations, while the AIV is responsible for vector operations. AIC and AIV are separated without a direct datapath, so ensuring their data interactions occur via the L2 cache is crucial when designing mixed kernels. Compared to GPUs, NPUs have larger core granularity, making load balancing between cores even more critical. The Memory Transfer Engine (MTE) handles data movement. Whether AIC, AIV, or MTE, all data processing and transferring occur at the tile level. AI accelerators with tile-based programming models are gaining increasing attention in the community. However, this tile-based processing model encounters significant challenges when managing dynamic workloads that vary at the token granularity. Typical solutions involve complex padding/unpadding operations, which result in wasted computation and memory, leading to substantial overhead.

When a user’s request enters the system, it first passes through the APC module, which matches the incoming prompt against existing prompts in the K/V cache, enabling token-wise reuse. Any unmatched tokens are added to the scheduling queues. Consequently, the prompt length in the scheduling queues is the user’s input length minus the length of tokens already cached in the K/V cache. Since both the user’s input and the cached token lengths are dynamically variable, the prompt lengths in the scheduling queues become even more dynamic. The scheduling queues also include decode and speculative tokens from previous requests awaiting processing.

Our scheduling system selects a fixed-budget length of tokens from the scheduling queues to form chunks, which may include tokens from prefill, verify, or decode stages. To improve first-token latency, prefill requests are prioritized. If a user’s prefill prompt exceeds the budget length, it is split into smaller parts to ensure each scheduled chunk remains within the budget. To minimize interruptions caused by prefill on decode and maintain a stable Time Between Tokens (TBT), certain slots are reserved for decode and speculative tokens. Prefill-only chunks are scheduled only when the queue has no decode or speculative tokens.

As shown in Fig. 6(a), the composition of P/D/V stages within a single chunk is inherently unpredictable, with token counts for each stage varying arbitrarily and each stage having a distinct historical K/V length. Additionally, the total token count in a chunk may not always match the budgeted length and can fall below the budget under a low system load. To address the four levels of dynamism, our scheduling operates entirely at the granularity of individual tokens without concern for their origin.

While token-wise scheduling improves efficiency by reducing bubbles and optimizing resource utilization, the four levels of dynamism it introduces pose significant challenges for execution, especially on AI accelerators with tile-based programming models.

To address these challenges, we propose a dynamic decomposition mechanism that converts dynamic workloads into hardware-friendly, tile-based computational units. Using the Token-Table, each stage is logically decomposed into tile blocks. At the tile level, computation modules can process these blocks in parallel without distinguishing their P/D/V stage origin. Importantly, this tiling decomposition is purely logical, requiring no changes to the physical data layout.

In the following sections, we present a detailed analysis of how tiling decomposition is applied to Attention and Linear layers.

After decomposing the dynamic workloads from P/D/V mixed stages into fundamental tile units, it is necessary to reorder these tile units and generate a Task-Table to enhance performance. The Task-Table is responsible for scheduling these tile units onto the hardware, with each entry specifying a coreID and the list of tiles assigned to that core. Based on this Task-Table, Attention, and Linear can simply retrieve the corresponding tile units according to their coreID. This approach not only maximizes hardware efficiency but also simplifies the design of Attention and Linear kernels.

As illustrated in Fig. 12(a), the dimension $M$ in GEMM is intrinsically related to the number of tokens. In the case of Linear GEMM, the $M$ dimension corresponds to the cumulative token count across P/D/V stages, while in Attention GEMM, it is determined by the token count in each stage. In cube-based or tensor-based AI accelerators, matrix computations are typically constrained by a minimum tiling size (e.g., $16\times 16$ for cube cores). A common practice is to pad the $M$ dimension to align with a multiple of the tiling size. However, this dynamic padding introduces non-trivial memory overhead and degrades performance.

To mitigate these issues, we replace the physical padding in global memory with virtual padding on the chip, combined with the selective read and write mechanisms shown in Fig. 12(a). This approach allows for efficient handling of matrices with arbitrary shapes. Specifically, on-chip buffer allocations are made in tiling-size units to fully exploit the hardware’s computational potential. During data transfer from global memory to the on-chip buffer, selective read operations copy only the actual, non-padding data. Similarly, selective write operations ensure that only the non-padding outputs are written back to global memory.

Because the virtual padding regions do not interfere with the computational results of the non-padding regions, this approach guarantees the correctness of the matrix computation results. Moreover, by limiting computations to fixed-shape matrix multiplications, we can apply highly customized optimizations for these fixed shapes, achieving both high efficiency and flexibility for dynamic workloads.

As illustrated in Fig. 4(a), the dimensions of $N$ and $K$ in linear operations are determined by the model structure. These dimensions are typically multiples of hardware tile size, such as $16$, eliminating the need for additional padding.

For Attention GEMM ops, dimensions $N$ and $K$ are tied to the sequence length, which can take arbitrary values. In theory, padding would be required for these dimensions. However, this is naturally handled by the K/V cache’s block page structure, which stores and reads data in blocks aligned to multiples of tile size. As shown in Fig. 12(b), by reading entire blocks, the read length inherently conforms to the required hardware tile size. Furthermore, this padding does not affect the final Attention computation because we explicitly set the values in the padded regions to zero during $Softmax$ calculation. This ensures that the padded values are effectively excluded from the Attention score. Additionally, since the padded data is reduced during the Attention $score$ $\times$ $value$ operation, it does not influence the shape of the final output.

Therefore, matrix multiplication for arbitrary shapes can be efficiently supported without introducing additional padding overhead for the $N$ and $K$ dimensions.

Given our focus on fixed shapes, such as multiples of the tiling size, and operating under budget constraints, we narrow the range of token sizes we handle. While the token count can vary significantly across different workloads, we apply a smoothing technique that reduces the shapes involved to fixed, well-defined sizes. This enables us to perform offline customization and optimization for these specific shapes, particularly for optimizations that are difficult to implement statically, such as swizzling[15].

Swizzling is a technique that optimizes matrix multiplication performance by altering the task allocation order across AI cores, enhancing L2 cache hit rates. However, determining the optimal allocation strategy for swizzling is complex and cannot be easily derived through theoretical formulas. To address this challenge, we conduct offline profiling to explore possible access patterns and identify the most efficient inter-core distribution strategy.

These optimal configurations are stored for future use. During online execution, we translate the dynamic workload into the nearest supported fixed shape. Using this fixed shape, we retrieve the pre-optimized configuration from the offline profiling results. This approach ensures that the dynamic workload is handled with the best configuration, enabling optimal performance during execution. By leveraging static optimizations, we can fully exploit the hardware’s potential, even in the face of dynamic workloads.

For the $QKV$ input, its shape is $[tokens,heads+2\times kvHeads,headSize]$, where the outermost dimension corresponds to the tokens, and each token contains its $QKV$ fusion. The attention operation is performed in parallel along the Head dimension, which requires the Head dimension to be positioned as the outermost dimension. Moreover, we need to read the $Q$, $K$, and $V$ separately rather than the fused $QKV$ tensor. Typically, $Split$ and $Permute$ operations are introduced to reorient the tensor for efficient computation. After the attention computation, another $Permute$ operation is applied to transform the output $O$ back to the shape $[tokens,heads,headSize]$.

To reduce memory overhead, we fuse $Split$ and $Permute$ directly into GEMM computations shown in Fig. 13. For the $QK$ operation, tile-based and stride-based reads are employed to access the required data directly from the $QKV$ tensor, eliminating the need for separate $Split$ and $Permute$ steps. Similarly, for the $SV$ computation, tile-based and stride-based writes are used to store the results in a format that is directly aligned with the $OProj$. This integration removes explicit $Permute$ operations, allowing the $SV$ output to seamlessly flow into subsequent matrix multiplication operations.

We built an Ascend-native inference system based on vLLM[31], leveraging Ascend intrinsic to implement core modules such as SmoothGEMM, Meta-Attention, and other essential operators like normalization, activation, and embedding. These operators were exposed to the Python API via pybind11[7] and seamlessly replaced the corresponding GPU kernels in vLLM, enabling vLLM to support the Ascend NPU.

To reduce the overhead of frequent Python calls, we offloaded the entire model-forward process to C++, integrating the optimized operators into a single C++ function. This model-forward function is then exposed to vLLM for invocation, ensuring a streamlined and efficient execution path.

Additionally, we replaced the native vLLM scheduler with a token-wise scheduling strategy, integrating workload decomposition and computation task reordering to better support dynamic workloads. We redesigned the speculative decoding framework in vLLM, enabling token tree construction and metadata generation for Meta-Attention. These enhancements allow us to implement tree-based speculative decoding algorithms, such as Lookahead Decoding[46], further improving inference performance.

In this section, we evaluate the performance of the attention kernel under dynamic workloads typically encountered in real-world systems. The comparison targets are PromptFlashAttention (PFA) [12] and IncreFlashAttention (IFA) [11] from torch-npu 2.1[6].

In practical systems, the input length from users can vary arbitrarily, ranging from $0$ to the maximum length supported by the model. For long sequences, to minimize the impact of prefill on decode and to optimize sequence parallelization, we typically adopt the Chunked Prefill strategy, which imposes a constraint on the maximum chunk length, such as $4k$. In real-world scenarios, lengths smaller than the chunk size may also be encountered. To evaluate performance across different conditions, we assess the LLMs with input lengths ranging from $1$ to the chunk size ($4096$).

We compare the performance of linear operators ($QKV$, $OProj$, GateUp, and $Down$) using shapes derived from Llama2-7B and Qwen2-7B with TP=$1$. As displayed in Fig. 17, the results indicate that SmoothGEMM outperforms torch-npu linear by an average of $14.6\%$, demonstrating superior performance across nearly all tested shapes. Moreover, the performance remains stable, with ideal MFU typically achieved for sequence sizes above $1k$.