perf(qwen3-a8w8): skip padded attention rows - #830
Conversation
📝 WalkthroughWalkthroughDecode entry points now accept an ChangesActive batch decode
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Sequence Diagram(s)sequenceDiagram
participant decode_fwd
participant _decode_layer
participant FAWorkTable
participant online_softmax
decode_fwd->>_decode_layer: pass active_batch
_decode_layer->>FAWorkTable: build work for active sequences
_decode_layer->>online_softmax: process active batch work
online_softmax->>_decode_layer: write outputs and pad inactive rows
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
models/qwen3/14b/decode_layer_a8w8.py (1)
966-966: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a partial-batch regression case.
The fixture always uses
active_batch = [BATCH], so the new active-row scheduling and padding paths are untested. Add at least anactive_batch = [1]or[2]case and verify that inactive output rows receive the row-0 attention result.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@models/qwen3/14b/decode_layer_a8w8.py` at line 966, Add a partial-batch regression case alongside the fixture using active_batch set to [1] or [2], exercise the active-row scheduling and padding paths, and assert that every inactive output row receives the row-0 attention result while retaining the existing full-batch coverage.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@models/qwen3/14b/decode_layer_a8w8.py`:
- Line 966: Add a partial-batch regression case alongside the fixture using
active_batch set to [1] or [2], exercise the active-row scheduling and padding
paths, and assert that every inactive output row receives the row-0 attention
result while retaining the existing full-batch coverage.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: e7c74854-c629-4c0c-86ff-78082af244c3
📒 Files selected for processing (1)
models/qwen3/14b/decode_layer_a8w8.py
2 participants
Summary
Optimize Qwen3-14B A8W8 decode for partially filled batches by making only the sequence-dependent attention work follow the real request count.
This PR:
active_batchscalar input to the fused Qwen3-14B A8W8 decode kernel;BATCH=16tensor layout and downstream kernel ABI.The main target is single-request and low-occupancy long-context decode, where processing replicated padding rows caused redundant KV-cache reads and a steep TPOT increase as the sequence grew.
Motivation
The A8W8 projection and MLP kernels are tiled around a fixed physical batch of 16. Serving therefore keeps model tensors shaped as 16 rows even when only one or two requests are active.
Before this change, attention also treated all 16 physical rows as real work. For a single request, the serving path replicated row-0 sequence metadata into the remaining 15 padding rows, so the decode kernel repeatedly:
This overhead grows with context length because the duplicated work includes sequence-length-dependent KV-cache traffic. Short decode could still look acceptable while long generation degraded substantially.
Changing the entire model to physical
BATCH=1is not safe: output projection, RMSNorm, MLP tiling, buffer layouts, and dependency structure still rely on the 16-row contract. This PR instead makes only attention active-batch-aware and leaves the fixed-batch downstream path unchanged.Design
Before
For
active_batch=1, 15 padding rows repeated the row-0 attention/KV work.After
This removes redundant long-context attention work without changing the tensor shapes consumed by output projection and MLP.
Implementation details
active_batchkernel ABIA scalar
active_batchtensor is added to:_decode_layer;_decode_layer_test_entry;decode_fwdentry.The companion serving path passes the actual number of active request rows for each decode step. The compile/test dummy input uses
active_batch=BATCHso the full physical-batch path remains covered.FlashAttention work construction
The FA work builder previously iterated over the compile-time
BATCHconstant:It now reads the runtime scalar and iterates only over:
The number of attention parts for each real row continues to depend on its actual sequence length. Padding rows no longer create FA work-table entries or trigger duplicate KV-cache reads.
Online-softmax work
The previous online-softmax work count was fixed at:
It is now computed at runtime as:
Only active rows participate in attention-part combination and context output generation.
Padding-row materialization
Downstream output projection and MLP still consume a physical
[16, hidden_size]tensor. After computing row 0 attention output, the kernel copies it into rows[active_batch, BATCH).This step is required to preserve the existing padding-row contract. It avoids leaving undefined padding data for fixed-batch consumers while keeping all user-visible active rows independently computed.
ABI and integration
This kernel change requires the matching serving-side
active_batchplumbing in pypto-serving PR #48.The caller must provide:
The PR intentionally does not change:
The optimized Gate/Up and BF16 KV data path from pypto-lib PR #810 remains unchanged.
Performance
Short deterministic gate
The strict 48-token test produced the exact golden token sequence and measured:
Long-context trend
A chat-template run generated 779 tokens, ended naturally with EOS, and measured:
A historical comparable long-chat run measured approximately
61.4 ms/token. Because the token sequence and EOS length are not identical, this comparison is trend evidence rather than a strict same-input causal percentage. The important signal is that TPOT remains much flatter as the context grows after duplicate padding-row attention work is removed.Why the gain grows with sequence length
Projection and MLP cost is largely independent of context length, while attention KV traffic grows with the number of cached tokens and attention parts. Removing up to 15 duplicated padding rows therefore has a larger effect during long decode than during the earliest decode steps.
Validation
Compile
a2a3simfrontend compile passed.active_batchargument.Accuracy and runtime gates
[9370, 99893].Batch semantics
The validation covers both:
active_batch=1, which exercises padding-row replication and removes 15 rows of duplicate attention work;active_batch=2, which verifies that multiple real request rows remain independently correct before the remaining padding rows are filled.Known boundary
A forced raw 1024-token completion can still exhaust the current runtime ring heap with
HEAP_RING_DEADLOCKbefore generation completes.An isolated unchanged-main run fails at the same capacity boundary with the same runtime classification, so this is not introduced by the active-batch attention change. This PR reduces attention work but does not change runtime ring admission, heap sizing, or task-window capacity.
The patch also does not make projection or MLP dynamically batched. Their physical 16-row ABI is deliberately retained; only sequence-length-dependent attention work is restricted to active request rows.
Change scope
Changed file:
Related work