# modules_benchmark

**Repository Path**: zksainx/modules_benchmark

## Basic Information

- **Project Name**: modules_benchmark
- **Description**: LLM inference, modules benchmark
- **Primary Language**: Unknown
- **License**: MIT
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2025-07-08
- **Last Updated**: 2026-05-12

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# Modules Benchmark

LLM 模块组件性能分析工具集，支持 GEMM、Attention、Normalization、Activation、MoE 等核心算子的 profiling。

## Overview

支持的模块:
- **GEMM** - Linear 层操作 (BF16/FP16/FP8)
- **Attention** - FlashInfer / vLLM / Flash Attention
- **Normalization** - RMSNorm / Fused Add RMSNorm
- **Positional Embeddings** - RoPE (Rotary Positional Embedding)
- **Activations** - SiLU / Fused SiLU and Mul
- **LM Head** - Vocabulary projection
- **Logits Processing** - Temperature / Softmax / TopP / TopK / Sampling
- **MoE** - Fused Mixture of Experts (Triton kernels)

## Quick Start

```bash
# 1. Environment Setup
cd modules
bash init.sh

# 2. Create output directory
mkdir -p /data/output/modules

# 3. Run benchmarks
python run_gemm.py
```

---

## GEMM Benchmarks

### Standard GEMM (`run_gemm.py`)

使用 `nn.Linear` 测试 GEMM 操作 (BF16/FP16)。

```bash
python run_gemm.py \
  --dir_path /data/output/gemm \
  --nk_file ../dim/tiny_gemm_N_K.csv \
  --seq_file ../dim/seq_len.csv \
  --num_iterations 5 \
  --params_dtype bfloat16 \
  --device cuda
```

| Parameter | Default | Description |
|-----------|---------|-------------|
| `--dir_path` | required | Output directory for traces |
| `--nk_file` | `../dim/tiny_gemm_N_K.csv` | CSV with (N, K) pairs |
| `--seq_file` | `../dim/seq_len.csv` | CSV with sequence lengths (M) |
| `--num_iterations` | 5 | Profiling iterations |
| `--params_dtype` | bfloat16 | Data type (bfloat16/float16/fp8) |
| `--device` | cuda | Device |

### FP8 GEMM (`run_fp8_gemm.py`)

使用 vLLM `cutlass_scaled_mm` 测试 FP8 量化 GEMM。

```bash
python run_fp8_gemm.py \
  --dir_path /data/output/fp8_gemm \
  --nk_file ../dim/tiny_gemm_N_K.csv \
  --seq_file ../dim/seq_len.csv \
  --num_iterations 5 \
  --output_dtype bfloat16 \
  --device cuda
```

| Parameter | Default | Description |
|-----------|---------|-------------|
| `--output_dtype` | bfloat16 | Output data type (bfloat16/float16) |

> **Note:** FP8 需要 CUDA compute capability 8.9+ (H100, H200)

### Multi-GPU GEMM (`run_gemm_multi_gpu.py`)

跨多 GPU 测试 GEMM 操作。

```bash
python run_gemm_multi_gpu.py
```

---

## LM Head Benchmark (`run_lmhead.py`)

测试 LM Head 层 (vocabulary projection)。

```bash
python run_lmhead.py \
  --dir_path /data/output/lmhead \
  --vocab_file ../dim/tiny_vocab_size.csv \
  --batch_file ../dim/batch_size.csv \
  --num_iterations 5 \
  --params_dtype bfloat16 \
  --device cuda
```

| Parameter | Default | Description |
|-----------|---------|-------------|
| `--vocab_file` | `../dim/tiny_vocab_size.csv` | CSV with (vocab_size, hidden_size) |
| `--batch_file` | `../dim/batch_size.csv` | CSV with batch sizes |

---

## Normalization Benchmarks

### FlashInfer RMSNorm (`run_flashinfer_rmsnorm.py`)

测试 FlashInfer 的 RMSNorm 和 Fused Add RMSNorm。

