Improve batched gemm swizzle

[dzoba]

Proposed changes

Improve batched GEMM performance by enabling swizzle for operations with large tile grids.

Modified swizzle heuristic in steel_matmul_regular_axpby to use swizzle_log=1 when:

• batch_size_out > 1 (batched operations)
• tm >= 8 && tn >= 8 (large tile grids, i.e., M >= 512 and N >= 512 with 64x64 blocks)

This improves cache efficiency for batched matrix multiplications commonly used in LLM training (attention projections, FFN layers).

Benchmark Results

Tested on Apple Silicon M4 Max:

Shape (B, M, N, K)	Before	After	Improvement
(16, 1024, 1024, 1024) float32	-7.6% vs PyTorch	-0.3%	+7.3 pp
(4, 1024, 1024, 4096) float32	-17.6%	+8.1%	+25.7 pp
(4, 1024, 4096, 1024) float32	-9.9%	+20.3%	+30.2 pp

Non-batched operations (batch=1) are unchanged as they don't benefit from swizzle.

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[awni]

@jagrit06 what do you think about this? Results look quite nice to me.