perf(qwen3-14b TQ decode): rotated-space attention, drop per-step KV …

[sunghajung6688]

The TQ decode dequant unrotated the entire KV cache every step (O(seq_len x head_dim^2) per step per layer), making long-sequence decode exceed the stream-sync timeout. Move the inverse rotation out of the per-row dequant into the cheap single-token Q (q.R) and the single attention output (.R^T): the paged cache stays rotated (K.R / V.R), and QK = (q.R).(K.R)^T = q.K^T (R orthogonal cancels). Precision unchanged.

• turboquant_kv.py: add turboquant_kv_dequant_chunk_rotated (unpack -> gather -> renorm -> scale, no xrot^T; emits two HALF-width rotated buffers, since assembling gather tiles into a full-width buffer is a non-mat tmov ptoas rejects).
• qwen3_14b_decode_tq.py: rotated-space attention -- q_rotate (q.R, tiled M=Q_HEAD_PAD=16 for Vec/UB), split dequant into lo/hi halves, QK = qk_lo.klo^T + qk_hi.khi^T, SV emits oi_lo/oi_hi, split online softmax, output unrotate ctx_lo.rot_lo^T + ctx_hi.rot_hi^T.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: cc574840-702e-4511-b224-29742a138764

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Changes

Rotated attention decode

Layer / File(s)	Summary
Rotated KV dequantization models/qwen3/14b/turboquant_kv.py	Adds rotated-space dequantization that emits separate low/high BF16 halves without applying the final unrotation.
Rotated QK computation models/qwen3/14b/qwen3_14b_decode_tq_draft.py	Rotates padded Q rows and computes QK scores as the sum of matmuls against separately dequantized K halves.
Split value accumulation and writeback models/qwen3/14b/qwen3_14b_decode_tq_draft.py	Tracks low/high value accumulation through online softmax, normalizes both halves, and performs one final unrotation during context writeback.

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

Sequence Diagram(s)

sequenceDiagram
  participant DecodeLayer
  participant RotatedDequant
  participant QKAndSV
  participant OnlineSoftmax
  participant Writeback
  DecodeLayer->>DecodeLayer: rotate padded Q by R
  DecodeLayer->>RotatedDequant: dequantize K and V into low/high halves
  RotatedDequant->>QKAndSV: rotated K halves
  QKAndSV->>OnlineSoftmax: summed QK scores and split value products
  OnlineSoftmax->>Writeback: normalized low/high context halves
  Writeback->>DecodeLayer: unrotate by R^T

Loading

Poem

I’m a bunny with buffers, hopping in pairs,
Low and high halves now travel upstairs.
Q turns by R, K and V stay near,
Softmax keeps both paths crystal clear.
One final hop unrotates the cheer!

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title matches the main change: rotated-space attention for Qwen3-14B TurboQuant decode to avoid per-step KV unrotation.
Description check	✅ Passed	The description directly matches the implemented rotated-space KV decode optimization and attention path changes.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
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.

---

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.}

[coderabbitai]

Actionable comments posted: 1

🤖 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.

Inline comments:
In `@models/qwen3/14b/qwen3_14b_decode_tq_draft.py`:
- Around line 518-532: Update the attention writeback block around ctx_unrot and
attn_out so only the first Q_HEAD_BATCH query rows are flattened and assembled.
Slice or otherwise restrict ctx_unrot before reshape, preserving the existing
unrotation while preventing padded rows from being written.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: a77419c9-a227-44a0-83c8-f1d13805aace

📥 Commits

Reviewing files that changed from the base of the PR and between c7ef3ca and fa7ccfd.

📒 Files selected for processing (2)

• models/qwen3/14b/qwen3_14b_decode_tq_draft.py
• models/qwen3/14b/turboquant_kv.py

[coderabbitai]

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Write back only the non-padded query rows.

ctx_unrot contains Q_HEAD_PAD rows, but this group owns only Q_HEAD_BATCH heads. Flattening every padded row can overwrite the following group’s output or exceed attn_out for the final group.