```bash
python run_flashinfer_rmsnorm.py \
  --dir_path /data/output/rmsnorm \
  --hidden_size_file ../dim/tiny_hidden_size.csv \
  --seq_file ../dim/seq_len.csv \
  --num_iterations 5 \
  --params_dtype bfloat16 \
  --device cuda
```

**Operations:**
- `rmsnorm`: Standard RMSNorm
- `fused_add_rmsnorm`: Fused residual addition + RMSNorm

### vLLM RMSNorm (`run_vllm_rmsnorm.py`)

测试 vLLM 的 RMSNorm 实现。

```bash
python run_vllm_rmsnorm.py \
  --dir_path /data/output/vllm_rmsnorm \
  --hidden_size_file ../dim/tiny_hidden_size.csv \
  --seq_file ../dim/seq_len.csv \
  --num_iterations 5 \
  --params_dtype bfloat16 \
  --device cuda
```

---

## Positional Embedding Benchmarks

### FlashInfer RoPE (`run_flashinfer_rope.py`)

测试 FlashInfer 的 RoPE (预计算 cos/sin cache)。

```bash
python run_flashinfer_rope.py \
  --dir_path /data/output/rope \
  --qkv_file ../dim/tiny_attn_qkv.csv \
  --seq_file ../dim/seq_len.csv \
  --num_iterations 5 \
  --params_dtype bfloat16 \
  --device cuda
```

| Parameter | Default | Description |
|-----------|---------|-------------|
| `--qkv_file` | `../dim/tiny_attn_qkv.csv` | CSV with (q_heads, kv_heads, head_dim) |

### vLLM RoPE (`run_vllm_rope.py`)

测试 vLLM 的 RoPE 实现。

```bash
python run_vllm_rope.py
```

---

## Activation Benchmarks

### FlashInfer Activation (`run_flashinfer_activation.py`)

测试 Fused SiLU and Mul: `silu(input[..., :hidden_size]) * input[..., hidden_size:]`

```bash
python run_flashinfer_activation.py \
  --dir_path /data/output/activation \
  --activation_file ../dim/tiny_activation_dim.csv \
  --seq_file ../dim/seq_len.csv \
  --num_iterations 5 \
  --params_dtype bfloat16 \
  --device cuda
```

| Parameter | Default | Description |
|-----------|---------|-------------|
| `--activation_file` | `../dim/tiny_activation_dim.csv` | CSV with input_size (= 2 * hidden_size) |

### vLLM Activation (`run_vllm_activation.py`)

测试 vLLM 的 SiluAndMul 层。

```bash
python run_vllm_activation.py \
  --dir_path /data/output/vllm_activation \
  --activation_file ../dim/tiny_activation_dim.csv \
  --seq_file ../dim/seq_len.csv \
  --num_iterations 5 \
  --params_dtype bfloat16 \
  --device cuda
```

---

## Logits Processing Benchmark (`run_flashinfer_logits_processor.py`)

测试完整的 token sampling pipeline。

```bash
python run_flashinfer_logits_processor.py \
  --dir_path /data/output/logits_processor \
  --vocab_file ../dim/tiny_vocab_size.csv \
  --batch_file ../dim/batch_size.csv \
  --num_iterations 5 \
  --params_dtype bfloat16 \
  --device cuda
```

**Pipeline Stages:**
1. **Temperature** - Scale logits by temperature
2. **Softmax** - Convert logits to probabilities
3. **TopK** - Apply top-k filtering
4. **TopP** - Apply nucleus (top-p) sampling
5. **Sample** - Sample tokens from distribution

---

## Attention Benchmarks

### FlashInfer Variable Length (`run_flashinfer_attn_varlen_control.py`)

测试 batch prefill (ragged tensor) 和 decode (paged KV cache)。

```bash
python run_flashinfer_attn_varlen_control.py \
  --dir_path /data/output/attention \
  --qkv_file ../dim/tiny_attn_qkv.csv \
  --seq_file ../dim/seq_len.csv \
  --num_iterations 10 \
  --params_dtype bfloat16 \
  --device cuda
```

