$ show MuLabPKU/TransArch
MuLabPKU/TransArch is a collection of research projects focused on efficient attention mechanisms for large language models (LLMs), particularly variants of Multi-Head Latent Attention (MLA). The repository provides software artifacts for converting standard attention to hardware-friendly forms (e.g., GQA→MLA via TransMLA, MLA→GQLA via per-group PCA, cross-layer SVD pruning via CLOVER), custom GPU kernels for sparse and FP8 attention (LISA, MISA), and integration patches for production inference frameworks like vLLM and SGLang. Each sub-project is self-contained, with its own entry points, configuration, and evaluation scripts, but they share common calibration utilities and model conversion patterns. The architecture is modular: model conversion pipelines (e.g., `converter.py`, `compression.py`) ingest pre-trained weights and calibration data to produce transformed models, which are then served via patched vLLM or SGLang instances. Custom CUDA/Triton kernels (e.g., LISA's linear and threshold scorers, MISA's tilelang-based sparse attention) are compiled and registered as custom ops. The OpenPWS sub-project provides a rigorous reproduction framework for benchmark experiments, with hash-verified inputs and provenance tracking. While the primary character is application-oriented (delivering reusable software for LLM inference optimization), the repo also contains research evaluation scripts, benchmark results (JSONL files), and calibration datasets, giving it a hybrid nature. The code is research-grade: functional and innovative, but with limited test coverage, hardware lock-in (SM100, Blackwell), and dependency on specific framework versions. Several sub-projects are marked as 'code coming soon' or are preprints, indicating ongoing development.
The LLM identified 5 files worth highlighting for reuse:
Converts GQA (Grouped-Query Attention) models to MLA (Multi-Head Latent Attention) by applying partial RoPE and low-rank QKV decomposition. The conversion is calibration-aware, using a small dataset to fit the low-rank matrices. The module is self-contained, with clear input/output (model checkpoint, config) and minimal dependencies on internal modules, making it reusable for any GQA-to-MLA migration.
Implements per-group PCA compression of the KV cache for MLA models, reducing the latent dimension while preserving accuracy. The pipeline collects covariance statistics from calibration data, performs eigenvalue decomposition, and absorbs the transformation into weights. The modular design (calibration → covariance → PCA → absorption) allows easy adaptation to other attention architectures.
A Triton-based FP8 linear scoring kernel that performs head collapse, scan, and top-k selection for sparse attention. It is designed for SM100 (B200/GB200) GPUs and demonstrates advanced use of Triton primitives for high-throughput attention scoring. The kernel is isolated and can be extracted as a standalone component for other sparse attention implementations.
A high-performance CUDA kernel for threshold-based attention scan on SM100 GPUs. It implements a prefix scan with compaction and recovery, using CUDA graphs for dynamic control flow. The kernel is self-contained with clear input/output interfaces and demonstrates advanced GPU programming techniques (warp-level primitives, ring buffer emission).
Generates final tables and plots for the OpenPWS reproduction study, with built-in validation checks (hash verification, consistency checks). It processes per-example results from code and retrieval experiments, computes Wilson confidence intervals, and exports evidence for README tables. The module is reusable for any benchmark reporting pipeline that requires integrity verification.