add batched_syrk triton kernel

[Rudin6]

No description provided.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

The PR adds a batched Triton SYRK kernel and integrates it into batched Newton–Schulz iterations for CUDA tensors.

• Adds batched tensor-descriptor setup, triangular tile mapping, autotuning, and input validation.
• Selects the batched SYRK path for eligible three-dimensional CUDA inputs while retaining CPU and alignment fallbacks.
• Adds benchmarks and correctness coverage for batched and distributed behavior.

Confidence Score: 4/5

The PR should not merge until the outstanding CUDA device-to-host synchronizations in the benchmark are addressed in accordance with the repository rule.

The benchmark still converts CUDA reduction results to Python via .item() at several correctness checks, so the previously reported synchronization issue remains present.

Files Needing Attention: benchmarks/bench_batched_syrk.py

Important Files Changed

Filename	Overview
emerging_optimizers/triton_kernels/batched_syrk.py	Introduces the batched BF16 SYRK wrapper and Triton kernel, including descriptor setup, triangular scheduling, validation, and autotuning.
emerging_optimizers/orthogonalized_optimizers/muon_utils.py	Routes eligible batched CUDA Newton–Schulz inputs through the new SYRK implementation while preserving GEMM fallbacks.
tests/test_muon_utils.py	Adds batched SYRK equivalence coverage and verifies the integrated Newton–Schulz path against GEMM.
benchmarks/bench_batched_syrk.py	Adds latency, accuracy, symmetry, accumulation, and transposed-input benchmark cases.

Sequence Diagram

sequenceDiagram
    participant NS as newton_schulz
    participant Step as batched_newton_schulz_step_tsyrk
    participant SYRK as batched_tsyrk_ex
    participant Kernel as Triton batched_syrk kernel
    NS->>Step: Batched BF16 CUDA tensor
    Step->>SYRK: "X @ X.mT"
    SYRK->>Kernel: Launch lower-triangular tile grid
    Kernel-->>SYRK: Symmetric Gram matrices
    SYRK-->>Step: A
    Step->>SYRK: "c*A@A + b*A (quintic path)"
    Step-->>NS: Updated batch via baddbmm

Loading

_{Reviews (2): Last reviewed commit: "Merge branch 'main' into main" | Re-trigger Greptile}

[greptile-apps]

CUDA reductions force host synchronization

When the benchmark reports correctness on CUDA, these .item() calls synchronously transfer reduction results to Python, stalling queued GPU work and perturbing benchmark execution; the same pattern occurs in the symmetry checks on lines 68 and 136.

Rule Used: Always flag device-to-host synchronizations (e.g.,... (source)

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[skyw]

Should be fine in benchmark if not in the timed region.

[greptile-apps]

Fair point — correctness checks that run outside the timed region don't skew the benchmark measurements. I'll stand down on this one.

Shall I remember this as a rule: "Do not flag .item() calls or other device-to-host synchronizations in benchmark files when they are used for correctness/validation checks outside the timed region"?

[skyw]

@gdengk could you help review the triton code? I can review the rest.