A multi-level kv cache operation primitive extension system and management method implemented in a triton framework

Abstract

This invention provides a multi-level KV cache operation primitive extension system and management method implemented within the Triton framework. It addresses the collaborative scheduling of L1 video memory, L2 distributed DRAM, and L3 native KV storage. The system designs loading, storage, and management primitives, and carries hierarchical, prefetching, hotness, snapshot, and transaction control semantics at the TritonIR ​​layer. Cross-chip optimization and machine instruction generation are completed at the LinalgIR / MLIR layer. Combined with a multi-level cache scheduler, it achieves cross-level migration, transaction control and rollback based on a four-dimensional hotness matrix, as well as execution feedback closed-loop updates. This scheme enables unified programming control, dynamic scheduling, and heterogeneous chip adaptation for multi-level KV caches, improving cache management efficiency and execution reliability in large-model inference scenarios.

Description

Technical Field

[0001] This invention relates to the field of large language model inference optimization and compiler extension technology, specifically to a multi-level KV cache operation primitive extension system and management method implemented within the Triton framework, which is suitable for accelerating large model inference in heterogeneous chip scenarios and can realize unified programming control and storage-computation co-optimization of multi-level cache. Background Technology

[0002] (I) Current Status of Triton Compiler Technology The Triton compiler is an open-source compiler for heterogeneous computing. Its core advantage lies in its layered architecture: "Python layer API - TritonIR ​​intermediate representation - LinalgIR (MLIR) optimization - hardware instruction generation." This architecture enables "write once, compile anywhere" cross-chip adaptation, significantly reducing the development cost of heterogeneous operators. The core process is as follows: upper-layer applications call Triton primitives (such as tl.load and tl.store) through the Python layer to describe the computational logic. This is then transformed into an intermediate representation by TritonIR, optimized by MLIR Passes, and finally generates native machine instructions adapted to different chips such as NVIDIA GPUs and Huawei Ascend NPUs.

[0003] (II) Current Status of Multi-Level Key-Value Caching Technology In large language model inference, the storage and scheduling of the KV cache (key-value vectors of historical tokens) directly affects inference efficiency. To balance performance, capacity, and cost, the industry generally adopts a multi-level cache architecture of L1-L2-L3: L1 cache (GPU memory) stores active data to ensure low-latency access; L2 cache (distributed DRAM) serves as an overflow layer to achieve cross-node data sharing; and L3 cache (native KV storage, built on NVMe SSDs) stores low-frequency data to provide large-capacity expansion. However, existing multi-level KV cache scheduling schemes lack a unified programming control method and require complex adaptation between upper-layer frameworks (such as PyTorch) and lower-layer hardware drivers, resulting in low coordination efficiency and difficulties in cross-chip adaptation.

[0004] In existing technologies, Triton primitives only support basic memory / video memory access and do not adapt to the hierarchical storage characteristics of multi-level KV caches, resulting in the following core defects: 1. Lack of multi-level cache control capabilities: Triton's native primitives cannot recognize L1-L2-L3 storage levels, requiring upper-layer frameworks to encapsulate hierarchical scheduling logic, which increases development costs and reduces collaboration efficiency; 2. High cost of cross-chip adaptation: The multi-level cache architecture of different chips varies greatly. Existing solutions require custom development of scheduling code for each chip, resulting in an "M×N" adaptation dilemma (M is the number of models, and N is the chip type). 3. Insufficient storage-computation synergy: Native primitives do not deeply integrate caching operations with computation operations, making it impossible to optimize hidden data migration latency through prefetching, parallel execution, etc. 4. Lack of cache lifecycle management: It lacks primitive-level cache isolation, snapshot, and eviction capabilities, which can easily lead to data conflicts in high-concurrency scenarios and make it impossible to quickly recover the cache after a session is interrupted.

[0005] In summary, existing technologies lack Triton operation primitives and extension schemes that are compatible with multi-level KV caches, making it impossible to achieve unified programming control and efficient cross-chip adaptation of multi-level caches, which restricts the improvement of large model inference efficiency. Summary of the Invention

[0006] To overcome the above problems, the present invention aims to provide a multi-level KV cache operation primitive extension system and management method implemented within the Triton framework. By designing three types of core primitives (multi-level storage-aware loading primitives, multi-level storage-aware storage primitives, and multi-level cache lifecycle management primitives) within the Triton system, and matching them with optimized passes for the TritonIR / LinalgIR / MLIR layers and a built-in multi-level cache scheduler, the collaboration between primitives, compiler passes, and underlying instructions is realized, thereby achieving multi-level KV cache management with cross-chip, hierarchical awareness, hot-driven, and transaction-guaranteed features.

[0007] To solve the above-mentioned technical problems, the present invention adopts the following technical solution.

[0008] An extension system of multi-level key-value cache operation primitives implemented within the Triton framework, comprising: The primitive set is used to support the loading, storage and full lifecycle management of L1 video memory, L2 distributed DRAM and L3 native KV memory. The primitive set includes at least loading primitives for multi-level KV data loading, storage primitives for multi-level KV data storage, and management primitives for cache isolation, snapshots, migration control, eviction and resource statistics. The TritonIR ​​semantic bearer module is used to map the loading primitives, storage primitives and management primitives to corresponding intermediate representation instructions, and to carry hierarchical, prefetching, hotness, snapshot and transaction control semantics. The compilation optimization module is used to perform at least one optimization process in the TritonIR ​​layer, including prefetching and computational parallelism, heat analysis and hierarchical adjustment, and instruction conflict detection, and to complete cross-chip optimization and machine instruction generation in the LinalgIR / MLIR layer. A multi-level cache scheduler is used to generate primitive invocation strategies based on scheduling requirements and heat matrix, and to coordinate primitive execution, status feedback and exception handling. The execution adaptation module is used to send the machine instructions to the execution paths corresponding to L1 video memory, L2 distributed DRAM and L3 native KV storage respectively, so as to complete multi-level KV cache operation.

[0009] Furthermore, the multi-level cache scheduler includes at least a demand parsing module, a heat perception module, a primitive scheduling strategy generation module, and an execution coordination module; wherein, the demand parsing module is used to receive at least one scheduling parameter among request ID, data access mode, priority, and sequence length, and generate structured scheduling requirements.

[0010] Furthermore, the heat perception module is used to construct and update a heat matrix based on primitive execution feedback. The heat matrix includes at least data address, access frequency, storage level, and request priority, and is used to drive cross-level migration of data between L1 video memory, L2 distributed DRAM, and L3 native KV storage.

[0011] Furthermore, the management primitives are used to perform transaction control on at least one composite operation among cross-level migration, snapshot creation, and snapshot recovery, and trigger rollback in the event of execution exceptions or verification failures to ensure data consistency.

[0012] Furthermore, the compilation optimization module performs at least one of the following processes in the TritonIR ​​layer: instruction conflict detection and prefetch and computation parallel optimization. Instruction conflict detection is used to analyze the access address or execution cycle of primitive instructions and to avoid or sequence potential read-write conflicts. Prefetch and computation parallel optimization is used to rearrange the timing of prefetch instructions and computation instructions to hide data migration delays.

[0013] Furthermore, the compilation optimization module at the LinalgIR / MLIR layer completes the adaptation to different chip platforms or different storage paths through dedicated optimization passes for loading primitives, storage primitives, and management primitives, and performs at least one of the following processes during data loading or data storage: data type conversion, compression, and integrity verification.

