Compare commits
10
Commits
c74176c065
...
bfa907d528
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
bfa907d528 | ||
|
|
1f74c3cfa2 | ||
|
|
09125567ea | ||
|
|
c539c8d02d | ||
|
|
c695836d81 | ||
|
|
20b8a77605 | ||
|
|
369f696a01 | ||
|
|
7bc53e7430 | ||
|
|
c34f67e39c | ||
|
|
593a9605b1 |
@@ -1,6 +1,8 @@
|
||||
# SAGE + Skip-Softmax — API Design & Test Plan
|
||||
|
||||
# Part 1 — API Design
|
||||
**Target: merge by 2026-07-24 for the v0.26.0 release.**
|
||||
|
||||
## Part 1 — API Design
|
||||
|
||||
*How the user uses the trtllm-gen backend.*
|
||||
|
||||
@@ -44,12 +46,14 @@ class TrtllmGenConfig: # mirrors TRT-LLM SkipSoftmaxAttentionCon
|
||||
disabled_until_timestep: float = 0.0 # keep initial (noisy) steps dense; per-model preset
|
||||
```
|
||||
|
||||
| CLI flag | env var | default | meaning |
|
||||
|----------|---------|:-------:|---------|
|
||||
| `--diffusion-attention-backend trtllm-gen` | `DIFFUSION_ATTENTION_BACKEND=trtllm-gen` | — | select backend (**BF16 dense**) |
|
||||
| `--trtllm-gen-sage` | `TRTLLM_GEN_SAGE=1` | **off** | FP8-SAGE |
|
||||
| `--trtllm-gen-target-sparsity <float>` | `TRTLLM_GEN_TARGET_SPARSITY` | unset (no skip) | skip amount; needs calibrated ckpt |
|
||||
| `--trtllm-gen-disabled-until-timestep <float>` | `TRTLLM_GEN_DISABLED_UNTIL_TIMESTEP` | per-model preset | keep initial noisy steps dense (fidelity guard); normalized denoise t 1→0 |
|
||||
| CLI flag | default | meaning |
|
||||
|----------|:-------:|---------|
|
||||
| `--diffusion-attention-backend trtllm-gen` | — | BF16 dense |
|
||||
| `--trtllm-gen-sage` | off | FP8-SAGE |
|
||||
| `--trtllm-gen-target-sparsity <float>` | unset | skip amount (needs calibrated ckpt) |
|
||||
| `--trtllm-gen-disabled-until-timestep <float>` | preset | keep initial steps dense (t 1→0) |
|
||||
|
||||
Env var = flag in UPPER_SNAKE (e.g. `TRTLLM_GEN_SAGE=1`).
|
||||
|
||||
**Maps to upstream:**
|
||||
|
||||
@@ -59,11 +63,15 @@ class TrtllmGenConfig: # mirrors TRT-LLM SkipSoftmaxAttentionCon
|
||||
| `target_sparsity` | `threshold_scale_factor` (via calibration) | `skip_softmax_threshold_scale_factor` |
|
||||
| `disabled_until_timestep` | — (DiT-only) | host gate |
|
||||
|
||||
Versions: **FlashInfer 0.6.14 · ModelOpt 0.45.0** (runtime deps); TRT-LLM = naming reference (blog16, no runtime dep).
|
||||
|
||||
---
|
||||
|
||||
# Part 2 — Testing Plan
|
||||
## Part 2 — Testing Plan
|
||||
|
||||
## Models
|
||||
Hardware: **SM103 (B300)**. Versions: **vLLM-Omni v0.25.0 · FlashInfer 0.6.14 · ModelOpt 0.45.0**.
|
||||
|
||||
### Models
|
||||
|
||||
| Model | HF checkpoint | params | shape |
|
||||
|-------|---------------|--------|-------|
|
||||
@@ -73,9 +81,7 @@ class TrtllmGenConfig: # mirrors TRT-LLM SkipSoftmaxAttentionCon
|
||||
|
||||
---
|
||||
|
||||
## ModelOpt skip calibration
|
||||
|
||||
Sweep `target_sparsity`, fit `factor = a·exp(b·target_sparsity)` per `config_group` → `config.json`. Steps 2–3 fix `target_sparsity`, sweep **D = `disabled_until_timestep`** (1.00 = off → 0.94 = most aggressive). D is a runtime gate, not calibrated.
|
||||
### ModelOpt skip calibration
|
||||
|
||||
**Setup** (per model)
|
||||
|
||||
@@ -103,7 +109,7 @@ Sweep `target_sparsity`, fit `factor = a·exp(b·target_sparsity)` per `config_g
|
||||
|
||||
---
|
||||
|
||||
## Step 1 — SAGE (FP8)
|
||||
### Step 1 — SAGE (FP8)
|
||||
|
||||
FP8-SAGE only, no calibration. trtllm-gen vs current backend: LPIPS parity + speedup.
|
||||
|
||||
@@ -116,11 +122,11 @@ FP8-SAGE only, no calibration. trtllm-gen vs current backend: LPIPS parity + spe
|
||||
| Cosmos 3 | BF16 dense (ref) | 0 (ref) | 1.00× | |
|
||||
| Cosmos 3 | FP8-SAGE | | | |
|
||||
|
||||
At BF16 dense: **trtgen == current backend** (LPIPS ≈ 0).
|
||||
At BF16 dense: **trtllm-gen == current backend** (LPIPS ≈ 0).
|
||||
|
||||
---
|
||||
|
||||
## Step 2 — Skip-Softmax
|
||||
### Step 2 — Skip-Softmax
|
||||
|
||||
Skip on BF16 attention (no SAGE). Fixed `target_sparsity` (→ factor); sweep **D = `disabled_until_timestep`**.
|
||||
|
||||
@@ -138,7 +144,7 @@ Skip on BF16 attention (no SAGE). Fixed `target_sparsity` (→ factor); sweep **
|
||||
|
||||
---
|
||||
|
||||
## Step 3 — SAGE + Skip
|
||||
### Step 3 — SAGE + Skip
|
||||
|
||||
FP8-SAGE + Skip. Same fixed `target_sparsity`; sweep **D = `disabled_until_timestep`**.
|
||||
|
||||
@@ -156,7 +162,7 @@ FP8-SAGE + Skip. Same fixed `target_sparsity`; sweep **D = `disabled_until_times
|
||||
|
||||
---
|
||||
|
||||
## Step 4 — cross-backend comparison
|
||||
### Step 4 — cross-backend comparison
|
||||
|
||||
Fix one model + shape + seed, swap only the attention backend. Speedup vs SDPA (= BF16 dense). Run per model (Wan / Hunyuan / Cosmos).
|
||||
|
||||
@@ -167,9 +173,3 @@ Fix one model + shape + seed, swap only the attention backend. Speedup vs SDPA (
|
||||
| **trtllm-gen** | FP8-SAGE | | | |
|
||||
| **trtllm-gen** | FP8-SAGE+Skip | | | |
|
||||
|
||||
---
|
||||
|
||||
## Notes
|
||||
|
||||
- SAGE = no calibration; Skip = per-model calibration.
|
||||
- LPIPS measured vs **BF16 dense**; perf on the **target Blackwell SKU under real concurrency**.
|
||||
Reference in new issue
Block a user