[gpt-oss] attention decode optimizations

[sraizada-tt]

Added padding to o_proj weights and bias to ensure tile-aligned dimensions in CCL operations, avoiding expensive untilize-pad-tilize cycles

https://github.com/tenstorrent/tt-metal/actions/runs/21719380426

[Copilot]

Pull request overview

This PR introduces optimizations for GPT attention operations, focusing on reducing overhead in tensor parallelism (TP) scenarios through padding for tile alignment and a fused QK RoPE kernel for decode mode.

Changes:

• Added padding to o_proj weights and bias to ensure tile-aligned dimensions in CCL operations, avoiding expensive untilize-pad-tilize cycles
• Implemented fused QK RoPE operation and fused KV cache update for decode mode when batch size ≤ 32
• Added slice operation after allreduce to remove padding and restore original hidden dimensions

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.

File	Description
models/demos/gpt_oss/tt/attention/weights.py	Adds padding logic to o_proj weight and bias for tile alignment in TP operations; updates cache keys to reflect padding
models/demos/gpt_oss/tt/attention/operations.py	Adds slice operation after allreduce to remove padding added to support tile-aligned CCL
models/demos/gpt_oss/tt/attention/decode.py	Implements fused QK RoPE optimization for batch_size ≤ 32; adds reshape logic for padded dimensions before allreduce
models/demos/gpt_oss/tt/attention/init.py	Pre-creates fused transformation matrix for fused QK RoPE to avoid host writes during trace