perf(qwen3-a8w8): optimize decode scheduling and data paths

[vegetabledoww]

Summary

Reduce Qwen3-14B A8W8 decode TPOT from the historical ~56 ms baseline to ~42 ms/token (about 25% faster) by optimizing both task scheduling and the decode data path.

Related to #665.

What changed

Decode scheduling

• Batch activation quantization across decode rows to reduce fine-grained loop and dispatch overhead.
• Increase the output-projection N tile size, reducing its logical task count.
• Remove Gate/Up multi-K-split seed, atomic accumulation, and fan-in overhead.
• Use two-way K parallelism for Down Projection with independent partial accumulators.

Gate/Up INT8 path

• Keep checkpoint Gate/Up weights in their native INT8 representation.
• Dynamically quantize post-attention RMSNorm activations once per row.
• Run INT8 × INT8 → INT32 Gate/Up matmuls and apply activation/weight scales during dequantization.
• Keep prefill compatible with the shared INT8 weights by dequantizing each weight tile to BF16.

BF16 KV cache

• Change K/V cache storage from INT8 plus FP32 row scales to BF16.
• Remove per-token K/V cache quantization from decode and prefill.
• Remove per-row K/V dequantization loops from attention.
• Remove the now-unused KV scale tensors from the kernel interfaces.

Performance

Measured with Qwen3-14B-W8A8 using a 48-token deterministic generation test:

Stage	Decode TPOT
Historical baseline after the earlier A8W8 path	~56 ms/token
Scheduling optimizations	~48 ms/token
Scheduling + INT8 Gate/Up + BF16 KV cache	~42 ms/token

Recent device-0 integration runs measured 42.0-42.4 ms/token. A strict same-environment experiment measured 46.1 → 41.7 ms/token for the data-path changes alone.

Test configuration:

• Prompt: 介绍一下北京故宫
• max_new_tokens=48
• max_seq_len=256
• temperature=0
• top_p=1

The generated output remained coherent with no repeated-output regression.

Validation

• Python syntax and diff checks passed.
• All newly introduced FP32-to-BF16 narrowing sites explicitly use mode="rint" to match RNE semantics.
• The one-layer A8W8 decode DSL compile test passed and generated 45 functions.
• Nine targeted A8W8 serving integration tests passed.
• End-to-end 48-token generation passed on A2/A3 hardware.
• The final branch contains one commit and modifies only the Qwen3-14B A8W8 decode/prefill kernels.

Integration note

This kernel ABI requires the matching serving-side update to:

• allocate BF16 KV caches;
• retain Gate/Up INT8 weights and their scales;
• remove KV scale cache arguments.

[coderabbitai]

📝 Walkthrough

Walkthrough

The Qwen3 14B A8W8 prefill and decode paths now store BF16 KV caches, dequantize INT8 gate/up weights with FP32 scales, and use revised decode quantization and split down-projection accumulation flows.

Changes

Qwen3 A8W8 cache and MLP dataflow

Layer / File(s)	Summary
BF16 cache and interface contracts models/qwen3/14b/decode_layer_a8w8.py, models/qwen3/14b/prefill_fwd_a8w8.py	Decode and prefill signatures replace KV-cache scale tensors with gate/up weight scales. RoPE writes and attention loads now use BF16 KV caches, with corresponding call-site and test-input updates.
Decode activation and gate/up quantization models/qwen3/14b/decode_layer_a8w8.py, models/qwen3/14b/prefill_fwd_a8w8.py	Decode activation quantization uses row scales, while MLP activations and INT8 gate/up weights are combined through INT32 matmul accumulators and FP32 dequantization scales.
Split down projection accumulation models/qwen3/14b/decode_layer_a8w8.py	Down projection work is split across K partitions, stored as FP32 partials, accumulated, and cast with the residual output.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Input
  participant prefill_layer
  participant KVCache
  participant Attention
  participant DecodeMLP
  Input->>prefill_layer: BF16 KV cache and INT8 gate/up weights
  prefill_layer->>KVCache: write rotated BF16 K/V
  Attention->>KVCache: load BF16 K/V tiles
  Attention->>DecodeMLP: pass post-attention normalized activations
  DecodeMLP->>DecodeMLP: quantize activations and apply gate/up scales