**Attention Modes:**
- **Ragged Prefill**: Variable length sequences using ragged tensors
- **Paged Decode**: Paged KV cache for decode phase

### FlashInfer Fixed Length (`run_flashinfer_attn_fixed.py`)

测试固定序列长度的 attention。

```bash
python run_flashinfer_attn_fixed.py
```

### FlashInfer Ragged Only (`run_flashinfer_ragged_only.py`)

仅测试 ragged prefill attention。

```bash
python run_flashinfer_ragged_only.py
```

### Flash Attention (`run_flash_attn.py`)

测试标准 Flash Attention 实现。

```bash
python run_flash_attn.py
```

### vLLM Flash Attention (`run_vllm_flash_attn.py`)

测试 vLLM 的 Flash Attention 实现。

```bash
python run_vllm_flash_attn.py
```

---

## MoE (Mixture of Experts) Benchmarks

### Fused MoE (`run_fused_moe.py`)

测试 Fused MoE 操作 (Triton kernel, round-robin expert assignment)。

```bash
python run_fused_moe.py \
  --dir_path /data/output/moe \
  --moe_file ../dim/tiny_moe.csv \
  --seq_file ../dim/tiny_seq_len.csv \
  --num_iterations 5 \
  --params_dtype bfloat16 \
  --device cuda \
  --smem 100
```

**Features:**
- Triton-based fused MoE kernels
- Round-robin expert assignment (ideal load balancing)
- Configurable block sizes and expert counts

### Fused MoE DSE (`run_fused_moe_dse.py`)

MoE Design Space Exploration，遍历不同 kernel 配置。

```bash
python run_fused_moe_dse.py \
  --dir_path /data/output/moe_dse \
  --num_iterations 5 \
  --params_dtype bfloat16 \
  --device cuda
```

**Explores:**
- `num_warps`: [4, 8]
- `num_stages`: [2, 3, 4]
- Block size configurations
- Expert count variations

### Fused MoE NCU (`run_fused_moe_ncu.py`)

专用于 NCU profiling 的 MoE kernel runner，使用 GEMM 风格的 MNK 输入。

```bash
# 基本用法 (使用默认 kernel config)
python run_fused_moe_ncu.py --M 1024 --N 4096 --K 1024

# 自定义 kernel config
python run_fused_moe_ncu.py --M 1024 --N 4096 --K 1024 \
    --BLOCK_SIZE_M 128 --BLOCK_SIZE_N 128 --BLOCK_SIZE_K 64 \
    --GROUP_SIZE_M 16 --num_warps 8 --num_stages 4 \
    --iterations 10
```

| Parameter | Default | Description |
|-----------|---------|-------------|
| `--M` | required | Batch size (rows) |
| `--N` | required | Output dimension (columns) |
| `--K` | required | Hidden/input dimension |
| `--BLOCK_SIZE_M` | 64 | Block size for M dimension |
| `--BLOCK_SIZE_N` | 64 | Block size for N dimension |
| `--BLOCK_SIZE_K` | 64 | Block size for K dimension |
| `--GROUP_SIZE_M` | 8 | Group size for M dimension |
| `--num_warps` | 4 | Number of warps |
| `--num_stages` | 3 | Number of pipeline stages |
| `--iterations` | 1 | Number of iterations |
| `--device` | cuda | Device |

**Internal Implementation:**
- 使用 E=1, topk=1 将 GEMM 格式转换为 MOE 格式
- A: `[M, K]`, B: `[K, N]` → B_moe: `[1, N, K]`
- Output C: `[M, 1, N]`

---

## Configuration Files

所有配置文件位于 `dim/` 目录。

### General Configuration

| File | Description |
|------|-------------|
| `batch_size.csv` | Batch sizes |
| `seq_len.csv` | Sequence lengths (M values) |

### Module-Specific Configuration

