AI-Hypercomputer / AI-Hypercomputer/maxtext

Qwen3-Next: sequential GDN paths (lax.scan recurrence, solve_triangular) bound TPU step time

オープン
#4,347 コメント 1 件 リアクション 0 件 担当者 1 名 @parambole が担当を希望しています GitHub で見る
主要言語
Python
スター
2.4k
フォーク
608
平均マージ
2日 19時間
マージ済み PR(30日)
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

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。