[0014] Furthermore, the multi-level cache scheduler is also used to trigger a heat-driven cross-level migration strategy based on heat changes in the heat matrix, so that high-heat data is preferentially residing in the low-latency layer, and low-heat data is migrated to the high-capacity layer.

[0015] In addition, this invention also discloses a multi-level key-value cache management method implemented within the Triton framework, comprising the following steps: S11: Receive KV cache scheduling requests through the application interface adaptation layer, and perform parameter verification and scheduling preprocessing; S12: Call the loading primitive, storage primitive, and management primitive, and encapsulate the primitive parameters into an intermediate representation structure that can be recognized by the TritonIR ​​layer; S13: Perform at least one optimization process in the TritonIR ​​layer, including prefetching and computational parallelism, heat analysis and hierarchical adjustment, and instruction conflict detection; S14: Perform cross-chip optimization and generate target hardware machine instructions at the LinalgIR / MLIR layer; S15: Based on the machine instructions, drive the L1 video memory, L2 distributed DRAM and L3 native KV storage to complete the corresponding cache operations, and feed the execution results back to the scheduler; S16: Update the heat matrix, scheduling policy and resource status based on execution feedback, and perform rollback or rescheduling in case of abnormal situations.

[0016] Furthermore, in step S16, a heat matrix containing data address, access frequency, storage level and request priority is constructed or updated, and cross-level migration is triggered based on the heat matrix; during cross-level migration, snapshot creation or snapshot restoration, if any abnormality or verification failure occurs in any link, a rollback is triggered through management primitives to ensure data consistency.

[0017] Furthermore, the present invention also discloses a computer-readable storage medium storing a computer program thereon, wherein the computer program, when executed by a processor, implements the management method described in the multi-level KV cache implemented within the Triton framework.

[0018] The multi-level KV cache operation primitive extension system and management method implemented in the Triton framework of this invention have the following advantages: 1) Primitives are deeply integrated with the compiler, enabling scheduling decisions to be directly mapped to hardware-level machine instructions, reducing instruction adaptation overhead and improving execution efficiency; 2) The hierarchical perception and heat-driven dynamic migration and prefetching mechanism significantly reduce the average access latency and improve the hit rate, adapting to large model long context reasoning scenarios. 3) Cross-chip and cross-level differentiated optimization (NVLink / RDMA / DMA / NVMe) improves the overall throughput and resource utilization of heterogeneous systems; 4) The combination of transactional and static analysis ensures the consistency and reliability of complex operations such as migration and snapshots; 5) Good scalability: When adding new hardware or new compression / quantization methods, only the corresponding kv-*-tl-opt implementation needs to be added to the LinalgIR layer for quick adaptation. Attached Figure Description

[0019] Figure 1 This is a schematic diagram of the overall architecture of the Triton operation primitive extension and optimization method for multi-level KV caching in this invention. It illustrates the collaborative relationship between the three primitives (kv_load_tl, kv_store_tl, kv_cache_ctrl_tl), the Triton compilation extension layer (TritonIR, LinalgIR / MLIR Passes), and the execution adaptation layer (video memory controller, RDMA / NIC, NVMe), as well as the built-in multi-level cache scheduler structure.

[0020] Figure 2 This is a schematic diagram of the primitive-based collaborative scheduling architecture, illustrating the interaction between the scheduler's four modules (requirement parsing, heat perception, primitive scheduling strategy generation module, and execution coordination) and the Triton underlying language and compiler Passes.

[0021] Figure 3 This diagram illustrates the technical architecture of the primitive collaborative scheduling process, showing a five-step closed-loop process (application requirements → primitive invocation → TritonIR ​​Passes → LinalgIR / MLIR adaptation → hardware execution and feedback) and the inputs and outputs of each step.

[0022] Figure 4 This document illustrates the semantic carrying and extension of the three core primitives at different implementation levels (Python API, TritonIR ​​instructions, LinalgIR / MLIR implementations), including instruction format examples and annotations of common parameter fields.

[0023] Figure 5 This diagram illustrates the transactional migration and rollback process, showing the two-phase commit process for cross-level migration, the command conflict detection trigger point, the rollback path for verification failure, and the status feedback callback path.

[0024] Figure 6 The flowchart of the multi-level KV cache management method of the present invention shows the processing from receiving scheduling requirements from the application interface adaptation layer to primitive calls, TritonIR ​​optimization, LinalgIR / MLIR cross-chip optimization, hardware execution, and feedback updates. Detailed Implementation

[0025] To better illustrate the present invention, the following detailed description is provided in conjunction with specific embodiments. However, these specific embodiments are merely for illustrative purposes and are not intended to limit the scope of the invention.

[0026] This invention discloses a multi-level KV cache operation primitive extension system and management method implemented within the Triton framework. It addresses the collaborative scheduling of L1 video memory, L2 distributed DRAM, and L3 native KV storage, designing loading, storage, and management primitives. The system carries hierarchical, prefetching, hotness, snapshot, and transaction control semantics at the TritonIR ​​layer, and performs cross-chip optimization and machine instruction generation at the LinalgIR / MLIR layer. Combined with a multi-level cache scheduler, it achieves cross-level migration, transaction control and rollback based on a four-dimensional hotness matrix, as well as execution feedback closed-loop updates. This scheme enables unified programming control, dynamic scheduling, and heterogeneous chip adaptation for multi-level KV caches, improving cache management efficiency and execution reliability in large-model inference scenarios.

[0027] like Figure 1 As shown, the overall technical architecture of this invention, from top to bottom, includes an upper-layer application, an application interface adaptation layer, a Triton compilation extension core layer, a three-level cache execution adaptation layer, and underlying hardware and L1-L2-L3 multi-level storage layers. The upper-layer application initiates KV cache scheduling requests, and the application interface adaptation layer receives and standardizes these scheduling requests. The Triton compilation extension core layer includes a multi-level cache scheduler and Triton compilation extension modules, used to organize primitive calls, carry semantics, optimize compilation, and coordinate execution. The three-level cache execution adaptation layer distributes the compiled target instructions to the corresponding execution paths of L1 video memory, L2 distributed DRAM, and L3 native KV storage. The underlying hardware actually performs data loading, migration, storage, snapshotting, and rollback operations. Thus, this invention achieves end-to-end collaboration from application requirement input, primitive calls, compilation optimization to multi-level storage execution.

[0028] like Figure 2 As shown, the multi-level cache scheduler includes a requirement parsing module, a heat perception module, a primitive scheduling strategy generation module, and an execution coordination module. The requirement parsing module receives scheduling parameters such as request ID, access mode, priority, and sequence length from the application interface adaptation layer and generates structured scheduling requirements. The heat perception module maintains a heat matrix based on primitive execution feedback, identifying changes in storage heat for different data across L1, L2, and L3. The primitive scheduling strategy generation module generates the calling order and parameter configuration for three types of core primitives based on the structured scheduling requirements and the heat matrix. The execution coordination module links the optimization passes of the TritonIR ​​layer and the LinalgIR / MLIR layer, coordinating primitive execution timing, conflict avoidance, and state feedback to form a closed-loop collaborative architecture of primitive scheduling, compilation optimization, and execution feedback.

