# gpu_agent_pro

**Repository Path**: Dongjialin0404/gpu_agent_pro

## Basic Information

- **Project Name**: gpu_agent_pro
- **Description**: 2025年秋季 国科大《GPU架构与编程》——大作业二（GPU Agent / MUSA S4000）
- **Primary Language**: Python
- **License**: MulanPSL-2.0
- **Default Branch**: main
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 1
- **Created**: 2025-12-23
- **Last Updated**: 2026-01-25

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# 2025年秋季 国科大《GPU架构与编程》——大作业二（GPU Agent / MUSA S4000）

本仓库为 2025 年秋季国科大《GPU架构与编程》课程大作业二实现与复现说明。本项目实现一个离线可部署的 GPU Agent：通过 FastAPI 提供推理服务，并提供本地网页 UI（/ui）进行可视化交互，支持 single/batch 两种推理模式，支持在 UI 中输入 @code 进入代码模式。

课程：2025年秋季 国科大《GPU架构与编程》
作业：大作业二（GPU Agent）
作者：董家麟
学院：前沿交叉科学学院

--------------------------------------------------------------------------------
1. 已验证运行环境
--------------------------------------------------------------------------------
GPU：摩尔线程 MUSA S4000
Python：3.10（建议 3.10.x）
PyTorch：2.2.0
MUSA Toolkit：3.1.0
Transformers：离线加载（TRANSFORMERS_OFFLINE=1 / HF_HUB_OFFLINE=1）
服务框架：FastAPI + Uvicorn

说明：
- 服务启动时强制检测 MUSA。若 torch_musa 无法导入或 torch.musa.is_available() 为 false，会报错退出。
- 默认启用离线模式：TRANSFORMERS_OFFLINE=1、HF_HUB_OFFLINE=1。

--------------------------------------------------------------------------------
2. 主要功能
--------------------------------------------------------------------------------
2.1 UI 可视化
- 启动服务后浏览器打开：http://127.0.0.1:8080/ui
- 支持 single/batch 模式切换
- 输入 @code 启动代码模式（输出更偏完整可复制代码，并保留换行/缩进）
- 支持流式输出（打字机效果），UI 通过 /predict_stream 使用 NDJSON 协议

2.2 HTTP 接口
- GET  /               返回 {"status":"batch"}（评测侧可用于探活）
- POST /predict         一次性返回（评测/脚本接口）
- POST /predict_stream  流式返回（UI 接口，NDJSON）
- GET  /ui              返回本地网页 UI
- GET  /ui_caps         返回 UI 能力声明（supports_stream/code_mode 等）

--------------------------------------------------------------------------------
3. 目录结构与相对路径约定（建议保持不变）
--------------------------------------------------------------------------------
推荐保持如下结构（与代码默认路径一致）：

.
├── serve.py                  # FastAPI 服务入口（推理 + KB 检索 + UI）
├── download_model.py         # 下载/准备基座模型脚本（先运行）
├── ref_gpu.jsonl             # 本地知识库（jsonl，每行一个 QA）
├── local-model/
│   ├── base/                 # 基座模型目录（需包含 config.json）
│   └── lora/                 # LoRA 目录（可选，存在 adapter_config.json 则自动合并）
└── web/
    └── index.html            # UI 页面（/ui）

默认读取路径（可用环境变量覆盖）：
- BASE_MODEL_DIR：默认 ./local-model/base
- LORA_DIR：默认 ./local-model/lora
- REF_FILE：默认 ./ref_gpu.jsonl

重要说明：
- 请在仓库根目录启动服务（与 serve.py 同级），以保证相对路径可复现。
- 若你把模型放在其它目录，可设置 BASE_MODEL_DIR 环境变量指向含 config.json 的目录。

--------------------------------------------------------------------------------
4. 步骤
--------------------------------------------------------------------------------
4.1 下载/准备基座模型（必须先做）
在仓库根目录执行：

python download_model.py

下载完成后请确认：
- ./local-model/base/config.json 存在
（或你设置的 BASE_MODEL_DIR 目录下存在 config.json）

4.2 启动服务
在仓库根目录执行：

python serve.py

默认监听：
- 0.0.0.0:8080

启动成功后终端会输出类似：
- [Info] MUSA Device: ...
- [Info] Loading tokenizer...
- [Info] Loading model (MUSA fp16)...
- [Info] Server Ready.

4.3 打开 UI 页面
浏览器访问：

http://127.0.0.1:8080/ui

--------------------------------------------------------------------------------
5. UI 使用说明
--------------------------------------------------------------------------------
5.1 Single 模式
- 直接输入问题，例如：什么是数据并行性

5.2 Batch 模式
- 在 UI 切换到 Batch 模式后，可一次输入多条问题进行批处理推理。

5.3 代码模式（@code）
在 UI 输入框以 @code 开头，例如：

@code 写一个 CUDA 矩阵乘法的代码

代码模式特点：
- 输出更偏“完整可复制的代码”
- 保留换行与缩进（不会把多行代码压成一行）
- 输出长度上限更高（受 max_new_tokens 与内部 cap 控制）

--------------------------------------------------------------------------------
6. API 调用示例（命令行可直接执行）
--------------------------------------------------------------------------------
6.1 /predict（评测/脚本接口）


返回格式：
- 单条：{"response":"..."}
- 批量：{"response":["...","..."]}

6.2 /predict_stream（UI 使用）
- UI 默认调用该接口以实现



--------------------------------------------------------------------------------
7. 可调参数（环境变量）
--------------------------------------------------------------------------------
模型与 KB：
- BASE_MODEL_DIR（默认 ./local-model/base）
- LORA_DIR（默认 ./local-model/lora）
- REF_FILE（默认 ./ref_gpu.jsonl）

性能与长度：
- HARD_MAX_BATCH（默认 96）
- MAX_PREFILL_TOKENS（默认 42000）
- MAX_INPUT_LEN（默认 448）
- PAD_TO_MULTIPLE（默认 16）
- MAX_NEW_CAP（默认 1024）

KB 阈值：
- KB_THRESHOLD（默认 0.16）
- HIGH_CONF_THRESH（默认 0.64）
- MID_CONF_THRESH（默认 0.40）

UI 流式：
- STREAM_YIELD_EVERY（默认 8）
- STREAM_FLUSH_INTERVAL（默认 0.05）
- STREAM_MAX_NEW_DEFAULT（默认 256）

--------------------------------------------------------------------------------
8. 致谢
--------------------------------------------------------------------------------
感谢国科大《GPU架构与编程》课程与助教团队；感谢摩尔线程 MUSA 平台提供算力与软件栈支持。