企业私有化 AI Agent 推理引擎并发优化：PagedAttention 显存分页与连续批处理（Continuous Batching）实战

/**
 * 企业级大模型推理引擎 PagedAttention 显存物理块管理器与连续批处理调度器
 * 模拟实现非连续物理显存分配、逻辑块表（Block Table）维护与迭代级（Iteration-level）动态调度
 */

/** 物理显存块接口定义 */
export interface PhysicalBlock {
  /** 物理块全局唯一标识符 */
  blockId: number;
  /** 引用计数（用于 Copy-on-Write 与多分支共享） */
  refCount: number;
  /** 当前物理块实际已存放的 Token 数量 */
  tokenCount: number;
}

/** 智能体推理请求状态枚举 */
export type RequestPhase = 'WAITING' | 'RUNNING' | 'FINISHED';

/** 推理请求上下文接口 */
export interface InferenceRequest {
  /** 业务唯一请求标识符 */
  requestId: string;
  /** 输入 Prompt 文本 */
  prompt: string;
  /** 输入已编码的 Token 序列 */
  promptTokens: number[];
  /** 已生成的输出 Token 序列 */
  outputTokens: number[];
  /** 用户设定的最大生成 Token 上限 */
  maxTokens: number;
  /** 当前运行阶段 */
  phase: RequestPhase;
  /** 属于该请求的逻辑块到物理块映射表（Block Table） */
  blockTable: number[];
  /** 请求进入队列的时间戳 */
  arrivedAt: number;
  /** 首字生成耗时（TTFT，毫秒） */
  ttftMs?: number;
}

/** 调度器核心配置选项 */
export interface SchedulerOptions {
  /** 每个 Block 容纳的 Token 数量（建议 16 或 32） */
  blockSize: number;
  /** GPU 预留显存池中可用的物理 Block 总数 */
  totalPhysicalBlocks: number;
  /** 单次迭代最大允许并发运行的序列数量 */
  maxNumSeqs: number;
}

/**
 * PagedAttention 物理显存块管理器（Block Manager）
 */
export class PhysicalBlockManager {
  private blockSize: number;
  private freeBlocks: PhysicalBlock[] = [];
  private allocatedBlocks: Map<number, PhysicalBlock> = new Map();

  constructor(totalBlocks: number, blockSize: number) {
    this.blockSize = blockSize;
    for (let i = 0; i < totalBlocks; i++) {
      this.freeBlocks.push({ blockId: i, refCount: 0, tokenCount: 0 });
    }
  }

  /** 获取当前剩余空闲物理块数量 */
  public getFreeBlockCount(): number {
    return this.freeBlocks.length;
  }

  /** 分配一个空闲物理块 */
  public allocateBlock(): PhysicalBlock {
    const block = this.freeBlocks.pop();
    if (!block) {
      throw new Error('GPU 显存物理块已耗尽（Out of Memory）');
    }
    block.refCount = 1;
    block.tokenCount = 0;
    this.allocatedBlocks.set(block.blockId, block);
    return block;
  }

  /** 释放并回收物理块 */
  public freeBlock(blockId: number): void {
    const block = this.allocatedBlocks.get(blockId);
    if (!block) return;

    block.refCount -= 1;
    if (block.refCount <= 0) {
      block.refCount = 0;
      block.tokenCount = 0;
      this.allocatedBlocks.delete(blockId);
      this.freeBlocks.push(block);
    }
  }

  /** 为指定请求追加一个 Token 并动态分配或扩展物理块 */
  public appendToken(request: InferenceRequest): void {
    const totalTokens = request.promptTokens.length + request.outputTokens.length;
    const neededBlocks = Math.ceil(totalTokens / this.blockSize);

    // 如果当前块表中拥有的物理块不足以容纳新增 Token，则动态申请新块
    if (request.blockTable.length < neededBlocks) {
      const newBlock = this.allocateBlock();
      request.blockTable.push(newBlock.blockId);
    }

    // 更新最新物理块中的 Token 计数
    const lastBlockId = request.blockTable[request.blockTable.length - 1];
    const lastBlock = this.allocatedBlocks.get(lastBlockId);
    if (lastBlock) {
      const currentInBlock = totalTokens % this.blockSize;
      lastBlock.tokenCount = currentInBlock === 0 ? this.blockSize : currentInBlock;
    }
  }