[0029] like Figure 3As shown, the primitive collaborative scheduling process of this invention includes five stages: application requirement input and scheduling preprocessing, Python layer primitive call and low-level instruction encapsulation, TritonIR ​​layer semantic transformation and Pass fusion optimization, LinalgIR / MLIR layer cross-chip optimization and machine instruction generation, hardware execution and status feedback. The upper-layer application first submits a KV cache scheduling requirement through the application interface adaptation layer. Then, the scheduler calls three types of primitives: tl.kv_load_tl, tl.kv_store_tl, and tl.kv_cache_ctrl_tl, and encapsulates the primitive parameters into an intermediate representation structure. Afterward, optimization processing such as prefetching parallelism, heat analysis, and conflict detection is completed at the TritonIR ​​layer. Then, the LinalgIR / MLIR layer generates execution instructions adapted to the target hardware. Finally, the execution adaptation layer drives L1, L2, and L3 to complete the corresponding cache operations and feeds back the execution status to the scheduler to update the heat matrix and trigger subsequent scheduling or rollback processing.

[0030] The scheduling logic for each level is implemented accordingly: L1 (video memory) focuses on low latency and high frequency access: it utilizes direct memory access instructions in the underlying language, combined with prefetching and computational parallel passes in the IR layer, to ensure that prefetching can mask video memory transfer latency.

[0031] L2 (Distributed DRAM) focuses on cross-node sharing and overflow buffering: it optimizes bandwidth and concurrent transmission by utilizing RDMA / NIC encapsulation instructions and cross-node migration optimization passes of the LinalgIR layer.

[0032] L3 (native KV storage, such as NVMe SSD) focuses on massive low-frequency storage and persistence: it encapsulates the NVMe parallel IO interface and LinalgIR's persistent storage optimization pass through the underlying language, supporting parallel IO, asynchronous commit and batch compression / decompression.

[0033] Primitive Cooperative Scheduling Process (Five-Step Closed Loop): Step 1 (Application Requirement Input and Scheduling Preprocessing): The upper-layer application submits a scheduling request through the adaptation layer, and the scheduler's requirement parsing module performs preliminary scheduling planning in conjunction with the heat matrix; Step 2 (Python layer primitive calls and underlying instruction encapsulation): The scheduler calls the Python API of the three primitives. The underlying language encapsulates the primitive parameters into an intermediate representation that TritonIR ​​can recognize, while performing type and validity checks. Step 3 (TritonIR ​​layer semantic transformation and Passes fusion optimization): The intermediate representation is converted into TritonIR ​​instructions and a series of Passes are executed (such as prefetch scheduling Pass, heat analysis Pass, instruction conflict detection Pass) to generate an instruction flow that integrates computing and storage logic; Step 4 (Cross-chip optimization and machine instruction generation at the LinalgIR layer): After TritonIR ​​is mapped to LinalgIR, the kv-*-tl-opt series MLIR Pass completes cross-chip adaptation (NVLink / RDMA / DMA / NVMe) and generates native machine instructions for the target chip; Step 5 (Hardware Execution and Status Feedback): Machine instructions are submitted by the execution adaptation layer to each hardware controller (graphics memory controller, RDMA / NIC, NVMe controller). Execution results and runtime statistics are fed back to the scheduler through the underlying language to update the heat matrix and trigger subsequent strategies. A rollback transaction is triggered when an anomaly occurs.

[0034] Consistency and reliability guarantees: Transactional management: A transaction mechanism is adopted for complex operations such as cross-level migration and snapshot creation / recovery. The transaction ID and status flag of tl.kv_cache_ctrl_tl are used to ensure atomicity (two-phase commit / rollback).

[0035] Instruction conflict detection Pass: Static analysis of access addresses and execution cycles to predict potential conflicts and resolve them through underlying language atomic instructions or serialization strategies.

[0036] Verification and Rollback: Storage operations are integrated with verification (CRC / SHA). If verification fails or an error occurs, rollback logic is triggered and the abnormal status is reported.

[0037] like Figure 4As shown, the three core primitives in this invention correspond to different levels of unified semantic carrying and implementation extension relationships. Among them, `tl.kv_load_tl` is used to implement multi-level KV data loading. In the Python API layer, it receives parameters such as level identifier, data address, data type, prefetch switch, and request ID. In the TritonIR ​​layer, it forms corresponding semantic instructions to carry control information such as prefetching, level, and request isolation. In the LinalgIR / MLIR layer, it is further mapped to the underlying execution logic adapted to the target chip and target storage path. `tl.kv_store_tl` is used to implement multi-level KV data writing. In the Python API layer, it receives parameters such as target level, data popularity, compression mode, and request ID. In the TritonIR ​​layer, it carries popularity-driven and write-back control semantics. In the LinalgIR / MLIR layer, it completes differentiated compression, storage write-back, and integrity verification. `tl.kv_cache_ctrl_tl` is used to implement full lifecycle management such as cache isolation, snapshot creation, snapshot recovery, migration control, eviction, and resource statistics. In each implementation layer, it carries management parameters, transaction control semantics, and the generation of underlying management instructions.

[0038] thus, Figure 4 This illustrates the unified mapping relationship between the three types of primitives from the upper-level interface to the intermediate representation and then to the lower-level implementation.

[0039] like Figure 5 As shown, when the system performs cross-level migration based on heat analysis results or resource status, the management primitive first initiates migration control and establishes a corresponding transaction context. Then, the loading primitive reads the target KV data from the source level, and the storage primitive writes the data to the target level. During the migration process, the compilation optimization layer performs conflict detection on relevant instructions and performs consistency checks on the migration results. When the check passes, the migration results are committed and the data level status is updated. When a conflict, exception, or check failure occurs, a rollback process is triggered, incomplete writes in the target level are undone, and the source level status is restored. Simultaneously, the failure result is fed back to the scheduler for rescheduling.

[0040] Therefore, this invention can ensure consistency and reliability in complex operations such as multi-level cache migration and snapshot recovery.

[0041] like Figure 6 As shown, the method flow of the present invention includes the following steps: S11: Receive KV cache scheduling requests through the application interface adaptation layer and complete parameter verification and scheduling preprocessing; S12: Call the three types of core primitives and encapsulate the primitive parameters into an intermediate representation structure that can be recognized by the TritonIR ​​layer; S13: Perform optimization passes such as prefetch parallelism, heat analysis, hierarchy adjustment and conflict detection in the TritonIR ​​layer; S14: Perform cross-chip optimization and generate target hardware machine instructions at the LinalgIR / MLIR layer; S15: Drives L1 video memory, L2 distributed DRAM and L3 native KV storage to complete the corresponding loading, migration, storage or management operations, and feeds back the execution results to the scheduler; S16: Update the heat matrix, scheduling policy and resource status based on execution feedback, and trigger rollback or rescheduling if necessary.

[0042] Through the above steps, this invention achieves unified control, dynamic scheduling, and execution feedback closed loop for multi-level KV cache in heterogeneous chip environments.

[0043] Innovative Triton operation primitive system design: The design incorporates three core primitives, covering the entire process of multi-level key-value cache loading, storage, and lifecycle management, with built-in hierarchical awareness, heat-driven operation, and cross-chip adaptation capabilities. (1) Multi-level storage-aware loading primitive (tl.kv_load_tl) Functionality: Extends the Triton native tl.load primitive to support precise loading of KV data from any level from L1 to L2 to L3, integrates hierarchical-aware prefetching and zero-copy transmission logic, and enables parallel loading and computation operations.

