Hacktoberfest 2026: những issue maintainer đã đánh dấu cho tháng Mười, đang mở và phù hợp người mới. Xem issue Hacktoberfest

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

Đang mở
  • Có gì không ổn?
#783 0 bình luận 0 reaction 0 người được giao Xem trên GitHub

Maintainer thường phản hồi trong vòng 1 ngày

Chưa có ai nhận issue này.

Đánh giá

Độ khó
4/5
Thời gian dự kiến
3-5 ngày
Mức phù hợp với người mới
55/100
Loại issue
Tính năng
Độ rõ ràng
Đặc tả rõ ràng
Mức độ hoạt động
Sôi nổi
Công nghệ
python, pytorch
Lĩnh vực
machine-learning, performance

Hướng nghiên cứu

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.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

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.

Ngôn ngữ chính
Python
Star
14.2k
Fork
1.4k
Merge trung bình
21 giờ 17 phút
Pull request đã merge (30 ngày)
97

Chuẩn bị môi trường

Dự án này không cung cấp dev container, Dockerfile hay hướng dẫn đóng góp, nên bạn cần tự thiết lập môi trường: hãy bắt đầu từ README và xem hướng dẫn đóng góp lần đầu của chúng tôi để biết các bước chung.

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của FlashML-org/FreeToken

Tất cả issue của FlashML-org/FreeToken

Issue tương tự

Thêm issue về Python

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.