perf(attn_res): remove Blackwell 16K token limit - #1219
Conversation
Signed-off-by: chenht2022 <chenht2022@gmail.com>
149264f to
b01a19c
Compare
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 149264f221
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # AMD Gluon currently specializes Kimi K3's H=7168 fused-output-norm path. | ||
| _MAX_BLACKWELL_TOKENS = 16384 | ||
| _MAX_BLACKWELL_TOKENS = ((1 << 31) - 1) // 7168 |
There was a problem hiding this comment.
This line newly routes Blackwell inputs with more than 16,384 tokens into pointer arithmetic that was previously unreachable, but the Blackwell CUDA tests in test/thirdparty/test_attn_res_online_v2.py only exercise T=1, and this commit adds no automated boundary or long-sequence test. Add coverage at least just above 16,384 and at a size such as 65,536 where the source-stride multiplication would overflow without the new widening, so future changes cannot silently restore either the fallback or incorrect addressing.
AGENTS.md reference: AGENTS.md:L11-L14
Useful? React with 👍 / 👎.
Signed-off-by: chenht2022 <chenht2022@gmail.com>
2 participants
Summary
Remove the obsolete 16K Blackwell AttnRes limit so long prefills use the fused CUDA kernel instead of the PyTorch fallback, while keeping long-sequence address handling safe.
Validation
On 8x B300 with Kimi-K3 NVFP4, attention TP=8, MoE TP=8, and no prefix-cache hits:
5.316s/10.375s -> 1.338s/2.674s(3.97x/3.88x); throughput:6.16K/5.78K -> 24.50K/22.44Ktok/s.264.9 -> 220.6 GiB/GPU(44.3 GiB saved).238.3 GiB/GPUpeak.max_abs=0atT=16,385,32,768, and65,536.