fix(triton): scope AMD gfx950 compiler knobs to GEMM kernel launches

[kyle-256]

Description

The AMD Triton knobs (use_async_copy / scalarize_packed_fops /
use_block_pingpong) were set process-globally once (guarded by a
_KNOBS_SET flag) and never reverted. After any turbo GEMM / grouped-GEMM ran,
these knobs leaked into unrelated Triton kernels running in the same
process (e.g. the training loop): they raise shared-memory (LDS) usage — fine
for a standalone op, but enough to OOM during training — and regress some other
Triton kernels.

These knobs are compile-time only: they affect codegen and only need to be
active while a kernel is JIT-compiled on its first launch (Triton caches the
compiled binary, so later launches are unaffected). This PR keeps the knob
flipping inside the GEMM / grouped-GEMM entry points but wraps every entry
point with a new scoped_amd_knobs decorator that snapshots the knob state on
entry and restores it on exit.

Net effect: the default state is now off; turbo turns these knobs on only
while its own GEMM / grouped-GEMM kernels compile, so nothing else in the
process inherits them.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• primus_turbo/triton/utils/triton_knobs_helper.py: drop the one-shot
  _KNOBS_SET global guard (knobs are now re-applied per call and restored by
  the surrounding scope); add scoped_amd_triton_knobs() context manager that
  snapshots/restores the AMD knobs (with triton.knobs.amd and env-var
  fallback), and a scoped_amd_knobs decorator wrapping it.
• Decorate all 16 GEMM / grouped-GEMM entry points with @scoped_amd_knobs:
  • primus_turbo/triton/gemm/gemm_kernel.py (1)
  • primus_turbo/triton/gemm/gemm_fp8_kernel.py (3)
  • primus_turbo/triton/grouped_gemm/grouped_gemm_kernel.py (2)
  • primus_turbo/triton/grouped_gemm/grouped_gemm_fp4_kernel.py (2)
  • primus_turbo/triton/grouped_gemm/grouped_gemm_fp8_kernel.py (8)
• The existing gfx942 layout-dependent _set_amd_knobs(...) calls are kept
  verbatim; they are now scoped (restored on function exit) too.

Checklist:

• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[Copilot]

Pull request overview

This PR fixes AMD Triton compiler knob leakage by scoping gfx950-specific knob toggles to only the GEMM / grouped-GEMM entry points that need them, snapshotting and restoring the previous knob state on exit so unrelated Triton kernels in the same process aren’t affected.

Changes:

• Reworks the gfx950 knob helper to remove the one-shot global guard and adds a scoped_amd_triton_knobs() context manager plus @scoped_amd_knobs decorator for snapshot/restore semantics.
• Wraps GEMM and grouped-GEMM public entry points with @scoped_amd_knobs so knob changes only apply during first-time JIT compilation / launch.
• Keeps existing layout-dependent _set_amd_knobs(...) behavior, but now implicitly scoped via the decorator on the entry points.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
primus_turbo/triton/utils/triton_knobs_helper.py	Introduces scoped snapshot/restore for AMD Triton knobs and removes process-global one-shot knob enabling.
primus_turbo/triton/gemm/gemm_kernel.py	Decorates the BF16 GEMM entry point to scope AMD knob changes.
primus_turbo/triton/gemm/gemm_fp8_kernel.py	Decorates FP8 GEMM entry points (tensorwise/rowwise/blockwise) to scope AMD knob changes.
primus_turbo/triton/grouped_gemm/grouped_gemm_kernel.py	Decorates grouped GEMM entry points (forward + variable-K backward) to scope AMD knob changes.
primus_turbo/triton/grouped_gemm/grouped_gemm_fp4_kernel.py	Decorates grouped MXFP4 entry points to scope AMD knob changes.
primus_turbo/triton/grouped_gemm/grouped_gemm_fp8_kernel.py	Decorates grouped FP8/MXFP8 entry points to scope AMD knob changes.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

Comments suppressed due to low confidence (1)

primus_turbo/triton/grouped_gemm/grouped_gemm_fp8_kernel.py:1699

• Similar to gemm_fp8_kernel.py, this file removed the previous non-gfx950 knob toggling (_set_amd_knobs) for the blockwise entry points, but the PR description states the gfx942 layout-dependent knob logic was kept. Please either update the PR description to match the new behavior or restore the non-gfx950 knob handling under the new scoped approach if it is still needed for performance/compatibility.

@scoped_amd_knobs
def grouped_gemm_fp8_blockwise_triton_kernel(
    a: torch.Tensor,
    b: torch.Tensor,
    a_scales: torch.Tensor,

[RuibinCheung]

LGTM

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

[RuibinCheung]

LGTM