[0044] Three-level extension implementation: Python layer API encapsulation: Defines a standardized interface with parameters including storage level identifier (0=L1, 1=L2, 2=L3), data address, data type, prefetch switch, request ID (for cache isolation), etc., with the interface format as follows: `tl.kv_load_tl( addr: tl.pointer, # Memory address of KV data at the target level)` level: tl.int32, # Storage level identifier: 0=L1 (video memory), 1=L2 (distributed DRAM), 2=L3 (native key-value storage) dtype: tl.dtype, # Data type (e.g., float16, bfloat16, int4) prefetch: tl.bool = False, # Prefetch switch: True = Preload data needed for subsequent calculations prefetch_addr: tl.pointer = None, # Address to prefetch data (only valid when prefetch=True) req_id: tl.int64 = 0 # Request ID: used for cache isolation to avoid data conflicts from multiple requests) -> tl.tensor # Returns the loaded KV tensor (a tensor concatenated with key and value) The TritonIR ​​layer semantics carry-over: The new triton.ir.op.kv_load_tl instruction encapsulates semantic information such as storage level, prefetch configuration, and request ID, enabling integration with attention computation IR; the new "prefetch and computation parallel" IR Pass embeds prefetch instructions into the idle period of attention computation, hiding migration latency.

[0045] Command format example: `%0=triton.ir.op.kv_load_tl%addr,%level{prefetch=true,prefetch_addr =%prefetch_addr, req_id = %req_id} : (!triton.pointer <f16>, i32) ->!triton.tensor <f16>` LinalgIR layer optimization implemented: A new kv-load-tl-opt MLIR Pass was added, achieving differentiated optimization based on the characteristics of the target chip: NVLink zero-copy transfer instructions were generated for NVIDIA GPUs, AscendCDMA instructions were called for Huawei Ascend NPUs, and NVMe SSD parallel IO instructions were adapted for L3 native KV storage; at the same time, data type quantization conversion (such as Int4 / FP8) was completed to ensure loading efficiency.

[0046] (2) Multi-level storage-aware storage primitive (tl.kv_store_tl) Features: Extends Triton's native tl.store primitive to support hierarchical storage of KV data to L1-L2-L3 based on data popularity, integrates differentiated compression logic, optimizes storage efficiency, and supports dynamic updates of popularity.

[0047] Three-level extension implementation: Python layer API encapsulation: Interface parameters include storage level, data to be stored, popularity marker, compression mode, request ID, etc., and support dynamic adjustment of storage level based on popularity.

[0048] The interface format is as follows: `tl.kv_store_tl(data: tl.tensor, # KV tensor to be stored` addr: tl.pointer, # Target level memory address level: tl.int32, # Target storage level: 0=L1, 1=L2, 2=L3 hot: tl.int32, # Data popularity flag: 0 = low frequency (stored in L3), 1 = medium to high frequency (stored in L2), 2 = high frequency (stored in L1) compress_mode: tl.int32 = 0, # Compression mode: 0 = no compression, 1 = FP8 quantization, 2 = Int4 quantization req_id: tl.int64 = 0 # Request ID: used for cache isolation) -> None` TritonIR ​​layer semantic support: A new `triton.ir.op.kv_store_tl` instruction has been added, encapsulating semantic information such as popularity tags and compression modes; a new "Popularity Analysis and Layer Adjustment" IR Pass has been added, dynamically updating popularity tags based on data access frequency and triggering cross-layer migration. Instruction format example: `triton.ir.op.kv_store_tl %data, %addr, %level {hot = %hot, compress_mode = %compress_mode, req_id = %req_id} : (!triton.tensor <f16>, !triton.pointer <f16>, i32) ->()` LinalgIR layer optimization implemented: The new kv-store-tl-opt MLIR Pass enables differentiated compression optimization (L1 lightweight compression, L2 / L3 deep compression), cross-chip storage instruction adaptation (interfacing with different chip storage controller interfaces), and storage integrity verification (integrated CRC check logic).

[0049] (3) Multi-level cache lifecycle management primitive (tl.kv_cache_ctrl_tl) Features: New dedicated primitives for full lifecycle management, covering core capabilities such as cache isolation, snapshot creation, snapshot recovery, cache eviction, and resource statistics, solving the problems of request interference and session interruption recovery in high-concurrency scenarios.

[0050] Three-level extension implementation: The Python layer API adopts a flexible design of "operation type + parameters" to support unified invocation of management functions such as cache isolation, snapshot creation, snapshot restoration, cache eviction, resource statistics, migration start, migration commit and rollback.