Proposed fix

                 ctx_unrot = pl.add(
                     pl.matmul(pl.cast(ctx_lo, target_type=pl.BF16), unrot_r_lo, b_trans=True, out_dtype=pl.FP32),
                     pl.matmul(pl.cast(ctx_hi, target_type=pl.BF16), unrot_r_hi, b_trans=True, out_dtype=pl.FP32),
                 )
+                ctx_valid = pl.slice(
+                    ctx_unrot, [Q_HEAD_BATCH, HEAD_DIM], [0, 0],
+                )
                 ctx_flat_bf16 = pl.cast(
-                    pl.reshape(ctx_unrot, [1, Q_HEAD_PAD * HEAD_DIM]),
+                    pl.reshape(ctx_valid, [1, Q_HEAD_BATCH * HEAD_DIM]),
                     target_type=pl.BF16,
                 )

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	# Finalize: ctx = oi / li (rotated), unrotate once (·R^T), write to attn_out.
	with pl.at(level=pl.Level.CORE_GROUP, name_hint="attention_writeback"):
	ctx = pl.row_expand_div(oi, li)
	ctx_lo = pl.row_expand_div(oi_lo, li)
	ctx_hi = pl.row_expand_div(oi_hi, li)
	unrot_r_lo = pl.slice(rot_slice, [HEAD_DIM, HALF_DIM], [0, 0])
	unrot_r_hi = pl.slice(rot_slice, [HEAD_DIM, HALF_DIM], [0, HALF_DIM])
	ctx_unrot = pl.add(
	pl.matmul(pl.cast(ctx_lo, target_type=pl.BF16), unrot_r_lo, b_trans=True, out_dtype=pl.FP32),
	pl.matmul(pl.cast(ctx_hi, target_type=pl.BF16), unrot_r_hi, b_trans=True, out_dtype=pl.FP32),
	)
	ctx_flat_bf16 = pl.cast(
	pl.reshape(ctx, [1, Q_HEAD_PAD * HEAD_DIM]),
	pl.reshape(ctx_unrot, [1, Q_HEAD_PAD * HEAD_DIM]),
	target_type=pl.BF16,
	)
	attn_out = pl.assemble(attn_out, ctx_flat_bf16, [b, q_base * HEAD_DIM])
	# Finalize: ctx = oi / li (rotated), unrotate once (·R^T), write to attn_out.
	with pl.at(level=pl.Level.CORE_GROUP, name_hint="attention_writeback"):
	ctx_lo = pl.row_expand_div(oi_lo, li)
	ctx_hi = pl.row_expand_div(oi_hi, li)
	unrot_r_lo = pl.slice(rot_slice, [HEAD_DIM, HALF_DIM], [0, 0])
	unrot_r_hi = pl.slice(rot_slice, [HEAD_DIM, HALF_DIM], [0, HALF_DIM])
	ctx_unrot = pl.add(
	pl.matmul(pl.cast(ctx_lo, target_type=pl.BF16), unrot_r_lo, b_trans=True, out_dtype=pl.FP32),
	pl.matmul(pl.cast(ctx_hi, target_type=pl.BF16), unrot_r_hi, b_trans=True, out_dtype=pl.FP32),
	)
	ctx_valid = pl.slice(
	ctx_unrot, [Q_HEAD_BATCH, HEAD_DIM], [0, 0],
	)
	ctx_flat_bf16 = pl.cast(
	pl.reshape(ctx_valid, [1, Q_HEAD_BATCH * HEAD_DIM]),
	target_type=pl.BF16,
	)
	attn_out = pl.assemble(attn_out, ctx_flat_bf16, [b, q_base * HEAD_DIM])

🤖 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/qwen3_14b_decode_tq_draft.py` around lines 518 - 532, Update
the attention writeback block around ctx_unrot and attn_out so only the first
Q_HEAD_BATCH query rows are flattened and assembled. Slice or otherwise restrict
ctx_unrot before reshape, preserving the existing unrotation while preventing
padded rows from being written.