[Diffusion] Move hf kernels diffusion cuda kernels skills to SGLD (#20001)
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Xiaoyu Zhang
@@ -5,22 +5,40 @@ description: Index for SGLang Diffusion kernel development skills.
|
||||
|
||||
# Diffusion Kernel Skills
|
||||
|
||||
## Rule: Follow User Kernel Language Preference
|
||||
|
||||
If the user explicitly states a preference for **Triton** or **CUDA**, follow that preference when implementing and optimizing kernels (even if the other option could work). Do not “pick for convenience”.
|
||||
|
||||
## Directory Layout
|
||||
|
||||
```
|
||||
python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/
|
||||
├── SKILL.md
|
||||
├── add-triton-kernel.md
|
||||
├── add-cuda-kernel.md
|
||||
├── diffusion-benchmark-and-profile.md
|
||||
├── nsight-profiler.md
|
||||
└── use-efficient-diffusion-kernels.md
|
||||
├── use-efficient-diffusion-kernels.md
|
||||
├── references/
|
||||
│ ├── kernel-templates.md # Copy-paste CUDA kernel templates (sglang JIT style)
|
||||
│ ├── troubleshooting.md # Build/perf/integration issues & fixes
|
||||
│ ├── h100-optimization-guide.md # H100 (sm_90) deep dive
|
||||
│ ├── a100-optimization-guide.md # A100 (sm_80) deep dive
|
||||
│ └── t4-optimization-guide.md # T4 (sm_75, FP16 only) deep dive
|
||||
└── scripts/
|
||||
├── bench_diffusion_rmsnorm.py # RMSNorm micro-benchmark vs PyTorch
|
||||
└── bench_diffusion_denoise.py # End-to-end denoise benchmark (sglang generate)
|
||||
```
|
||||
|
||||
## Index
|
||||
|
||||
- [add-triton-kernel.md](./add-triton-kernel.md)
|
||||
|
||||
Step-by-step guide for adding a new Triton kernel to SGLang Diffusion's `jit_kernel` module, including authoring, autotune, `torch.compile` compatibility, integration, and tests.
|
||||
Step-by-step guide for adding a new Triton kernel to SGLang Diffusion's `jit_kernel/diffusion/triton/` module, including authoring, autotune, `torch.compile` compatibility, integration, and tests. Use for fused elementwise ops, norm variants, RoPE variants, or when NPU/CPU fallback is needed.
|
||||
|
||||
- [add-cuda-kernel.md](./add-cuda-kernel.md)
|
||||
|
||||
Step-by-step guide for adding a JIT CUDA kernel. CUDA source goes in `jit_kernel/csrc/diffusion/<op>.cuh`; Python wrapper at `jit_kernel/diffusion/<op>.py`. Uses SGLang's JIT compilation system (`load_jit`, `cache_once`) and internal abstractions (`TensorMatcher`, `device::AlignedVector`, `host::LaunchKernel`, `device::warp::reduce_sum`). Use for bandwidth-bound reductions (RMSNorm, LayerNorm) or ops needing fine-grained vectorization and shared memory control. Adapted from [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels).
|
||||
|
||||
- [use-efficient-diffusion-kernels.md](./use-efficient-diffusion-kernels.md)
|
||||
|
||||
@@ -28,8 +46,23 @@ python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/
|
||||
|
||||
- [diffusion-benchmark-and-profile.md](./diffusion-benchmark-and-profile.md)
|
||||
|
||||
End-to-end benchmarking and profiling guide for SGLang Diffusion models, including denoise latency measurement, per-layer breakdown, and regression tracking.
|
||||
Denoise-stage benchmark and profiling guide for SGLang Diffusion models. Three profiling levels: Level 1 (torch.profiler — kernel time ranking), Level 2 (nsys — category breakdown), Level 3 (ncu — per-kernel bandwidth/occupancy/roofline analysis). **ncu is critical for kernel optimization** — always use it when writing or tuning custom kernels to verify hardware saturation.
|
||||
|
||||
- [nsight-profiler.md](./nsight-profiler.md)
|
||||
|
||||
Advanced profiling skill for NVIDIA Nsight Systems / Nsight Compute: collecting traces, reading reports, and interpreting kernel-level performance metrics.
|
||||
|
||||
## References (GPU optimization guides, templates, troubleshooting)
|
||||
|
||||
Loaded by `add-cuda-kernel.md`. Adapted from [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels).
|
||||
|
||||
- [references/kernel-templates.md](references/kernel-templates.md) — copy-paste ready sglang JIT CUDA templates: element-wise (SiLU), row-reduction (RMSNorm), fused AdaLN, Python wrapper, test, benchmark
|
||||
- [references/troubleshooting.md](references/troubleshooting.md) — build errors, performance issues, torch.compile compatibility, kernel injection pitfalls
|
||||
- [references/h100-optimization-guide.md](references/h100-optimization-guide.md) — H100 (sm_90): AlignedVector benchmarks, warp reductions, occupancy, TMA, PDL
|
||||
- [references/a100-optimization-guide.md](references/a100-optimization-guide.md) — A100 (sm_80): cp.async, TF32, 2:4 sparsity, H100→A100 migration checklist
|
||||
- [references/t4-optimization-guide.md](references/t4-optimization-guide.md) — T4 (sm_75): FP16 only, 320 GB/s bandwidth, 64 KB shared mem, 16 GB memory management
|
||||
|
||||
## Scripts (runnable benchmarks)
|
||||
|
||||
- [scripts/bench_diffusion_rmsnorm.py](scripts/bench_diffusion_rmsnorm.py) — RMSNorm micro-benchmark: JIT CUDA vs PyTorch, correctness check, bandwidth efficiency analysis
|
||||
- [scripts/bench_diffusion_denoise.py](scripts/bench_diffusion_denoise.py) — end-to-end denoise benchmark via `sglang generate`, baseline vs custom kernels comparison table
|
||||
|
||||
@@ -0,0 +1,542 @@
|
||||
---
|
||||
name: add-cuda-kernel
|
||||
description: Step-by-step guide for adding a new JIT CUDA kernel to SGLang Diffusion. CUDA source files go in jit_kernel/csrc/diffusion/<op>.cuh; Python wrapper at jit_kernel/diffusion/<op>.py. Use when implementing optimized CUDA kernels for diffusion model operators (RMSNorm, RoPE, AdaLN, GEGLU, etc.) on NVIDIA GPUs (H100, A100). Covers kernel authoring with sglang abstractions, JIT compilation, Python wrapper, integration into the denoise stage, and benchmarking. Adapted from https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels.
|
||||
---
|
||||
|
||||
# Adding a CUDA Kernel to SGLang Diffusion (JIT Style)
|
||||
|
||||
> **Origin**: This skill is adapted from the [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels), rewritten to follow SGLang's JIT compilation system and internal abstractions.
|
||||
>
|
||||
> **Extended references** (in this directory's `references/` and `scripts/`):
|
||||
> - [references/kernel-templates.md](references/kernel-templates.md) — copy-paste ready templates for element-wise, row-reduction (RMSNorm), fused AdaLN
|
||||
> - [references/troubleshooting.md](references/troubleshooting.md) — build errors, perf issues, integration pitfalls
|
||||
> - [references/h100-optimization-guide.md](references/h100-optimization-guide.md) — H100 (sm_90) deep dive
|
||||
> - [references/a100-optimization-guide.md](references/a100-optimization-guide.md) — A100 (sm_80) deep dive
|
||||
> - [references/t4-optimization-guide.md](references/t4-optimization-guide.md) — T4 (sm_75, FP16 only) deep dive
|
||||
> - [scripts/bench_diffusion_rmsnorm.py](scripts/bench_diffusion_rmsnorm.py) — RMSNorm micro-benchmark vs PyTorch
|
||||
> - [scripts/bench_diffusion_denoise.py](scripts/bench_diffusion_denoise.py) — end-to-end denoise benchmark with/without kernels
|
||||
|
||||
## When to Use CUDA vs Triton
|
||||
|
||||
| Scenario | Use |
|
||||
|----------|-----|
|
||||
| Fused elementwise / norm variants / RoPE | **Triton** (`add-triton-kernel.md`) — faster iteration |
|
||||
| Bandwidth-bound reduction (RMSNorm, LayerNorm) requiring max vectorization | **CUDA** — full control over `__nv_bfloat162` / `float4` vectorization |
|
||||
| Attention pattern or tile-based ops needing shared memory tuning | **CUDA** — warp-level primitives, shared memory layout |
|
||||
| Prototype or NPU/CPU fallback needed | **Triton** — portable across backends |
|
||||
|
||||
For most diffusion-model elementwise ops, **start with Triton**. Switch to CUDA when profiling shows Triton can't reach hardware bandwidth limits.
|
||||
|
||||
## Directory Layout
|
||||
|
||||
```
|
||||
python/sglang/jit_kernel/
|
||||
├── csrc/
|
||||
│ ├── diffusion/ # JIT CUDA source files for diffusion kernels (this skill)
|
||||
│ │ ├── timestep_embedding.cuh # existing example
|
||||
│ │ ├── rmsnorm.cuh # NEW: add here
|
||||
│ │ └── adaln.cuh # NEW: add here
|
||||
│ └── elementwise/ # shared JIT CUDA csrc (non-diffusion)
|
||||
├── diffusion/
|
||||
│ ├── triton/ # Triton kernels (scale_shift, norm, rope, ...)
|
||||
│ ├── cutedsl/ # CuTe DSL kernels
|
||||
│ └── rmsnorm.py # NEW: CUDA JIT Python wrapper (add here)
|
||||
├── timestep_embedding.py # existing CUDA diffusion kernel Python wrapper (legacy)
|
||||
```
|
||||
|
||||
New diffusion CUDA kernel source files go into `python/sglang/jit_kernel/csrc/diffusion/<op_name>.cuh`.
|
||||
The Python wrapper goes at `python/sglang/jit_kernel/diffusion/<op_name>.py`
|
||||
(inside `diffusion/`, alongside the `triton/` and `cutedsl/` subdirectories).
|
||||
|
||||
---
|
||||
|
||||
## SGLang Kernel Abstractions (Required)
|
||||
|
||||
Always use these — do **not** use raw CUDA primitives directly.
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/tensor.h> // TensorMatcher, SymbolicSize, SymbolicDevice
|
||||
#include <sgl_kernel/type.cuh> // fp16_t, bf16_t, fp32_t, dtype_trait, packed_t
|
||||
#include <sgl_kernel/utils.h> // RuntimeCheck, div_ceil
|
||||
#include <sgl_kernel/utils.cuh> // LaunchKernel, SGL_DEVICE, type aliases
|
||||
#include <sgl_kernel/vec.cuh> // AlignedVector<T, N> — 128-bit vector loads
|
||||
#include <sgl_kernel/warp.cuh> // warp::reduce_sum, warp::reduce_max
|
||||
#include <sgl_kernel/math.cuh> // device::math::rsqrt, sqrt, ...
|
||||
#include <sgl_kernel/tile.cuh> // tile::Memory (strided access pattern)
|
||||
```
|
||||
|
||||
Key types: `fp16_t` = `__half`, `bf16_t` = `__nv_bfloat16`, `fp32_t` = `float`.
|
||||
Packed variants: `fp16x2_t`, `bf16x2_t`. Use `packed_t<T>` for the 2-element alias.
|
||||
|
||||
---
|
||||
|
||||
## Step 1: Write the CUDA Kernel
|
||||
|
||||
Create `python/sglang/jit_kernel/csrc/diffusion/rmsnorm.cuh` (RMSNorm as example).
|
||||
|
||||
### 1a. Vectorized RMSNorm Kernel
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/type.cuh>
|
||||
#include <sgl_kernel/utils.h>
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
#include <sgl_kernel/vec.cuh>
|
||||
#include <sgl_kernel/warp.cuh>
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
|
||||
namespace {
|
||||
|
||||
// ---------------------------------------------------------------
|
||||
// RMSNorm kernel: y = x / rms(x) * weight
|
||||
// T = fp16_t | bf16_t | fp32_t
|
||||
// kVecN = vectorized elements per load (8 for fp16/bf16, 4 for fp32)
|
||||
// ---------------------------------------------------------------
|
||||
template <typename T, int kVecN>
|
||||
__global__ void rmsnorm_kernel(
|
||||
T* __restrict__ dst,
|
||||
const T* __restrict__ src,
|
||||
const T* __restrict__ weight, // may be nullptr if no affine weight
|
||||
uint32_t hidden_size,
|
||||
uint32_t n_vecs, // hidden_size / kVecN
|
||||
float eps)
|
||||
{
|
||||
using vec_t = device::AlignedVector<T, kVecN>;
|
||||
|
||||
const uint32_t row = blockIdx.x;
|
||||
const T* row_src = src + row * hidden_size;
|
||||
T* row_dst = dst + row * hidden_size;
|
||||
|
||||
// --- Pass 1: accumulate sum of squares (vectorized) ---
|
||||
float sum_sq = 0.f;
|
||||
for (uint32_t vi = threadIdx.x; vi < n_vecs; vi += blockDim.x) {
|
||||
vec_t v;
|
||||
v.load(row_src, vi);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecN; ++i) {
|
||||
float val = static_cast<float>(v[i]);
|
||||
sum_sq += val * val;
|
||||
}
|
||||
}
|
||||
|
||||
// --- Warp reduction ---
|
||||
sum_sq = device::warp::reduce_sum<float>(sum_sq);
|
||||
|
||||
// --- Block reduction via shared memory ---
|
||||
__shared__ float smem[32];
|
||||
if (threadIdx.x % 32 == 0) {
|
||||
smem[threadIdx.x / 32] = sum_sq;
|
||||
}
|
||||
__syncthreads();
|
||||
if (threadIdx.x < 32) {
|
||||
sum_sq = (threadIdx.x < blockDim.x / 32) ? smem[threadIdx.x] : 0.f;
|
||||
sum_sq = device::warp::reduce_sum<float>(sum_sq);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const float rms_inv = device::math::rsqrt<float>(sum_sq / static_cast<float>(hidden_size) + eps);
|
||||
|
||||
// --- Pass 2: normalize + apply weight (vectorized) ---
|
||||
for (uint32_t vi = threadIdx.x; vi < n_vecs; vi += blockDim.x) {
|
||||
vec_t v_in, v_w, v_out;
|
||||
v_in.load(row_src, vi);
|
||||
if (weight != nullptr) {
|
||||
v_w.load(weight, vi);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecN; ++i) {
|
||||
float val = static_cast<float>(v_in[i]) * rms_inv;
|
||||
if (weight != nullptr) {
|
||||
val *= static_cast<float>(v_w[i]);
|
||||
}
|
||||
v_out[i] = static_cast<T>(val);
|
||||
}
|
||||
v_out.store(row_dst, vi);
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------
|
||||
// Launcher
|
||||
// ---------------------------------------------------------------
|
||||
template <typename T>
|
||||
void rmsnorm(
|
||||
tvm::ffi::TensorView dst,
|
||||
tvm::ffi::TensorView src,
|
||||
tvm::ffi::TensorView weight, // pass empty / nullptr for no-weight case
|
||||
float eps)
|
||||
{
|
||||
using namespace host;
|
||||
|
||||
// Validate
|
||||
SymbolicSize B{"batch_tokens"}, H{"hidden_size"};
|
||||
SymbolicDevice device;
|
||||
device.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({B, H})
|
||||
.with_dtype<T>()
|
||||
.with_device(device)
|
||||
.verify(dst)
|
||||
.verify(src);
|
||||
|
||||
const uint32_t num_rows = static_cast<uint32_t>(B.unwrap());
|
||||
const uint32_t hidden = static_cast<uint32_t>(H.unwrap());
|
||||
const DLDevice dev = device.unwrap();
|
||||
|
||||
RuntimeCheck(hidden % (16 / sizeof(T)) == 0,
|
||||
"rmsnorm: hidden_size must be divisible by vector width, got ", hidden);
|
||||
|
||||
constexpr int kVecN = 16 / sizeof(T); // 128-bit vector: 8×fp16/bf16, 4×fp32
|
||||
const uint32_t n_vecs = hidden / kVecN;
|
||||
|
||||
// Thread count: enough warps to cover n_vecs, max 512 threads
|
||||
uint32_t threads = std::min(n_vecs, 512u);
|
||||
threads = (threads + 31) / 32 * 32; // round up to warp boundary
|
||||
|
||||
const T* w_ptr = (weight.data_ptr() != nullptr)
|
||||
? static_cast<const T*>(weight.data_ptr()) : nullptr;
|
||||
|
||||
LaunchKernel(num_rows, threads, dev)(
|
||||
rmsnorm_kernel<T, kVecN>,
|
||||
static_cast<T*>(dst.data_ptr()),
|
||||
static_cast<const T*>(src.data_ptr()),
|
||||
w_ptr,
|
||||
hidden,
|
||||
n_vecs,
|
||||
eps);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Step 2: Python Wrapper
|
||||
|
||||
Create `python/sglang/jit_kernel/diffusion/rmsnorm.py`:
|
||||
|
||||
```python
|
||||
from __future__ import annotations
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_rmsnorm_module(dtype: torch.dtype) -> Module:
|
||||
args = make_cpp_args(dtype)
|
||||
return load_jit(
|
||||
"diffusion_rmsnorm",
|
||||
*args,
|
||||
cuda_files=["diffusion/rmsnorm.cuh"], # relative to csrc/
|
||||
cuda_wrappers=[("rmsnorm", f"rmsnorm<{args}>")],
|
||||
extra_cuda_cflags=["-O3", "--use_fast_math"],
|
||||
)
|
||||
|
||||
|
||||
def diffusion_rmsnorm(
|
||||
src: torch.Tensor,
|
||||
weight: torch.Tensor | None = None,
|
||||
eps: float = 1e-6,
|
||||
out: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
RMSNorm for diffusion DiT layers.
|
||||
|
||||
y = x / rms(x) * weight (weight=None → no affine scaling)
|
||||
|
||||
Supported dtypes: float16, bfloat16, float32.
|
||||
hidden_size must be divisible by 8 (fp16/bf16) or 4 (fp32).
|
||||
"""
|
||||
assert src.is_cuda, "src must be a CUDA tensor"
|
||||
assert src.dtype in (torch.float16, torch.bfloat16, torch.float32)
|
||||
|
||||
if out is None:
|
||||
out = torch.empty_like(src)
|
||||
|
||||
# Pass a zero-sized tensor when weight is absent (launcher checks data_ptr == nullptr)
|
||||
w = weight if weight is not None else torch.empty(0, dtype=src.dtype, device=src.device)
|
||||
|
||||
module = _jit_rmsnorm_module(src.dtype)
|
||||
module.rmsnorm(out, src, w, eps)
|
||||
return out
|
||||
```
|
||||
|
||||
**Key rules for the wrapper:**
|
||||
- Use `cache_once` — never `functools.lru_cache` (breaks `torch.compile`)
|
||||
- First arg(s) to `load_jit` form the unique build cache key
|
||||
- `cuda_files` are relative to `python/sglang/jit_kernel/csrc/`
|
||||
- `cuda_wrappers`: `(python_name, cpp_template_instantiation)`
|
||||
|
||||
---
|
||||
|
||||
## Step 3: Integrate into Denoising Stage
|
||||
|
||||
The kernel replaces a slow operator inside the DiT forward pass. Find the correct module in:
|
||||
|
||||
```
|
||||
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
|
||||
python/sglang/multimodal_gen/runtime/models/dits/<model>.py
|
||||
```
|
||||
|
||||
**Pattern — monkey-patch the DiT block's RMSNorm:**
|
||||
|
||||
```python
|
||||
from sglang.jit_kernel.diffusion.rmsnorm import diffusion_rmsnorm
|
||||
|
||||
def _patch_rmsnorm(model: torch.nn.Module) -> None:
|
||||
for name, module in model.named_modules():
|
||||
cls_name = type(module).__name__
|
||||
if cls_name in ("RMSNorm", "LlamaRMSNorm") or "RMSNorm" in cls_name:
|
||||
eps = getattr(module, "eps", getattr(module, "variance_epsilon", 1e-6))
|
||||
has_weight = hasattr(module, "weight") and module.weight is not None
|
||||
|
||||
if has_weight:
|
||||
def _make_fwd(mod, epsilon):
|
||||
def forward(x):
|
||||
return diffusion_rmsnorm(x, weight=mod.weight, eps=epsilon)
|
||||
return forward
|
||||
module.forward = _make_fwd(module, eps)
|
||||
else:
|
||||
def _make_fwd_noweight(epsilon):
|
||||
def forward(x):
|
||||
return diffusion_rmsnorm(x, weight=None, eps=epsilon)
|
||||
return forward
|
||||
module.forward = _make_fwd_noweight(eps)
|
||||
```
|
||||
|
||||
**Critical:** inject kernels **before** `torch.compile` and before any CPU offload is enabled.
|
||||
|
||||
---
|
||||
|
||||
## Step 4: Key Kernel Patterns Reference
|
||||
|
||||
### Diffusion-Specific Operators
|
||||
|
||||
| Operator | Kernel Pattern | Notes |
|
||||
|----------|---------------|-------|
|
||||
| **RMSNorm** | 2-pass row reduction + vectorized normalize | Weight may be `None` (`elementwise_affine=False`) |
|
||||
| **AdaLN modulation** | `y = norm(x) * (1 + scale) + shift` | Fuse norm + scale + shift in one pass |
|
||||
| **RoPE 3D** | Read `(t, h, w)` cos/sin tables, apply to `(q, k)` | Layout: `[batch, t*h*w, heads, head_dim]` |
|
||||
| **GEGLU** | Split last dim → `gate * silu(linear)` | Input `[B, L, 2*H]` → output `[B, L, H]` |
|
||||
| **SiLU gate** | `out = a * sigmoid(a)` fused | Avoid separate elementwise ops |
|
||||
|
||||
### Vectorized Memory Access
|
||||
|
||||
```cpp
|
||||
// BF16: 8 elements × 2 bytes = 16 bytes per vector load (AlignedVector<bf16_t, 8>)
|
||||
// FP16: 8 elements × 2 bytes = 16 bytes (AlignedVector<fp16_t, 8>)
|
||||
// FP32: 4 elements × 4 bytes = 16 bytes (AlignedVector<fp32_t, 4>)
|
||||
constexpr int kVecN = 16 / sizeof(T);
|
||||
using vec_t = device::AlignedVector<T, kVecN>;
|
||||
```
|
||||
|
||||
### Warp / Block Reductions
|
||||
|
||||
```cpp
|
||||
// Warp reduction (within 32 threads)
|
||||
float result = device::warp::reduce_sum<float>(partial);
|
||||
|
||||
// Block reduction via shared memory (see rmsnorm example above)
|
||||
__shared__ float smem[32];
|
||||
// ... write warp-leaders into smem, sync, reduce again
|
||||
```
|
||||
|
||||
### Thread Configuration
|
||||
|
||||
```cpp
|
||||
// Element-wise (RoPE, GEGLU, SiLU): simple 1D grid
|
||||
constexpr uint32_t kBlock = 256;
|
||||
uint32_t grid = host::div_ceil(total_elements, kBlock);
|
||||
LaunchKernel(grid, kBlock, dev)(kernel, ...);
|
||||
|
||||
// Row reduction (RMSNorm, LayerNorm): one block per row
|
||||
uint32_t threads = std::min(hidden_size / kVecN, 512u);
|
||||
threads = (threads + 31) / 32 * 32;
|
||||
LaunchKernel(num_rows, threads, dev)(kernel, ...);
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Step 5: GPU Architecture Targets
|
||||
|
||||
| GPU | Compute Cap | Memory BW | BF16 | Key Note |
|
||||
|-----|------------|-----------|------|----------|
|
||||
| H100 | sm_90 | 3.35 TB/s | Yes | Primary target; 132 SMs, 192 KB shared mem/SM |
|
||||
| A100 | sm_80 | 2.0 TB/s | Yes | 108 SMs, 164 KB shared mem/SM |
|
||||
| T4 | sm_75 | 320 GB/s | **No** | FP16 only; no `__nv_bfloat16` |
|
||||
|
||||
If kernel requires SM90+ features (e.g., TMA, wgmma), raise a clear error:
|
||||
|
||||
```python
|
||||
if torch.cuda.get_device_capability()[0] < 9:
|
||||
raise RuntimeError("This kernel requires SM90 (H100/Hopper) or later")
|
||||
```
|
||||
|
||||
**Grid sizing for H100** (132 SMs): aim for grid multiples of 132 for good occupancy.
|
||||
|
||||
---
|
||||
|
||||
## Step 6: Tests
|
||||
|
||||
Create `python/sglang/jit_kernel/tests/test_diffusion_rmsnorm.py`:
|
||||
|
||||
```python
|
||||
import pytest
|
||||
import torch
|
||||
from sglang.jit_kernel.diffusion.rmsnorm import diffusion_rmsnorm
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
|
||||
@pytest.mark.parametrize("shape", [(1, 2048), (4, 3072), (16, 4096)])
|
||||
@pytest.mark.parametrize("has_weight", [True, False])
|
||||
def test_rmsnorm_correctness(dtype, shape, has_weight):
|
||||
batch, hidden = shape
|
||||
src = torch.randn(batch, hidden, dtype=dtype, device="cuda")
|
||||
weight = torch.randn(hidden, dtype=dtype, device="cuda") if has_weight else None
|
||||
|
||||
out_jit = diffusion_rmsnorm(src, weight=weight, eps=1e-6)
|
||||
|
||||
# Reference: torch.nn.functional
|
||||
ref = torch.nn.functional.rms_norm(
|
||||
src.float(), (hidden,), weight.float() if weight is not None else None, eps=1e-6
|
||||
).to(dtype)
|
||||
|
||||
tol = {"rtol": 1e-2, "atol": 1e-2} if dtype != torch.float32 else {"rtol": 1e-5, "atol": 1e-6}
|
||||
torch.testing.assert_close(out_jit, ref, **tol)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v", "-s"])
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Step 7: Benchmark
|
||||
|
||||
Create `python/sglang/jit_kernel/benchmark/bench_diffusion_rmsnorm.py`:
|
||||
|
||||
```python
|
||||
import torch
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import DEFAULT_DEVICE, DEFAULT_DTYPE, run_benchmark
|
||||
from sglang.jit_kernel.diffusion.rmsnorm import diffusion_rmsnorm
|
||||
|
||||
SHAPES = [(4096, 2048), (4096, 3072), (4096, 4096)]
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["hidden"],
|
||||
x_vals=[s[1] for s in SHAPES],
|
||||
line_arg="provider",
|
||||
line_vals=["jit_cuda", "torch"],
|
||||
line_names=["SGLang JIT CUDA", "PyTorch rms_norm"],
|
||||
styles=[("blue", "-"), ("red", "--")],
|
||||
ylabel="us",
|
||||
plot_name="diffusion-rmsnorm",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(hidden: int, provider: str):
|
||||
src = torch.randn(4096, hidden, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE)
|
||||
w = torch.ones(hidden, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE)
|
||||
|
||||
if provider == "jit_cuda":
|
||||
fn = lambda: diffusion_rmsnorm(src, weight=w, eps=1e-6)
|
||||
else:
|
||||
fn = lambda: torch.nn.functional.rms_norm(src, (hidden,), w, eps=1e-6)
|
||||
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark.run(print_data=True)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Step 8: Profile with Nsight Compute (required)
|
||||
|
||||
After correctness + benchmarking, you must collect **Nsight Compute (ncu)** data to validate:
|
||||
|
||||
- Whether the kernel reaches reasonable bandwidth/throughput (avoid false positives where it is “faster” but under-utilizes hardware)
|
||||
- Whether there are clear occupancy / register / shared memory limiters
|
||||
|
||||
Use the canonical docs in this directory (do not duplicate CLI details across multiple skills):
|
||||
|
||||
- `diffusion-benchmark-and-profile.md` → Step 3.5 (ncu workflow, including CUDA graph profiling)
|
||||
- `nsight-profiler.md` (metrics interpretation: bandwidth / occupancy / roofline / stall reasons)
|
||||
|
||||
---
|
||||
|
||||
## Common Pitfalls
|
||||
|
||||
| Issue | Fix |
|
||||
|-------|-----|
|
||||
| `RMSNorm weight is None` | Use `type(module).__name__` check; pass `None` weight explicitly |
|
||||
| `isinstance(m, torch.nn.RMSNorm)` misses diffusers variants | Use `"RMSNorm" in type(m).__name__` |
|
||||
| Kernel patched after `torch.compile` | Inject **before** any compile call |
|
||||
| Kernel patched after `enable_model_cpu_offload()` | Inject **before** CPU offload |
|
||||
| `hidden_size` not divisible by `kVecN` | Add `RuntimeCheck(hidden % kVecN == 0, ...)` in launcher |
|
||||
| `torch.compile` fails with custom CUDA kernel | Register as `@torch.library.custom_op` or use Triton instead |
|
||||
| T4 GPU with BF16 kernel | Gate on compute capability; T4 is `sm_75`, no native BF16 |
|
||||
|
||||
---
|
||||
|
||||
## Summary of Files
|
||||
|
||||
```
|
||||
python/sglang/jit_kernel/csrc/diffusion/
|
||||
└── rmsnorm.cuh # NEW: JIT CUDA kernel source
|
||||
|
||||
python/sglang/jit_kernel/diffusion/
|
||||
└── rmsnorm.py # NEW: Python wrapper + load_jit
|
||||
|
||||
python/sglang/jit_kernel/tests/
|
||||
└── test_diffusion_rmsnorm.py # NEW: correctness tests
|
||||
|
||||
python/sglang/jit_kernel/benchmark/
|
||||
└── bench_diffusion_rmsnorm.py # NEW: benchmark
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## References
|
||||
|
||||
### This Skill's Extended Docs (references/ and scripts/)
|
||||
|
||||
| File | Contents |
|
||||
|------|----------|
|
||||
| [references/kernel-templates.md](references/kernel-templates.md) | Copy-paste templates: element-wise, RMSNorm, AdaLN, Python wrapper, test, benchmark |
|
||||
| [references/troubleshooting.md](references/troubleshooting.md) | Build errors, perf issues, torch.compile compatibility, debugging checklist |
|
||||
| [references/h100-optimization-guide.md](references/h100-optimization-guide.md) | H100 (sm_90): memory hierarchy, warp reductions, occupancy, vectorization benchmarks |
|
||||
| [references/a100-optimization-guide.md](references/a100-optimization-guide.md) | A100 (sm_80): cp.async, TF32, 2:4 sparsity, H100→A100 migration checklist |
|
||||
| [references/t4-optimization-guide.md](references/t4-optimization-guide.md) | T4 (sm_75): FP16 only, low bandwidth, tile size limits, memory constraints |
|
||||
| [scripts/bench_diffusion_rmsnorm.py](scripts/bench_diffusion_rmsnorm.py) | Micro-benchmark: JIT CUDA RMSNorm vs PyTorch, correctness check, bandwidth analysis |
|
||||
| [scripts/bench_diffusion_denoise.py](scripts/bench_diffusion_denoise.py) | End-to-end: `sglang generate` baseline vs custom kernels, comparison table |
|
||||
|
||||
### SGLang Internals
|
||||
|
||||
- **JIT system**: `add-jit-kernel` skill (`sglang/.claude/skills/add-jit-kernel/SKILL.md`)
|
||||
- **JIT utils**: `python/sglang/jit_kernel/utils.py` — `cache_once`, `load_jit`, `make_cpp_args`
|
||||
- **Abstractions**: `python/sglang/jit_kernel/include/sgl_kernel/` — `tensor.h`, `utils.cuh`, `vec.cuh`, `warp.cuh`, `math.cuh`, `tile.cuh`
|
||||
- **Real csrc examples**: `python/sglang/jit_kernel/csrc/elementwise/rmsnorm.cuh`, `python/sglang/jit_kernel/csrc/elementwise/qknorm.cuh`
|
||||
|
||||
### Other Diffusion Kernel Skills (this directory)
|
||||
|
||||
- **Triton alternative**: `add-triton-kernel.md` — prefer Triton unless bandwidth analysis shows CUDA needed
|
||||
- **Existing fused kernels**: `use-efficient-diffusion-kernels.md` — check here first before writing new kernels
|
||||
- **Profiling**: `diffusion-benchmark-and-profile.md` — workflow to identify bottleneck before implementing
|
||||
- **Nsight Compute deep dive**: `nsight-profiler.md` — full guide: occupancy analysis, roofline model, warp efficiency, kernel comparison
|
||||
|
||||
### External
|
||||
|
||||
- [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels) — original source adapted for this skill
|
||||
@@ -366,101 +366,14 @@ python python/sglang/jit_kernel/benchmark/bench_<op_name>.py
|
||||
|
||||
---
|
||||
|
||||
## Step 6: Profile with Nsight Compute (optional but recommended)
|
||||
## Step 6: Profile with Nsight Compute (required for optimization work)
|
||||
|
||||
After the kernel passes correctness tests, use **`nsight-profiler.md`** to measure its hardware-level efficiency. This step requires **Nsight Compute (`ncu`)** to be installed.
|
||||
After correctness tests, you must use **ncu (Nsight Compute)** to validate hardware efficiency (bandwidth/throughput/occupancy/bottleneck type).
|
||||
|
||||
### Dependency Check
|
||||
To avoid duplicating ncu CLI details across multiple skills, this skill does not repeat command flags. Follow the canonical docs:
|
||||
|
||||
Before profiling, verify `ncu` is available:
|
||||
|
||||
```bash
|
||||
ncu --version # must print a version string, e.g. "NVIDIA Nsight Compute 2024.1.0"
|
||||
```
|
||||
|
||||
If `ncu` is missing, install it via the CUDA Toolkit package or the standalone [Nsight Compute installer](https://developer.nvidia.com/nsight-compute).
|
||||
|
||||
### Quick Kernel Profile
|
||||
|
||||
```bash
|
||||
# Profile the Triton kernel during the benchmark script
|
||||
# --kernel-name: match the Triton-mangled name (check with nsys first if unsure)
|
||||
ncu --set full \
|
||||
-o /tmp/triton_<op_name> \
|
||||
python python/sglang/jit_kernel/benchmark/bench_<op_name>.py
|
||||
```
|
||||
|
||||
To profile only the Triton kernel (skip PyTorch reference and warmup launches), add `--launch-skip N --launch-count M`:
|
||||
|
||||
```bash
|
||||
# Skip first 2 launches (warmup), capture 3 kernel invocations
|
||||
ncu --set full --launch-skip 2 --launch-count 3 \
|
||||
-o /tmp/triton_<op_name> \
|
||||
python python/sglang/jit_kernel/benchmark/bench_<op_name>.py
|
||||
```
|
||||
|
||||
### Key Metrics to Check for a Triton Kernel
|
||||
|
||||
| Metric | Healthy Range | Action if Low |
|
||||
|--------|--------------|---------------|
|
||||
| **Achieved Occupancy** | ≥ 50% | Reduce register usage or shared memory; try smaller block sizes |
|
||||
| **Memory Throughput** | ≥ 70% of peak BW | Check for non-coalesced access (pass contiguous strides) |
|
||||
| **Compute Throughput** | ≥ 50% of peak | Increase arithmetic intensity; fuse more ops per load |
|
||||
| **Warp Efficiency (No Stall)** | ≥ 60% | Reduce branch divergence; avoid `tl.atomic_*` on hot paths |
|
||||
| **L1/L2 Hit Rate** | L2 ≥ 40% | Reorder loads for locality; check broadcast patterns |
|
||||
|
||||
### CLI-Only Analysis Workflow
|
||||
|
||||
```bash
|
||||
# 1. Collect profile (--csv is NOT valid here; only -o to save .ncu-rep)
|
||||
ncu --set full \
|
||||
--launch-skip 2 --launch-count 1 \
|
||||
-o /tmp/prof_<op_name> \
|
||||
python python/sglang/jit_kernel/benchmark/bench_<op_name>.py
|
||||
|
||||
# 2. Export key metrics to CSV from the saved .ncu-rep (--csv is valid here)
|
||||
ncu --import /tmp/prof_<op_name>.ncu-rep \
|
||||
--page details --csv \
|
||||
> /tmp/prof_<op_name>_details.csv
|
||||
|
||||
# 3. Quick summary — top bottleneck sections
|
||||
python3 - << 'EOF'
|
||||
import csv, sys
|
||||
rows = list(csv.DictReader(open("/tmp/prof_<op_name>_details.csv")))
|
||||
# print section names and their achieved % of peak
|
||||
for r in rows:
|
||||
name = r.get("Metric Name", "")
|
||||
val = r.get("Metric Value", "")
|
||||
if any(k in name for k in ["sol", "Occupancy", "Throughput"]):
|
||||
print(f"{name:60s} {val}")
|
||||
EOF
|
||||
```
|
||||
|
||||
### Compare Two Kernel Versions
|
||||
|
||||
```bash
|
||||
# Profile baseline
|
||||
ncu --set full --launch-skip 2 --launch-count 1 \
|
||||
-o /tmp/baseline python .../bench_<op_name>.py
|
||||
|
||||
# Profile optimized version (after your changes)
|
||||
ncu --set full --launch-skip 2 --launch-count 1 \
|
||||
-o /tmp/optimized python .../bench_<op_name>.py
|
||||
|
||||
# Diff (CSV, no GUI)
|
||||
ncu --import /tmp/baseline.ncu-rep \
|
||||
--import /tmp/optimized.ncu-rep \
|
||||
--page diff --csv > /tmp/diff_<op_name>.csv
|
||||
|
||||
python3 -c "
|
||||
import csv
|
||||
rows = list(csv.DictReader(open('/tmp/diff_<op_name>.csv')))
|
||||
for r in rows[:30]:
|
||||
print(r.get('Metric Name','')[:55], r.get('Baseline',''), '->', r.get('Comparison',''))
|
||||
"
|
||||
```
|
||||
|
||||
> For a complete guide to Nsight Compute metrics, occupancy analysis, roofline model interpretation, and warp efficiency optimization, refer to **`nsight-profiler.md`** in this directory.
|
||||
- `diffusion-benchmark-and-profile.md` → Step 3.5 (ncu workflow, including CUDA graph profiling)
|
||||
- `nsight-profiler.md` (metrics interpretation: bandwidth / occupancy / roofline / warp stalls)
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -1,547 +1,232 @@
|
||||
---
|
||||
name: diffusion-benchmark-and-profile
|
||||
description: End-to-end benchmark and per-layer kernel profiling guide for SGLang Diffusion models. Use when measuring SGLang Diffusion generation performance, running latency benchmarks across Qwen-Image, FLUX, Z-Image-Turbo, and Wan2.2 models, profiling DiT layer kernel breakdown with torch.profiler or nsys+gputrc2graph.py, investigating performance bottlenecks, or tracking performance regressions. Always verify output correctness before and after any optimization.
|
||||
description: Denoise-stage benchmark and per-layer kernel profiling guide for SGLang Diffusion models. Use when measuring denoising latency, profiling DiT kernel breakdown with torch.profiler or nsys+gputrc2graph.py, investigating performance bottlenecks, or optimizing with custom Triton/CUDA kernels. Always verify output correctness before and after any optimization.
|
||||
---
|
||||
|
||||
# SGLang Diffusion Benchmark and Profile Guide
|
||||
|
||||
**Overview**
|
||||
This skill covers how to run end-to-end benchmarks for SGLang Diffusion (`sglang generate`) across a standard set of models, profile per-layer kernel execution inside the DiT, and use those results to continuously improve performance. These metrics collectively reflect the overall performance of the current SGLang Diffusion release.
|
||||
|
||||
> **Correctness First**: Any performance optimization must be validated for correctness before being considered complete. Faster but incorrect output is not an improvement. Always compare generated images/videos against a reference baseline before and after any change.
|
||||
|
||||
**Benchmarking / Perf Report (Reuse Existing Skill)**
|
||||
For end-to-end latency measurement, JSON dump format, before/after comparison, and extracting a single number for reports, reuse the `diffusion-perf` skill:
|
||||
|
||||
- `../diffusion-perf/SKILL.md`
|
||||
|
||||
This document focuses on:
|
||||
- The standard diffusion model suite commands (what to run)
|
||||
- When and how to do per-layer / per-kernel profiling (how to analyze)
|
||||
|
||||
**Primary Metric: Denoise Latency**
|
||||
The most important latency signal is the **denoising loop latency** — the total time spent running the DiT forward pass across all inference steps. This is the dominant cost in every diffusion model and the main target for optimization. End-to-end latency (including VAE decode and text encoding) is also recorded as a secondary metric, but denoising latency is the key indicator of DiT model performance.
|
||||
The denoising loop latency — total DiT forward pass time across all inference steps — is the dominant cost (>80% of end-to-end) and the **sole optimization target** for kernel work. End-to-end latency is recorded as a secondary check only.
|
||||
|
||||
> **Correctness First**: Faster but incorrect output is not an improvement. Always compare generated images/videos against a reference baseline before and after any change.
|
||||
|
||||
---
|
||||
|
||||
## Prerequisites
|
||||
|
||||
### Tool Dependency Check
|
||||
|
||||
Before running any benchmark or profiling command, verify that all required tools are available. Run the following check script:
|
||||
|
||||
```bash
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
PASS=0; FAIL=0
|
||||
|
||||
check() {
|
||||
local label=$1; shift
|
||||
if "$@" &>/dev/null; then
|
||||
echo " [OK] $label"
|
||||
((PASS++)) || true
|
||||
else
|
||||
echo " [MISS] $label"
|
||||
((FAIL++)) || true
|
||||
fi
|
||||
}
|
||||
|
||||
echo "=== SGLang Diffusion Benchmark Prerequisites ==="
|
||||
|
||||
# Core runtime
|
||||
check "sglang CLI" python3 -c "import sglang"
|
||||
check "torch" python3 -c "import torch; assert torch.cuda.is_available()"
|
||||
check "CUDA available" python3 -c "import torch; torch.zeros(1).cuda()"
|
||||
|
||||
# Profiling (torch.profiler is built into torch — no extra install needed)
|
||||
check "torch.profiler" python3 -c "import torch.profiler"
|
||||
|
||||
# nsys (optional, for Level 2 profiling)
|
||||
check "nsys in PATH" which nsys
|
||||
|
||||
# gputrc2graph.py dependencies
|
||||
# The tool lives in the sglang repo at examples/profiler/nsys_profile_tools/gputrc2graph.py
|
||||
# Set SGLANG_REPO to your local sglang repo root, e.g.:
|
||||
# export SGLANG_REPO=/workspace/sglang
|
||||
SGLANG_REPO="${SGLANG_REPO:-$(python3 -c "import sglang, os; print(os.path.abspath(os.path.join(os.path.dirname(sglang.__file__), '../../..')))" 2>/dev/null || echo "")}"
|
||||
GPUTRC="${SGLANG_REPO}/examples/profiler/nsys_profile_tools/gputrc2graph.py"
|
||||
check "gputrc2graph.py exists (set SGLANG_REPO if missing)" test -f "$GPUTRC"
|
||||
check "pandas" python3 -c "import pandas"
|
||||
check "regex" python3 -c "import regex"
|
||||
check "plotly (optional)" python3 -c "import plotly"
|
||||
|
||||
echo ""
|
||||
echo "Result: $PASS passed, $FAIL missing"
|
||||
if [ "$FAIL" -gt 0 ]; then
|
||||
echo ""
|
||||
echo "Install missing dependencies:"
|
||||
echo " pip install pandas regex plotly # for gputrc2graph.py"
|
||||
echo " # nsys: install NVIDIA Nsight Systems from https://developer.nvidia.com/nsight-systems"
|
||||
fi
|
||||
# Quick dependency check
|
||||
check() { "$@" &>/dev/null && echo "[OK] $1" || echo "[MISS] $1"; }
|
||||
check "sglang" python3 -c "import sglang"
|
||||
check "torch+CUDA" python3 -c "import torch; assert torch.cuda.is_available()"
|
||||
check "torch.profiler" python3 -c "import torch.profiler"
|
||||
check "nsys (Level 2)" which nsys
|
||||
check "pandas" python3 -c "import pandas"
|
||||
check "plotly" python3 -c "import plotly"
|
||||
```
|
||||
|
||||
**Minimum required for benchmarking**: `sglang`, `torch` with CUDA.
|
||||
**Additional for Level 1 profiling**: `torch.profiler` (bundled with torch, always available).
|
||||
**Additional for Level 2 profiling**: `nsys` on PATH, `gputrc2graph.py` present, `pandas`, `regex`.
|
||||
`plotly` is only needed to generate the HTML chart; `result.csv` is generated regardless.
|
||||
|
||||
---
|
||||
|
||||
### Download Required Input Images
|
||||
|
||||
Some models (image editing / image-guided video generation) require input images:
|
||||
**Minimum for benchmarking**: `sglang`, `torch` with CUDA.
|
||||
**Level 1 profiling**: `torch.profiler` (bundled with torch).
|
||||
**Level 2 profiling**: `nsys`, `pandas`, `plotly` + `gputrc2graph.py` from the sglang repo.
|
||||
|
||||
Download input images required by some models:
|
||||
```bash
|
||||
mkdir -p /workspace/gen_benchmark/figs
|
||||
cd /workspace/gen_benchmark/figs
|
||||
|
||||
wget -O cat.png https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png
|
||||
wget -O astronaut.jpg https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg
|
||||
wget -O /workspace/gen_benchmark/figs/cat.png \
|
||||
https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png
|
||||
wget -O /workspace/gen_benchmark/figs/astronaut.jpg \
|
||||
https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Standard Benchmark Model Suite
|
||||
## Benchmark Commands
|
||||
|
||||
All commands include `--warmup` (pre-warm torch.compile) and `--enable-torch-compile` to reflect real production deployment performance. Timing is measured after warmup completes.
|
||||
All commands include `--warmup` and `--enable-torch-compile` for real production performance. Add `--perf-dump-path <file>.json` for machine-readable output.
|
||||
|
||||
If you want a machine-readable perf dump for comparison/reporting, add `--perf-dump-path <file>.json` to any command below, then follow `diffusion-perf` (`../diffusion-perf/SKILL.md`) to interpret and compare the dumps.
|
||||
### Perf dump & before/after compare
|
||||
|
||||
### 1. Qwen/Qwen-Image-2512 (Text-to-Image, single GPU)
|
||||
For every benchmark run, always write a perf dump JSON:
|
||||
|
||||
**Task**: Text-to-Image
|
||||
**Resolution**: 1024×1024, 50 steps
|
||||
```bash
|
||||
sglang generate ... --warmup --perf-dump-path <result>.json
|
||||
```
|
||||
|
||||
Before/after comparison (outputs a Markdown table suitable for PR descriptions):
|
||||
|
||||
```bash
|
||||
# Baseline (on main branch or before changes)
|
||||
sglang generate ... --warmup --perf-dump-path baseline.json
|
||||
|
||||
# New (after changes)
|
||||
sglang generate ... --warmup --perf-dump-path new.json
|
||||
|
||||
python python/sglang/multimodal_gen/benchmarks/compare_perf.py baseline.json new.json
|
||||
```
|
||||
|
||||
### Qwen-Image-2512 (1024×1024, 50 steps)
|
||||
```bash
|
||||
sglang generate \
|
||||
--model-path=Qwen/Qwen-Image-2512 \
|
||||
--prompt="A futuristic cyberpunk city at night, neon lights reflecting on wet streets, highly detailed, 8k" \
|
||||
'--negative-prompt= ' \
|
||||
--width=1024 \
|
||||
--height=1024 \
|
||||
--num-inference-steps=50 \
|
||||
--guidance-scale=4.0 \
|
||||
--seed=42 \
|
||||
--save-output \
|
||||
--enable-torch-compile \
|
||||
--warmup \
|
||||
--dit-cpu-offload false \
|
||||
--text-encoder-cpu-offload false
|
||||
--width=1024 --height=1024 --num-inference-steps=50 --guidance-scale=4.0 \
|
||||
--seed=42 --save-output --enable-torch-compile --warmup \
|
||||
--dit-cpu-offload false --text-encoder-cpu-offload false
|
||||
```
|
||||
|
||||
**Key metrics**: denoise latency (s, primary), end-to-end latency (s/image), peak GPU memory (GB)
|
||||
|
||||
---
|
||||
|
||||
### 2. Qwen/Qwen-Image-Edit-2511 (Image Editing, single GPU)
|
||||
|
||||
**Task**: Text-guided Image Editing
|
||||
**Prerequisite**: `cat.png` (see Prerequisites)
|
||||
**Resolution**: 1024×1024, 50 steps
|
||||
|
||||
### Qwen-Image-Edit-2511 (image editing, 1024×1024, 50 steps)
|
||||
```bash
|
||||
sglang generate \
|
||||
--model-path=Qwen/Qwen-Image-Edit-2511 \
|
||||
'--prompt=Transform into anime style' \
|
||||
'--negative-prompt= ' \
|
||||
'--prompt=Transform into anime style' '--negative-prompt= ' \
|
||||
--image-path=/workspace/gen_benchmark/figs/cat.png \
|
||||
--width=1024 \
|
||||
--height=1024 \
|
||||
--num-inference-steps=50 \
|
||||
--guidance-scale=4.0 \
|
||||
--seed=42 \
|
||||
--save-output \
|
||||
--enable-torch-compile \
|
||||
--warmup \
|
||||
--dit-cpu-offload false \
|
||||
--text-encoder-cpu-offload false
|
||||
--width=1024 --height=1024 --num-inference-steps=50 --guidance-scale=4.0 \
|
||||
--seed=42 --save-output --enable-torch-compile --warmup \
|
||||
--dit-cpu-offload false --text-encoder-cpu-offload false
|
||||
```
|
||||
|
||||
**Key metrics**: denoise latency (s, primary), end-to-end latency (s/image), peak GPU memory (GB)
|
||||
|
||||
---
|
||||
|
||||
### 3. black-forest-labs/FLUX.1-dev (Text-to-Image, single GPU)
|
||||
|
||||
**Task**: Text-to-Image
|
||||
**Resolution**: 1024×1024, 50 steps
|
||||
|
||||
### FLUX.1-dev (1024×1024, 50 steps)
|
||||
```bash
|
||||
sglang generate \
|
||||
--model-path=black-forest-labs/FLUX.1-dev \
|
||||
--prompt="A futuristic cyberpunk city at night, neon lights reflecting on wet streets, highly detailed, 8k" \
|
||||
--width=1024 \
|
||||
--height=1024 \
|
||||
--num-inference-steps=50 \
|
||||
--guidance-scale=4.0 \
|
||||
--seed=42 \
|
||||
--save-output \
|
||||
--warmup \
|
||||
--enable-torch-compile
|
||||
--width=1024 --height=1024 --num-inference-steps=50 --guidance-scale=4.0 \
|
||||
--seed=42 --save-output --enable-torch-compile --warmup
|
||||
```
|
||||
|
||||
**Key metrics**: denoise latency (s, primary), end-to-end latency (s/image), peak GPU memory (GB)
|
||||
|
||||
---
|
||||
|
||||
### 4. black-forest-labs/FLUX.2-dev (Text-to-Image, single GPU)
|
||||
|
||||
**Task**: Text-to-Image
|
||||
**Resolution**: 1024×1024
|
||||
|
||||
### FLUX.2-dev (1024×1024)
|
||||
```bash
|
||||
sglang generate \
|
||||
--model-path black-forest-labs/FLUX.2-dev \
|
||||
--prompt "A Logo With Bold Large Text: SGL Diffusion" \
|
||||
--width=1024 \
|
||||
--height=1024 \
|
||||
--dit-layerwise-offload false \
|
||||
--enable-torch-compile \
|
||||
--warmup \
|
||||
--dit-cpu-offload false \
|
||||
--text-encoder-cpu-offload true \
|
||||
--vae-cpu-offload false
|
||||
--width=1024 --height=1024 \
|
||||
--dit-layerwise-offload false --enable-torch-compile --warmup \
|
||||
--dit-cpu-offload false --text-encoder-cpu-offload true --vae-cpu-offload false
|
||||
```
|
||||
|
||||
**Key metrics**: denoise latency (s, primary), end-to-end latency (s/image), peak GPU memory (GB)
|
||||
|
||||
---
|
||||
|
||||
### 5. Tongyi-MAI/Z-Image-Turbo (Turbo Text-to-Image, single GPU)
|
||||
|
||||
**Task**: Text-to-Image (few-step turbo mode, guidance=0)
|
||||
**Resolution**: 1024×1024, **9 steps**
|
||||
|
||||
### Z-Image-Turbo (1024×1024, 9 steps)
|
||||
```bash
|
||||
sglang generate \
|
||||
--model-path=Tongyi-MAI/Z-Image-Turbo \
|
||||
--log-level=info \
|
||||
--prompt='A fantasy landscape with mountains and a river, detailed, vibrant colors' \
|
||||
--width=1024 \
|
||||
--height=1024 \
|
||||
--num-inference-steps=9 \
|
||||
--guidance-scale=0.0 \
|
||||
--seed=42 \
|
||||
--save-output \
|
||||
--enable-torch-compile \
|
||||
--warmup \
|
||||
--dit-cpu-offload false \
|
||||
--text-encoder-cpu-offload false
|
||||
--width=1024 --height=1024 --num-inference-steps=9 --guidance-scale=0.0 \
|
||||
--seed=42 --save-output --enable-torch-compile --warmup \
|
||||
--dit-cpu-offload false --text-encoder-cpu-offload false
|
||||
```
|
||||
|
||||
**Key metrics**: denoise latency (s, primary), end-to-end latency (s/image), peak GPU memory (GB)
|
||||
|
||||
---
|
||||
|
||||
### 6. Wan-AI/Wan2.2-T2V-A14B-Diffusers 720P (Text-to-Video, 8 GPUs)
|
||||
|
||||
**Task**: Text-to-Video
|
||||
**Resolution**: 720P, 81 frames, 40 steps
|
||||
**Parallelism**: 8 GPUs, Ulysses degree=4, CFG parallel + layerwise offload
|
||||
|
||||
### Wan2.2-T2V-A14B 720P (8 GPUs, 81 frames, 40 steps)
|
||||
```bash
|
||||
sglang generate \
|
||||
--model-path=Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--log-level=info \
|
||||
--prompt="A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon. The kitchen is cozy, with sunlight streaming through the window." \
|
||||
--negative-prompt=" " \
|
||||
--720p \
|
||||
--num-inference-steps=40 \
|
||||
--num-frames=81 \
|
||||
--guidance-scale=5.0 \
|
||||
--seed=42 \
|
||||
--save-output \
|
||||
--num-gpus=8 \
|
||||
--enable-cfg-parallel \
|
||||
--ulysses-degree=4 \
|
||||
--dit-layerwise-offload true \
|
||||
--dit-cpu-offload false \
|
||||
--vae-cpu-offload false \
|
||||
--text-encoder-cpu-offload true \
|
||||
--warmup \
|
||||
--enable-torch-compile
|
||||
--prompt="A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon." \
|
||||
--negative-prompt=" " --720p --num-inference-steps=40 --num-frames=81 \
|
||||
--guidance-scale=5.0 --seed=42 --save-output \
|
||||
--num-gpus=8 --enable-cfg-parallel --ulysses-degree=4 \
|
||||
--dit-layerwise-offload true --dit-cpu-offload false \
|
||||
--vae-cpu-offload false --text-encoder-cpu-offload true \
|
||||
--warmup --enable-torch-compile
|
||||
```
|
||||
|
||||
**Key metrics**: denoise latency (s, primary), total end-to-end latency (s/video), peak GPU memory per device (GB)
|
||||
|
||||
---
|
||||
|
||||
### 7. Wan-AI/Wan2.2-TI2V-5B-Diffusers 720P (Text-Image-to-Video, single GPU)
|
||||
|
||||
**Task**: Text-Image-to-Video
|
||||
**Prerequisite**: `astronaut.jpg` (see Prerequisites)
|
||||
**Resolution**: 720P, 81 frames, 50 steps
|
||||
|
||||
### Wan2.2-TI2V-5B 720P (single GPU, 81 frames, 50 steps)
|
||||
```bash
|
||||
sglang generate \
|
||||
--model-path Wan-AI/Wan2.2-TI2V-5B-Diffusers \
|
||||
--log-level info \
|
||||
--warmup \
|
||||
--dit-layerwise-offload false \
|
||||
--dit-cpu-offload false \
|
||||
--vae-cpu-offload false \
|
||||
--text-encoder-cpu-offload false \
|
||||
--enable-torch-compile \
|
||||
--prompt "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot." \
|
||||
--negative-prompt "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" \
|
||||
--prompt "An astronaut hatching from an egg, on the surface of the moon..." \
|
||||
--negative-prompt "Bright tones, overexposed, static, blurred details..." \
|
||||
--image-path=/workspace/gen_benchmark/figs/astronaut.jpg \
|
||||
--num-frames 81 \
|
||||
--720p \
|
||||
--num-inference-steps 50 \
|
||||
--guidance-scale 5.0 \
|
||||
--seed 42 \
|
||||
--save-output
|
||||
--num-frames 81 --720p --num-inference-steps 50 --guidance-scale 5.0 \
|
||||
--seed 42 --save-output \
|
||||
--dit-layerwise-offload false --dit-cpu-offload false \
|
||||
--vae-cpu-offload false --text-encoder-cpu-offload false \
|
||||
--enable-torch-compile --warmup
|
||||
```
|
||||
|
||||
**Key metrics**: denoise latency (s, primary), total end-to-end latency (s/video), peak GPU memory (GB)
|
||||
**Key metrics** (all models): denoise latency ★, end-to-end latency, peak GPU memory.
|
||||
|
||||
---
|
||||
|
||||
## Result Recording Format
|
||||
## Performance Bottleneck Workflow
|
||||
|
||||
For recording, extracting, and comparing benchmark results across versions, reuse `diffusion-perf` (`../diffusion-perf/SKILL.md`):
|
||||
### Step 1: Identify the Slow DiT Operation
|
||||
|
||||
- Use `--perf-dump-path` to generate a JSON perf dump
|
||||
- Use `python python/sglang/multimodal_gen/benchmarks/compare_perf.py baseline.json new.json` to generate a Markdown table for PR reports
|
||||
Add `--log-level=info` and observe:
|
||||
- **Denoise loop latency** ★ — primary target
|
||||
- Per-step DiT latency — denoise ÷ steps
|
||||
|
||||
---
|
||||
|
||||
## Performance Bottleneck Investigation Workflow
|
||||
|
||||
### Step 1: Identify the Slow Stage
|
||||
|
||||
Add `--log-level=info` (or `debug`) and observe the following stages in order of priority:
|
||||
|
||||
- **Denoise loop latency** ★ — total time across all DiT forward passes; this is the primary optimization target and usually accounts for >80% of end-to-end latency
|
||||
- **Per-step DiT forward latency** — denoise latency divided by number of steps; useful for pinpointing per-step overhead
|
||||
- VAE decode latency — runs once after the denoising loop; significant for video models
|
||||
- Text encoder encoding latency — runs once before the denoising loop
|
||||
- Warmup / torch.compile compilation time — excluded from all reported latency numbers
|
||||
|
||||
Focus optimization effort on denoise latency first. Improvements to VAE or text encoder only matter after denoise is already well-optimized.
|
||||
|
||||
### Step 2: Cross-reference Kernel Optimization Skill
|
||||
|
||||
After identifying the slow stage, refer to `use-efficient-diffusion-kernels.md`:
|
||||
- Slow attention → check whether attention backend is FlashAttention (FA3/FA4)
|
||||
- Slow AdaLN modulation → verify `LayerNormScaleShift` / `fuse_scale_shift_kernel` is active
|
||||
- Slow RMSNorm → verify `sgl_kernel.rmsnorm` / `fused_add_rmsnorm` is hit
|
||||
- Slow RoPE → check FlashInfer inplace RoPE or Triton RoPE fallback
|
||||
- Slow QK Norm → verify `fused_inplace_qknorm` path; confirm `head_dim` is in the supported list (`64, 128, 256, 512, 1024`)
|
||||
|
||||
If no existing fused kernel covers the slow operation (e.g., a new elementwise fusion opportunity, a norm variant, or a custom DiT sub-op), implement a new Triton kernel using **`add-triton-kernel.md`**. That skill covers the full workflow: kernel authoring, autotune, `torch.compile` compatibility, NPU fallback, layer integration, and tests.
|
||||
|
||||
### Step 3: Check torch.compile Coverage
|
||||
### Step 2: Profile with torch.profiler (Level 1)
|
||||
|
||||
```bash
|
||||
TORCH_COMPILE_DEBUG=1 sglang generate ...
|
||||
```
|
||||
|
||||
Key things to watch:
|
||||
- Dynamic shape changes trigger recompilation — fix resolution and frame count when benchmarking
|
||||
- Conditional branches containing `tensor.item()` cause graph breaks and must be rewritten
|
||||
|
||||
### Step 4: Multi-GPU Parallel Efficiency (Wan2.2-T2V-A14B)
|
||||
|
||||
- Verify `--ulysses-degree` evenly divides `--num-gpus`
|
||||
- Confirm `--enable-cfg-parallel` is active (requires `guidance_scale > 1`)
|
||||
- Use `torch.distributed` profiling or `nsys` to identify communication bottlenecks
|
||||
- `--dit-layerwise-offload true` introduces CPU↔GPU transfer overhead; only enable when memory-constrained
|
||||
|
||||
### Step 5: Offload Strategy Trade-offs
|
||||
|
||||
| Offload Flag | Effect | Cost |
|
||||
|--------------------------------|--------------------------------|---------------------------------|
|
||||
| `--dit-cpu-offload true` | Move all DiT weights to CPU | Significantly increases per-step latency |
|
||||
| `--dit-layerwise-offload true` | Load DiT layer-by-layer on demand | Moderate latency, saves GPU memory |
|
||||
| `--text-encoder-cpu-offload` | Move text encoder to CPU | Only affects encoding phase (once per run) |
|
||||
| `--vae-cpu-offload` | Move VAE to CPU | Affects decode phase only |
|
||||
|
||||
When establishing a GPU performance baseline, disable all offloading (`false`).
|
||||
|
||||
---
|
||||
|
||||
## Per-Layer Kernel Profiling
|
||||
|
||||
When you have any **model / operator optimization** requirement (e.g., a new fused kernel, a torch.compile graph-break fix, a suspected regression, or a perf goal you need to justify/validate), you should proactively profile to understand **which layer and which kernel** inside the DiT is responsible. Two profiling levels are used, from coarse to fine:
|
||||
|
||||
```
|
||||
Level 1 — torch.profiler (built-in --profile flag) → per-PyTorch-op breakdown, per DiT layer
|
||||
Level 2 — nsys + gputrc2graph.py → CUDA kernel category breakdown, CPU vs GPU time
|
||||
```
|
||||
|
||||
Use Level 1 to find the slow DiT layer and sub-component; use Level 2 to get the full CUDA kernel category distribution and confirm where GPU time is spent.
|
||||
|
||||
---
|
||||
|
||||
### Level 1: torch.profiler — Built-in `--profile` Flag
|
||||
|
||||
SGLang Diffusion has a built-in profiler (`SGLDiffusionProfiler` in `python/sglang/multimodal_gen/runtime/utils/profiler.py`) that is activated via CLI flags. No driver script needed.
|
||||
|
||||
**Key flags:**
|
||||
|
||||
| Flag | Default | Description |
|
||||
|------|---------|-------------|
|
||||
| `--profile` | off | Enable torch profiler for the denoising stage |
|
||||
| `--num-profiled-timesteps N` | 5 | Profile N denoising steps (use `-1` for all steps) |
|
||||
| `--profile-all-stages` | off | Also profile text encoding and VAE decode stages |
|
||||
|
||||
**Output:** gzipped Chrome trace JSON at `$SGLANG_TORCH_PROFILER_DIR/{uuid4}-{mode}-global-rank0.trace.json.gz` (default dir: `./logs`).
|
||||
The `{uuid4}` part is a randomly generated request ID. Use `ls -t` or `ls -1` to find the latest trace file.
|
||||
|
||||
**Example — profile 3 denoising steps of FLUX.1-dev:**
|
||||
|
||||
```bash
|
||||
SGLANG_TORCH_PROFILER_DIR=/workspace/gen_benchmark/profiles \
|
||||
sglang generate \
|
||||
--model-path=black-forest-labs/FLUX.1-dev \
|
||||
--prompt="A futuristic cyberpunk city at night, neon lights reflecting on wet streets, highly detailed, 8k" \
|
||||
--width=1024 --height=1024 \
|
||||
--num-inference-steps=50 \
|
||||
--guidance-scale=4.0 \
|
||||
--seed=42 \
|
||||
--enable-torch-compile \
|
||||
--warmup \
|
||||
--profile \
|
||||
--num-profiled-timesteps 3
|
||||
```
|
||||
|
||||
**Example — profile all pipeline stages (text encoder + denoise + VAE):**
|
||||
|
||||
```bash
|
||||
SGLANG_TORCH_PROFILER_DIR=/workspace/gen_benchmark/profiles \
|
||||
SGLANG_TORCH_PROFILER_DIR=/workspace/profiles \
|
||||
sglang generate \
|
||||
--model-path=black-forest-labs/FLUX.1-dev \
|
||||
--prompt="A futuristic cyberpunk city at night" \
|
||||
--width=1024 --height=1024 \
|
||||
--num-inference-steps=5 \
|
||||
--seed=42 \
|
||||
--enable-torch-compile \
|
||||
--warmup \
|
||||
--profile \
|
||||
--profile-all-stages
|
||||
--width=1024 --height=1024 --num-inference-steps=50 \
|
||||
--seed=42 --enable-torch-compile --warmup \
|
||||
--profile --num-profiled-timesteps 3
|
||||
```
|
||||
|
||||
**Reading the trace without a UI:**
|
||||
|
||||
The output is a `.trace.json.gz` file (Chrome trace format). Parse it with Python to get a ranked op table — no browser needed:
|
||||
|
||||
Parse the trace without a browser:
|
||||
```python
|
||||
import gzip, json, collections, glob, os
|
||||
|
||||
# The request_id is a uuid4, so the filename is e.g.:
|
||||
# 550e8400-e29b-41d4-a716-446655440000-3_steps-global-rank0.trace.json.gz
|
||||
# Find the latest trace file automatically:
|
||||
log_dir = os.environ.get("SGLANG_TORCH_PROFILER_DIR", "./logs")
|
||||
traces = sorted(glob.glob(f"{log_dir}/*.trace.json.gz"), key=os.path.getmtime, reverse=True)
|
||||
assert traces, f"No trace files found in {log_dir}"
|
||||
trace_path = traces[0]
|
||||
print(f"Reading: {trace_path}")
|
||||
trace_path = sorted(glob.glob(f"{log_dir}/*.trace.json.gz"), key=os.path.getmtime, reverse=True)[0]
|
||||
with gzip.open(trace_path, "rb") as f:
|
||||
data = json.loads(f.read())
|
||||
|
||||
events = data.get("traceEvents", [])
|
||||
|
||||
# Collect CUDA kernel durations per op name
|
||||
cuda_ops = collections.defaultdict(lambda: {"total_us": 0, "count": 0})
|
||||
for e in events:
|
||||
for e in data.get("traceEvents", []):
|
||||
if e.get("cat") in ("kernel", "gpu_memcpy") and "dur" in e:
|
||||
name = e.get("name", "unknown")
|
||||
cuda_ops[name]["total_us"] += e["dur"]
|
||||
cuda_ops[name]["count"] += 1
|
||||
cuda_ops[e.get("name","unknown")]["total_us"] += e["dur"]
|
||||
cuda_ops[e.get("name","unknown")]["count"] += 1
|
||||
|
||||
# Print top 40 by total CUDA time
|
||||
sorted_ops = sorted(cuda_ops.items(), key=lambda x: x[1]["total_us"], reverse=True)
|
||||
print(f"{'Kernel Name':<80} {'Total (ms)':>12} {'Count':>8}")
|
||||
print("-" * 102)
|
||||
for name, stats in sorted_ops[:40]:
|
||||
print(f"{name:<80} {stats['total_us']/1000:>12.3f} {stats['count']:>8}")
|
||||
print(f"{'Kernel':<80} {'Total(ms)':>10} {'Count':>6}")
|
||||
for name, s in sorted(cuda_ops.items(), key=lambda x: -x[1]["total_us"])[:30]:
|
||||
print(f"{name:<80} {s['total_us']/1000:>10.3f} {s['count']:>6}")
|
||||
```
|
||||
|
||||
**Add `record_function` scopes for per-layer attribution:**
|
||||
|
||||
The built-in profiler captures all PyTorch ops, but adding named scopes makes it easy to attribute costs to specific DiT layers. Add these directly in the DiT block `forward()`:
|
||||
|
||||
Add `record_function` scopes in the DiT block for per-layer attribution:
|
||||
```python
|
||||
import torch
|
||||
|
||||
# Inside DiT transformer block forward() — e.g., FluxTransformerBlock, QwenDiTBlock
|
||||
with torch.profiler.record_function(f"dit_block_{block_idx}.norm"):
|
||||
x = self.norm1(x, scale, shift)
|
||||
with torch.profiler.record_function(f"dit_block_{block_idx}.attn"):
|
||||
with torch.profiler.record_function(f"dit_block_{idx}.attn"):
|
||||
x = self.attn(x)
|
||||
with torch.profiler.record_function(f"dit_block_{block_idx}.mlp"):
|
||||
x = self.mlp(x)
|
||||
with torch.profiler.record_function(f"dit_block_{idx}.norm"):
|
||||
x = self.norm(x)
|
||||
```
|
||||
|
||||
These scopes appear as named spans in the trace and in `key_averages()` output.
|
||||
**Expected dominant kernels per DiT sub-component:**
|
||||
|
||||
**Key ops to watch per DiT sub-component:**
|
||||
| Sub-component | Expected kernel |
|
||||
|--------------|-----------------|
|
||||
| QKV / output / MLP projections | `cutlass_gemm` / `ampere_*_gemm` |
|
||||
| Attention | `flash_attn_fwd` / `fmha_*` (FA3/FA4) |
|
||||
| AdaLN modulation | `fuse_scale_shift_kernel` |
|
||||
| RMSNorm / LayerNorm | `sgl_kernel_rmsnorm` / Triton norm |
|
||||
| SiLU gate | `vectorized_elementwise_kernel` |
|
||||
| RoPE | `apply_rotary_embedding` (Triton) |
|
||||
| QK Norm | `fused_inplace_qknorm` (JIT) |
|
||||
|
||||
| Sub-component | Expected dominant kernel |
|
||||
|----------------------|-------------------------------------------------------------|
|
||||
| QKV projection | `cutlass_gemm` / `ampere_*_gemm` |
|
||||
| Attention compute | `flash_attn_fwd` / `fmha_*` (FA3/FA4) |
|
||||
| Output projection | `cutlass_gemm` |
|
||||
| AdaLN modulation | `fuse_scale_shift_kernel` / `fuse_scale_shift_gate_*` |
|
||||
| RMSNorm / LayerNorm | `sgl_kernel_rmsnorm` / `fused_add_rmsnorm` / Triton norm |
|
||||
| MLP fc1 / fc2 | `cutlass_gemm` |
|
||||
| SiLU gate | `vectorized_elementwise_kernel` / `silu_and_mul` |
|
||||
| RoPE | `apply_rotary_embedding` (Triton) / FlashInfer inplace |
|
||||
| QK Norm | `fused_inplace_qknorm` (JIT) or fallback `rmsnorm` |
|
||||
|
||||
---
|
||||
|
||||
### Level 2: nsys + gputrc2graph.py — CUDA Kernel Category Breakdown
|
||||
|
||||
SGLang ships a CLI analysis tool at `examples/profiler/nsys_profile_tools/gputrc2graph.py` ([PR #9314](https://github.com/sgl-project/sglang/pull/9314)). It processes `.nsys-rep` files and outputs:
|
||||
- `result.csv` — kernel-to-category mapping with elapsed time per category (read directly on server)
|
||||
- `result.html` — stacked bar chart (for environments with a browser)
|
||||
|
||||
**Step 1: Install dependencies**
|
||||
### Step 3: Deep CUDA Kernel Breakdown (Level 2 — nsys)
|
||||
|
||||
```bash
|
||||
pip install pandas regex plotly
|
||||
```
|
||||
|
||||
**Step 2: Collect nsys trace**
|
||||
|
||||
Run two passes: one with profiling (to get the trace) and one without (to measure the true runtime for CPU time calculation):
|
||||
|
||||
```bash
|
||||
# Pass A — with profiling (skip warmup using --delay)
|
||||
# Estimate DELAY as: warmup time + a few seconds buffer
|
||||
nsys profile \
|
||||
-t cuda \
|
||||
-o /workspace/gen_benchmark/profiles/flux_dev \
|
||||
-f true \
|
||||
--trace-fork-before-exec=true \
|
||||
--delay 120 \
|
||||
--duration 60 \
|
||||
# Pass A — collect nsys trace (skip warmup with --delay)
|
||||
nsys profile -t cuda -o /workspace/profiles/flux_dev -f true \
|
||||
--trace-fork-before-exec=true --delay 120 --duration 60 \
|
||||
sglang generate \
|
||||
--model-path=black-forest-labs/FLUX.1-dev \
|
||||
--prompt="A futuristic cyberpunk city at night, neon lights reflecting on wet streets, highly detailed, 8k" \
|
||||
--width=1024 --height=1024 \
|
||||
--num-inference-steps=50 \
|
||||
--guidance-scale=4.0 \
|
||||
--seed=42 \
|
||||
--enable-torch-compile \
|
||||
--warmup
|
||||
--prompt="A futuristic cyberpunk city at night" \
|
||||
--width=1024 --height=1024 --num-inference-steps=50 \
|
||||
--seed=42 --enable-torch-compile --warmup
|
||||
|
||||
# Pass B — without profiling, record total wall-clock time
|
||||
time sglang generate \
|
||||
--model-path=black-forest-labs/FLUX.1-dev \
|
||||
--prompt="A futuristic cyberpunk city at night, neon lights reflecting on wet streets, highly detailed, 8k" \
|
||||
--width=1024 --height=1024 --num-inference-steps=50 \
|
||||
--guidance-scale=4.0 --seed=42 \
|
||||
# Pass B — measure wall-clock time without profiling
|
||||
time sglang generate --model-path=black-forest-labs/FLUX.1-dev \
|
||||
--width=1024 --height=1024 --num-inference-steps=50 --seed=42 \
|
||||
--enable-torch-compile --warmup
|
||||
# Record the elapsed seconds as ELAPSED_SEC
|
||||
# Record ELAPSED_SEC from Pass B
|
||||
```
|
||||
|
||||
**Step 3: Add a diffusion kernel classification JSON**
|
||||
|
||||
`gputrc2graph.py` loads all `.json` files in its directory. Create a diffusion-specific one at `examples/profiler/nsys_profile_tools/sglang_diffusion_engine_model.json`:
|
||||
|
||||
Create classification JSON at `examples/profiler/nsys_profile_tools/sglang_diffusion_engine_model.json`:
|
||||
```json
|
||||
{
|
||||
"sglang": {
|
||||
@@ -561,29 +246,18 @@ time sglang generate \
|
||||
}
|
||||
```
|
||||
|
||||
**Step 4: Run analysis (CLI-only, no UI)**
|
||||
|
||||
Run analysis:
|
||||
```bash
|
||||
cd examples/profiler/nsys_profile_tools
|
||||
|
||||
python3 gputrc2graph.py \
|
||||
--in_file /workspace/gen_benchmark/profiles/flux_dev.nsys-rep,sglang,diffusion,ELAPSED_SEC \
|
||||
--out_dir /workspace/gen_benchmark/profiles/analysis \
|
||||
--in_file /workspace/profiles/flux_dev.nsys-rep,sglang,diffusion,ELAPSED_SEC \
|
||||
--out_dir /workspace/profiles/analysis \
|
||||
--title "FLUX.1-dev denoise kernel breakdown"
|
||||
```
|
||||
|
||||
Replace `ELAPSED_SEC` with the wall-clock seconds measured in Pass B (e.g. `132` if `time sglang generate ...` reported `real 2m12s`). The tool uses this value to compute CPU (non-GPU) idle time = `ELAPSED_SEC - total_GPU_sec`. Passing `0` is also valid but will use the nsys-measured elapsed time, which may inflate non-GPU time if profiling overhead is significant.
|
||||
|
||||
**Step 5: Read results without a UI**
|
||||
|
||||
```bash
|
||||
# Read the CSV directly — sorted by elapsed time, largest first
|
||||
cat /workspace/gen_benchmark/profiles/analysis/result.csv | column -t -s,
|
||||
|
||||
# Or with Python for a clean summary
|
||||
# Read results
|
||||
python3 - << 'EOF'
|
||||
import pandas as pd
|
||||
df = pd.read_csv("/workspace/gen_benchmark/profiles/analysis/result.csv")
|
||||
df = pd.read_csv("/workspace/profiles/analysis/result.csv")
|
||||
summary = df.groupby("Category")["Elapsed Time (sec)"].sum().sort_values(ascending=False)
|
||||
total = summary.sum()
|
||||
for cat, sec in summary.items():
|
||||
@@ -591,94 +265,214 @@ for cat, sec in summary.items():
|
||||
EOF
|
||||
```
|
||||
|
||||
**What to look for in the category breakdown:**
|
||||
**What the category breakdown tells you:**
|
||||
|
||||
| Category shows high time | Investigation |
|
||||
|--------------------------|---------------|
|
||||
| `gemm` dominant | Check QKV / output / MLP projections; verify tensor parallelism is active |
|
||||
| `attn` dominant | Verify FA3/FA4 is active; check sequence length and head_dim |
|
||||
| `adaln_modulation` unexpectedly high | Verify fused `fuse_scale_shift_kernel` path is used |
|
||||
| `norm` high | Verify `sgl_kernel_rmsnorm` / CuTe DSL fused path; check D alignment |
|
||||
| `nccl_comm` high | Multi-GPU: check Ulysses degree; consider reducing TP degree |
|
||||
| `triton_kernel` high | Identify which Triton kernel; check if a CuTe DSL or sgl-kernel replacement exists |
|
||||
| `non-gpu-H_D_memops` high | CPU↔GPU copy detected; check for accidental offload or `.cpu()` calls mid-denoising |
|
||||
| `CPU(non-GPU)` high | Python dispatch overhead; check for graph breaks in torch.compile |
|
||||
| Category high | Investigation |
|
||||
|--------------|---------------|
|
||||
| `gemm` dominant | Check tensor parallelism; QKV/MLP bottleneck |
|
||||
| `attn` dominant | Verify FA3/FA4 is active |
|
||||
| `adaln_modulation` high | Verify fused `fuse_scale_shift_kernel` is used |
|
||||
| `norm` high | Verify `sgl_kernel_rmsnorm` / CuTe DSL path; check D alignment |
|
||||
| `nccl_comm` high | Multi-GPU: tune Ulysses degree |
|
||||
| `triton_kernel` high | Identify which Triton kernel; consider CUDA replacement |
|
||||
| `non-gpu-H_D_memops` high | Accidental CPU offload or `.cpu()` calls mid-denoising |
|
||||
| `CPU(non-GPU)` high | Python dispatch overhead / torch.compile graph breaks |
|
||||
|
||||
**Comparing two versions:**
|
||||
### Step 3.5: Per-Kernel Deep Analysis (Level 3 — ncu)
|
||||
|
||||
The 4th field in `--in_file` is `elapsed_nonprofiled_sec` — the **total wall-clock time (seconds) of the same run without profiling** (measured via `time sglang generate ...` in Pass B above). The tool uses this to calculate how much time was spent on CPU (non-GPU) work. Replace `120` and `118` with the actual measured seconds for each run:
|
||||
**CRITICAL**: `ncu` (Nsight Compute) is the essential tool for kernel-level optimization. While nsys and torch.profiler tell you **which** kernels are slow, only ncu tells you **why** — memory bandwidth utilization, compute throughput, occupancy limiters, warp stall reasons, and roofline position. **Always use ncu when optimizing or writing custom kernels.**
|
||||
|
||||
#### When to use ncu
|
||||
|
||||
- After writing a new Triton or CUDA kernel — verify it saturates hardware bandwidth
|
||||
- When a kernel shows up as a top bottleneck in Level 1/2 profiling
|
||||
- When comparing your fused kernel vs PyTorch baseline or torch.compile output
|
||||
- When tuning Triton autotune configs (block sizes, num_warps)
|
||||
|
||||
#### Basic ncu workflow
|
||||
|
||||
```bash
|
||||
python3 gputrc2graph.py \
|
||||
--in_file \
|
||||
before.nsys-rep,sglang,diffusion,120 \ # replace 120 with Pass B elapsed seconds for "before"
|
||||
after.nsys-rep,sglang,diffusion,118 \ # replace 118 with Pass B elapsed seconds for "after"
|
||||
--out_dir /workspace/gen_benchmark/profiles/compare \
|
||||
--title "FLUX.1-dev before vs after optimization"
|
||||
# 1. Profile a specific kernel by name (skip warmup launches, collect 3 invocations)
|
||||
ncu --kernel-name "_fused_gated_residual_add_kernel" \
|
||||
--launch-skip 10 --launch-count 3 \
|
||||
--set full \
|
||||
-o /workspace/ncu_reports/gated_residual \
|
||||
sglang generate \
|
||||
--model-path=black-forest-labs/FLUX.1-dev \
|
||||
--prompt="test" --width=1024 --height=1024 \
|
||||
--num-inference-steps=5 --seed=42
|
||||
|
||||
# Then read the comparison CSV
|
||||
python3 - << 'EOF'
|
||||
import pandas as pd
|
||||
df = pd.read_csv("/workspace/gen_benchmark/profiles/compare/result.csv")
|
||||
pivot = df.pivot_table(values="Elapsed Time (sec)", index="Category",
|
||||
columns="Model_Engine", aggfunc="sum").round(3)
|
||||
pivot["delta_s"] = pivot.iloc[:, 1] - pivot.iloc[:, 0]
|
||||
pivot["delta_%"] = (pivot["delta_s"] / pivot.iloc[:, 0] * 100).round(1)
|
||||
print(pivot.sort_values("delta_s"))
|
||||
EOF
|
||||
# 2. Profile all kernels in a short run (use few steps to limit time)
|
||||
ncu --launch-skip 50 --launch-count 200 \
|
||||
--set full \
|
||||
-o /workspace/ncu_reports/all_kernels \
|
||||
sglang generate \
|
||||
--model-path=black-forest-labs/FLUX.1-dev \
|
||||
--prompt="test" --width=1024 --height=1024 \
|
||||
--num-inference-steps=3 --seed=42
|
||||
|
||||
# 3. For CUDA graph mode, use --graph-profiling=node to profile inside the graph
|
||||
ncu --graph-profiling node \
|
||||
--kernel-name "_fused_gated_residual_add_kernel" \
|
||||
--launch-skip 5 --launch-count 3 \
|
||||
--set full \
|
||||
-o /workspace/ncu_reports/gated_residual_cudagraph \
|
||||
sglang generate \
|
||||
--model-path=black-forest-labs/FLUX.1-dev \
|
||||
--prompt="test" --width=1024 --height=1024 \
|
||||
--num-inference-steps=5 --seed=42 \
|
||||
--enable-piecewise-cuda-graph
|
||||
```
|
||||
|
||||
#### Reading ncu results (CLI, no GUI needed)
|
||||
|
||||
```bash
|
||||
# Summary of all profiled kernels
|
||||
ncu --import /workspace/ncu_reports/gated_residual.ncu-rep --page raw --csv 2>/dev/null | head -50
|
||||
|
||||
# Key metrics to extract:
|
||||
ncu --import /workspace/ncu_reports/gated_residual.ncu-rep \
|
||||
--page details --csv 2>/dev/null | python3 -c "
|
||||
import csv, sys
|
||||
reader = csv.DictReader(sys.stdin)
|
||||
key_metrics = [
|
||||
'gpu__time_duration.avg', # Kernel duration
|
||||
'sm__throughput.avg.pct_of_peak_sustained_elapsed', # SM utilization
|
||||
'dram__throughput.avg.pct_of_peak_sustained_elapsed', # DRAM bandwidth util
|
||||
'l1tex__throughput.avg.pct_of_peak_sustained_elapsed', # L1 throughput
|
||||
'sm__warps_active.avg.pct_of_peak_sustained_active', # Achieved occupancy
|
||||
'launch__occupancy_limit_registers', # Occupancy limiter
|
||||
'launch__occupancy_limit_shared_mem',
|
||||
]
|
||||
for row in reader:
|
||||
name = row.get('Metric Name', '')
|
||||
if any(m in name for m in key_metrics):
|
||||
print(f'{name:<60} {row.get(\"Metric Value\",\"\")}')
|
||||
"
|
||||
```
|
||||
|
||||
#### Interpreting ncu results for kernel optimization
|
||||
|
||||
| Metric | Good | Action if bad |
|
||||
|--------|------|--------------|
|
||||
| DRAM throughput > 80% peak | Memory-bound, near optimal | Already saturating HBM — fuse with adjacent ops to reduce total memory traffic |
|
||||
| DRAM throughput < 50% peak | Not saturating memory bandwidth | Check coalescing, increase vector width, tune BLOCK sizes |
|
||||
| SM throughput > 60% peak | Compute-bound, near optimal | Reduce arithmetic, use faster instructions (e.g., FMA) |
|
||||
| SM throughput < 30% peak | Underutilized compute | Increase occupancy, reduce warp stalls, check instruction mix |
|
||||
| Achieved occupancy > 50% | Acceptable for most kernels | — |
|
||||
| Achieved occupancy < 25% | Too few active warps | Reduce register pressure or shared memory; increase block size |
|
||||
|
||||
#### Comparing before/after with ncu
|
||||
|
||||
```bash
|
||||
# Profile baseline kernel
|
||||
ncu --kernel-name "vectorized_elementwise_kernel" \
|
||||
--launch-skip 10 --launch-count 3 --set full \
|
||||
-o /workspace/ncu_reports/baseline ./program
|
||||
|
||||
# Profile optimized kernel
|
||||
ncu --kernel-name "_fused_gated_residual_add_kernel" \
|
||||
--launch-skip 10 --launch-count 3 --set full \
|
||||
-o /workspace/ncu_reports/optimized ./program
|
||||
|
||||
# Compare key metrics
|
||||
for report in baseline optimized; do
|
||||
echo "=== $report ==="
|
||||
ncu --import /workspace/ncu_reports/${report}.ncu-rep \
|
||||
--page details --csv 2>/dev/null | grep -E "time_duration|throughput.*pct|occupancy"
|
||||
done
|
||||
```
|
||||
|
||||
**Decision rule after ncu analysis:**
|
||||
- Kernel already at >80% DRAM bandwidth → fuse with neighbors to reduce total traffic
|
||||
- Kernel at <50% DRAM bandwidth → tune block sizes, fix coalescing, increase vectorization
|
||||
- Kernel compute-bound (SM util high, DRAM low) → reduce FLOPs or switch to a faster algorithm
|
||||
- Low occupancy → reduce registers (simplify kernel) or increase block size in autotune configs
|
||||
|
||||
### Step 4: Apply Kernel Optimization
|
||||
|
||||
After pinpointing the slow op, choose the right tool:
|
||||
|
||||
| Scenario | Skill to use |
|
||||
|----------|-------------|
|
||||
| New fused elementwise, norm variant, RoPE variant | **`add-triton-kernel.md`** — Triton JIT, faster iteration, NPU fallback |
|
||||
| Bandwidth-bound reduction (RMSNorm) needing max vectorization | **`add-cuda-kernel.md`** — CUDA JIT with `AlignedVector`, warp reductions |
|
||||
| Attention or tile-based op needing shared memory tuning | **`add-cuda-kernel.md`** — full control over CUDA primitives |
|
||||
| Slow op already covered by existing fused kernel | **`use-efficient-diffusion-kernels.md`** — check constraints & enable |
|
||||
|
||||
**Quick decision rule**: start with Triton. Switch to CUDA JIT only when profiling shows Triton can't saturate hardware bandwidth.
|
||||
|
||||
Both kernel types use SGLang's JIT compilation:
|
||||
- **Triton**: `python/sglang/jit_kernel/diffusion/triton/<op>.py`
|
||||
- **CUDA JIT**: `python/sglang/jit_kernel/csrc/diffusion/<op>.cuh` + wrapper `python/sglang/jit_kernel/diffusion/<op>.py`
|
||||
|
||||
### Step 5: torch.compile Coverage
|
||||
|
||||
```bash
|
||||
TORCH_COMPILE_DEBUG=1 sglang generate ...
|
||||
```
|
||||
- Dynamic shape changes trigger recompilation → fix resolution and frame count when benchmarking
|
||||
- `tensor.item()` in conditional branches causes graph breaks → rewrite as tensor ops
|
||||
|
||||
### Step 6: Multi-GPU Efficiency (Wan2.2-T2V-A14B)
|
||||
|
||||
- Verify `--ulysses-degree` evenly divides `--num-gpus`
|
||||
- Confirm `--enable-cfg-parallel` is active (requires `guidance_scale > 1`)
|
||||
- `--dit-layerwise-offload true` introduces CPU↔GPU transfer overhead; disable when memory permits
|
||||
|
||||
---
|
||||
|
||||
### Profiling Workflow Summary
|
||||
## Optimization Workflow Summary
|
||||
|
||||
```
|
||||
0. CORRECTNESS BASELINE
|
||||
0. BASELINE
|
||||
sglang generate --seed=42 --save-output → save reference images/videos
|
||||
(do this before any optimization; compare against it after every change)
|
||||
↓
|
||||
1. Run benchmark → establish baseline / measure impact
|
||||
1. BENCHMARK
|
||||
Run benchmark commands above → record denoise latency baseline
|
||||
↓
|
||||
2. sglang generate --profile --num-profiled-timesteps 3
|
||||
→ parse .trace.json.gz → rank ops by self_cuda_time_total
|
||||
→ identify slow DiT layer / sub-component (norm / attn / mlp / rope)
|
||||
2. LEVEL 1 PROFILE (torch.profiler)
|
||||
--profile --num-profiled-timesteps 3
|
||||
→ parse .trace.json.gz → rank ops by CUDA time
|
||||
→ identify slow DiT layer (norm / attn / mlp / rope / adaln)
|
||||
↓
|
||||
3. nsys profile + gputrc2graph.py → result.csv
|
||||
→ read category breakdown (gemm / attn / adaln / norm / nccl / cpu-overhead)
|
||||
3. LEVEL 2 PROFILE (nsys + gputrc2graph.py)
|
||||
→ result.csv category breakdown (gemm / attn / adaln / norm / triton / cpu)
|
||||
→ confirm where GPU time is concentrated
|
||||
↓
|
||||
4. Cross-reference use-efficient-diffusion-kernels.md
|
||||
→ apply existing fused kernel if available
|
||||
→ if no existing kernel covers the case, use add-triton-kernel.md
|
||||
to implement and integrate a new Triton kernel
|
||||
4. LEVEL 3 PROFILE (ncu — per-kernel deep analysis) ★ CRITICAL
|
||||
→ ncu --set full on target kernel(s)
|
||||
→ extract DRAM bandwidth util, SM throughput, achieved occupancy
|
||||
→ determine if kernel is memory-bound, compute-bound, or latency-bound
|
||||
→ for CUDA graph: use --graph-profiling node
|
||||
↓
|
||||
5. VERIFY CORRECTNESS FIRST
|
||||
sglang generate --seed=42 --save-output → diff against reference baseline
|
||||
If output differs beyond tolerance → reject the optimization regardless of speedup
|
||||
5. KERNEL OPTIMIZATION
|
||||
Existing fused kernel? → use-efficient-diffusion-kernels.md
|
||||
New Triton kernel? → add-triton-kernel.md
|
||||
New CUDA JIT kernel? → add-cuda-kernel.md
|
||||
After writing kernel → ncu again to verify bandwidth/occupancy ★
|
||||
↓
|
||||
6. Re-run benchmark → verify denoise latency improvement; no regression elsewhere
|
||||
6. VERIFY CORRECTNESS
|
||||
sglang generate --seed=42 --save-output → diff against reference
|
||||
If output differs beyond tolerance → reject optimization
|
||||
↓
|
||||
7. RE-BENCHMARK
|
||||
Verify denoise latency improvement; no regression on other models
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Continuous Improvement Checklist
|
||||
|
||||
Before merging any PR that affects diffusion performance, run the full benchmark suite and compare.
|
||||
|
||||
> **Rule: correctness gates performance.** A PR that improves latency but changes output is not acceptable. Correctness checks must pass before performance numbers are even considered.
|
||||
|
||||
### Correctness Checks (must pass first)
|
||||
## Checklist Before Merging
|
||||
|
||||
### Correctness (must pass first)
|
||||
- [ ] Reference outputs collected with `--seed=42 --save-output` **before** any change
|
||||
- [ ] After change: regenerate with identical args and compare against reference
|
||||
- [ ] After change: regenerate with identical args and compare
|
||||
- [ ] No visible quality degradation in generated images / videos
|
||||
- [ ] For numerical changes: pixel-level diff or PSNR/SSIM within agreed tolerance
|
||||
- [ ] Correctness verified on **all 7 benchmark models**, not just the model being optimized
|
||||
- [ ] Correctness verified on all benchmark models
|
||||
|
||||
### Performance Checks (only after correctness passes)
|
||||
|
||||
- [ ] All 7 model benchmarks executed; denoise latency (★), end-to-end latency, and peak memory recorded
|
||||
- [ ] No regression in denoise latency vs. previous baseline (allow ±2% variance)
|
||||
- [ ] New optimization shows measurable improvement in denoise latency on at least 2 models
|
||||
- [ ] No new graph breaks introduced (verify via `torch._dynamo` logs)
|
||||
### Performance (only after correctness passes)
|
||||
- [ ] All benchmark models executed; denoise latency ★, end-to-end, peak memory recorded
|
||||
- [ ] No regression in denoise latency vs. previous baseline (±2% tolerance)
|
||||
- [ ] New kernel shows measurable improvement on at least 2 models
|
||||
- [ ] No new torch.compile graph breaks introduced
|
||||
- [ ] Results reproducible with all offloads disabled and fixed `--seed=42`
|
||||
|
||||
@@ -54,8 +54,9 @@ nsys profile --gpu-metrics-device=all -o report ./cuda_program
|
||||
# Profile specific duration
|
||||
nsys profile -d 10 -o report ./cuda_program
|
||||
|
||||
# Export to multiple formats
|
||||
nsys export -t sqlite,json report.nsys-rep
|
||||
# Export to multiple formats (one type per command)
|
||||
nsys export -t sqlite report.nsys-rep
|
||||
nsys export -t json report.nsys-rep
|
||||
|
||||
# Generate summary statistics
|
||||
nsys stats report.nsys-rep
|
||||
@@ -179,14 +180,23 @@ ncu --set full -o baseline ./program_v1
|
||||
# Step 2: Profile optimized version
|
||||
ncu --set full -o optimized ./program_v2
|
||||
|
||||
# Step 3: Generate comparison report (CLI, no GUI needed)
|
||||
# Both --import flags required; --page diff generates a side-by-side diff
|
||||
ncu --import baseline.ncu-rep \
|
||||
--import optimized.ncu-rep \
|
||||
--page diff --csv > comparison.csv
|
||||
```
|
||||
# Step 3: Export both profiles to CSV, then compare with Python (no GUI needed)
|
||||
# Note: --import can only be specified once; --page diff is not a valid page value.
|
||||
ncu --import baseline.ncu-rep --page details --csv > baseline_details.csv
|
||||
ncu --import optimized.ncu-rep --page details --csv > optimized_details.csv
|
||||
|
||||
> **Note**: `ncu --diff` (the old single-flag syntax) was removed in Nsight Compute 2022.x. Always use two `--import` flags with `--page diff` for comparisons.
|
||||
python3 -c "
|
||||
import csv
|
||||
def load(p):
|
||||
return {r.get('Metric Name',''): r.get('Metric Value','')
|
||||
for r in csv.DictReader(open(p))}
|
||||
b = load('baseline_details.csv')
|
||||
o = load('optimized_details.csv')
|
||||
for k in sorted(set(b) | set(o)):
|
||||
bv, ov = b.get(k,''), o.get(k,'')
|
||||
if bv != ov:
|
||||
print(f'{k[:55]:<55} {bv} -> {ov}')
|
||||
"
|
||||
|
||||
### 8. Performance Recommendations
|
||||
|
||||
|
||||
@@ -0,0 +1,283 @@
|
||||
# A100 GPU Optimization Guide — SGLang Diffusion JIT Kernels
|
||||
|
||||
Deep dive into A100-specific optimizations for diffusion model CUDA kernels in SGLang's JIT system.
|
||||
|
||||
> **Adapted from**: [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels)
|
||||
|
||||
---
|
||||
|
||||
## A100 Ampere Architecture Overview
|
||||
|
||||
| Component | A100 40GB | A100 80GB | Notes |
|
||||
|-----------|-----------|-----------|-------|
|
||||
| Compute Capability | sm_80 | sm_80 | Use `"-arch=sm_80"` in `extra_cuda_cflags` |
|
||||
| SMs | 108 | 108 | Grid: aim for multiples of 108 |
|
||||
| Shared Memory | 164 KB/SM | 164 KB/SM | Configurable: 48/96/164 KB |
|
||||
| L2 Cache | 40 MB | 40 MB | Less than H100 (50 MB) |
|
||||
| Memory Bandwidth | 1.55 TB/s | 2.0 TB/s | HBM2e |
|
||||
| Max Threads/SM | 2048 | 2048 | Same as H100 |
|
||||
| Tensor Cores | 3rd gen | 3rd gen | FP16, BF16, TF32, INT8, INT4 |
|
||||
|
||||
### A100 vs H100 Comparison
|
||||
|
||||
| Feature | A100 | H100 | Impact on JIT Kernels |
|
||||
|---------|------|------|-----------------------|
|
||||
| Memory BW | 2.0 TB/s | 3.35 TB/s | H100 ~67% faster for memory-bound ops |
|
||||
| SMs | 108 | 132 | Adjust persistent kernel grid sizing |
|
||||
| Shared Mem/SM | 164 KB | 192 KB | Reduce max tile sizes on A100 |
|
||||
| L2 Cache | 40 MB | 50 MB | Attention tile reuse still works well |
|
||||
| TMA | No | Yes | Can't use `cp.async.bulk` on A100 |
|
||||
| FP8 | No | Yes | Use FP16/BF16 only on A100 |
|
||||
|
||||
---
|
||||
|
||||
## Memory Access Optimization
|
||||
|
||||
Same coalescing and vectorization rules as H100; lower bandwidth makes them even more critical.
|
||||
|
||||
### `AlignedVector` Vectorization (same pattern as H100)
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/vec.cuh>
|
||||
|
||||
constexpr int kVecN = 16 / sizeof(T); // 8 for bf16/fp16, 4 for fp32
|
||||
using vec_t = device::AlignedVector<T, kVecN>;
|
||||
|
||||
vec_t v;
|
||||
v.load(src, vi);
|
||||
// ... process elements ...
|
||||
v.store(dst, vi);
|
||||
```
|
||||
|
||||
**Expected A100 performance (BF16 RMSNorm):**
|
||||
|
||||
| Implementation | A100 (ms) | H100 (ms) | A100 Speedup |
|
||||
|:---|:---:|:---:|:---:|
|
||||
| Scalar loads | ~0.10 | 0.065 | 1.00x |
|
||||
| `AlignedVector<bf16_t, 8>` | ~0.03 | 0.019 | ~3x |
|
||||
|
||||
**Target bandwidth**: 30–40% of A100's 2.0 TB/s = 600–800 GB/s.
|
||||
|
||||
### Shared Memory Configuration
|
||||
|
||||
```cpp
|
||||
// A100 max: 164 KB/SM
|
||||
cudaFuncSetAttribute(
|
||||
your_kernel,
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize,
|
||||
164 * 1024 // 164 KB max on A100
|
||||
);
|
||||
```
|
||||
|
||||
Attention tile sizes for A100:
|
||||
|
||||
```
|
||||
BLOCK_SIZE_M = 128 (Q block)
|
||||
BLOCK_SIZE_N = 64 (K,V block)
|
||||
Tile = 128×64×2 = 16 KB (FP16) — fits in 164 KB shared mem
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Occupancy Tuning
|
||||
|
||||
**Grid sizing for A100 (108 SMs):**
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/runtime.cuh>
|
||||
|
||||
// Cap blocks to SM × occupancy (same pattern as H100)
|
||||
static const uint32_t max_occ = host::runtime::get_blocks_per_sm(kernel, kBlockSize);
|
||||
static const uint32_t num_sm = host::runtime::get_sm_count(device.unwrap().device_id);
|
||||
const uint32_t num_blocks = std::min(num_sm * max_occ, host::div_ceil(n, kBlockSize));
|
||||
```
|
||||
|
||||
**Recommended block sizes (same as H100):**
|
||||
|
||||
| Kernel Type | Threads/Block | Notes |
|
||||
|-------------|---------------|-------|
|
||||
| Element-wise | 256 | High occupancy |
|
||||
| Row reduction | 512 | Full reduction per row |
|
||||
| Tiled/attention | 256 | Balance shared mem |
|
||||
|
||||
---
|
||||
|
||||
## A100-Specific Features
|
||||
|
||||
### Async Memory Copy (sm_80)
|
||||
|
||||
A100 introduced `cp.async` for overlapping compute and memory. Use this in custom kernels for prefetching:
|
||||
|
||||
```cuda
|
||||
#if __CUDA_ARCH__ >= 800
|
||||
// Async copy from global to shared (A100+)
|
||||
__pipeline_memcpy_async(smem_ptr, global_ptr, bytes);
|
||||
__pipeline_commit();
|
||||
__pipeline_wait_prior(0);
|
||||
#endif
|
||||
```
|
||||
|
||||
### TF32 Mode (A100 specific)
|
||||
|
||||
Enables FP32-range with FP16-like throughput for GEMM. Enable in Python:
|
||||
|
||||
```python
|
||||
# Enable TF32 for matmuls (A100+)
|
||||
torch.backends.cuda.matmul.allow_tf32 = True
|
||||
torch.backends.cudnn.allow_tf32 = True
|
||||
```
|
||||
|
||||
TF32 is automatic for FP32 GEMMs via cuBLAS — no kernel changes needed.
|
||||
|
||||
### Structural Sparsity (2:4)
|
||||
|
||||
A100 tensor cores support 50% structured sparsity:
|
||||
|
||||
```python
|
||||
from torch.sparse import to_sparse_semi_structured
|
||||
sparse_weight = to_sparse_semi_structured(dense_weight)
|
||||
# ~2x GEMM speedup for matmul with sparse weight
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## JIT Compilation for A100
|
||||
|
||||
```python
|
||||
return load_jit(
|
||||
"my_kernel",
|
||||
*args,
|
||||
cuda_files=["diffusion/my_kernel.cuh"],
|
||||
cuda_wrappers=[("my_kernel", f"my_kernel<{args}>")],
|
||||
extra_cuda_cflags=[
|
||||
"-O3",
|
||||
"--use_fast_math",
|
||||
"-arch=sm_80", # A100 only; omit for multi-arch
|
||||
],
|
||||
)
|
||||
```
|
||||
|
||||
**Multi-arch (A100 + H100):**
|
||||
|
||||
```python
|
||||
extra_cuda_cflags=[
|
||||
"-O3", "--use_fast_math",
|
||||
"-gencode=arch=compute_80,code=sm_80", # A100
|
||||
"-gencode=arch=compute_90,code=sm_90", # H100
|
||||
]
|
||||
```
|
||||
|
||||
Runtime arch guard (in Python wrapper):
|
||||
|
||||
```python
|
||||
cap = torch.cuda.get_device_capability()
|
||||
if cap < (8, 0):
|
||||
raise RuntimeError(f"This kernel requires sm_80 (A100) or later, got sm_{cap[0]}{cap[1]}")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## H100 → A100 Migration Checklist
|
||||
|
||||
When porting an H100-optimized kernel to A100:
|
||||
|
||||
| Item | H100 | A100 | Change Required |
|
||||
|------|------|------|-----------------|
|
||||
| Shared memory | 192 KB | 164 KB | Reduce `cudaFuncSetAttribute` size |
|
||||
| Grid sizing | ×132 SMs | ×108 SMs | `get_sm_count()` handles automatically |
|
||||
| TMA bulk copy | Available | **Not available** | Remove `cp.async.bulk`; use standard `__pipeline_memcpy_async` |
|
||||
| FP8 | Available | **Not available** | Fall back to FP16/BF16 |
|
||||
| PDL | Supported | Supported | `.enable_pdl(true)` works on sm_80 |
|
||||
| Warp shuffles | Same | Same | No changes |
|
||||
| `AlignedVector` | Same | Same | No changes |
|
||||
|
||||
**Conditional compilation:**
|
||||
|
||||
```cuda
|
||||
#if __CUDA_ARCH__ >= 900
|
||||
// H100-only: TMA, FP8, thread block clusters
|
||||
#define USE_TMA 1
|
||||
#elif __CUDA_ARCH__ >= 800
|
||||
// A100: cp.async, TF32, 2:4 sparsity
|
||||
#define USE_ASYNC_COPY 1
|
||||
#endif
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Precision Notes
|
||||
|
||||
| Type | Available on A100 | Notes |
|
||||
|------|-------------------|-------|
|
||||
| FP16 | Yes | Good, watch overflow in attention |
|
||||
| BF16 | Yes | Preferred for training and inference |
|
||||
| TF32 | Yes (A100 specific) | Auto for FP32 GEMMs |
|
||||
| FP8 | **No** | H100 only |
|
||||
|
||||
---
|
||||
|
||||
## Performance Profiling
|
||||
|
||||
### NVIDIA Nsight Systems (nsys)
|
||||
|
||||
```bash
|
||||
nsys profile -o a100_profile python scripts/bench_diffusion_rmsnorm.py
|
||||
|
||||
# Key metrics to watch:
|
||||
# - Kernel duration
|
||||
# - Memory transfer time
|
||||
# - GPU idle time
|
||||
# - Stream utilization
|
||||
```
|
||||
|
||||
### NVIDIA Nsight Compute (ncu)
|
||||
|
||||
```bash
|
||||
# Full metrics
|
||||
ncu --set full -o a100_metrics.ncu-rep \
|
||||
python scripts/bench_diffusion_rmsnorm.py
|
||||
|
||||
# Specific metrics for bandwidth / occupancy checks
|
||||
ncu --metrics sm__throughput.avg.pct_of_peak_sustained_elapsed,\
|
||||
dram__throughput.avg.pct_of_peak_sustained_elapsed \
|
||||
python scripts/bench_diffusion_rmsnorm.py
|
||||
|
||||
# Key metrics for A100 diffusion kernels:
|
||||
# - Achieved occupancy (sm__warps_active.avg.pct_of_peak_sustained_active)
|
||||
# - Memory throughput (dram__throughput.avg.pct_of_peak_sustained_elapsed)
|
||||
# → Target: 30–40% of 2.0 TB/s (600–800 GB/s) for vectorized kernels
|
||||
# - Compute throughput (sm__throughput.avg.pct_of_peak_sustained_elapsed)
|
||||
# - Warp stall reasons (smsp__warp_issue_stalled_*.avg.pct_of_peak_sustained_active)
|
||||
# - Kernel time (gpu__time_duration.avg)
|
||||
```
|
||||
|
||||
### Common A100 Performance Issues
|
||||
|
||||
1. **Memory bound below target**: `dram__throughput` < 30%
|
||||
- Fix: Use `AlignedVector<bf16_t, 8>` (128-bit vector loads)
|
||||
|
||||
2. **Low occupancy**: Grid too small for 108 SMs
|
||||
- Fix: Use `runtime::get_sm_count()` persistent kernel pattern
|
||||
|
||||
3. **No TF32 for FP32 GEMMs**: torch.backends.cuda.matmul.allow_tf32 not set
|
||||
- Fix: `torch.backends.cuda.matmul.allow_tf32 = True`
|
||||
|
||||
---
|
||||
|
||||
## Best Practices Summary (A100)
|
||||
|
||||
1. **Bandwidth**: Even more critical than H100 — profile with `ncu` first
|
||||
2. **Vectorization**: `AlignedVector<bf16_t, 8>` gives ~3x over scalar
|
||||
3. **TF32**: Enable for any FP32 matmul workload
|
||||
4. **Shared memory**: Cap at 164 KB; use `cudaFuncSetAttribute`
|
||||
5. **Grid sizing**: Multiples of 108 SMs via `runtime::get_sm_count`
|
||||
6. **cp.async**: Use for prefetching in tiled kernels
|
||||
7. **Multi-arch**: Build for both `sm_80` and `sm_90` to support both GPUs
|
||||
8. **Same abstractions**: `AlignedVector`, `TensorMatcher`, `LaunchKernel` work identically
|
||||
|
||||
## Reference Benchmark Results (A100 80GB, BF16)
|
||||
|
||||
| Kernel | Shape | A100 (ms) | H100 (ms) | H100 Speedup |
|
||||
|--------|-------|-----------|-----------|--------------|
|
||||
| RMSNorm | [2, 1024, 2048] | ~0.08 | 0.054 | 1.5x |
|
||||
| GEGLU | [2, 1024, 4096] | ~0.05 | 0.030 | 1.7x |
|
||||
@@ -0,0 +1,364 @@
|
||||
# H100 GPU Optimization Guide — SGLang Diffusion JIT Kernels
|
||||
|
||||
Deep dive into H100-specific optimizations for diffusion model CUDA kernels, written for SGLang's JIT kernel system.
|
||||
|
||||
> **Adapted from**: [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels)
|
||||
|
||||
---
|
||||
|
||||
## H100 Hopper Architecture Overview
|
||||
|
||||
| Component | Specification | Optimization Implication |
|
||||
|-----------|---------------|--------------------------|
|
||||
| Compute Capability | sm_90 | Use `extra_cuda_cflags=["-arch=sm_90"]` in `load_jit` |
|
||||
| SMs | 132 | Grid: aim for multiples of 132 |
|
||||
| Shared Memory | 192 KB/SM | Configurable: 96/144/192 KB |
|
||||
| L2 Cache | 50 MB | Tile K,V of attention to fit in L2 |
|
||||
| Memory Bandwidth | 3.35 TB/s | BF16 vectorized: achieves ~38% (~1.27 TB/s) |
|
||||
| Max Threads/SM | 2048 | Max 16 blocks of 128 threads per SM |
|
||||
| Warp Size | 32 | All reductions use `warp::reduce_sum` |
|
||||
| Registers | 64K 32-bit/SM | 255 per thread max |
|
||||
|
||||
### New Hopper Features (sm_90+)
|
||||
|
||||
1. **Thread Block Clusters** — groups cooperating via Distributed Shared Memory
|
||||
2. **TMA (Tensor Memory Accelerator)** — hardware-accelerated bulk copies
|
||||
3. **FP8 support** — native 8-bit floating point in tensor cores
|
||||
4. **PDL (Programmatic Dependent Launch)** — enable with `.enable_pdl(true)` in `LaunchKernel`
|
||||
|
||||
Gate sm_90+ features with a runtime check before calling `load_jit`:
|
||||
|
||||
```python
|
||||
if torch.cuda.get_device_capability()[0] < 9:
|
||||
raise RuntimeError("This kernel requires H100 (sm_90+)")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Memory Hierarchy Optimization
|
||||
|
||||
### Coalesced Global Memory Access
|
||||
|
||||
```cpp
|
||||
// GOOD: threads read consecutive addresses → 128-byte transaction per warp
|
||||
uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
fp16_t val = src[idx];
|
||||
|
||||
// BAD: strided access → multiple transactions, lower effective bandwidth
|
||||
uint32_t idx = threadIdx.x * stride; // avoid stride > 1
|
||||
```
|
||||
|
||||
**Transaction sizes**: 32 bytes minimum, 128 bytes optimal (full warp, FP32).
|
||||
|
||||
### Vectorized Memory Access with `AlignedVector`
|
||||
|
||||
SGLang's `AlignedVector<T, N>` provides 128-bit (16-byte) vector loads. Always use this instead of raw pointer reinterprets.
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/vec.cuh>
|
||||
|
||||
// 16 bytes per load: 8×bf16_t, 8×fp16_t, or 4×fp32_t
|
||||
constexpr int kVecN = 16 / sizeof(T);
|
||||
using vec_t = device::AlignedVector<T, kVecN>;
|
||||
|
||||
// Load
|
||||
vec_t v;
|
||||
v.load(src, vi); // loads src[vi * kVecN .. vi * kVecN + kVecN - 1]
|
||||
|
||||
// Process
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecN; ++i) {
|
||||
float val = static_cast<float>(v[i]);
|
||||
// ... compute ...
|
||||
v[i] = static_cast<T>(result);
|
||||
}
|
||||
|
||||
// Store
|
||||
v.store(dst, vi);
|
||||
```
|
||||
|
||||
**RMSNorm benchmark (H100 80GB, BF16):**
|
||||
|
||||
| Implementation | Time (ms) | Speedup |
|
||||
|:---|:---:|:---:|
|
||||
| Scalar loads | 0.065 | 1.00x |
|
||||
| `AlignedVector<bf16_t, 8>` | 0.019 | **3.37x** |
|
||||
|
||||
Bandwidth achieved: **~38% of 3.35 TB/s** = 1.27 TB/s.
|
||||
|
||||
### L2 Cache Utilization (50 MB)
|
||||
|
||||
For attention, tile K and V so they stay in L2 while Q iterates:
|
||||
|
||||
```
|
||||
BLOCK_SIZE_M = 128 (Q block)
|
||||
BLOCK_SIZE_N = 64 (K,V block)
|
||||
With head_dim=64: tile = 128×64×2 = 16 KB (FP16), multiple tiles fit in L2
|
||||
```
|
||||
|
||||
### Shared Memory Configuration
|
||||
|
||||
Request max shared memory for attention kernels:
|
||||
|
||||
```cpp
|
||||
// In launcher (after selecting kernel function pointer):
|
||||
cudaFuncSetAttribute(
|
||||
your_kernel,
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize,
|
||||
192 * 1024 // 192 KB max on H100
|
||||
);
|
||||
```
|
||||
|
||||
Shared memory has 32 banks (4 bytes/bank). Avoid conflicts with padding:
|
||||
|
||||
```cpp
|
||||
__shared__ float data[32][33]; // 33 instead of 32 → no bank conflict
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Warp & CTA Reductions (SGLang Abstractions)
|
||||
|
||||
Use `sgl_kernel/warp.cuh` and `sgl_kernel/cta.cuh` — never raw `__shfl_xor_sync`.
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/warp.cuh>
|
||||
#include <sgl_kernel/cta.cuh>
|
||||
|
||||
// Warp-level sum (uses __shfl_xor_sync internally)
|
||||
float result = device::warp::reduce_sum<float>(partial);
|
||||
|
||||
// Warp-level max
|
||||
float mx = device::warp::reduce_max<float>(val);
|
||||
|
||||
// CTA-wide max via shared memory
|
||||
__shared__ float smem[32];
|
||||
device::cta::reduce_max<float>(val, smem, -1e38f);
|
||||
// smem[0] holds the result after __syncthreads()
|
||||
```
|
||||
|
||||
**Block reduction pattern for RMSNorm:**
|
||||
|
||||
```cpp
|
||||
// 1. Warp reduction
|
||||
sum_sq = device::warp::reduce_sum<float>(sum_sq);
|
||||
|
||||
// 2. Write warp leaders to smem
|
||||
__shared__ float smem_r[32];
|
||||
if (threadIdx.x % 32 == 0) smem_r[threadIdx.x / 32] = sum_sq;
|
||||
__syncthreads();
|
||||
|
||||
// 3. Final warp reduction over warp leaders
|
||||
if (threadIdx.x < 32) {
|
||||
sum_sq = (threadIdx.x < blockDim.x / 32) ? smem_r[threadIdx.x] : 0.f;
|
||||
sum_sq = device::warp::reduce_sum<float>(sum_sq);
|
||||
}
|
||||
__syncthreads();
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Occupancy Tuning
|
||||
|
||||
```
|
||||
Occupancy = Active Warps per SM / Max Warps per SM (64)
|
||||
|
||||
Limiting factors on H100:
|
||||
1. Registers: 65536 / (threads_per_block × regs_per_thread)
|
||||
2. Shared Memory: 192 KB / smem_per_block
|
||||
3. Threads: 2048 / threads_per_block
|
||||
```
|
||||
|
||||
**Recommended block sizes:**
|
||||
|
||||
| Kernel Type | Threads/Block | Warps | Reasoning |
|
||||
|-------------|---------------|-------|-----------|
|
||||
| Element-wise (RoPE, GEGLU) | 256 | 8 | High occupancy, simple |
|
||||
| Row reduction (RMSNorm, LayerNorm) | 256–512 | 8–16 | Enough threads for full reduction |
|
||||
| Tiled (attention) | 256 | 8 | Balance shared mem and registers |
|
||||
|
||||
**Persistent kernel pattern** (cap grid to SM × occupancy):
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/runtime.cuh>
|
||||
|
||||
static const uint32_t max_occ = host::runtime::get_blocks_per_sm(kernel, kBlockSize);
|
||||
static const uint32_t num_sm = host::runtime::get_sm_count(device.unwrap().device_id);
|
||||
const uint32_t num_blocks = std::min(num_sm * max_occ, host::div_ceil(n, kBlockSize));
|
||||
host::LaunchKernel(num_blocks, kBlockSize, device.unwrap())(kernel, params);
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Precision and Numerical Stability
|
||||
|
||||
| Type | Exponent Bits | Mantissa Bits | Range | Use Case |
|
||||
|------|--------------|---------------|-------|----------|
|
||||
| FP16 | 5 | 10 | ±65504 | Inference; attention score overflow risk |
|
||||
| BF16 | 8 | 7 | ±3.39×10³⁸ | Training/inference preferred; safer for attn |
|
||||
| FP32 | 8 | 23 | ±3.39×10³⁸ | Accumulation only |
|
||||
|
||||
**Mixed precision pattern** (always accumulate in FP32):
|
||||
|
||||
```cpp
|
||||
// Input via AlignedVector
|
||||
vec_t v;
|
||||
v.load(src, vi);
|
||||
float acc = 0.f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecN; ++i) {
|
||||
float val = static_cast<float>(v[i]); // promote to FP32
|
||||
acc += val * val;
|
||||
}
|
||||
// Output
|
||||
v[i] = static_cast<T>(fp32_result); // demote back
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Diffusion-Specific Patterns
|
||||
|
||||
### DiT Block Operators
|
||||
|
||||
| Operator | Pattern | Key Constraint |
|
||||
|----------|---------|----------------|
|
||||
| **RMSNorm** | 2-pass row reduction | weight may be `None` |
|
||||
| **AdaLN** | `norm(x) * (1 + scale) + shift` | fuse norm+scale+shift |
|
||||
| **RoPE 3D** | `[B, t*h*w, heads, head_dim]` | layout: `seq = t*h*w` |
|
||||
| **GEGLU** | `gelu(gate) * value`, input `[B,L,2H]` | don't use for LTX-Video (uses GELU) |
|
||||
| **SiLU gate** | `x * sigmoid(x)` | fuse with MLP linear |
|
||||
|
||||
### Online Softmax (for custom attention)
|
||||
|
||||
```cuda
|
||||
// Numerically stable without materializing full [seq×seq] score matrix
|
||||
float row_max = -INFINITY, row_sum = 0.f;
|
||||
for each K block:
|
||||
compute local_scores
|
||||
new_max = max(row_max, max(local_scores))
|
||||
rescale = exp(row_max - new_max)
|
||||
row_sum = row_sum * rescale + sum(exp(local_scores - new_max))
|
||||
out_acc = out_acc * rescale + softmax(local_scores) @ V_block
|
||||
row_max = new_max
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Profiling and Debugging
|
||||
|
||||
### NVIDIA Nsight Systems (nsys)
|
||||
|
||||
System-wide profiling to see kernel durations, memory transfers, and GPU idle time:
|
||||
|
||||
```bash
|
||||
nsys profile -o profile_report python scripts/bench_diffusion_rmsnorm.py
|
||||
|
||||
# Key metrics to watch:
|
||||
# - Kernel duration
|
||||
# - Memory transfer time
|
||||
# - GPU idle time
|
||||
# - Stream utilization
|
||||
```
|
||||
|
||||
For end-to-end denoise profiling via `sglang generate`, see `diffusion-benchmark-and-profile.md` (Level 2: nsys + gputrc2graph.py).
|
||||
|
||||
### NVIDIA Nsight Compute (ncu)
|
||||
|
||||
Detailed per-kernel analysis for tuning individual JIT CUDA kernels:
|
||||
|
||||
```bash
|
||||
# Full metrics — use when you need everything (slow)
|
||||
ncu --set full -o metrics.ncu-rep \
|
||||
python scripts/bench_diffusion_rmsnorm.py
|
||||
|
||||
# Specific metrics — use for targeted bandwidth / occupancy checks
|
||||
ncu --metrics sm__throughput.avg.pct_of_peak_sustained_elapsed,\
|
||||
dram__throughput.avg.pct_of_peak_sustained_elapsed \
|
||||
python scripts/bench_diffusion_rmsnorm.py
|
||||
|
||||
# Key metrics for diffusion JIT kernels:
|
||||
# - Achieved occupancy (sm__warps_active.avg.pct_of_peak_sustained_active)
|
||||
# - Memory throughput (dram__throughput.avg.pct_of_peak_sustained_elapsed)
|
||||
# - Compute throughput (sm__throughput.avg.pct_of_peak_sustained_elapsed)
|
||||
# - Warp stall reasons (smsp__warp_issue_stalled_*.avg.pct_of_peak_sustained_active)
|
||||
# - L1 cache hit rate (l1tex__t_requests_pipe_lsu_mem_global_op_ld.sum)
|
||||
```
|
||||
|
||||
### Common Performance Issues
|
||||
|
||||
1. **Low occupancy**: Too many registers or shared memory per block
|
||||
- Check: `--ptxas-options=-v` in `extra_cuda_cflags` to see register count
|
||||
- Fix: Reduce `--maxrregcount=N`; use smaller block size
|
||||
|
||||
2. **Memory bound, low bandwidth**: Achieved < 30% of 3.35 TB/s
|
||||
- Check: `dram__throughput.avg.pct_of_peak_sustained_elapsed`
|
||||
- Fix: Switch to `AlignedVector<T, 16/sizeof(T)>` for 128-bit vector loads
|
||||
|
||||
3. **Shared memory bank conflicts**: `l1tex__data_bank_conflicts_pipe_lmem_op_st.sum` is high
|
||||
- Fix: Add padding — `__shared__ float data[32][33]`
|
||||
|
||||
4. **Warp divergence**: Conditional branches splitting warps
|
||||
- Check: `smsp__warp_issue_stalled_branch.avg.pct_of_peak_sustained_active`
|
||||
- Fix: Restructure so elements with identical branches are in the same warp
|
||||
|
||||
5. **Too many small kernels**: High kernel launch overhead
|
||||
- Fix: Fuse operations (e.g., norm + scale + shift → AdaLN in one kernel)
|
||||
|
||||
---
|
||||
|
||||
## JIT Compilation Notes
|
||||
|
||||
SGLang's JIT compiles kernels on first use via `load_jit`. For H100-specific flags:
|
||||
|
||||
```python
|
||||
return load_jit(
|
||||
"my_kernel",
|
||||
*args,
|
||||
cuda_files=["diffusion/my_kernel.cuh"],
|
||||
cuda_wrappers=[("my_kernel", f"my_kernel<{args}>")],
|
||||
extra_cuda_cflags=[
|
||||
"-O3",
|
||||
"--use_fast_math",
|
||||
"-arch=sm_90", # H100 only; omit for multi-arch
|
||||
"--ptxas-options=-v", # Remove after tuning
|
||||
],
|
||||
)
|
||||
```
|
||||
|
||||
For multi-arch (H100 + A100):
|
||||
|
||||
```python
|
||||
extra_cuda_cflags=[
|
||||
"-O3",
|
||||
"--use_fast_math",
|
||||
"-gencode=arch=compute_80,code=sm_80", # A100
|
||||
"-gencode=arch=compute_90,code=sm_90", # H100
|
||||
]
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Best Practices Summary
|
||||
|
||||
1. **Memory access**: Coalesce writes, align to 128-byte boundaries
|
||||
2. **Vectorization**: Use `AlignedVector<T, 16/sizeof(T)>` for all element-wise loads/stores
|
||||
3. **Reductions**: Use `warp::reduce_sum/max`, then shared memory pattern above
|
||||
4. **Precision**: BF16 for I/O, FP32 for accumulation; use `static_cast<float>`
|
||||
5. **Block size**: 256 threads default; 512 for reductions; tune with `runtime::get_blocks_per_sm`
|
||||
6. **Grid sizing**: Multiples of 132 SMs; use persistent kernel pattern for small N
|
||||
7. **Shared memory**: Add padding (`[32][33]`) to avoid bank conflicts
|
||||
8. **Profile**: Run `ncu` before claiming a speedup; check dram throughput %
|
||||
9. **Fuse**: Combine norm + scale + shift into a single pass to reduce memory traffic
|
||||
10. **Abstractions**: Always use `TensorMatcher`, `AlignedVector`, `LaunchKernel` — never raw CUDA
|
||||
|
||||
## Reference Benchmark Results (H100 80GB, BF16)
|
||||
|
||||
| Kernel | Shape | Time (ms) |
|
||||
|--------|-------|-----------|
|
||||
| RMSNorm | [2, 1024, 2048] | 0.054 |
|
||||
| GEGLU | [2, 1024, 4096] → [2, 1024, 2048] | 0.030 |
|
||||
| RoPE 3D | [2, 480, 8, 64] | 1.670 |
|
||||
| RMSNorm vectorized | [1, 1024, 2048] | 0.019 |
|
||||
| RMSNorm vectorized | [4, 4096, 3072] | 0.157 |
|
||||
|
||||
> See `kernel-templates.md` for copy-paste ready sglang JIT kernel implementations.
|
||||
@@ -0,0 +1,569 @@
|
||||
# CUDA Kernel Templates — SGLang Diffusion JIT Style
|
||||
|
||||
Copy-paste ready templates for JIT CUDA kernels in `python/sglang/jit_kernel/csrc/diffusion/`.
|
||||
All templates use SGLang's internal abstractions; no raw CUDA headers needed.
|
||||
|
||||
> **Adapted from**: [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels)
|
||||
|
||||
---
|
||||
|
||||
## Prerequisite: Standard Includes
|
||||
|
||||
Every kernel file in `csrc/diffusion/` starts with:
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/tensor.h> // TensorMatcher, SymbolicSize, SymbolicDevice
|
||||
#include <sgl_kernel/type.cuh> // fp16_t, bf16_t, fp32_t, dtype_trait, packed_t
|
||||
#include <sgl_kernel/utils.h> // RuntimeCheck, Panic, div_ceil
|
||||
#include <sgl_kernel/utils.cuh> // LaunchKernel, SGL_DEVICE, type aliases
|
||||
#include <sgl_kernel/vec.cuh> // AlignedVector<T, N>
|
||||
#include <sgl_kernel/warp.cuh> // warp::reduce_sum, warp::reduce_max
|
||||
#include <sgl_kernel/math.cuh> // device::math::rsqrt, sqrt, ...
|
||||
#include <sgl_kernel/tile.cuh> // tile::Memory (strided access pattern)
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
```
|
||||
|
||||
**Key type aliases** (from `utils.cuh`):
|
||||
- `fp16_t` = `__half`, `fp16x2_t` = `__half2`
|
||||
- `bf16_t` = `__nv_bfloat16`, `bf16x2_t` = `__nv_bfloat162`
|
||||
- `fp32_t` = `float`, `fp32x2_t` = `float2`
|
||||
- `SGL_DEVICE` = `__forceinline__ __device__`
|
||||
|
||||
---
|
||||
|
||||
## Template 1: Element-wise Operation
|
||||
|
||||
Use for ops that process elements independently: RoPE, SiLU, GEGLU, scale+bias.
|
||||
|
||||
### `.cuh` file: `csrc/diffusion/silu_gate.cuh`
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/type.cuh>
|
||||
#include <sgl_kernel/utils.h>
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
#include <sgl_kernel/vec.cuh>
|
||||
#include <sgl_kernel/math.cuh>
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
|
||||
namespace {
|
||||
|
||||
// SiLU gate: out[i] = x[i] * sigmoid(x[i])
|
||||
// Input layout: [B, L, hidden]
|
||||
template <typename T, int kVecN>
|
||||
__global__ void silu_gate_kernel(
|
||||
T* __restrict__ dst,
|
||||
const T* __restrict__ src,
|
||||
uint32_t n_vecs,
|
||||
uint32_t n_remainder,
|
||||
uint32_t n_total)
|
||||
{
|
||||
using vec_t = device::AlignedVector<T, kVecN>;
|
||||
|
||||
const uint32_t stride = blockDim.x * gridDim.x;
|
||||
|
||||
// --- vectorized body ---
|
||||
for (uint32_t vi = blockIdx.x * blockDim.x + threadIdx.x; vi < n_vecs; vi += stride) {
|
||||
vec_t v;
|
||||
v.load(src, vi);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecN; ++i) {
|
||||
float val = static_cast<float>(v[i]);
|
||||
float sig = 1.f / (1.f + device::math::exp<float>(-val));
|
||||
v[i] = static_cast<T>(val * sig);
|
||||
}
|
||||
v.store(dst, vi);
|
||||
}
|
||||
|
||||
// --- scalar tail (for sizes not divisible by kVecN) ---
|
||||
const uint32_t base = n_vecs * kVecN;
|
||||
for (uint32_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n_remainder; i += stride) {
|
||||
float val = static_cast<float>(src[base + i]);
|
||||
float sig = 1.f / (1.f + device::math::exp<float>(-val));
|
||||
dst[base + i] = static_cast<T>(val * sig);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void silu_gate(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) {
|
||||
using namespace host;
|
||||
|
||||
SymbolicSize N{"num_elements"};
|
||||
SymbolicDevice device;
|
||||
device.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({N})
|
||||
.with_dtype<T>()
|
||||
.with_device(device)
|
||||
.verify(dst)
|
||||
.verify(src);
|
||||
|
||||
const uint32_t n = static_cast<uint32_t>(N.unwrap());
|
||||
const DLDevice dev = device.unwrap();
|
||||
RuntimeCheck(n > 0, "silu_gate: num_elements must be > 0");
|
||||
|
||||
constexpr int kVecN = 16 / sizeof(T); // 128-bit vector load
|
||||
const uint32_t n_vecs = n / kVecN;
|
||||
const uint32_t n_rem = n % kVecN;
|
||||
|
||||
constexpr uint32_t kBlock = 256;
|
||||
const uint32_t grid = div_ceil(std::max(n_vecs, n_rem), kBlock);
|
||||
|
||||
LaunchKernel(grid, kBlock, dev)(
|
||||
silu_gate_kernel<T, kVecN>,
|
||||
static_cast<T*>(dst.data_ptr()),
|
||||
static_cast<const T*>(src.data_ptr()),
|
||||
n_vecs, n_rem, n);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
```
|
||||
|
||||
### Python wrapper: `diffusion/silu_gate.py`
|
||||
|
||||
```python
|
||||
from __future__ import annotations
|
||||
import torch
|
||||
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
|
||||
|
||||
@cache_once
|
||||
def _jit_silu_gate_module(dtype: torch.dtype):
|
||||
args = make_cpp_args(dtype)
|
||||
return load_jit(
|
||||
"diffusion_silu_gate",
|
||||
*args,
|
||||
cuda_files=["diffusion/silu_gate.cuh"],
|
||||
cuda_wrappers=[("silu_gate", f"silu_gate<{args}>")],
|
||||
extra_cuda_cflags=["-O3", "--use_fast_math"],
|
||||
)
|
||||
|
||||
def diffusion_silu_gate(src: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
|
||||
assert src.is_cuda and src.dtype in (torch.float16, torch.bfloat16, torch.float32)
|
||||
if out is None:
|
||||
out = torch.empty_like(src)
|
||||
module = _jit_silu_gate_module(src.dtype)
|
||||
module.silu_gate(out, src)
|
||||
return out
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Template 2: Row-wise Reduction (RMSNorm / LayerNorm)
|
||||
|
||||
Use for ops that reduce across the last dimension of each row.
|
||||
|
||||
### `.cuh` file: `csrc/diffusion/rmsnorm.cuh`
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/type.cuh>
|
||||
#include <sgl_kernel/utils.h>
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
#include <sgl_kernel/vec.cuh>
|
||||
#include <sgl_kernel/warp.cuh>
|
||||
#include <sgl_kernel/math.cuh>
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
|
||||
namespace {
|
||||
|
||||
// RMSNorm: y = x / rms(x) * weight
|
||||
// One block per row; vectorized loads/stores; warp + shared-mem reduction
|
||||
template <typename T, int kVecN>
|
||||
__global__ void rmsnorm_kernel(
|
||||
T* __restrict__ dst,
|
||||
const T* __restrict__ src,
|
||||
const T* __restrict__ weight, // nullptr if no affine weight
|
||||
uint32_t hidden,
|
||||
uint32_t n_vecs,
|
||||
float eps)
|
||||
{
|
||||
using vec_t = device::AlignedVector<T, kVecN>;
|
||||
|
||||
const uint32_t row = blockIdx.x;
|
||||
const T* row_src = src + row * hidden;
|
||||
T* row_dst = dst + row * hidden;
|
||||
|
||||
// Pass 1: sum of squares
|
||||
float sum_sq = 0.f;
|
||||
for (uint32_t vi = threadIdx.x; vi < n_vecs; vi += blockDim.x) {
|
||||
vec_t v;
|
||||
v.load(row_src, vi);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecN; ++i) {
|
||||
float val = static_cast<float>(v[i]);
|
||||
sum_sq += val * val;
|
||||
}
|
||||
}
|
||||
|
||||
// Warp + block reduction
|
||||
sum_sq = device::warp::reduce_sum<float>(sum_sq);
|
||||
__shared__ float smem[32];
|
||||
if (threadIdx.x % 32 == 0) smem[threadIdx.x / 32] = sum_sq;
|
||||
__syncthreads();
|
||||
if (threadIdx.x < 32) {
|
||||
sum_sq = (threadIdx.x < blockDim.x / 32) ? smem[threadIdx.x] : 0.f;
|
||||
sum_sq = device::warp::reduce_sum<float>(sum_sq);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const float rms_inv = device::math::rsqrt<float>(sum_sq / static_cast<float>(hidden) + eps);
|
||||
|
||||
// Pass 2: normalize + optional weight
|
||||
for (uint32_t vi = threadIdx.x; vi < n_vecs; vi += blockDim.x) {
|
||||
vec_t v_in, v_out;
|
||||
v_in.load(row_src, vi);
|
||||
if (weight != nullptr) {
|
||||
vec_t v_w;
|
||||
v_w.load(weight, vi);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecN; ++i)
|
||||
v_out[i] = static_cast<T>(static_cast<float>(v_in[i]) * rms_inv
|
||||
* static_cast<float>(v_w[i]));
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecN; ++i)
|
||||
v_out[i] = static_cast<T>(static_cast<float>(v_in[i]) * rms_inv);
|
||||
}
|
||||
v_out.store(row_dst, vi);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void rmsnorm(
|
||||
tvm::ffi::TensorView dst,
|
||||
tvm::ffi::TensorView src,
|
||||
tvm::ffi::TensorView weight, // data_ptr == nullptr → no weight
|
||||
float eps)
|
||||
{
|
||||
using namespace host;
|
||||
|
||||
SymbolicSize B{"batch_tokens"}, H{"hidden_size"};
|
||||
SymbolicDevice device;
|
||||
device.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({B, H})
|
||||
.with_dtype<T>()
|
||||
.with_device(device)
|
||||
.verify(dst)
|
||||
.verify(src);
|
||||
|
||||
const uint32_t num_rows = static_cast<uint32_t>(B.unwrap());
|
||||
const uint32_t hidden = static_cast<uint32_t>(H.unwrap());
|
||||
const DLDevice dev = device.unwrap();
|
||||
|
||||
constexpr int kVecN = 16 / sizeof(T);
|
||||
RuntimeCheck(hidden % kVecN == 0,
|
||||
"rmsnorm: hidden_size (", hidden, ") must be divisible by ", kVecN);
|
||||
const uint32_t n_vecs = hidden / kVecN;
|
||||
|
||||
uint32_t threads = std::min(n_vecs, 512u);
|
||||
threads = (threads + 31) / 32 * 32;
|
||||
|
||||
const T* w_ptr = (weight.data_ptr() != nullptr)
|
||||
? static_cast<const T*>(weight.data_ptr()) : nullptr;
|
||||
|
||||
LaunchKernel(num_rows, threads, dev)(
|
||||
rmsnorm_kernel<T, kVecN>,
|
||||
static_cast<T*>(dst.data_ptr()),
|
||||
static_cast<const T*>(src.data_ptr()),
|
||||
w_ptr, hidden, n_vecs, eps);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Template 3: Fused Row-Reduction + Element-wise (AdaLN)
|
||||
|
||||
Combines RMSNorm + AdaLN modulation into one pass: `y = norm(x) * (1 + scale) + shift`.
|
||||
|
||||
### `.cuh` file: `csrc/diffusion/adaln.cuh`
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/type.cuh>
|
||||
#include <sgl_kernel/utils.h>
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
#include <sgl_kernel/vec.cuh>
|
||||
#include <sgl_kernel/warp.cuh>
|
||||
#include <sgl_kernel/math.cuh>
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
|
||||
namespace {
|
||||
|
||||
// AdaLN: y = norm(x) * (1 + scale) + shift
|
||||
// scale, shift: [batch, hidden] (one per row)
|
||||
template <typename T, int kVecN>
|
||||
__global__ void adaln_kernel(
|
||||
T* __restrict__ dst,
|
||||
const T* __restrict__ src,
|
||||
const T* __restrict__ weight,
|
||||
const T* __restrict__ scale,
|
||||
const T* __restrict__ shift,
|
||||
uint32_t hidden,
|
||||
uint32_t n_vecs,
|
||||
float eps)
|
||||
{
|
||||
using vec_t = device::AlignedVector<T, kVecN>;
|
||||
|
||||
const uint32_t row = blockIdx.x;
|
||||
const T* row_src = src + row * hidden;
|
||||
const T* row_scale = scale + row * hidden;
|
||||
const T* row_shift = shift + row * hidden;
|
||||
T* row_dst = dst + row * hidden;
|
||||
|
||||
// Pass 1: compute RMS
|
||||
float sum_sq = 0.f;
|
||||
for (uint32_t vi = threadIdx.x; vi < n_vecs; vi += blockDim.x) {
|
||||
vec_t v;
|
||||
v.load(row_src, vi);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecN; ++i) {
|
||||
float val = static_cast<float>(v[i]);
|
||||
sum_sq += val * val;
|
||||
}
|
||||
}
|
||||
sum_sq = device::warp::reduce_sum<float>(sum_sq);
|
||||
__shared__ float smem[32];
|
||||
if (threadIdx.x % 32 == 0) smem[threadIdx.x / 32] = sum_sq;
|
||||
__syncthreads();
|
||||
if (threadIdx.x < 32) {
|
||||
sum_sq = (threadIdx.x < blockDim.x / 32) ? smem[threadIdx.x] : 0.f;
|
||||
sum_sq = device::warp::reduce_sum<float>(sum_sq);
|
||||
}
|
||||
__syncthreads();
|
||||
const float rms_inv = device::math::rsqrt<float>(sum_sq / static_cast<float>(hidden) + eps);
|
||||
|
||||
// Pass 2: normalize + modulate
|
||||
for (uint32_t vi = threadIdx.x; vi < n_vecs; vi += blockDim.x) {
|
||||
vec_t v_in, v_w, v_sc, v_sh, v_out;
|
||||
v_in.load(row_src, vi);
|
||||
v_w.load(weight, vi);
|
||||
v_sc.load(row_scale, vi);
|
||||
v_sh.load(row_shift, vi);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecN; ++i) {
|
||||
float x = static_cast<float>(v_in[i]) * rms_inv * static_cast<float>(v_w[i]);
|
||||
float sc = static_cast<float>(v_sc[i]);
|
||||
float sh = static_cast<float>(v_sh[i]);
|
||||
v_out[i] = static_cast<T>(x * (1.f + sc) + sh);
|
||||
}
|
||||
v_out.store(row_dst, vi);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void adaln(
|
||||
tvm::ffi::TensorView dst,
|
||||
tvm::ffi::TensorView src,
|
||||
tvm::ffi::TensorView weight,
|
||||
tvm::ffi::TensorView scale,
|
||||
tvm::ffi::TensorView shift,
|
||||
float eps)
|
||||
{
|
||||
using namespace host;
|
||||
|
||||
SymbolicSize B{"batch_tokens"}, H{"hidden_size"};
|
||||
SymbolicDevice device;
|
||||
device.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({B, H})
|
||||
.with_dtype<T>()
|
||||
.with_device(device)
|
||||
.verify(dst).verify(src).verify(weight).verify(scale).verify(shift);
|
||||
|
||||
const uint32_t num_rows = static_cast<uint32_t>(B.unwrap());
|
||||
const uint32_t hidden = static_cast<uint32_t>(H.unwrap());
|
||||
const DLDevice dev = device.unwrap();
|
||||
|
||||
constexpr int kVecN = 16 / sizeof(T);
|
||||
RuntimeCheck(hidden % kVecN == 0, "adaln: hidden_size must be divisible by ", kVecN);
|
||||
const uint32_t n_vecs = hidden / kVecN;
|
||||
|
||||
uint32_t threads = std::min(n_vecs, 512u);
|
||||
threads = (threads + 31) / 32 * 32;
|
||||
|
||||
LaunchKernel(num_rows, threads, dev)(
|
||||
adaln_kernel<T, kVecN>,
|
||||
static_cast<T*>(dst.data_ptr()),
|
||||
static_cast<const T*>(src.data_ptr()),
|
||||
static_cast<const T*>(weight.data_ptr()),
|
||||
static_cast<const T*>(scale.data_ptr()),
|
||||
static_cast<const T*>(shift.data_ptr()),
|
||||
hidden, n_vecs, eps);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Template 4: Python Wrapper (generic pattern)
|
||||
|
||||
File location: `python/sglang/jit_kernel/diffusion/<op>.py`
|
||||
|
||||
```python
|
||||
from __future__ import annotations
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_module(dtype: torch.dtype) -> Module:
|
||||
"""Cache key: dtype (and any other template params you need)."""
|
||||
args = make_cpp_args(dtype)
|
||||
return load_jit(
|
||||
"diffusion_your_op", # unique build cache key
|
||||
*args,
|
||||
cuda_files=["diffusion/your_op.cuh"], # relative to csrc/
|
||||
cuda_wrappers=[("your_op", f"your_op<{args}>")],
|
||||
extra_cuda_cflags=["-O3", "--use_fast_math"],
|
||||
)
|
||||
|
||||
|
||||
def diffusion_your_op(
|
||||
src: torch.Tensor,
|
||||
out: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Your op description.
|
||||
|
||||
Supported dtypes: float16, bfloat16, float32.
|
||||
"""
|
||||
assert src.is_cuda, "src must be a CUDA tensor"
|
||||
assert src.dtype in (torch.float16, torch.bfloat16, torch.float32), (
|
||||
f"Unsupported dtype {src.dtype}"
|
||||
)
|
||||
if out is None:
|
||||
out = torch.empty_like(src)
|
||||
|
||||
module = _jit_module(src.dtype)
|
||||
module.your_op(out, src)
|
||||
return out
|
||||
```
|
||||
|
||||
**`make_cpp_args` conversion table:**
|
||||
|
||||
| `torch.dtype` | C++ type |
|
||||
|---------------|----------|
|
||||
| `torch.float16` | `fp16_t` |
|
||||
| `torch.bfloat16` | `bf16_t` |
|
||||
| `torch.float32` | `fp32_t` |
|
||||
|
||||
---
|
||||
|
||||
## Template 5: Correctness Test
|
||||
|
||||
```python
|
||||
# python/sglang/jit_kernel/tests/test_diffusion_<op>.py
|
||||
import pytest
|
||||
import torch
|
||||
from sglang.jit_kernel.diffusion.<op> import diffusion_<op>
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
|
||||
@pytest.mark.parametrize("shape", [(1, 2048), (4, 3072), (16, 4096)])
|
||||
def test_<op>_correctness(dtype, shape):
|
||||
src = torch.randn(*shape, dtype=dtype, device="cuda")
|
||||
|
||||
out_jit = diffusion_<op>(src)
|
||||
ref = reference_<op>(src.float()).to(dtype) # reference in fp32
|
||||
|
||||
tol = {"rtol": 1e-2, "atol": 1e-2} if dtype != torch.float32 else {"rtol": 1e-5, "atol": 1e-6}
|
||||
torch.testing.assert_close(out_jit, ref, **tol)
|
||||
|
||||
|
||||
def test_<op>_out_param():
|
||||
src = torch.randn(1024, 2048, dtype=torch.bfloat16, device="cuda")
|
||||
out = torch.empty_like(src)
|
||||
result = diffusion_<op>(src, out=out)
|
||||
assert result is out
|
||||
|
||||
|
||||
def test_<op>_cpu_error():
|
||||
src = torch.randn(128, dtype=torch.float16) # CPU tensor
|
||||
with pytest.raises(AssertionError):
|
||||
diffusion_<op>(src)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v", "-s"])
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Template 6: Benchmark
|
||||
|
||||
```python
|
||||
# python/sglang/jit_kernel/benchmark/bench_diffusion_<op>.py
|
||||
import torch
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import DEFAULT_DEVICE, DEFAULT_DTYPE, run_benchmark
|
||||
from sglang.jit_kernel.diffusion.<op> import diffusion_<op>
|
||||
|
||||
SHAPES = [(4096, 2048), (4096, 3072), (4096, 4096)]
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["hidden"],
|
||||
x_vals=[s[1] for s in SHAPES],
|
||||
line_arg="provider",
|
||||
line_vals=["jit_cuda", "torch"],
|
||||
line_names=["SGLang JIT CUDA", "PyTorch"],
|
||||
styles=[("blue", "-"), ("red", "--")],
|
||||
ylabel="us",
|
||||
plot_name="diffusion-<op>",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(hidden: int, provider: str):
|
||||
src = torch.randn(4096, hidden, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE)
|
||||
|
||||
if provider == "jit_cuda":
|
||||
fn = lambda: diffusion_<op>(src)
|
||||
else:
|
||||
fn = lambda: reference_<op>(src) # torch baseline
|
||||
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark.run(print_data=True)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Summary of New Files per Kernel
|
||||
|
||||
```
|
||||
python/sglang/jit_kernel/csrc/diffusion/
|
||||
└── <op>.cuh # CUDA kernel + launcher
|
||||
|
||||
python/sglang/jit_kernel/diffusion/
|
||||
└── <op>.py # Python wrapper (load_jit + cache_once)
|
||||
|
||||
python/sglang/jit_kernel/tests/
|
||||
└── test_diffusion_<op>.py # correctness tests
|
||||
|
||||
python/sglang/jit_kernel/benchmark/
|
||||
└── bench_diffusion_<op>.py # triton.testing benchmark
|
||||
```
|
||||
|
||||
> See `scripts/bench_diffusion_rmsnorm.py` and `scripts/bench_diffusion_denoise.py` for full runnable examples.
|
||||
@@ -0,0 +1,335 @@
|
||||
# T4 GPU Optimization Guide — SGLang Diffusion JIT Kernels
|
||||
|
||||
T4 is a Turing architecture GPU (GCP n1+T4, AWS g4dn) commonly used for cloud inference.
|
||||
Its key constraint for diffusion kernels: **no BF16 support** — FP16 only.
|
||||
|
||||
> **Adapted from**: [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels)
|
||||
|
||||
---
|
||||
|
||||
## T4 Turing Architecture Overview
|
||||
|
||||
| Component | T4 | A100 | H100 |
|
||||
|-----------|-----|------|------|
|
||||
| Compute Capability | sm_75 | sm_80 | sm_90 |
|
||||
| SMs | 40 | 108 | 132 |
|
||||
| Shared Memory/SM | **64 KB** | 164 KB | 192 KB |
|
||||
| L2 Cache | 4 MB | 40 MB | 50 MB |
|
||||
| Memory Bandwidth | **320 GB/s** | 2.0 TB/s | 3.35 TB/s |
|
||||
| Memory | 16 GB GDDR6 | 40–80 GB HBM2e | 80 GB HBM3 |
|
||||
| Max Threads/SM | **1024** | 2048 | 2048 |
|
||||
| BF16 Support | **No** | Yes | Yes |
|
||||
|
||||
### Critical T4 Constraints
|
||||
|
||||
1. **No BFloat16** — must use FP16 everywhere
|
||||
2. **320 GB/s bandwidth** — ~10x lower than H100; vectorization is critical
|
||||
3. **16 GB memory** — limits model size; use offloading
|
||||
4. **64 KB shared memory/SM** — smaller attention tiles
|
||||
5. **Max 1024 threads/SM** — half of A100/H100; affects occupancy calculations
|
||||
|
||||
---
|
||||
|
||||
## No BF16: Always Use FP16
|
||||
|
||||
This is the most impactful constraint. **Never use `bf16_t` or `__nv_bfloat16` on T4.**
|
||||
|
||||
**Python wrapper guard:**
|
||||
|
||||
```python
|
||||
import torch
|
||||
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
|
||||
|
||||
@cache_once
|
||||
def _jit_rmsnorm_module(dtype: torch.dtype):
|
||||
# T4 (sm_75) does not support BF16
|
||||
cap = torch.cuda.get_device_capability()
|
||||
if cap < (8, 0) and dtype == torch.bfloat16:
|
||||
raise RuntimeError(
|
||||
f"T4 (sm_75) does not support BF16. Use torch.float16 instead. "
|
||||
f"Got dtype={dtype}"
|
||||
)
|
||||
args = make_cpp_args(dtype)
|
||||
return load_jit(
|
||||
"diffusion_rmsnorm",
|
||||
*args,
|
||||
cuda_files=["diffusion/rmsnorm.cuh"],
|
||||
cuda_wrappers=[("rmsnorm", f"rmsnorm<{args}>")],
|
||||
)
|
||||
```
|
||||
|
||||
**Conditional type in kernel:**
|
||||
|
||||
```cuda
|
||||
#if __CUDA_ARCH__ >= 800
|
||||
// A100/H100: BF16 available
|
||||
using DefaultHalf = bf16_t;
|
||||
#else
|
||||
// T4/Turing: FP16 only
|
||||
using DefaultHalf = fp16_t;
|
||||
#endif
|
||||
```
|
||||
|
||||
**Runtime detection helper:**
|
||||
|
||||
```python
|
||||
def get_diffusion_dtype() -> torch.dtype:
|
||||
"""Return the appropriate half-precision dtype for the current GPU."""
|
||||
cap = torch.cuda.get_device_capability()
|
||||
if cap >= (8, 0):
|
||||
return torch.bfloat16 # A100/H100: prefer BF16
|
||||
else:
|
||||
return torch.float16 # T4/older: FP16 only
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Memory Access Optimization
|
||||
|
||||
With only 320 GB/s, **vectorization is more critical on T4 than on A100/H100**.
|
||||
|
||||
### `AlignedVector` (same abstraction, FP16 only)
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/vec.cuh>
|
||||
|
||||
// On T4, T must be fp16_t or fp32_t (NOT bf16_t)
|
||||
constexpr int kVecN = 16 / sizeof(T); // 8 for fp16, 4 for fp32
|
||||
using vec_t = device::AlignedVector<T, kVecN>;
|
||||
```
|
||||
|
||||
**Target bandwidth**: 40–50% of T4's 320 GB/s = 128–160 GB/s.
|
||||
|
||||
### Increase Arithmetic Intensity
|
||||
|
||||
With low bandwidth, fusing ops saves more on T4 than on H100:
|
||||
|
||||
```cpp
|
||||
// BAD on T4: separate passes → 2× memory traffic
|
||||
output1[i] = input[i] * scale; // pass 1
|
||||
output2[i] = output1[i] + bias; // pass 2
|
||||
|
||||
// GOOD: fuse → single memory read, single write
|
||||
float val = static_cast<float>(v[i]);
|
||||
val = val * scale + bias;
|
||||
val = device::math::max<float>(val, 0.f); // ReLU
|
||||
v[i] = static_cast<T>(val);
|
||||
```
|
||||
|
||||
### Expected T4 Performance
|
||||
|
||||
| Kernel | T4 (ms) | A100 (ms) | H100 (ms) | T4 vs H100 |
|
||||
|--------|---------|-----------|-----------|------------|
|
||||
| RMSNorm [2, 1024, 2048] | ~0.5 | ~0.08 | 0.054 | ~9x slower |
|
||||
| GEGLU [2, 1024, 4096] | ~0.3 | ~0.05 | 0.030 | ~10x slower |
|
||||
|
||||
---
|
||||
|
||||
## Shared Memory Configuration
|
||||
|
||||
T4 max: **64 KB/SM**. Use smaller tiles vs A100/H100.
|
||||
|
||||
```cpp
|
||||
// T4: request max shared memory (64 KB)
|
||||
cudaFuncSetAttribute(
|
||||
your_kernel,
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize,
|
||||
64 * 1024
|
||||
);
|
||||
```
|
||||
|
||||
**Attention tile sizes for T4** (halved vs H100):
|
||||
|
||||
```
|
||||
H100/A100: BLOCK_SIZE_M = 128, BLOCK_SIZE_N = 64
|
||||
T4: BLOCK_SIZE_M = 64, BLOCK_SIZE_N = 32 ← reduced for 64 KB limit
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Occupancy Tuning
|
||||
|
||||
T4 max: **1024 threads/SM** (vs 2048 on A100/H100). This halves max occupancy for a given block size.
|
||||
|
||||
**Block sizes for T4:**
|
||||
|
||||
| Kernel Type | Threads/Block | Notes |
|
||||
|-------------|---------------|-------|
|
||||
| Element-wise | 256 | Same as H100 |
|
||||
| Row reduction | 256–512 | Avoid > 512 to fit multiple blocks/SM |
|
||||
| Tiled/attention | 128–256 | Small tiles due to 64 KB shared mem |
|
||||
|
||||
**Grid sizing for T4 (40 SMs)** — `runtime::get_sm_count` handles this automatically:
|
||||
|
||||
```cpp
|
||||
// get_sm_count() returns 40 on T4, 108 on A100, 132 on H100
|
||||
const uint32_t num_sm = host::runtime::get_sm_count(device.unwrap().device_id);
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Numerical Stability with FP16
|
||||
|
||||
FP16 has a smaller dynamic range (±65504) vs BF16 (±3.39×10³⁸). Watch for overflow in attention:
|
||||
|
||||
```cuda
|
||||
// Scale attention scores to prevent FP16 overflow
|
||||
float scale_factor = 1.0f / sqrtf(static_cast<float>(head_dim));
|
||||
// For very long sequences on T4, may need additional scaling:
|
||||
// if (score * scale_factor > 65000.f) { /* clamp */ }
|
||||
```
|
||||
|
||||
Always accumulate in FP32:
|
||||
|
||||
```cpp
|
||||
float acc = 0.f; // FP32 accumulation
|
||||
for (uint32_t vi = threadIdx.x; vi < n_vecs; vi += blockDim.x) {
|
||||
vec_t v;
|
||||
v.load(src, vi);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecN; ++i) {
|
||||
float val = static_cast<float>(v[i]); // fp16 → fp32
|
||||
acc += val * val;
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Memory Management for 16 GB
|
||||
|
||||
T4's 16 GB requires careful planning for large diffusion models.
|
||||
|
||||
**sglang generate flags for T4:**
|
||||
|
||||
```bash
|
||||
# Enable CPU offloading to fit within 16 GB
|
||||
sglang generate \
|
||||
--model-path=black-forest-labs/FLUX.1-dev \
|
||||
--dit-cpu-offload true \ # DiT weights to CPU
|
||||
--text-encoder-cpu-offload true \
|
||||
--vae-cpu-offload true \
|
||||
--width=512 --height=512 \ # Reduce resolution
|
||||
--num-inference-steps=20 \ # Fewer steps
|
||||
--seed=42
|
||||
```
|
||||
|
||||
**Resolution recommendations for T4:**
|
||||
|
||||
| Model | H100/A100 | T4 |
|
||||
|-------|-----------|-----|
|
||||
| FLUX.1-dev | 1024×1024 | 512×512 |
|
||||
| Wan2.2-TI2V-5B | 720P | 480P |
|
||||
| FLUX.2-dev | 1024×1024 | 512×512 |
|
||||
|
||||
---
|
||||
|
||||
## JIT Compilation for T4
|
||||
|
||||
```python
|
||||
return load_jit(
|
||||
"my_kernel",
|
||||
*args,
|
||||
cuda_files=["diffusion/my_kernel.cuh"],
|
||||
cuda_wrappers=[("my_kernel", f"my_kernel<{args}>")],
|
||||
extra_cuda_cflags=[
|
||||
"-O3",
|
||||
"--use_fast_math",
|
||||
"-arch=sm_75", # T4 only; omit for multi-arch
|
||||
],
|
||||
)
|
||||
```
|
||||
|
||||
**Multi-arch (T4 + A100 + H100):**
|
||||
|
||||
```python
|
||||
extra_cuda_cflags=[
|
||||
"-O3", "--use_fast_math",
|
||||
"-gencode=arch=compute_75,code=sm_75", # T4
|
||||
"-gencode=arch=compute_80,code=sm_80", # A100
|
||||
"-gencode=arch=compute_90,code=sm_90", # H100
|
||||
]
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## H100/A100 → T4 Migration Checklist
|
||||
|
||||
| Item | H100/A100 | T4 | Action |
|
||||
|------|-----------|-----|--------|
|
||||
| BF16 | Available | **Not available** | Replace `bf16_t` with `fp16_t`; guard in Python wrapper |
|
||||
| Shared memory | 164–192 KB | **64 KB** | Halve tile sizes |
|
||||
| Grid sizing | ×108/132 SMs | ×40 SMs | `get_sm_count()` auto-handles |
|
||||
| Max threads/SM | 2048 | **1024** | Don't exceed 512 threads/block |
|
||||
| Memory | 40–80 GB | **16 GB** | Enable CPU offloading |
|
||||
| cp.async | Available | No (Turing has limited async) | Remove async copy patterns |
|
||||
| `AlignedVector` | Same | Same | No changes |
|
||||
| `warp::reduce_sum` | Same | Same | No changes |
|
||||
|
||||
---
|
||||
|
||||
## Performance Profiling
|
||||
|
||||
### NVIDIA Nsight Systems (nsys)
|
||||
|
||||
```bash
|
||||
nsys profile -o t4_profile python scripts/bench_diffusion_rmsnorm.py
|
||||
|
||||
# Key metrics to watch:
|
||||
# - Kernel duration
|
||||
# - Memory transfer time
|
||||
# - GPU idle time
|
||||
# - Stream utilization
|
||||
```
|
||||
|
||||
### NVIDIA Nsight Compute (ncu)
|
||||
|
||||
```bash
|
||||
# Full metrics
|
||||
ncu --set full -o t4_metrics.ncu-rep \
|
||||
python scripts/bench_diffusion_rmsnorm.py
|
||||
|
||||
# Specific metrics — T4 is memory-bound; focus on dram throughput
|
||||
ncu --metrics sm__throughput.avg.pct_of_peak_sustained_elapsed,\
|
||||
dram__throughput.avg.pct_of_peak_sustained_elapsed \
|
||||
python scripts/bench_diffusion_rmsnorm.py
|
||||
|
||||
# Key metrics for T4 diffusion kernels:
|
||||
# - Memory throughput (dram__throughput.avg.pct_of_peak_sustained_elapsed)
|
||||
# → Target: 40–50% of 320 GB/s (128–160 GB/s) for vectorized kernels
|
||||
# - SM utilization (sm__throughput.avg.pct_of_peak_sustained_elapsed)
|
||||
# → Target high with only 40 SMs
|
||||
# - Achieved occupancy (sm__warps_active.avg.pct_of_peak_sustained_active)
|
||||
# → Max 1024 threads/SM on T4 — block size ≤ 512 for decent occupancy
|
||||
# - Warp stall reasons (smsp__warp_issue_stalled_*.avg.pct_of_peak_sustained_active)
|
||||
```
|
||||
|
||||
### Common T4 Bottlenecks
|
||||
|
||||
1. **Memory Bandwidth** — 320 GB/s is the primary limit; if `dram__throughput` < 40% → use `AlignedVector`
|
||||
2. **Limited Memory** — 16 GB; enable `--dit-cpu-offload`/`--vae-cpu-offload` as needed
|
||||
3. **No BF16** — guard in Python wrapper; FP16 overflow risk in long-sequence attention
|
||||
4. **Smaller tiles** — 64 KB shared memory; reduce `BLOCK_SIZE_M/N` vs H100
|
||||
|
||||
---
|
||||
|
||||
## Best Practices Summary (T4)
|
||||
|
||||
1. **No BF16**: Guard in Python wrapper, raise clear error
|
||||
2. **Vectorization**: Even more critical at 320 GB/s — always use `AlignedVector`
|
||||
3. **Tile sizes**: 64 KB shared memory limit → halve BLOCK_SIZE vs H100
|
||||
4. **Block size**: Max 512 threads/block for decent occupancy (max 1024 threads/SM)
|
||||
5. **Grid sizing**: 40 SMs — `runtime::get_sm_count()` auto-handles
|
||||
6. **FP32 accumulation**: Always accumulate in FP32 to avoid FP16 overflow
|
||||
7. **Memory**: Plan for 16 GB; use `--dit-cpu-offload`/`--vae-cpu-offload` as needed
|
||||
8. **Fuse more**: Low bandwidth makes kernel fusion more impactful than on H100
|
||||
9. **Multi-arch build**: Always build for `sm_75,sm_80,sm_90` together
|
||||
|
||||
## T4 Cloud Instance Quick Reference
|
||||
|
||||
| Provider | Instance | Notes |
|
||||
|----------|----------|-------|
|
||||
| GCP | n1-standard-4 + T4 | Most common inference setup |
|
||||
| AWS | g4dn.xlarge | 1× T4, 16 GB |
|
||||
| AWS | g4dn.12xlarge | 4× T4, 64 GB total |
|
||||
| Azure | NC4as T4 v3 | 1× T4 |
|
||||
@@ -0,0 +1,328 @@
|
||||
# Troubleshooting Guide — SGLang Diffusion JIT CUDA Kernels
|
||||
|
||||
Common issues and solutions when writing and integrating JIT CUDA kernels for SGLang Diffusion.
|
||||
|
||||
> **Adapted from**: [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels)
|
||||
|
||||
---
|
||||
|
||||
## Build / Compile Issues
|
||||
|
||||
### 1. JIT compilation fails: "No such file or directory"
|
||||
|
||||
**Problem:** `load_jit` cannot find your `.cuh` file.
|
||||
|
||||
```
|
||||
FileNotFoundError: .../jit_kernel/csrc/diffusion/your_op.cuh not found
|
||||
```
|
||||
|
||||
**Fix:** Ensure the file is under `python/sglang/jit_kernel/csrc/diffusion/`. The path passed to `cuda_files` is relative to `csrc/`:
|
||||
|
||||
```python
|
||||
# CORRECT — file lives at csrc/diffusion/your_op.cuh
|
||||
load_jit(..., cuda_files=["diffusion/your_op.cuh"])
|
||||
# resolves to: python/sglang/jit_kernel/csrc/diffusion/your_op.cuh
|
||||
|
||||
# ALSO CORRECT — absolute path (pathlib replaces the csrc/ prefix)
|
||||
load_jit(..., cuda_files=["/full/absolute/path/to/your_op.cuh"])
|
||||
```
|
||||
|
||||
### 2. Type conversion errors (FP16/BF16)
|
||||
|
||||
**Problem:** Implicit FP16/BF16 conversion fails because PyTorch compiles with `-D__CUDA_NO_HALF_OPERATORS__`:
|
||||
|
||||
```
|
||||
error: no suitable conversion function from "__half" to "float" exists
|
||||
```
|
||||
|
||||
**Fix:** SGLang's `static_cast<float>` works because `fp16_t` and `bf16_t` are typedef'd with proper conversion operators. Always use explicit casts:
|
||||
|
||||
```cpp
|
||||
// CORRECT — explicit cast
|
||||
float val = static_cast<float>(v[i]); // fp16_t / bf16_t → float
|
||||
v[i] = static_cast<T>(fp32_result); // float → T
|
||||
|
||||
// WRONG — implicit conversion (disabled by PyTorch build flags)
|
||||
float val = v[i]; // compile error
|
||||
v[i] = fp32_result; // compile error
|
||||
```
|
||||
|
||||
If you need the raw intrinsics for packed types:
|
||||
```cpp
|
||||
// bf16x2_t → two floats
|
||||
bf16x2_t packed = ...;
|
||||
float v0 = __bfloat162float(packed.x);
|
||||
float v1 = __bfloat162float(packed.y);
|
||||
```
|
||||
|
||||
### 3. Template instantiation explodes / slow first compile
|
||||
|
||||
**Problem:** Many template combinations makes the first JIT compile very slow.
|
||||
|
||||
**Fix:** Reduce template argument combinations. Move compile-time constants to runtime if they don't affect performance critically:
|
||||
|
||||
```cpp
|
||||
// Fewer template args = fewer instantiations
|
||||
template <typename T> // only dtype varies
|
||||
void my_op(tvm::ffi::TensorView dst, tvm::ffi::TensorView src, int block_size);
|
||||
```
|
||||
|
||||
### 4. SM check: kernel requires sm_90 but device is sm_80
|
||||
|
||||
**Problem:** Kernel uses H100-only features on A100.
|
||||
|
||||
**Fix:** Add a Python guard before calling `load_jit`:
|
||||
|
||||
```python
|
||||
cap = torch.cuda.get_device_capability()
|
||||
if cap[0] < 9:
|
||||
raise RuntimeError(
|
||||
f"This kernel requires H100 (sm_90+). "
|
||||
f"Got compute capability {cap[0]}.{cap[1]}. "
|
||||
f"Use the Triton fallback instead: diffusion_triton_<op>()"
|
||||
)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Performance Issues
|
||||
|
||||
### 5. Kernel is slower than Triton / PyTorch baseline
|
||||
|
||||
**Steps to diagnose:**
|
||||
|
||||
1. Check dtype: are you using `bf16_t` on T4? (T4 has no BF16 — silently falls back to slow emulation)
|
||||
2. Check vectorization: is `hidden_size` divisible by `kVecN = 16/sizeof(T)` (8 for bf16, 4 for fp32)?
|
||||
3. Profile with `ncu`:
|
||||
```bash
|
||||
ncu --set full --csv -o metrics.csv \
|
||||
python -c "from sglang.jit_kernel.diffusion.rmsnorm import diffusion_rmsnorm; ..."
|
||||
```
|
||||
Look at `dram__throughput.avg.pct_of_peak_sustained_elapsed` — if < 30%, check coalescing.
|
||||
|
||||
4. Check occupancy: run with `--ptxas-options=-v` in `extra_cuda_cflags` to see register usage.
|
||||
|
||||
### 6. Shared memory bank conflicts
|
||||
|
||||
**Problem:** `ncu` reports high `l1tex__data_bank_conflicts_pipe_lmem_op_st.sum`.
|
||||
|
||||
**Fix:** Add padding to shared memory arrays:
|
||||
|
||||
```cpp
|
||||
// Conflict (all threads hit same bank when stride=32)
|
||||
__shared__ float data[32][32];
|
||||
|
||||
// Fixed with padding
|
||||
__shared__ float data[32][33]; // 33 instead of 32
|
||||
```
|
||||
|
||||
### 7. Low occupancy from too many registers
|
||||
|
||||
**Problem:** `nvcc --ptxas-options=-v` shows high register count; occupancy < 25%.
|
||||
|
||||
**Fix:** Add `--maxrregcount=N` to limit registers:
|
||||
|
||||
```python
|
||||
extra_cuda_cflags=["-O3", "--use_fast_math", "--maxrregcount=64"]
|
||||
```
|
||||
|
||||
Reduces registers per thread at the cost of possible register spilling to local memory.
|
||||
|
||||
---
|
||||
|
||||
## Integration Issues
|
||||
|
||||
### 8. RMSNorm weight is None (`elementwise_affine=False`)
|
||||
|
||||
**Problem:**
|
||||
```
|
||||
AttributeError: 'NoneType' object has no attribute 'data_ptr'
|
||||
```
|
||||
|
||||
**Root Cause:** DiT transformer blocks often use `RMSNorm(dim, elementwise_affine=False)` — no learnable weight.
|
||||
|
||||
**Fix in Python wrapper:** pass an empty tensor when weight is absent; the kernel launcher checks `data_ptr == nullptr`:
|
||||
|
||||
```python
|
||||
w = weight if weight is not None else torch.empty(0, dtype=src.dtype, device=src.device)
|
||||
module.rmsnorm(out, src, w, eps)
|
||||
```
|
||||
|
||||
**Fix in `.cuh` launcher:**
|
||||
|
||||
```cpp
|
||||
const T* w_ptr = (weight.data_ptr() != nullptr)
|
||||
? static_cast<const T*>(weight.data_ptr()) : nullptr;
|
||||
// ... pass w_ptr to kernel ...
|
||||
```
|
||||
|
||||
**Fix in module patching:**
|
||||
|
||||
```python
|
||||
has_weight = hasattr(module, "weight") and module.weight is not None
|
||||
if has_weight:
|
||||
def _fwd(mod, eps):
|
||||
def forward(x): return diffusion_rmsnorm(x, weight=mod.weight, eps=eps)
|
||||
return forward
|
||||
module.forward = _fwd(module, module.eps)
|
||||
else:
|
||||
def _fwd_noweight(eps):
|
||||
def forward(x): return diffusion_rmsnorm(x, weight=None, eps=eps)
|
||||
return forward
|
||||
module.forward = _fwd_noweight(module.eps)
|
||||
```
|
||||
|
||||
### 9. `isinstance(module, torch.nn.RMSNorm)` misses diffusion variants
|
||||
|
||||
**Problem:** Patching doesn't apply because diffusers / sglang diffusion models define their own `RMSNorm` class that is **not** a subclass of `torch.nn.RMSNorm`.
|
||||
|
||||
**Fix:** Match by class name string:
|
||||
|
||||
```python
|
||||
# WRONG — misses diffusers/sglang RMSNorm
|
||||
if isinstance(module, torch.nn.RMSNorm):
|
||||
|
||||
# CORRECT — catches all variants
|
||||
if type(module).__name__ == "RMSNorm":
|
||||
# or for broader matching:
|
||||
if "RMSNorm" in type(module).__name__:
|
||||
```
|
||||
|
||||
### 10. Kernel patching doesn't persist after CPU offloading
|
||||
|
||||
**Problem:** After calling `pipe.enable_model_cpu_offload()`, patched modules revert.
|
||||
|
||||
**Fix:** Always inject **after** moving to CUDA, **before** enabling any offloading:
|
||||
|
||||
```python
|
||||
pipe = load_pipeline(...)
|
||||
pipe.to("cuda") # 1. Move to CUDA
|
||||
inject_optimized_kernels(pipe) # 2. Patch modules
|
||||
pipe.enable_model_cpu_offload() # 3. Now safe to enable offloading
|
||||
```
|
||||
|
||||
### 11. Kernel patched after `torch.compile`
|
||||
|
||||
**Problem:** Module is already compiled; patching its `forward` after compilation has no effect.
|
||||
|
||||
**Fix:** Apply patches **before** any `torch.compile` call:
|
||||
|
||||
```python
|
||||
inject_optimized_kernels(pipe) # FIRST: patch
|
||||
pipe.transformer = torch.compile(...) # SECOND: compile
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## `torch.compile` Compatibility
|
||||
|
||||
### 12. Custom CUDA kernel causes graph break
|
||||
|
||||
**Problem:**
|
||||
```
|
||||
torch._dynamo.exc.Unsupported: Attempted to call function marked as skipped
|
||||
```
|
||||
or:
|
||||
```
|
||||
torch._dynamo.exc.TorchRuntimeError: Cannot access data pointer of Tensor (FakeTensor)
|
||||
```
|
||||
|
||||
**Root Cause:** `torch.compile` traces with "fake tensors" that have no real data. Any kernel that calls `.data_ptr()` during tracing fails.
|
||||
|
||||
**Options:**
|
||||
|
||||
**Option A (simplest):** Don't use `torch.compile` with CUDA JIT kernels — use Triton instead:
|
||||
```python
|
||||
# Triton kernels are torch.compile compatible
|
||||
from sglang.jit_kernel.diffusion.triton.norm import fused_rmsnorm
|
||||
```
|
||||
|
||||
**Option B:** Register as a `@torch.library.custom_op` (advanced):
|
||||
```python
|
||||
import torch
|
||||
|
||||
@torch.library.custom_op("diffusion_jit::rmsnorm", mutates_args={"out"})
|
||||
def _rmsnorm_op(out: torch.Tensor, src: torch.Tensor,
|
||||
weight: torch.Tensor, eps: float) -> None:
|
||||
module = _jit_rmsnorm_module(src.dtype)
|
||||
module.rmsnorm(out, src, weight, eps)
|
||||
|
||||
@_rmsnorm_op.register_fake
|
||||
def _(out, src, weight, eps):
|
||||
pass # no shape changes; output already allocated in 'out'
|
||||
```
|
||||
|
||||
**Performance trade-off:**
|
||||
|
||||
| Approach | Speedup (denoise) | torch.compile | Notes |
|
||||
|----------|-------------------|---------------|-------|
|
||||
| CUDA JIT kernel | best | Yes (via `torch.library.custom_op`) | Performance-optimal regardless of whether `torch.compile` is enabled; use `custom_op` + `register_fake` for compile compatibility |
|
||||
| Triton kernel | good | Yes | Use when you need faster iteration/portability, or when you do not have a well-tuned CUDA kernel yet |
|
||||
| Triton + compile | good | Yes | Use for end-to-end `torch.compile` integration convenience; typically slower than a well-tuned CUDA kernel |
|
||||
|
||||
### 13. Unstable benchmark results from JIT timing
|
||||
|
||||
**Problem:** First few runs are slow due to JIT compilation; timing is noisy.
|
||||
|
||||
**Fix:** Use `triton.testing.do_bench` / `run_benchmark` which use CUDA-graph-based timing automatically. Always do a warmup run first:
|
||||
|
||||
```python
|
||||
# Pre-compile by running once before timing
|
||||
diffusion_rmsnorm(dummy_src, weight=dummy_w, eps=1e-6)
|
||||
torch.cuda.synchronize()
|
||||
# Now time
|
||||
result = run_benchmark(lambda: diffusion_rmsnorm(src, weight=w, eps=1e-6))
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Debugging Checklist
|
||||
|
||||
```bash
|
||||
# 1. Verify CUDA device and compute capability
|
||||
python -c "import torch; print(torch.cuda.get_device_name(), torch.cuda.get_device_capability())"
|
||||
|
||||
# 2. Force synchronous CUDA execution to get real error location
|
||||
CUDA_LAUNCH_BLOCKING=1 python scripts/bench_diffusion_rmsnorm.py
|
||||
|
||||
# 3. Run memory sanitizer to catch illegal accesses
|
||||
compute-sanitizer --tool memcheck python scripts/bench_diffusion_rmsnorm.py
|
||||
|
||||
# 4. Check register and shared memory usage
|
||||
# Add to extra_cuda_cflags: "--ptxas-options=-v"
|
||||
|
||||
# 5a. Kernel-level profiling — full metrics
|
||||
ncu --set full -o metrics.ncu-rep \
|
||||
python scripts/bench_diffusion_rmsnorm.py
|
||||
|
||||
# 5b. Kernel-level profiling — targeted bandwidth + occupancy check
|
||||
ncu --metrics sm__throughput.avg.pct_of_peak_sustained_elapsed,\
|
||||
dram__throughput.avg.pct_of_peak_sustained_elapsed \
|
||||
python scripts/bench_diffusion_rmsnorm.py
|
||||
|
||||
# Key metrics to interpret:
|
||||
# - sm__throughput : compute utilization % of peak
|
||||
# - dram__throughput: memory bandwidth % of peak (target ≥ 30% on H100/A100)
|
||||
# - smsp__warp_issue_stalled_*: warp stall breakdown (memory_dependency / math_pipe)
|
||||
|
||||
# 6. System-level profiling (per-op breakdown inside sglang generate)
|
||||
nsys profile -o denoise_profile \
|
||||
sglang generate --model-path=black-forest-labs/FLUX.1-dev \
|
||||
--width=1024 --height=1024 --num-inference-steps=50 \
|
||||
--seed=42 --enable-torch-compile --warmup
|
||||
|
||||
# 7. Verify a patched module produces correct output
|
||||
python - << 'EOF'
|
||||
import torch
|
||||
from sglang.jit_kernel.diffusion.rmsnorm import diffusion_rmsnorm
|
||||
|
||||
x = torch.randn(4, 2048, dtype=torch.bfloat16, device="cuda")
|
||||
w = torch.ones(2048, dtype=torch.bfloat16, device="cuda")
|
||||
|
||||
out_jit = diffusion_rmsnorm(x, weight=w, eps=1e-6)
|
||||
out_ref = torch.nn.functional.rms_norm(x.float(), (2048,), w.float(), eps=1e-6).to(torch.bfloat16)
|
||||
|
||||
max_diff = (out_jit - out_ref).abs().max().item()
|
||||
print(f"Max diff: {max_diff:.2e} ({'PASS' if max_diff < 0.02 else 'FAIL'})")
|
||||
EOF
|
||||
```
|
||||
@@ -0,0 +1,475 @@
|
||||
"""
|
||||
End-to-end denoise-stage benchmark for SGLang Diffusion with/without custom JIT CUDA kernels.
|
||||
|
||||
Measures denoise latency (primary metric ★) and peak GPU memory.
|
||||
All model configs are kept in exact sync with diffusion-benchmark-and-profile.md.
|
||||
|
||||
Adapted from: https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels
|
||||
|
||||
Usage:
|
||||
# Baseline — single model
|
||||
python scripts/bench_diffusion_denoise.py --model flux
|
||||
|
||||
# With custom JIT CUDA kernels
|
||||
python scripts/bench_diffusion_denoise.py --model flux --custom-kernels
|
||||
|
||||
# Side-by-side comparison
|
||||
python scripts/bench_diffusion_denoise.py --model flux --compare
|
||||
|
||||
# All 7 models, comparison
|
||||
python scripts/bench_diffusion_denoise.py --all --compare
|
||||
|
||||
Input images required for image-guided models:
|
||||
mkdir -p /workspace/gen_benchmark/figs
|
||||
wget -O /workspace/gen_benchmark/figs/cat.png \
|
||||
https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png
|
||||
wget -O /workspace/gen_benchmark/figs/astronaut.jpg \
|
||||
https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Model configs — kept in exact sync with diffusion-benchmark-and-profile.md
|
||||
# Each entry produces the same `sglang generate` command as shown in that doc.
|
||||
# ---------------------------------------------------------------------------
|
||||
MODELS = {
|
||||
# 1. Qwen/Qwen-Image-2512 — Text-to-Image, 1024×1024, 50 steps
|
||||
"qwen": {
|
||||
"path": "Qwen/Qwen-Image-2512",
|
||||
"prompt": "A futuristic cyberpunk city at night, neon lights reflecting on wet streets, highly detailed, 8k",
|
||||
"negative_prompt": " ",
|
||||
"extra_args": [
|
||||
"--width=1024",
|
||||
"--height=1024",
|
||||
"--num-inference-steps=50",
|
||||
"--guidance-scale=4.0",
|
||||
"--dit-cpu-offload",
|
||||
"false",
|
||||
"--text-encoder-cpu-offload",
|
||||
"false",
|
||||
],
|
||||
},
|
||||
# 2. Qwen/Qwen-Image-Edit-2511 — Image Editing, 1024×1024, 50 steps
|
||||
# Requires: /workspace/gen_benchmark/figs/cat.png
|
||||
"qwen-edit": {
|
||||
"path": "Qwen/Qwen-Image-Edit-2511",
|
||||
"prompt": "Transform into anime style",
|
||||
"negative_prompt": " ",
|
||||
"image_path": "/workspace/gen_benchmark/figs/cat.png",
|
||||
"extra_args": [
|
||||
"--width=1024",
|
||||
"--height=1024",
|
||||
"--num-inference-steps=50",
|
||||
"--guidance-scale=4.0",
|
||||
"--dit-cpu-offload",
|
||||
"false",
|
||||
"--text-encoder-cpu-offload",
|
||||
"false",
|
||||
],
|
||||
},
|
||||
# 3. black-forest-labs/FLUX.1-dev — Text-to-Image, 1024×1024, 50 steps
|
||||
"flux": {
|
||||
"path": "black-forest-labs/FLUX.1-dev",
|
||||
"prompt": "A futuristic cyberpunk city at night, neon lights reflecting on wet streets, highly detailed, 8k",
|
||||
"extra_args": [
|
||||
"--width=1024",
|
||||
"--height=1024",
|
||||
"--num-inference-steps=50",
|
||||
"--guidance-scale=4.0",
|
||||
],
|
||||
},
|
||||
# 4. black-forest-labs/FLUX.2-dev — Text-to-Image, 1024×1024
|
||||
"flux2": {
|
||||
"path": "black-forest-labs/FLUX.2-dev",
|
||||
"prompt": "A Logo With Bold Large Text: SGL Diffusion",
|
||||
"extra_args": [
|
||||
"--width=1024",
|
||||
"--height=1024",
|
||||
"--dit-layerwise-offload",
|
||||
"false",
|
||||
"--dit-cpu-offload",
|
||||
"false",
|
||||
"--text-encoder-cpu-offload",
|
||||
"true",
|
||||
"--vae-cpu-offload",
|
||||
"false",
|
||||
],
|
||||
},
|
||||
# 5. Tongyi-MAI/Z-Image-Turbo — Turbo Text-to-Image, 1024×1024, 9 steps
|
||||
"zimage": {
|
||||
"path": "Tongyi-MAI/Z-Image-Turbo",
|
||||
"prompt": "A fantasy landscape with mountains and a river, detailed, vibrant colors",
|
||||
"extra_args": [
|
||||
"--width=1024",
|
||||
"--height=1024",
|
||||
"--num-inference-steps=9",
|
||||
"--guidance-scale=0.0",
|
||||
"--dit-cpu-offload",
|
||||
"false",
|
||||
"--text-encoder-cpu-offload",
|
||||
"false",
|
||||
],
|
||||
},
|
||||
# 6. Wan-AI/Wan2.2-T2V-A14B-Diffusers — Text-to-Video, 720P, 8 GPUs, 81 frames, 40 steps
|
||||
"wan-t2v": {
|
||||
"path": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
"prompt": "A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon.",
|
||||
"negative_prompt": " ",
|
||||
"extra_args": [
|
||||
"--720p",
|
||||
"--num-inference-steps=40",
|
||||
"--num-frames=81",
|
||||
"--guidance-scale=5.0",
|
||||
"--num-gpus=8",
|
||||
"--enable-cfg-parallel",
|
||||
"--ulysses-degree=4",
|
||||
"--dit-layerwise-offload",
|
||||
"true",
|
||||
"--dit-cpu-offload",
|
||||
"false",
|
||||
"--vae-cpu-offload",
|
||||
"false",
|
||||
"--text-encoder-cpu-offload",
|
||||
"true",
|
||||
],
|
||||
},
|
||||
# 7. Wan-AI/Wan2.2-TI2V-5B-Diffusers — Text-Image-to-Video, 720P, 1 GPU, 81 frames, 50 steps
|
||||
# Requires: /workspace/gen_benchmark/figs/astronaut.jpg
|
||||
"wan-ti2v": {
|
||||
"path": "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
|
||||
"prompt": "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot.",
|
||||
"negative_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards",
|
||||
"image_path": "/workspace/gen_benchmark/figs/astronaut.jpg",
|
||||
"extra_args": [
|
||||
"--num-frames",
|
||||
"81",
|
||||
"--720p",
|
||||
"--num-inference-steps",
|
||||
"50",
|
||||
"--guidance-scale",
|
||||
"5.0",
|
||||
"--dit-layerwise-offload",
|
||||
"false",
|
||||
"--dit-cpu-offload",
|
||||
"false",
|
||||
"--vae-cpu-offload",
|
||||
"false",
|
||||
"--text-encoder-cpu-offload",
|
||||
"false",
|
||||
],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def build_sglang_cmd(
|
||||
model_key: str,
|
||||
use_custom_kernels: bool,
|
||||
perf_dump_path: Optional[str] = None,
|
||||
warmup: bool = True,
|
||||
torch_compile: bool = True,
|
||||
seed: int = 42,
|
||||
save_output: bool = True,
|
||||
) -> list[str]:
|
||||
"""
|
||||
Build the `sglang generate` command for the given model.
|
||||
Matches the commands in diffusion-benchmark-and-profile.md exactly.
|
||||
"""
|
||||
cfg = MODELS[model_key]
|
||||
|
||||
cmd = [
|
||||
"sglang",
|
||||
"generate",
|
||||
f"--model-path={cfg['path']}",
|
||||
f"--prompt={cfg['prompt']}",
|
||||
f"--seed={seed}",
|
||||
"--log-level=info",
|
||||
]
|
||||
|
||||
if "negative_prompt" in cfg:
|
||||
cmd.append(f"--negative-prompt={cfg['negative_prompt']}")
|
||||
|
||||
if "image_path" in cfg:
|
||||
cmd.append(f"--image-path={cfg['image_path']}")
|
||||
|
||||
cmd.extend(cfg["extra_args"])
|
||||
|
||||
if save_output:
|
||||
cmd.append("--save-output")
|
||||
if warmup:
|
||||
cmd.append("--warmup")
|
||||
if torch_compile:
|
||||
cmd.append("--enable-torch-compile")
|
||||
if perf_dump_path:
|
||||
cmd.extend(["--perf-dump-path", perf_dump_path])
|
||||
|
||||
return cmd
|
||||
|
||||
|
||||
def run_benchmark_once(
|
||||
model_key: str,
|
||||
use_custom_kernels: bool,
|
||||
output_dir: Path,
|
||||
warmup: bool = True,
|
||||
) -> dict:
|
||||
"""Run a single benchmark pass and return results dict."""
|
||||
label = "custom" if use_custom_kernels else "baseline"
|
||||
perf_path = output_dir / f"{model_key}_{label}.json"
|
||||
|
||||
cmd = build_sglang_cmd(
|
||||
model_key,
|
||||
use_custom_kernels=use_custom_kernels,
|
||||
perf_dump_path=str(perf_path),
|
||||
warmup=warmup,
|
||||
)
|
||||
|
||||
env = os.environ.copy()
|
||||
if use_custom_kernels:
|
||||
# NOTE: This env var is a convention for user-implemented kernel injection
|
||||
# logic. SGLang runtime does not read it by default — you must add handling
|
||||
# in your denoising stage or model code to check this var and apply patches.
|
||||
env["SGLANG_DIFFUSION_CUSTOM_CUDA_KERNELS"] = "1"
|
||||
|
||||
print(f"\n{'=' * 64}")
|
||||
print(f"[{label.upper()}] {model_key}")
|
||||
print(" " + " \\\n ".join(cmd))
|
||||
print()
|
||||
|
||||
t0 = time.time()
|
||||
result = subprocess.run(cmd, env=env, text=True)
|
||||
elapsed = time.time() - t0
|
||||
|
||||
if result.returncode != 0:
|
||||
print(f" ERROR: exit code {result.returncode}")
|
||||
return {"model": model_key, "label": label, "error": True, "elapsed_s": elapsed}
|
||||
|
||||
metrics = {"model": model_key, "label": label, "elapsed_s": elapsed, "error": False}
|
||||
if perf_path.exists():
|
||||
try:
|
||||
with open(perf_path) as f:
|
||||
perf = json.load(f)
|
||||
|
||||
# e2e latency: total_duration_ms (set by PerformanceLogger.dump_benchmark_report)
|
||||
total_ms = perf.get("total_duration_ms")
|
||||
metrics["e2e_latency_s"] = (
|
||||
float(total_ms) / 1000.0 if total_ms is not None else None
|
||||
)
|
||||
|
||||
# denoise latency: look in "steps" list for the "DenoisingStage" entry
|
||||
# steps = [{"name": "DenoisingStage", "duration_ms": 1234.5}, ...]
|
||||
denoise_latency_s = None
|
||||
for step in perf.get("steps", []):
|
||||
if (
|
||||
step.get("name") == "DenoisingStage"
|
||||
and step.get("duration_ms") is not None
|
||||
):
|
||||
denoise_latency_s = float(step["duration_ms"]) / 1000.0
|
||||
break
|
||||
|
||||
# fallback: sum all per-step durations from denoise_steps_ms
|
||||
# denoise_steps_ms = [{"step": 0, "duration_ms": 100.5}, ...]
|
||||
if denoise_latency_s is None:
|
||||
denoise_steps = perf.get("denoise_steps_ms", [])
|
||||
if denoise_steps:
|
||||
denoise_latency_s = (
|
||||
sum(s.get("duration_ms", 0.0) for s in denoise_steps) / 1000.0
|
||||
)
|
||||
metrics["denoise_latency_s"] = denoise_latency_s
|
||||
|
||||
# peak memory: max peak_reserved_mb across all memory checkpoints (→ GB)
|
||||
# memory_checkpoints = {"after_DenoisingStage": {"peak_reserved_mb": 12288.0, ...}}
|
||||
peak_memory_gb = None
|
||||
for snapshot in perf.get("memory_checkpoints", {}).values():
|
||||
peak_mb = snapshot.get("peak_reserved_mb")
|
||||
if peak_mb is not None:
|
||||
candidate = float(peak_mb) / 1024.0
|
||||
if peak_memory_gb is None or candidate > peak_memory_gb:
|
||||
peak_memory_gb = candidate
|
||||
metrics["peak_memory_gb"] = peak_memory_gb
|
||||
|
||||
except Exception as e:
|
||||
print(f" Warning: could not parse perf dump: {e}")
|
||||
|
||||
return metrics
|
||||
|
||||
|
||||
def print_results_table(results: list[dict]):
|
||||
"""Print baseline vs custom kernel comparison table."""
|
||||
print()
|
||||
print("=" * 80)
|
||||
print("BENCHMARK RESULTS — Denoise Latency (primary metric ★)")
|
||||
print("(Models and params match diffusion-benchmark-and-profile.md)")
|
||||
print("=" * 80)
|
||||
|
||||
by_model: dict[str, dict] = {}
|
||||
for r in results:
|
||||
by_model.setdefault(r["model"], {})[r["label"]] = r
|
||||
|
||||
print(
|
||||
f"{'Model':<16} {'Baseline(s)':>12} {'Custom(s)':>10} {'Speedup':>9} {'Peak Mem(GB)':>14}"
|
||||
)
|
||||
print("-" * 64)
|
||||
|
||||
for model_key in MODELS: # preserve order
|
||||
if model_key not in by_model:
|
||||
continue
|
||||
runs = by_model[model_key]
|
||||
base = runs.get("baseline", {})
|
||||
custom = runs.get("custom", {})
|
||||
|
||||
base_lat = base.get("denoise_latency_s")
|
||||
custom_lat = custom.get("denoise_latency_s")
|
||||
peak_mem = base.get("peak_memory_gb") or custom.get("peak_memory_gb")
|
||||
|
||||
speedup = f"{base_lat / custom_lat:.2f}x" if base_lat and custom_lat else "n/a"
|
||||
base_s = f"{base_lat:.2f}" if base_lat else "n/a"
|
||||
custom_s = f"{custom_lat:.2f}" if custom_lat else "n/a"
|
||||
mem_s = f"{peak_mem:.1f}" if isinstance(peak_mem, float) else "n/a"
|
||||
|
||||
print(f"{model_key:<16} {base_s:>12} {custom_s:>10} {speedup:>9} {mem_s:>14}")
|
||||
|
||||
print("-" * 64)
|
||||
print()
|
||||
print("★ Denoise latency = total DiT forward pass time across all inference steps.")
|
||||
print(
|
||||
" See diffusion-benchmark-and-profile.md for full Level 1/2 profiling workflow."
|
||||
)
|
||||
|
||||
|
||||
def inject_kernels_example():
|
||||
"""
|
||||
Show the kernel injection pattern used when SGLANG_DIFFUSION_CUSTOM_CUDA_KERNELS=1.
|
||||
After implementing add-cuda-kernel.md, this logic lives in denoising.py or
|
||||
the model's transformer.py — NOT in this script.
|
||||
|
||||
Call patch_rmsnorm(dit_model) BEFORE torch.compile and BEFORE any CPU offloading.
|
||||
"""
|
||||
import torch.nn as nn
|
||||
|
||||
try:
|
||||
from sglang.jit_kernel.diffusion.rmsnorm import diffusion_rmsnorm
|
||||
except ImportError:
|
||||
print(
|
||||
"diffusion.rmsnorm JIT kernel not available. "
|
||||
"Implement add-cuda-kernel.md first."
|
||||
)
|
||||
return
|
||||
|
||||
def patch_rmsnorm(model: nn.Module, verbose: bool = False) -> int:
|
||||
"""Monkey-patch all RMSNorm variants to use the JIT CUDA kernel."""
|
||||
patched = 0
|
||||
for name, module in model.named_modules():
|
||||
if "RMSNorm" not in type(module).__name__:
|
||||
continue
|
||||
eps = getattr(module, "eps", getattr(module, "variance_epsilon", 1e-6))
|
||||
has_weight = hasattr(module, "weight") and module.weight is not None
|
||||
|
||||
if has_weight:
|
||||
|
||||
def _make(mod, ep):
|
||||
def fwd(x):
|
||||
return diffusion_rmsnorm(x, weight=mod.weight, eps=ep)
|
||||
|
||||
return fwd
|
||||
|
||||
module.forward = _make(module, eps)
|
||||
else:
|
||||
|
||||
def _make_no_w(ep):
|
||||
def fwd(x):
|
||||
return diffusion_rmsnorm(x, weight=None, eps=ep)
|
||||
|
||||
return fwd
|
||||
|
||||
module.forward = _make_no_w(eps)
|
||||
|
||||
patched += 1
|
||||
if verbose:
|
||||
print(f" Patched: {name} (weight={has_weight})")
|
||||
return patched
|
||||
|
||||
return patch_rmsnorm
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="SGLang Diffusion denoise benchmark — baseline vs JIT CUDA kernels"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
choices=list(MODELS.keys()),
|
||||
help="Model to benchmark (default: flux)",
|
||||
)
|
||||
parser.add_argument("--all", action="store_true", help="Benchmark all 7 models")
|
||||
parser.add_argument(
|
||||
"--custom-kernels",
|
||||
action="store_true",
|
||||
help="Run with custom JIT CUDA kernels (SGLANG_DIFFUSION_CUSTOM_CUDA_KERNELS=1)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-custom-kernels",
|
||||
action="store_true",
|
||||
help="Run baseline (no custom kernels)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--compare",
|
||||
action="store_true",
|
||||
help="Run both baseline and custom, print comparison table",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
type=str,
|
||||
default="/workspace/gen_benchmark/bench_results",
|
||||
help="Directory for perf dump JSON files",
|
||||
)
|
||||
parser.add_argument("--no-warmup", action="store_true", help="Skip warmup")
|
||||
parser.add_argument(
|
||||
"--show-injection-example",
|
||||
action="store_true",
|
||||
help="Print kernel injection pattern and exit",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.show_injection_example:
|
||||
patch_fn = inject_kernels_example()
|
||||
if patch_fn:
|
||||
print(
|
||||
"patch_rmsnorm function defined. "
|
||||
"Call it on the DiT model before torch.compile and CPU offloading."
|
||||
)
|
||||
return
|
||||
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
warmup = not args.no_warmup
|
||||
|
||||
models_to_run = list(MODELS.keys()) if args.all else [args.model or "flux"]
|
||||
results = []
|
||||
|
||||
for model_key in models_to_run:
|
||||
if args.compare:
|
||||
results.append(run_benchmark_once(model_key, False, output_dir, warmup))
|
||||
results.append(run_benchmark_once(model_key, True, output_dir, warmup))
|
||||
elif args.custom_kernels:
|
||||
results.append(run_benchmark_once(model_key, True, output_dir, warmup))
|
||||
else:
|
||||
results.append(run_benchmark_once(model_key, False, output_dir, warmup))
|
||||
|
||||
if results:
|
||||
print_results_table(results)
|
||||
|
||||
print(f"Perf dump JSONs → {output_dir}")
|
||||
print(
|
||||
"Compare across runs: follow diffusion-benchmark-and-profile.md → Perf dump & before/after compare."
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,193 @@
|
||||
"""
|
||||
Micro-benchmark for the SGLang Diffusion JIT CUDA RMSNorm kernel.
|
||||
|
||||
Compares:
|
||||
1. SGLang JIT CUDA kernel (diffusion_rmsnorm)
|
||||
2. PyTorch baseline (torch.nn.functional.rms_norm)
|
||||
|
||||
Adapted from: https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels
|
||||
|
||||
Usage:
|
||||
python scripts/bench_diffusion_rmsnorm.py
|
||||
|
||||
Requirements:
|
||||
pip install triton # for triton.testing timing utilities
|
||||
# SGLang must be installed and CUDA available
|
||||
"""
|
||||
|
||||
import time
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Import the JIT CUDA kernel.
|
||||
# When you implement add-cuda-kernel.md, the file will be at:
|
||||
# python/sglang/jit_kernel/diffusion/rmsnorm.py
|
||||
# ---------------------------------------------------------------------------
|
||||
try:
|
||||
from sglang.jit_kernel.diffusion.rmsnorm import diffusion_rmsnorm
|
||||
|
||||
JIT_AVAILABLE = True
|
||||
except ImportError:
|
||||
JIT_AVAILABLE = False
|
||||
print(
|
||||
"WARNING: diffusion.rmsnorm JIT kernel not available. "
|
||||
"Run after implementing add-cuda-kernel.md."
|
||||
)
|
||||
|
||||
|
||||
def pytorch_rmsnorm(
|
||||
x: torch.Tensor,
|
||||
weight: torch.Tensor | None = None,
|
||||
eps: float = 1e-6,
|
||||
) -> torch.Tensor:
|
||||
"""Reference PyTorch implementation of RMSNorm."""
|
||||
hidden = x.shape[-1]
|
||||
return torch.nn.functional.rms_norm(
|
||||
x.float(), (hidden,), weight.float() if weight is not None else None, eps=eps
|
||||
).to(x.dtype)
|
||||
|
||||
|
||||
def benchmark_kernel(
|
||||
func,
|
||||
args,
|
||||
warmup: int = 20,
|
||||
iterations: int = 100,
|
||||
) -> Tuple[float, float]:
|
||||
"""Benchmark a kernel function. Returns (avg_ms, min_ms)."""
|
||||
for _ in range(warmup):
|
||||
func(*args)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
times = []
|
||||
for _ in range(iterations):
|
||||
torch.cuda.synchronize()
|
||||
t0 = time.perf_counter()
|
||||
func(*args)
|
||||
torch.cuda.synchronize()
|
||||
times.append((time.perf_counter() - t0) * 1000)
|
||||
|
||||
return sum(times) / len(times), min(times)
|
||||
|
||||
|
||||
def run_benchmark():
|
||||
print("=" * 72)
|
||||
print("SGLang Diffusion RMSNorm Micro-Benchmark: JIT CUDA vs PyTorch")
|
||||
print("=" * 72)
|
||||
print(f"Device: {torch.cuda.get_device_name(0)}")
|
||||
cap = torch.cuda.get_device_capability()
|
||||
print(f"Compute Capability: sm_{cap[0]}{cap[1]}")
|
||||
print()
|
||||
|
||||
if not JIT_AVAILABLE:
|
||||
print("Skipping JIT kernel benchmark (kernel not available).")
|
||||
return
|
||||
|
||||
# Determine dtype: T4 (sm_75) has no BF16
|
||||
dtype = torch.bfloat16 if cap >= (8, 0) else torch.float16
|
||||
print(f"Dtype: {dtype}")
|
||||
print()
|
||||
|
||||
# Typical DiT hidden sizes for sglang diffusion models:
|
||||
# FLUX.1-dev: hidden=3072
|
||||
# Qwen-Image: hidden=2048
|
||||
# Wan2.2: hidden=4096
|
||||
configs = [
|
||||
# (batch_tokens, hidden_size, has_weight)
|
||||
(1024, 2048, True), # Qwen-Image: 1 sample × 1024 tokens
|
||||
(4096, 2048, True), # Qwen-Image: larger batch
|
||||
(1024, 3072, True), # FLUX: 1 sample × 1024 tokens
|
||||
(4096, 3072, True), # FLUX: larger
|
||||
(4096, 4096, True), # Wan2.2
|
||||
(4096, 2048, False), # no-weight (elementwise_affine=False)
|
||||
(16384, 3072, True), # long sequence
|
||||
]
|
||||
|
||||
print(
|
||||
f"{'Config':<32} {'JIT(ms)':>10} {'PyTorch(ms)':>12} {'Speedup':>9} {'Weight'}"
|
||||
)
|
||||
print("-" * 72)
|
||||
|
||||
total_speedup = 0
|
||||
n = 0
|
||||
|
||||
for batch_tokens, hidden, has_weight in configs:
|
||||
x = torch.randn(batch_tokens, hidden, dtype=dtype, device="cuda")
|
||||
weight = torch.ones(hidden, dtype=dtype, device="cuda") if has_weight else None
|
||||
|
||||
jit_avg, _ = benchmark_kernel(
|
||||
diffusion_rmsnorm, (x, weight, 1e-6), warmup=20, iterations=100
|
||||
)
|
||||
pt_avg, _ = benchmark_kernel(
|
||||
pytorch_rmsnorm, (x, weight, 1e-6), warmup=20, iterations=100
|
||||
)
|
||||
|
||||
speedup = pt_avg / jit_avg
|
||||
total_speedup += speedup
|
||||
n += 1
|
||||
|
||||
w_str = "yes" if has_weight else "no "
|
||||
cfg = f"[{batch_tokens}×{hidden}]"
|
||||
print(f"{cfg:<32} {jit_avg:>10.3f} {pt_avg:>12.3f} {speedup:>8.2f}x {w_str}")
|
||||
|
||||
print("-" * 72)
|
||||
print(f"{'Average Speedup':>56} {total_speedup / n:.2f}x")
|
||||
print()
|
||||
|
||||
# -----------------------------------------------------------------------
|
||||
# Correctness check
|
||||
# -----------------------------------------------------------------------
|
||||
print("Correctness Check (BF16 tolerance 0.02):")
|
||||
x = torch.randn(4096, 3072, dtype=dtype, device="cuda")
|
||||
weight = torch.ones(3072, dtype=dtype, device="cuda")
|
||||
|
||||
out_jit = diffusion_rmsnorm(x, weight=weight, eps=1e-6)
|
||||
out_ref = pytorch_rmsnorm(x, weight=weight, eps=1e-6)
|
||||
|
||||
max_diff = (out_jit - out_ref).abs().max().item()
|
||||
rel_diff = ((out_jit - out_ref).abs() / (out_ref.abs() + 1e-8)).max().item()
|
||||
passed = max_diff < 0.02
|
||||
|
||||
print(f" Max absolute diff: {max_diff:.2e}")
|
||||
print(f" Max relative diff: {rel_diff:.2e}")
|
||||
print(f" Correctness: {'PASS ✓' if passed else 'FAIL ✗'}")
|
||||
print()
|
||||
|
||||
# -----------------------------------------------------------------------
|
||||
# Memory bandwidth analysis
|
||||
# -----------------------------------------------------------------------
|
||||
print("Memory Bandwidth Analysis:")
|
||||
bt, hid = 4096, 3072
|
||||
x = torch.randn(bt, hid, dtype=dtype, device="cuda")
|
||||
weight = torch.ones(hid, dtype=dtype, device="cuda")
|
||||
|
||||
bytes_per_elem = dtype.itemsize
|
||||
total_bytes = (
|
||||
bt * hid + hid + bt * hid
|
||||
) * bytes_per_elem # read x + read w + write out
|
||||
jit_avg, _ = benchmark_kernel(diffusion_rmsnorm, (x, weight, 1e-6))
|
||||
|
||||
bandwidth_gbps = (total_bytes / 1e9) / (jit_avg / 1000)
|
||||
theoretical_bw = {
|
||||
(9, 0): 3350, # H100: 3.35 TB/s
|
||||
(8, 0): 2000, # A100 80GB
|
||||
}.get(
|
||||
cap, 320
|
||||
) # T4: 320 GB/s
|
||||
efficiency = bandwidth_gbps / theoretical_bw * 100
|
||||
|
||||
print(f" Shape: [{bt} × {hid}] dtype: {dtype}")
|
||||
print(f" Total data: {total_bytes / 1e6:.1f} MB")
|
||||
print(f" Achieved: {bandwidth_gbps:.1f} GB/s")
|
||||
print(f" Theoretical ({torch.cuda.get_device_name(0)}): {theoretical_bw} GB/s")
|
||||
print(f" Bandwidth efficiency: {efficiency:.1f}%")
|
||||
print()
|
||||
print("Target: ≥ 30% efficiency (H100/A100), ≥ 40% (T4)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if not torch.cuda.is_available():
|
||||
print("CUDA not available.")
|
||||
else:
|
||||
run_benchmark()
|
||||
@@ -20,7 +20,7 @@ This skill focuses on SGLang Diffusion (`sglang.multimodal_gen`) kernel fusion p
|
||||
- `python/sglang/jit_kernel/norm.py`
|
||||
- `python/sglang/multimodal_gen/runtime/platforms/cuda.py`
|
||||
- `python/sglang/multimodal_gen/runtime/layers/attention/selector.py`
|
||||
- `docs/diffusion/performance/attention_backends.md`
|
||||
- `docs/diffusion/performance/attention_backends.md` (repo root)
|
||||
|
||||
**Core Fusion Patterns**
|
||||
|
||||
@@ -109,3 +109,4 @@ This skill focuses on SGLang Diffusion (`sglang.multimodal_gen`) kernel fusion p
|
||||
- Keep CuTe compile cache keys aligned to `(dtype, ndim, D)`.
|
||||
- Avoid implicit broadcasts that force hidden `contiguous()` copies.
|
||||
- Preserve NPU and ROCm fallback paths.
|
||||
- **Always verify with ncu** (`ncu --set full`) that the kernel achieves adequate memory bandwidth utilization (>70% of peak for bandwidth-bound ops) and occupancy (>50%). See `diffusion-benchmark-and-profile.md` Step 3.5 for the ncu workflow.
|
||||
|
||||
@@ -1,81 +1,18 @@
|
||||
---
|
||||
name: diffusion-perf
|
||||
description: Measure and compare sglang-diffusion performance. Use when benchmarking a model, comparing before/after performance, or generating a perf report for a PR.
|
||||
user-invocable: true
|
||||
description: Deprecated alias (merged into diffusion-kernel).
|
||||
user-invocable: false
|
||||
allowed-tools: Bash, Read
|
||||
argument-hint: <model-path> [--prompt "..."] [--baseline baseline.json]
|
||||
---
|
||||
|
||||
# Diffusion Performance Measurement
|
||||
|
||||
Measure sglang-diffusion e2e latency via `--perf-dump-path`, then extract or compare results from the JSON dump.
|
||||
This skill has been merged into the canonical docs under `diffusion-kernel`:
|
||||
|
||||
## JSON dump structure
|
||||
- `../diffusion-kernel/diffusion-benchmark-and-profile.md` → **Perf dump & before/after compare**
|
||||
|
||||
`--perf-dump-path` writes a JSON file with:
|
||||
Follow that document as the single source of truth:
|
||||
|
||||
```json
|
||||
{
|
||||
"total_duration_ms": 14959.11,
|
||||
"steps": [
|
||||
{"name": "TextEncodingStage", "duration_ms": 611.83},
|
||||
{"name": "DenoisingStage", "duration_ms": 14289.46}
|
||||
],
|
||||
"denoise_steps_ms": [
|
||||
{"step": 0, "duration_ms": 240.5},
|
||||
{"step": 1, "duration_ms": 279.1}
|
||||
],
|
||||
"commit_hash": "abc123",
|
||||
"timestamp": "...",
|
||||
"memory_checkpoints": {}
|
||||
}
|
||||
```
|
||||
|
||||
Key fields:
|
||||
- `total_duration_ms` — e2e walltime (warmup excluded when `--warmup` is used)
|
||||
- `steps` — per-stage breakdown
|
||||
- `denoise_steps_ms` — per denoising step timing
|
||||
|
||||
## Workflow
|
||||
|
||||
### 1. Single measurement
|
||||
|
||||
```bash
|
||||
sglang generate --model-path $MODEL --prompt "$PROMPT" --warmup --perf-dump-path result.json
|
||||
```
|
||||
|
||||
Then read `total_duration_ms` from `result.json`.
|
||||
|
||||
### 2. Before/after comparison
|
||||
|
||||
```bash
|
||||
# Baseline (on main branch or before changes)
|
||||
sglang generate --model-path $MODEL --prompt "$PROMPT" --warmup --perf-dump-path baseline.json
|
||||
|
||||
# New (after changes)
|
||||
sglang generate --model-path $MODEL --prompt "$PROMPT" --warmup --perf-dump-path new.json
|
||||
|
||||
# Compare — outputs a Markdown table suitable for PR descriptions
|
||||
python python/sglang/multimodal_gen/benchmarks/compare_perf.py baseline.json new.json
|
||||
```
|
||||
|
||||
### 3. Extracting a single number
|
||||
|
||||
To get e2e latency in seconds from a dump:
|
||||
|
||||
```bash
|
||||
python3 -c "import json; print(f\"{json.load(open('result.json'))['total_duration_ms']/1000:.2f}\")"
|
||||
```
|
||||
|
||||
## Arguments
|
||||
|
||||
If `$ARGUMENTS` is provided, parse it as:
|
||||
- First positional arg → `--model-path`
|
||||
- `--prompt "..."` → generation prompt (default: `"A curious raccoon"`)
|
||||
- `--baseline <file>` → if given, run comparison against this baseline file
|
||||
|
||||
## Notes
|
||||
|
||||
- Always use `--warmup` for accurate timing (excludes CUDA warmup from measurement).
|
||||
- Keep `--prompt` and all server/sampling args identical between baseline and new runs.
|
||||
- For PR descriptions, paste the output of `compare_perf.py` directly.
|
||||
- Always run `sglang generate ... --warmup --perf-dump-path <file>.json`
|
||||
- Use `python python/sglang/multimodal_gen/benchmarks/compare_perf.py <baseline.json> <new.json>` to generate a PR-ready comparison table
|
||||
|
||||
Reference in New Issue
Block a user