[0051] The interface format is as follows: `tl.kv_cache_ctrl_tl( op_type: tl.int32, # Operation type: 0=Cache isolation, 1=Create snapshot, 2=Restore snapshot, 3=Cache eviction, 4=Resource statistics req_id: tl.int64 = 0, # Request ID: used to identify the dedicated cache space snapshot_addr: tl.pointer = None, # Snapshot storage address (only valid when op_type=1 / 2) eviction_policy: tl.int32 = 0 # Eviction policy: 0 = priority-based LRU (default), 1 = FIFO ) -> tl.tensor # Returns the operation result: such as resource statistics, operation success indicator. Triton IR layer semantic support: A new `triton.ir.op.kv_cache_ctrl_tl` directive has been added, encapsulating semantic information such as operation type, request ID, and snapshot address; a new "Cache Isolation and Snapshot Management" IR Pass has been added, enabling efficient cache space allocation and snapshot data read / write based on request ID. Example of directive format: `%0 = triton.ir.op.kv_cache_ctrl_tl %op_type {req_id = %req_id,snapshot_addr = %snapshot_addr, eviction_policy = %eviction_policy} : (i32) ->!triton.tensor <i32>` LinalgIR layer optimization implemented: The new kv-cache-ctrl-tl-opt MLIR Pass enables hardware-level cache isolation (allocating dedicated space through chip memory isolation instructions), optimized parallel read / write of snapshot data, cross-chip eviction policy adaptation, and real-time resource statistics (collecting data such as cache utilization and hit rate at each level).

[0052] (II) Primitive Cooperative Scheduling Architecture (Implementation of L1-L2-L3 Multi-level KV Cache Scheduling Technology) The primitive-based collaborative scheduling architecture of this invention uses the Triton low-level language as the execution foundation and compiler Passes optimization as the efficiency core. Following the design principles of "multi-level cache hierarchy awareness, primitive instruction-driven, and dynamic hot-spot adaptation," it leverages the scheduler built into the core layer of the Triton compiler extension to achieve precise end-to-end scheduling of L1-L2-L3 multi-level key-value caches through the linkage of three core primitives. The core of the architecture relies on the Triton low-level language to achieve precise hardware-level control of primitive instructions. Through Passes optimization at the TritonIR / LinalgIR layer, it achieves deep collaboration between scheduling and computation, constructing a closed-loop driving mechanism of "primitive instruction generation (low-level language) - scheduling decision - Passes optimization - cache execution." The scheduler, as the core hub, parses application requirements and generates primitive calling strategies. The Triton low-level language ensures the hardware-adaptive execution of these strategies, while compiler Passes optimization improves the collaborative efficiency of scheduling and computation. These three elements work together to achieve efficient management of multi-level caches.

[0053] Figure 2 This is a schematic diagram of primitive collaborative scheduling architecture technology.

[0054] The specific technical implementation is as follows: (1) A "multi-level cache scheduler" is built into the Triton compilation extension core layer. As the core hub of primitive cooperative scheduling, it integrates four major functional modules. Each module is deeply dependent on the core capabilities of the Triton underlying language and the compiler Passes to achieve deep cooperation between scheduling decisions and primitive execution: The requirement parsing module receives requirement parameters (including request ID, sequence length, priority, data access mode, etc.) from the application interface adaptation layer. Through the standardized parsing interface encapsulated by the Triton underlying language, it generates structured scheduling requirements such as "load high-frequency KV data into L1 cache" and "migrate low-frequency data to L3 and create a snapshot". Its core advantage is the strong type validation feature of the underlying language interface, which can avoid scheduling anomalies caused by incorrect parameter formats and improve the reliability of requirement parsing.

[0055] The popularity awareness module: Based on the newly added "popularity analysis pass" in the TritonIR ​​layer, the core capability is realized. This pass tracks the data flow of primitive instructions (access instructions of tl.kv_load_tl and storage instructions of tl.kv_store_tl) generated by the Triton underlying language, counts the access frequency and period of data, and combines the resource statistics results fed back by tl.kv_cache_ctrl_tl to construct a four-dimensional popularity matrix of "data address-access frequency-storage level-request priority" and dynamically update the data popularity level. The technical principle is to use the data flow analysis capability of the compiler pass to complete the popularity collection without intruding on the primitive execution logic. The advantages are high collection efficiency and no additional interference to inference performance.

[0056] Primitive scheduling strategy generation module: Based on scheduling requirements and a heat matrix, it generates primitive invocation strategies (clearly defining primitive types, invocation order, and parameter configurations), and transforms the strategies into executable primitive instruction sequences through the Triton underlying language; for example, for high-frequency active data, it generates the invocation strategy "tl.kv_load_tl(level=0, prefetch=True)", and ensures a non-redundant conversion from strategy to instruction through the instruction concatenation feature of the underlying language; the core advantage is the native compatibility between the underlying language and the Triton compilation system, which can avoid instruction loss caused by cross-framework adaptation.

[0057] The execution coordination module relies on the TritonIR ​​layer's "Instruction Conflict Detection Pass" and the atomic operation instructions of the underlying language to achieve timing coordination and conflict resolution of primitive instructions. The Pass analyzes the memory access address and execution cycle of primitive instructions to predict and avoid conflict scenarios such as "the same data being operated on by tl.kv_load_tl and tl.kv_store_tl at the same time". The atomic operation instructions of the underlying language ensure the hardware-level reliable execution of the conflict avoidance logic. At the same time, through the execution status tracking capability of the compiler Pass, the execution results of primitives (cache hit status, migration progress) are fed back to the requirement parsing module to form a scheduling closed loop. The advantage is that the static analysis of the Pass + the dynamic execution coordination of the underlying language balances coordination efficiency and execution reliability.

[0058] (2) Based on the characteristics of different storage levels L1-L2-L3, precise control of hierarchical scheduling is achieved through a combination of "Triton low-level language custom primitive instructions + compiler Passes optimized execution logic". The core of the linkage between scheduling logic and primitives at each level depends on the collaborative support of the two, as follows: L1 cache (video memory) scheduling: focuses on "low-latency access to high-frequency active data", and relies on the direct memory access instructions of the Triton underlying language and the "prefetch and computation parallel pass" of the TritonIR ​​layer. Technical Implementation Details: ① The scheduler filters high-frequency data (heat level = 2) through a heat perception module, generates tl.kv_load_tl call instructions through the underlying language, configures level = 0 (L1) and prefetch = True, and the instructions directly connect to the GPU / NPU memory controller interface to ensure low-latency loading; ② Utilizing the instruction reordering capability of the "prefetch and computation parallelism Pass", prefetch instructions are embedded into the idle execution cycle of the attention computation operator. The Pass analyzes the cycle duty cycle of the computation instructions generated by the Triton underlying language and the prefetch instructions, dynamically adjusts the instruction execution timing, and achieves parallelism between storage prefetching and computation; ③ Through the cache isolation instruction tl.kv_cache_ctrl_tl (a hardware memory isolation interface encapsulated by the underlying language), dedicated memory space is allocated based on req_id to avoid conflicts from multiple requests; The core advantage is the hardware native instruction adaptability of the underlying language + the timing optimization capability of the Pass, which ensures low access latency and improves resource utilization.

[0059] L2 cache (distributed DRAM) scheduling: focuses on "cross-node sharing and overflow buffering of mid-to-high frequency data", and relies on the RDMA communication instruction encapsulation of Triton's underlying language and the "cross-node data migration optimization pass" of the LinalgIR layer. Technical Implementation Details: ① When L1 cache space is insufficient, the scheduler generates the `tl.kv_store_tl` instruction through the underlying language, configuring level=1, hot=1, and compress_mode=1. The instruction integrates FP8 quantization compression logic (a vector operation interface optimized by the underlying language) to improve storage efficiency; ② In cross-node sharing scenarios, the `tl.kv_load_tl` instruction uses native RDMA instructions encapsulated in the underlying language to read data across nodes. Combined with the "Cross-Node Data Migration Optimization Pass" fragmented transmission strategy, the Pass splits large-size KV data into fragments adapted to the RDMA transmission MTU, and achieves parallel transmission through asynchronous instructions in the underlying language, reducing migration latency; ③ The L2 utilization rate is periodically collected through the resource statistics instruction of `tl.kv_cache_ctrl_tl` (the underlying language interfaces with the chip monitoring interface). When it exceeds 80%, low-hot data is triggered to migrate to L3. The advantage is the efficient encapsulation of RDMA instructions in the underlying language + the fragmented optimization of the Pass, which balances cross-node transmission efficiency and reliability.

[0060] L3 cache (native KV storage) scheduling: focuses on "large-capacity storage and long-term reuse of low-frequency data", and relies on the NVMe SSD interface encapsulation of Triton underlying language and the "persistent storage optimization pass" of LinalgIR layer. Technical Implementation Details: ① For data with a heat level of 0, the `tl.kv_store_tl` command writes to the L3 cache through the underlying language's native NVMe IO interface. Combined with the sequential write optimization of the "Persistent Storage Optimization Pass," discrete KV data is organized into contiguous blocks, improving SSD write throughput. ② At the end of the session, a data snapshot is created using the `tl.kv_cache_ctrl_tl` snapshot command (an SSD snapshot interface encapsulated in the underlying language). During subsequent recovery, data pre-reading is achieved through the "Snapshot Fast Load Pass," which analyzes the access frequency of snapshot data and prioritizes pre-reading high-frequency segments to the L2 cache. ③ When L3 capacity is insufficient, low-priority data is filtered through the "Priority-Based LRU Eviction Pass," combined with efficient deletion commands from the underlying language to release resources. The core advantage is the native storage medium interface adaptation of the underlying language + the storage optimization capabilities of the Pass, achieving a balance between large-capacity storage and efficient reuse.

[0061] (3) To adapt to scenarios where data popularity changes dynamically and request priority is adjusted, the implementation of the dynamic scheduling strategy relies entirely on "the dynamic instruction generation capability of the Triton underlying language + the real-time optimization capability of the compiler Pass". The strategy is closed-loop through real-time updating of primitive parameters and combined invocation, as follows: The heat-driven cross-level migration strategy relies on the TritonIR ​​layer's "heat level triggering pass" and the dynamic command concatenation capability of the underlying language. Technical principle: The pass monitors changes in the heat matrix in real time. When the data heat level changes, it automatically triggers migration command generation logic—using the underlying language's command template mechanism to quickly concatenate the command sequence "tl.kv_load_tl (source level) → tl.kv_store_tl (target level)" without recompiling the entire operator. For example, when data is upgraded from hot=1 to hot=2, the pass triggers command sequence generation, and the underlying language quickly encapsulates the native commands for L2 loading and L1 storage, achieving efficient migration. The advantage is that the real-time triggering capability of the pass combined with the rapid command generation of the underlying language ensures fast migration response without interrupting the inference process.

[0062] To ensure the stability of multi-level cache scheduling, two main technical approaches are employed: atomic instructions in the Triton underlying language and consistency checks by the compiler pass. These ensure the consistency and reliability of primitive instruction execution, and the core implementation relies on the collaborative support of both. Primitive instruction transaction management: For combined primitive operations such as cross-level migration and snapshot creation-restore, a transaction management mechanism is adopted, achieving "atomic execution" through the status flag function of tl.kv_cache_ctrl_tl. For example, in the combined operation of "tl.kv_load_tl → tl.kv_store_tl" for migrating data from L1 to L2, the scheduler updates the data storage level flag only if both primitives execute successfully; if either primitive fails, a rollback operation is triggered through tl.kv_cache_ctrl_tl (such as deleting incomplete data in L2 and restoring the data state in L1).

[0063] (III) Primitive Co-scheduling Process (based on the Triton underlying language and the compiler Passes in a closed-loop process) This process follows a closed loop of "application requirements → primitive calls → compilation optimization → hardware execution → status feedback". The entire process relies on the Triton underlying language to achieve precise generation and execution of instructions. The compiler Passes completes deep collaboration between scheduling and computation and cross-chip adaptation, ultimately achieving efficient scheduling of L1-L2-L3 multi-level KV caches.

[0064] The specific process is divided into 5 core steps, each closely linked and progressively building upon the previous one: (1) Step 1: Application requirement input and scheduling preprocessing (scheduler core driver) Upper-layer applications (such as large language model inference frameworks) pass standardized KV cache scheduling requirements to the "multi-level cache scheduler" of the Triton compilation extension core layer through the application interface adaptation layer. These requirements include parameters such as request ID, data access mode, priority, and sequence length. The scheduler's requirement parsing module verifies the validity of the parameters through a strongly typed interface encapsulated in the Triton underlying language. Combined with the "four-dimensional heat matrix" (data address-access frequency-storage level-request priority) maintained by the heat awareness module, it generates structured scheduling instructions (such as "load high-frequency KV data into L1 cache and prefetch the next round of data" and "migrate low-frequency data to L3 and create a snapshot"), providing accurate basis for subsequent primitive calls.

[0065] (2) Step 2: Python layer primitive calls and low-level instruction encapsulation The scheduler invokes the three core primitives (tl.kv_load_tl / tl.kv_store_tl / tl.kv_cache_ctrl_tl) based on structured scheduling instructions. The Python layer encapsulates primitive parameters (such as storage level, hotness flags, and prefetch configuration) into an intermediate representation structure recognizable by the Triton IR layer using the Triton underlying language, while simultaneously performing parameter type conversion and secondary validity checks. The core advantage lies in the native compatibility between the underlying language and the Triton compilation system, avoiding instruction redundancy caused by cross-framework adaptation and ensuring lightweight primitive calls. For example, when calling tl.kv_load_tl, the underlying language encapsulates parameters such as "level=0 (L1)" and "prefetch=True" into a standardized instruction header, concatenating it with information such as the data address and request ID to form a complete primitive call instruction.

[0066] (3) Step 3: Semantic transformation of TritonIR ​​layer and optimization of Passes fusion The primitive call instructions generated by the Python layer are passed to the TritonIR ​​layer, where they are first converted into corresponding TritonIR ​​semantic instructions (such as `triton.ir.op.kv_load_tl` and `triton.ir.op.kv_store_tl`), completing the transformation from "primitive logic to intermediate representation". Subsequently, the scheduler calls a series of newly added optimization passes in the TritonIR ​​layer to complete the fusion optimization: ① "Prefetch and Computation Parallel Pass": analyzes the duty cycle of primitive instructions and attention computation instructions, embedding prefetch instructions into idle computation cycles; ② "Instruction Conflict Detection Pass": predicts and avoids read / write conflicts for the same data, ensuring execution order through instruction rearrangement; ③ "Hotness Level Trigger Pass": dynamically adjusts the storage level parameters of primitive instructions based on the heat matrix (e.g., changing the `tl.kv_store_tl` instruction level corresponding to data with upgraded heat to 0). After optimization, a TritonIR ​​instruction stream is generated that integrates storage operations and computation logic, achieving deep collaboration between scheduling and computation.

[0067] (4) Step 4: LinalgIR layer cross-chip optimization and machine instruction generation The optimized TritonIR ​​instruction stream is transformed into LinalgIR (MLIR standard intermediate representation). The LinalgIR layer completes cross-chip adaptation and low-level machine instruction generation through three dedicated MLIR Passes: ① kv-load-tl-opt Pass generates zero-copy transfer instructions (such as NVLink instructions and AscendC DMA instructions) adapted to the hardware characteristics of the target chip (NVIDIA GPU / Huawei Ascend NPU); ② kv-store-tl-opt Pass implements differentiated compression optimization (L1 lightweight compression, L2 / L3 deep compression) and storage integrity verification (CRC check); ③ kv-cache-ctrl-tl-opt Pass adapts to the chip's memory isolation and snapshot interface, generating hardware-level cache management instructions. Finally, the MLIR target code generation module transforms the optimized LinalgIR into the target chip's native machine instructions, shielding the hardware details of different chips.

[0068] (5) Step 5: Hardware execution and status feedback (closed-loop termination) The L3 cache execution adaptation layer receives native machine instructions generated by the LinalgIR layer, driving L1 (video memory), L2 (distributed DRAM), and L3 (native KV storage) to complete corresponding operations: L1 performs high-frequency data loading / storage through the video memory controller, L2 completes cross-node data migration through the RDMA interface, and L3 achieves low-frequency data persistence through the NVMe SSD controller. Simultaneously, the execution adaptation layer uses a status feedback interface encapsulated in the Triton low-level language to send operation results (such as cache hit rate, data migration progress, and instruction execution status) back to the scheduler's heat perception module and demand analysis module. The heat perception module updates the four-dimensional heat matrix accordingly, and the demand analysis module feeds the results back to the upper-layer application, completing the entire scheduling loop. If data verification fails or instruction execution anomalies occur, the scheduler will trigger a rollback mechanism (through the tl.kv_cache_ctrl_tl primitive) to ensure scheduling reliability.

[0069] The following is the relevant content of the application test case section.

[0070] Application Test Example 1 This embodiment builds a heterogeneous computing cluster based on NVIDIA GPU (H100) and Huawei Ascend NPU (Ascend 910B). The three-level cache configuration is as follows: L1 cache (H100 HBM memory, 80GB), L2 cache (distributed DRAM, 1TB cluster shared memory pool, built based on RDMA), and L3 cache (native KV storage, built based on NVMe SSD, capacity 10TB). The software environment is Triton compiler version 2.0, Python 3.9, CUDA 12.1, and Ascend Toolkit 8.0.RC1.

[0071] The implementation steps for the entire process based on the Triton underlying language and the Passes compiler are as follows: 1) Primitive extension implementation (three-level extension + underlying language / Passes development) Based on the extended interface of Triton compiler version 2.0, a three-level extension development of the three core primitives, namely "Python layer - TritonIR ​​layer - LinalgIR layer", was completed. The entire process relied on the Triton low-level language to implement instruction encapsulation, and the accompanying compiler Passes was developed simultaneously. The specific implementation is as follows: Python Layer: Primitive API Encapsulation and Low-Level Language Interface Binding Three new primitive functions have been added to the Triton Python API library (triton / language / core.py). These functions bind to underlying instructions through the `tl._register_op` interface of the Triton underlying language, enabling parameter parsing, type validation, and semantic encapsulation. The core code snippet is as follows: # Core logic for binding low-level language interfaces @tl._register_op("kv_load_tl") def kv_load_tl(addr: tl.pointer, level: tl.int32, dtype: tl.dtype, prefetch: tl.bool = False, prefetch_addr: tl.pointer = None, req_id: tl.int64= 0) ->tl.tensor: # Strong type checking (underlying language feature) tl.static_assert(level in [0,1,2], "level must be 0(L1),1(L2),2(L3)") tl.static_assert(addr.dtype == dtype, "addr dtype mismatch with datadtype") # Call the underlying language to generate primitive instructions return tl._op("kv_load_tl", addr, level, prefetch=prefetch, prefetch_addr=prefetch_addr, req_id=req_id, dtype=dtype)` TritonIR ​​Layer: Semantic Instruction Extension and IR Pass Development The TritonIR ​​instruction set was extended by adding three semantic instructions, `triton.ir.op.kv_load_tl`, `triton.ir.op.kv_store_tl`, and `triton.ir.op.kv_cache_ctrl_tl`, to `triton / ir / ops.py`, defining parameter constraints and type inference rules for these instructions. Simultaneously, the "Prefetch and Computation Parallelism Pass," "Instruction Conflict Detection Pass," and "Heat Analysis Pass" were developed and integrated into the TritonIR ​​optimization process (triton / ir / passes.py). Taking the "Prefetch and Computation Parallelism Pass" as an example, the core logic involves analyzing the cycle duty cycle of attention computation instructions and prefetch instructions through the underlying language's data flow tracing interface, thereby achieving instruction rearrangement and parallel embedding.