Loading

Possibly related PRs

• hw-native-sys/pypto-lib#760: Modifies the same fused-attention scheduling and K/V cache handling paths.
• hw-native-sys/pypto-lib#806: Modifies the same RoPE/KV-cache writeback logic.

Poem

A bunny hops through BF16 snow,
While INT8 gates make matrices glow.
K and V cache their softer hue,
Split sums gather, neat and true.
“Hop!” says the rabbit—“the scales now flow!”

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title accurately summarizes the main change: performance-focused optimization of Qwen3 A8W8 decode scheduling and data paths.
Description check	✅ Passed	The description matches the changeset and explains the decode, MLP, and KV-cache ABI updates in detail.

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[gemini-code-assist]

Code Review

This pull request refactors the Qwen3 14B decode and prefill layers to store BF16 keys and values directly in the KV cache, eliminating the need for cache scales. Additionally, it quantizes the MLP gate and up projections to INT8, introduces corresponding weight scales, and parallelizes the down projection using K-splits. The review feedback recommends specifying mode="rint" in all pl.cast calls converting FP32 to BF16 to ensure standard round-to-nearest-even behavior and avoid precision loss on Ascend hardware.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[coderabbitai]

🧹 Nitpick comments (1)

models/qwen3/14b/decode_layer_a8w8.py (1)

66-94: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Add coverage invariants for the split down-projection constants.

The new down-proj partitioning silently assumes DOWN_ON % DOWN_DUAL_PAIR == 0 and INTERMEDIATE % DOWN_K_SPLITS == 0. The writers in down_dual_proj cover DOWN_DUAL_ON * DOWN_DUAL_PAIR column tiles (DOWN_DUAL_ON = DOWN_ON // DOWN_DUAL_PAIR), while down_cast_residual reads all DOWN_ON tiles. If DOWN_ON were odd (or INTERMEDIATE not divisible by DOWN_K_SPLITS), the truncating integer division would leave trailing tiles/rows of down_acc_parts unwritten and produce silently wrong output. Both hold for current dims, but this matches the existing defensive assert style on Line 94.

🛡️ Proposed defensive asserts

-assert DOWN_TN == 256 and DOWN_TK == 512 and DOWN_K_SLICE % DOWN_DUAL_L0_K == 0
+assert DOWN_TN == 256 and DOWN_TK == 512 and DOWN_K_SLICE % DOWN_DUAL_L0_K == 0
+assert HIDDEN % (DOWN_DUAL_PAIR * DOWN_TN) == 0 and INTERMEDIATE % DOWN_K_SPLITS == 0

🤖 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` around lines 66 - 94, Extend the
existing constant-validation asserts near DOWN_TN and DOWN_TK to require
INTERMEDIATE divisibility by DOWN_K_SPLITS and DOWN_ON divisibility by
DOWN_DUAL_PAIR. Keep the checks tied to the down-projection constants used by
down_dual_proj and down_cast_residual, preserving the current assert style.

🤖 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`:
- Around line 66-94: Extend the existing constant-validation asserts near
DOWN_TN and DOWN_TK to require INTERMEDIATE divisibility by DOWN_K_SPLITS and
DOWN_ON divisibility by DOWN_DUAL_PAIR. Keep the checks tied to the
down-projection constants used by down_dual_proj and down_cast_residual,
preserving the current assert style.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: c71609d0-1b77-4c9c-91d3-98dd33c2714d

📥 Commits

Reviewing files that changed from the base of the PR and between 0999dba and 8b34511.

📒 Files selected for processing (2)

• models/qwen3/14b/decode_layer_a8w8.py
• models/qwen3/14b/prefill_fwd_a8w8.py