| File | Description |
|------|-------------|
| `tiny_gemm_N_K.csv` | GEMM (N, K) pairs |
| `tiny_vocab_size.csv` | (vocab_size, hidden_size) |
| `tiny_attn_qkv.csv` | (q_heads, kv_heads, head_dim) |
| `tiny_hidden_size.csv` | Hidden dimensions |
| `tiny_activation_dim.csv` | Activation dimensions (input_size = 2 * hidden_size) |
| `tiny_moe.csv` | MoE configurations |

### Example Formats

**tiny_gemm_N_K.csv:**
```csv
N,K
4096,4096
8192,4096
4096,11008
```

**tiny_attn_qkv.csv:**
```csv
q,kv,head_dim
32,32,128
32,8,128
40,40,128
```

**seq_len.csv:**
```csv
M
1
16
32
128
512
```

---

## Repository Structure

```
modules_benchmark/
├── modules/                              # Benchmark scripts
│   ├── run_gemm.py                       # Standard GEMM
│   ├── run_fp8_gemm.py                   # FP8 GEMM
│   ├── run_gemm_multi_gpu.py             # Multi-GPU GEMM
│   ├── run_lmhead.py                     # LM Head
│   ├── run_flashinfer_rmsnorm.py         # FlashInfer RMSNorm
│   ├── run_vllm_rmsnorm.py               # vLLM RMSNorm
│   ├── run_flashinfer_rope.py            # FlashInfer RoPE
│   ├── run_vllm_rope.py                  # vLLM RoPE
│   ├── run_flashinfer_activation.py      # FlashInfer SiluAndMul
│   ├── run_vllm_activation.py            # vLLM SiluAndMul
│   ├── run_flashinfer_logits_processor.py # Logits processing
│   ├── run_flashinfer_attn_varlen_control.py # Variable length attention
│   ├── run_flashinfer_attn_fixed.py      # Fixed length attention
│   ├── run_flashinfer_ragged_only.py     # Ragged prefill only
│   ├── run_flash_attn.py                 # Flash Attention
│   ├── run_vllm_flash_attn.py            # vLLM Flash Attention
│   ├── run_fused_moe.py                  # Fused MoE
│   ├── run_fused_moe_dse.py              # MoE DSE
│   ├── run_fused_moe_ncu.py              # MoE NCU profiling
│   ├── utils.py                          # Utilities
│   └── vllm_utils.py                     # vLLM utilities
├── dim/                                  # Dimension config CSVs
├── extract/                              # Data extraction utilities
│   ├── extract_gemm_data_multiprocess.py
│   ├── extract_flashinfer_attn_varlen.py
│   ├── extract_fused_add_rmsnorm.py
│   ├── extract_siluandmul.py
│   ├── extract_fp8_gemm.py
│   └── extract_moe_dse.py
├── sglang_benchmark/                     # SGLang benchmarks
│   └── bench_one_batch.py
└── config/                               # Other configs
```

---

## Common Parameters

所有 benchmark 脚本支持以下通用参数:

| Parameter | Default | Description |
|-----------|---------|-------------|
| `--dir_path` | varies | Output directory for profiling traces (JSON) |
| `--num_iterations` | 5 | Number of profiling iterations per configuration |
| `--params_dtype` | bfloat16 | Data type (bfloat16 / float16 / fp8) |
| `--device` | cuda | Device to run benchmarks |

---

## Output Format

所有 benchmarks 生成 Chrome trace 文件 (JSON format)，可通过以下方式查看:

1. **Chrome Tracing**: 在 Chrome/Edge 浏览器中打开 `chrome://tracing`
2. **Perfetto UI**: 访问 https://ui.perfetto.dev/

Trace 文件包含详细的 kernel 执行时间、内存操作等 profiling 信息。

---

## Notes

- 所有 benchmarks 使用 PyTorch Profiler + CUDA synchronization 确保准确计时
- Benchmarks 包含 warmup iterations 确保 GPU 状态稳定
- 结果自动保存到指定目录，文件名包含配置信息
- vLLM benchmarks 需要正确安装 vLLM
- FlashInfer benchmarks 需要正确安装 FlashInfer
- FP8 benchmarks 需要 CUDA compute capability 8.9+ (H100, H200)