[0072] LinalgIR Layer: Dedicated MLIR Pass Development and Cross-Chip Adaptation Three dedicated optimization passes (kv-load-tl-opt, kv-store-tl-opt, and kv-cache-ctrl-tl-opt) were developed based on the MLIR framework and integrated into Triton's LinalgIR optimization pipeline. Each pass achieves differentiated optimizations through a chip adaptation interface encapsulated in Triton's low-level language. Adaptation for NVIDIA GPUs: kv-load-tl-opt Pass generates NVLink zero-copy transfer instructions (`nvlink::load`), and kv-store-tl-opt Pass integrates CUDA's FP8 / Int4 quantization interface; Adapting to Huawei Ascend NPU: Generate DMA data transfer instructions (`aclrtMemcpyAsync`) through the AscendC interface of the underlying language, and connect the kv-cache-ctrl-tl-opt Pass to Ascend's memory isolation interface (`aclrtSetDeviceMemFlags`).

[0073] Application Test Example 2 The following is an example of an application call for ChatGLM3-6B inference + primitive cooperative scheduling: Using the attention computation operator of the large language model ChatGLM3-6B as a carrier, three core primitives are invoked to achieve dynamic scheduling of multi-level key-value caches (L1-L2-L3). Combined with the scheduler's heat-driven strategy, the entire process of "loading-computation-storage-snapshot" is completed. The core calling logic is as follows, and the relationship between key primitive calls and scheduling strategies in the code is annotated: Python import triton import triton.language as tl import torch # Global scheduler configuration (interfacing with the scheduler for compiling and extending the core layer of Triton) SCHEDULER_CONFIG = {"priority": 1, # High priority request (level 1)} "eviction_policy": 0, # LRU eviction policy "hot_threshold": 3# Marks 3 consecutive accesses as high-frequency data. @triton.jit def chatglm3_attention(q: tl.tensor, k_addr: tl.pointer, v_addr:tl.pointer, output: tl.pointer, seq_len: tl.int32, req_id: tl.int64): # Step 1: Call the management primitive to initialize the cache isolation space (the scheduler's cache isolation policy) tl.kv_cache_ctrl_tl(op_type=0,req_id=req_id, eviction_policy=SCHEDULER_CONFIG["eviction_policy"]) # Step 2: Call the load primitive to load high-frequency key-value data from L1 and enable prefetching (the scheduler's prefetching strategy). # The prefetched address contains the key-value data needed for the next round of calculation, which is predicted by the scheduler's heat perception module. k=tl.kv_load_tl(addr=k_addr,level=0,dtype=tl.float16,prefetch=True,prefetch_addr=k_addr + seq_len * 128,# 128 is the KV vector dimension req_id=req_id) v = tl.kv_load_tl( addr=v_addr, level=0, dtype=tl.float16,prefetch=True, prefetch_addr=v_addr + seq_len * 128,req_id=req_id ) # Step 3: Attention core computation (in parallel with prefetching, implemented with TritonIR ​​layer Pass optimization) attn = tl.dot(q, kT) attn = tl.softmax(attn) output_val = tl.dot(attn, v) tl.store(output, output_val) # Step 4: Call the storage primitive to migrate inactive KV data to L2 (the scheduler's hot migration strategy).

