Fuse reduce slices and squeezed epilogues, tile short reductions, and match attention output through view ops

[pfultz2]

Motivation

In int4 decode the matvecs are already fused reductions (dequantize, mul, reduce_sum), but two common patterns around them still fall out of the fusion:

• The gate_up projection produces one reduction whose output is sliced into the gate and up halves for SwiGLU. The slices sit between the reduction and the activation, so the SwiGLU runs as a separate pointwise kernel that re-reads the projection output.
• A pointwise that consumes a squeezed reduction output does not fuse as an epilogue, because the squeeze changes the rank the reduction module works in.

Short reductions with many outputs (blocked int4 weights with a small K) also leave each lane of a full workgroup with only a few elements, so per-workgroup overhead dominates and too few loads are in flight.

Separately, once the projection dot after attention is rewritten to a reduction, simplify_reshapes reshuffles the transpose/reshape epilogue after the second gemm (for example leaving an unsqueeze in it). fuse_attention anchored on the exact reshape(transpose(dot)) shape, so it stopped matching and the attention was left unfused.

Technical Details

fuse_reduce

• find_reduce_slice: when every consumer of a fused_reduce is a slice along the same non-reduced axis, the reduction is split into one reduce per slice. The slices are pushed onto the reduce inputs, the submodule is re-fused with its broadcasts narrowed to the sliced length, and any unit reshapes between the reduce and the slice are replayed on the new reduce. The slice consumers can then fuse with their own reduction (SwiGLU over the gate_up halves). It waits until any unpack_int4 feeding the reduce has been fused, since the slices move into the reduce inputs.
• find_reduce_squeeze_pointwise: a pointwise over a squeezed fused_reduce output is moved into the reduce space: the other inputs are unsqueezed instead and the squeeze is applied after the pointwise, so the existing reduce+pointwise fusion picks it up as an epilogue.

simplify_reshapes

• find_unary_shape_transforms now also treats a fused pointwise module with a single input as unary, so a squeeze or reshape can move across it. The module inputs are preserved when the instruction is re-inserted.

GPU fused reduce compiler

• find_short_reduce_tile: when no broadcast-based tile is found, and the tiled block size for the reduction is below a full workgroup while the outputs still fill the device, the last non-reduced axis is tiled two outputs per workgroup. This amortizes the per-workgroup overhead over two reductions and doubles the loads in flight per lane. The tile is only offered to the tuner (block_tile / block_batch), which benchmarks it against the plain block algorithm.

fuse_attention

• find_kv_cache_attention now ends its match at the second gemm instead of reshape(transpose(dot)). find_output_end walks the single-consumer view ops after the gemm (transpose, reshape, unsqueeze, squeeze) to find where the attention output ends, so the fusion works in whatever form the reshape simplifications left the epilogue.

Tests

• fuse_reduce: reduce_slice_pointwise, reduce_squeeze_pointwise
• simplify_reshapes: pointwise_module_reshape_unary
• fuse_attention: kv_cache_attention_unsqueezed_output
• verify: test_unpack_int4_dequant_reduce_slice_swiglu, test_unpack_int4_dequant_reduce_tiled

Changelog Category

Add a CHANGELOG.md entry for any option other than Not Applicable

• • Added: New functionality.
• • Changed: Changes to existing functionality.
• • Removed: Functionality or support that has been removed. (Compared to a previous release)
• • Optimized: Component performance that has been optimized or improved.
• • Resolved Issues: Known issues from a previous version that have been resolved.
• • Not Applicable: This PR is not to be included in the changelog.

Follow the LLVM AI Tool Use Policy for contributions using AI.

[Copilot]

🟡 Changes recommended

The new rewrites can fail on dynamic or input-driven slices, collide on generated module names, and mishandle tuple-valued pointwise modules.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Extends GPU fusion to handle sliced/squeezed reductions, short-reduction tiling, and attention output view chains.

Changes:

• Adds reduction slice and squeeze-pointwise fusion rewrites.
• Adds short-reduction tiling candidates.
• Broadens attention matching and adds regression coverage.

Review used a single-pass analysis without agent fan-out.

File summaries

File	Description
src/fuse_reduce.cpp	Adds slice and squeeze fusion rewrites.
src/simplify_reshapes.cpp	Supports unary fused-pointwise reshaping.
src/fuse_attention.cpp	Follows attention outputs through view operations.
src/targets/gpu/jit/reduce.cpp	Adds short-reduction tiling candidates.
test/fuse_reduce.cpp	Tests reduction fusion rewrites.
test/simplify_reshapes_test.cpp	Tests fused-pointwise reshaping.
test/fuse_attention.cpp	Tests unsqueezed attention outputs.
test/verify/test_unpack_int4_dequant_reduce_slice_swiglu.cpp	Verifies sliced int4 SwiGLU reductions.
test/verify/test_unpack_int4_dequant_reduce_tiled.cpp	Verifies tiled int4 reductions.

Review details

• Files reviewed: 9/9 changed files
• Comments generated: 5
• Review effort level: Balanced

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.