  /** 请求完成时释放其占用的全部物理块 */
  public freeRequest(request: InferenceRequest): void {
    for (const blockId of request.blockTable) {
      this.freeBlock(blockId);
    }
    request.blockTable = [];
  }
}

/**
 * 连续批处理迭代级调度引擎（Continuous Batching Scheduler）
 */
export class ContinuousBatchingEngine {
  private blockManager: PhysicalBlockManager;
  private maxNumSeqs: number;
  private blockSize: number;

  /** 等待队列 */
  private waitingQueue: InferenceRequest[] = [];
  /** 当前迭代正在执行的批次队列 */
  private runningBatch: InferenceRequest[] = [];
  /** 已完成的历史请求 */
  private completedRequests: InferenceRequest[] = [];

  constructor(options: SchedulerOptions) {
    this.blockSize = options.blockSize;
    this.maxNumSeqs = options.maxNumSeqs;
    this.blockManager = new PhysicalBlockManager(options.totalPhysicalBlocks, options.blockSize);
  }

  /** 提交新的智能体推理任务 */
  public submitRequest(requestId: string, prompt: string, promptTokens: number[], maxTokens: number = 64): void {
    const req: InferenceRequest = {
      requestId,
      prompt,
      promptTokens,
      outputTokens: [],
      maxTokens,
      phase: 'WAITING',
      blockTable: [],
      arrivedAt: Date.now()
    };
    this.waitingQueue.push(req);
  }

  /**
   * 执行单个迭代步（Iteration Step）
   * 返回当前迭代的批处理执行状态指标
   */
  public step(currentIteration: number): {
    runningCount: number;
    waitingCount: number;
    freeBlocks: number;
    newlyFinished: string[];
  } {
    // 阶段一：动态准入（Dynamic Admission）
    // 检查是否有足够的物理块容纳等待队列中的新请求
    while (this.waitingQueue.length > 0 && this.runningBatch.length < this.maxNumSeqs) {
      const candidate = this.waitingQueue[0];
      const initialBlocksNeeded = Math.ceil(candidate.promptTokens.length / this.blockSize);

      // 预留至少额外 2 个块作为解码缓冲安全边界
      if (this.blockManager.getFreeBlockCount() >= initialBlocksNeeded + 2) {
        this.waitingQueue.shift();
        candidate.phase = 'RUNNING';

        // 初始 Prefill 分配物理块并建立块表
        for (let i = 0; i < initialBlocksNeeded; i++) {
          const block = this.blockManager.allocateBlock();
          candidate.blockTable.push(block.blockId);
        }
        this.runningBatch.push(candidate);
      } else {
        // 显存水位紧张，停止准入新请求，等待已有请求释放
        break;
      }
    }

    const newlyFinishedIds: string[] = [];
    const remainingBatch: InferenceRequest[] = [];

    // 阶段二：逐序列迭代生成与 Token 追加
    for (const req of this.runningBatch) {
      // 模拟生成 1 个新 Token（此处用随机数代指）
      const nextToken = Math.floor(Math.random() * 1000) + 1;
      req.outputTokens.push(nextToken);

      // 记录首字延迟
      if (req.outputTokens.length === 1 && !req.ttftMs) {
        req.ttftMs = Date.now() - req.arrivedAt;
      }

      // 显存物理块动态感知与自增
      this.blockManager.appendToken(req);

      // 判定终止条件：到达最大长度或命中结束标记（此处模拟 5% 概率命中 EOS）
      const isEos = Math.random() < 0.05;
      const isMaxReached = req.outputTokens.length >= req.maxTokens;

      if (isEos || isMaxReached) {
        req.phase = 'FINISHED';
        // 即刻释放占用的物理显存块！
        this.blockManager.freeRequest(req);
        newlyFinishedIds.push(req.requestId);
        this.completedRequests.push(req);
      } else {
        remainingBatch.push(req);
      }
    }

    // 更新运行中批次（已完成请求立即剔除，不产生空闲气泡）
    this.runningBatch = remainingBatch;

    return {
      runningCount: this.runningBatch.length,
      waitingCount: this.waitingQueue.length,
      freeBlocks: this.blockManager.getFreeBlockCount(),
      newlyFinished: newlyFinishedIds
    };
  }
}