[0074] # The popularity marker is calculated in real time by the scheduler's popularity analysis pass (if the current data has not been accessed once consecutively, it is marked as medium to high frequency).

[0075] tl.kv_store_tl(data=k, addr=k_addr + seq_len * 128 * 2, level=1, hot=1, compress_mode=1, # FP8 quantization compression (L2 layer deep compression strategy) req_id=req_id) # Step 5: Invoke the management primitive to create a session snapshot (scheduler's snapshot management policy) tl.kv_cache_ctrl_tl(op_type=1, req_id=req_id,snapshot_addr=k_addr +seq_len * 128 * 4 # Snapshots are stored in the L3 cache) # Upper-layer PyTorch application calls (interfacing with the application interface adaptation layer) def torch_attention_wrapper(q, k, v, req_id): seq_len = q.shape[1] output = torch.empty_like(q) # Calling the Triton operator triggers the primitive cooperative scheduling process chatglm3_attention[(128,)](q=q, k_addr=k.data_ptr(), v_addr=v.data_ptr(), output=output.data_ptr(), seq_len=seq_len, req_id=req_id) return output ``` Compilation and execution process (end-to-end collaborative verification): After the upper-layer application calls the `torch_attention_wrapper` function, the Triton compiler initiates a full-chain process of "primitive call - compilation optimization - hardware execution". Each step is closely dependent on the Triton underlying language and the compiler Passes. The specific execution steps are as follows (divided into two scenarios: NVIDIA GPU and Huawei Ascend NPU): 1. Primitive call parsing and underlying instruction generation The Triton Python layer parses the primitive call parameters in `chatglm3_attention`, generates standardized primitive instructions through the `tl._op` interface of the underlying language, and simultaneously passes the priority, eviction policy, and other parameters in `SCHEDULER_CONFIG` to the scheduler of the Triton compilation extension core layer.

[0076] 2. Semantic transformation of TritonIR ​​layer and Passes optimization The generated primitive instructions are converted into TritonIR ​​semantic instructions, which then trigger the execution of the three major IR Passes in sequence: "Instruction Conflict Detection Pass": Analyze the memory access addresses of `tl.kv_load_tl` and `tl.kv_store_tl` to confirm that there are no read / write conflicts.

[0077] 3. "Prefetch and Computation Parallel Pass": Reorders the execution timing of prefetch instructions and attention computation instructions, embedding the prefetch operation into the idle cycle of `tl.dot`; 4. "Popularity Analysis Pass": Tracks the number of visits to key-value (KV) data, updates the popularity matrix, and provides a basis for the `hot=1` parameter of `tl.kv_store_tl`.

[0078] 5. LinalgIR layer cross-chip optimization and machine instruction generation After the TritonIR ​​instruction stream is converted to LinalgIR, the three dedicated MLIR passes initiate differentiated optimizations: For NVIDIA GPU scenarios: kv-load-tl-opt Pass generates NVLink zero-copy transfer instructions to directly load KV data from L1 memory into the GPU computing core; kv-store-tl-opt Pass calls the CUDA FP8 quantization interface to complete data compression and generate NVMe SSD write instructions. 6. Adapting to Huawei Ascend NPU scenarios: kv-load-tl-opt Pass generates AscendC DMA asynchronous transmission instructions to load data from L2 distributed DRAM via RDMA interface; kv-cache-ctrl-tl-opt Pass calls the Ascend memory isolation interface to allocate dedicated cache space based on `req_id`.

[0079] 7. Hardware Execution and State Inversion The L3 cache executes adaptation layer instructions received from the machine, driving L1-L2-L3 to complete the corresponding operations: NVIDIA GPU scenario: L1 memory performs high-frequency key-value data loading, L2 distributed DRAM stores high-frequency data, and L3 NVMe SSD stores snapshot data; 8. Huawei Ascend NPU scenario: L1-L2 data migration is completed through DMA, and the Ascend memory isolation mechanism ensures dedicated cache space.

[0080] The above description is merely a preferred embodiment of the present invention and is not intended to limit the present invention in any way. Therefore, any simple modifications, equivalent changes, and alterations made to the above embodiments based on the technical essence of the present invention without departing from the scope of the present invention shall still fall within the scope of the present invention.

Claims

1. A multi-level key-value cache operation primitive extension system implemented within the Triton framework, characterized in that, include: The primitive set is used to support the loading, storage and full lifecycle management of L1 video memory, L2 distributed DRAM and L3 native KV memory. The primitive set includes at least loading primitives for multi-level KV data loading, storage primitives for multi-level KV data storage, and management primitives for cache isolation, snapshots, migration control, eviction and resource statistics. The TritonIR ​​semantic bearer module is used to map the loading primitives, storage primitives and management primitives to corresponding intermediate representation instructions, and to carry hierarchical, prefetching, hotness, snapshot and transaction control semantics. The compilation optimization module is used to perform at least one optimization process in the TritonIR ​​layer, including prefetching and computational parallelism, heat analysis and hierarchical adjustment, and instruction conflict detection, and to complete cross-chip optimization and machine instruction generation in the LinalgIR / MLIR layer. A multi-level cache scheduler is used to generate primitive invocation strategies based on scheduling requirements and heat matrix, and to coordinate primitive execution, status feedback and exception handling. The execution adaptation module is used to send the machine instructions to the execution paths corresponding to L1 video memory, L2 distributed DRAM and L3 native KV storage respectively, so as to complete multi-level KV cache operation.

2. The system as described in claim 1, characterized in that: The multi-level cache scheduler includes at least a demand parsing module, a heat perception module, a primitive scheduling strategy generation module, and an execution coordination module; wherein, the demand parsing module is used to receive at least one scheduling parameter among request ID, data access mode, priority, and sequence length, and generate structured scheduling requirements.

3. The system as described in claim 2, characterized in that: The heat perception module is used to construct and update the heat matrix based on primitive execution feedback. The heat matrix includes at least data address, access frequency, storage level and request priority, and is used to drive the cross-level migration of data between L1 video memory, L2 distributed DRAM and L3 native KV storage.

4. The system as described in claim 1, characterized in that: The management primitives are used to perform transaction control on at least one of the composite operations of cross-level migration, snapshot creation, and snapshot recovery, and trigger rollback in the event of execution exceptions or verification failures to ensure data consistency.

5. The system as described in claim 1, characterized in that: The compilation optimization module performs at least one of the following processes in the TritonIR ​​layer: instruction conflict detection and prefetch and computation parallel optimization. Instruction conflict detection is used to analyze the access address or execution cycle of primitive instructions and to avoid or sequence potential read-write conflicts. Prefetch and computation parallel optimization is used to rearrange the timing of prefetch instructions and computation instructions to hide data migration delays.

6. The system as described in claim 1, characterized in that: The compilation optimization module, at the LinalgIR / MLIR layer, completes the adaptation to different chip platforms or different storage paths through dedicated optimization passes for loading primitives, storage primitives, and management primitives, and performs at least one of the following processes during data loading or data storage: data type conversion, compression, and integrity verification.

7. The system as described in claim 1 or 3, characterized in that: The multi-level cache scheduler is also used to trigger a heat-driven cross-level migration strategy based on heat changes in the heat matrix, so that high-heat data is preferentially residing in the low-latency layer and low-heat data is migrated to the high-capacity layer.

8. A multi-level key-value cache management method implemented within the Triton framework, characterized in that, Includes the following steps: S11: Receive KV cache scheduling requests through the application interface adaptation layer, and perform parameter verification and scheduling preprocessing; S12: Call the loading primitive, storage primitive, and management primitive, and encapsulate the primitive parameters into an intermediate representation structure that can be recognized by the TritonIR ​​layer; S13: Perform at least one optimization process in the TritonIR ​​layer, including prefetching and computational parallelism, heat analysis and hierarchical adjustment, and instruction conflict detection; S14: Perform cross-chip optimization and generate target hardware machine instructions at the LinalgIR / MLIR layer; S15: Based on the machine instructions, drive the L1 video memory, L2 distributed DRAM and L3 native KV storage to complete the corresponding cache operations, and feed the execution results back to the scheduler; S16: Update the heat matrix, scheduling policy and resource status based on execution feedback, and perform rollback or rescheduling in case of abnormal situations.

9. The method as described in claim 8, characterized in that: In step S16, a heat matrix containing data address, access frequency, storage level and request priority is constructed or updated, and cross-level migration is triggered based on the heat matrix; During cross-level migration, snapshot creation, or snapshot recovery, if any step encounters an anomaly or verification fails, a rollback is triggered through management primitives to ensure data consistency.

10. A computer-readable storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, it implements the method of claim 8 or 9.

Images