Stream-k mxfp4 kernel

[xiaohuguo2023]

• add stream-k mxfp4 kernel
• add unit and perf tests

[Copilot]

Pull request overview

This PR adds Stream‑K load balancing support for FP4 GEMM in tritonblas, along with an expanded FP4 test/benchmark suite and a pytest marker for performance benchmarks.

Changes:

• Add a new Triton kernel implementing Stream‑K for FP4 matmul and wire it into the matmul_fp4 dispatch.
• Extend FP4 tests to cover Stream‑K correctness/behavior and add Stream‑K performance benchmark tests.
• Register a performance pytest marker in pyproject.toml.

Reviewed changes

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

Show a summary per file

File	Description
include/tritonblas/matmul.py	Adds FP4 Stream‑K dispatch and scratch-buffer handling; updates Stream‑K buffer reuse logic.
include/tritonblas/kernels/fp4_streamk_gemm.py	New FP4 Stream‑K GEMM Triton kernel implementation.
include/tritonblas/kernels/__init__.py	Exports the new FP4 Stream‑K kernel symbol.
tests/test_matmul_fp4.py	Adds Stream‑K FP4 correctness tests and multiple Stream‑K benchmark tests.
pyproject.toml	Defines the performance pytest marker.

Comments suppressed due to low confidence (1)

include/tritonblas/matmul.py:201

• Global Stream-K scratch buffers (_global_locks / _global_P) are allocated once on the default CUDA device at import time. The reuse path here does not verify that a.device matches those buffers, which can crash or silently misbehave when matmul runs on a non-default GPU (e.g. cuda:1). Consider gating reuse on device match (or keeping a per-device cache of scratch buffers) and falling back to per-call allocation otherwise.

    # Reuse pre-allocated global buffers at runtime, but allocate fresh during
    # torch.compile tracing (FakeTensorMode cannot slice real tensors).
    if not torch.compiler.is_compiling() and grids <= MAX_SMS and block_size <= MAX_BLOCK_SIZE:
        locks = _global_locks[:grids]
        P = _global_P[:grids, :block_size]
    else:
        locks = torch.empty(grids, device=a.device, dtype=torch.uint8)
        P = torch.empty(grids, block_size, device=a.device, dtype=torch.float32)

---

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