AI-Hypercomputer / AI-Hypercomputer/maxtext
Qwen3-Next: sequential GDN paths (lax.scan recurrence, solve_triangular) bound TPU step time
- 主要語言
- Python
- 星號
- 2.4k
- 分支
- 608
- 平均合併
- 2 天 19 小時
- 30 天內合併 PR
- 158
描述
### Problem
In the Qwen3-Next Gated Delta Net (`jax_chunk_gated_delta_rule`), the two sequential computations dominate TPU step time:
- the inter-chunk recurrence runs as `lax.scan`, so the recurrent state round-trips HBM on every chunk step, and
- the UT-transform inverse `A = (I+S)^-1` uses `solve_triangular`, which substitutes row by row and barely uses the MXU.
Neither overlaps with other work, so they bound the step directly.
### Measurement
v5e-256, Qwen3-Next dense-8B config, seq 4096, bs 1, remat full (details in #4348):
| build | step time | tokens/s/device |
|---|---|---|
| main | 3.373s | 1,214 |
| #4348 (Pallas kernels) | **2.571s** | **1,593 (+31%)** |
Loss trajectories match to bf16 rounding level.
貢獻指南
研究方向
Start at the Qwen3-Next Gated Delta Net entry point jax_chunk_gated_delta_rule and inspect the lax.scan recurrence and solve_triangular inverse described in the issue. Compare the Pallas-kernel results referenced in #4348, then benchmark the v5e-256, dense-8B, sequence-4096, batch-1 configuration; done means lower step time with loss matching to bf16 rounding.
由索引模型根據 Issue 內容生成。
評估
- 技術堆疊
- python
- 領域
- machine-learning, performance
- Issue 類型
- 重構
- 難度
- 4/5
- 預估耗時
- 3-5 天
- 活躍度
- 活躍
- 描述清晰度
- 基本清楚
- 新手友好度
- 35/100