Hacktoberfest 2026：メンテナが10月に向けて印を付けた、オープンで初心者向けの issue。 Hacktoberfest の issue を見る

perf(qsa): prefill indexer scoring pads 4 heads to 16 per row; 16-row tiles are ~5x faster and bitwise identical

オープン
  • 何か問題がありますか？
#783 コメント 0 件 リアクション 0 件 担当者 0 名 GitHub で見る

メンテナーはふだん 1 日以内に返信

まだ誰も着手していません。

評価

難易度
4/5
見積もり時間
3〜5日
初心者へのやさしさ
55/100
issue の種類
機能追加
明瞭さ
明確に書かれている
活発さ
活発
技術スタック
python, pytorch
領域
machine-learning, performance

調査の方向性

Start with python/freetoken/kernel/triton/qsa/score.py and _qsa_mqa_paged_kernel, then trace its prefill and decode callers in attention/qsa_sparse.py. The proposed change needs a prefill tile map built from qo_indptr, reused across QSA layers, while decode and CUDA graph capture keep the current kernel. Done means matching score bits and qsa_block_topk selections on the listed ragged, partial, sliced, and long-context cases, with the reported prefill speedup.

索引モデルが issue の本文から書いたものです。

説明

What happened

In prefill, the QSA indexer score kernel (_qsa_mqa_paged_kernel) runs one query row per program and pads the 4 index heads to MAX_N = 16. So most of each tl.dot is padding. On Qwen3.8-Flash-Next it takes 196-200 ms of every 4096-token prefill chunk at a 249K context, 14.5-15% of the chunk on an RTX 5080.

Proposal

For prefill only, give each program 16 query rows of one request × 4 heads. That is N = 16 with no padding: one dot per head, summed over the heads in the same order as today.

  • Tiles are cut at request boundaries, using a tile map built from qo_indptr on the host and cached across the 12 QSA layers.
  • Decode and CUDA-graph capture keep the current kernel.

Each output is still a single K=128 dot followed by the same head-sum order, so the scores are bitwise identical. Checked as int32 bits on ragged multi-request batches, partial tiles, mid-forward row slices and a 249K slice. qsa_block_topk selects identical blocks. Changing the head-sum order makes the check fail, so it does detect order differences.

Measured on an RTX 5080 (sm_120) with production flags, 2 interleaved rounds:

current 16-row tiles
score kernel per 4K chunk at 249K 196-200 ms 38 ms
microbench, one 4096-row chunk at 249K context, per layer 18.47 ms 3.24 ms
cold TTFT 249K 85.5 s 83.5 s

The end-to-end gain here is small because this box streams MoE experts over PCIe and late chunks are copy-bound. On a setup where prefill is compute-bound, the score kernel is a larger share.

Cause (current main, cc56b6bd)

python/freetoken/kernel/triton/qsa/score.py: _qsa_mqa_paged_kernel launches one program per query row, with the head axis padded to MAX_N. It is called from attention/qsa_sparse.py for prefill and decode alike.

Environment

RTX 5080 16 GB (sm_120, driver 615.78.08), torch 2.14.1+cu130, Triton 3.8.0, nvidia/Qwen3.8-Flash-Next-NVFP4, main cc56b6bd.

主要言語
Python
スター
14.2k
フォーク
1.4k
平均マージ
21時間 39分
マージ済み PR（30日）
100

環境構築

このプロジェクトには開発コンテナ、Dockerfile、コントリビューションガイドがありません。まず README を読み、一般的な手順ははじめてのコントリビューションガイドを参照してください。

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

FlashML-org/FreeToken のほかの issue

FlashML-org/FreeToken の issue をすべて見る

似ている issue

Python の issue をもっと見る

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

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