[Qubitium]

TL;DR

Route the vendored HIP BF16T WMMA reference through the same tile and split-K dispatcher as NF12. A split-K plan can no longer launch grid.z > 1 while the BF16T kernel hardcodes kSplits=1, skips the partial plane, and feeds uninitialized workspace to the reducer.

1. The problem

The host plan selected split 4 for m=16,n=4096,k=4096, but the BF16T branch always instantiated kSplits=1. It wrote directly to out; the host still launched the reducer from partial, whose contents were uninitialized.

flowchart TD
    A["wmma_plan: splits 4"] --> B["grid.z = 4"]
    B --> C["BF16T kernel: kSplits = 1"]
    C --> D["write out directly"]
    B --> E["wmma_reduce_kernel"]
    E --> F["read uninitialized partial"]
    F --> X["corrupt output"]
    style X fill:#fee,stroke:#b22

Loading

2. The new flow

flowchart TD
    A["linear_wmma_impl<WeightSource>"] --> B["wmma_plan"]
    B --> C["NF12_AD_SWITCH for NF12 or BF16T"]
    C --> D["NF12_AD_DISPATCH: 1, 2, 4, or 8"]
    D --> E["wmma_kernel with matching kSplits"]
    E --> F{"splits == 1?"}
    F -->|"yes"| G["write BF16 output"]
    F -->|"no"| H["write FP32 partial"]
    H --> I["wmma_reduce_kernel"]
    style G fill:#efe,stroke:#282
    style I fill:#efe,stroke:#282

Loading

The same flow is documented in docs/nf12-hip-device-workspace.md.

Upstream

The implementation is already present in dphnAI/nf12 PR #3 commit 10a8825. This PR also adds upstream coverage on the same PR branch in commit 9fc6283:

• feat(hip): RDNA4 kernels and torch.ops.nf12 parity on ROCm dphnAI/nf12#3

Coverage

• Static launcher contract runs without ROCm and rejects the old hardcoded kSplits=1 BF16T path.
• Hardware tests cover the reachable split-2, split-4, and split-8 shapes:
  • 4096 x 4096, m=16
  • 4096 x 4096, m=64
  • 2048 x 4096, m=16
• Each compares BF16T and NF12 output bitwise.

Validation

TMPDIR=/dev/shm PYTHONDONTWRITEBYTECODE=1 /root/venv/bin/python -m pytest tests/test_nf12_hip_review_contracts.py tests/test_nf12_ad_codec.py -q -p no:cacheprovider -rs
17 passed, 90 skipped

ruff check tests/test_nf12_hip_review_contracts.py
All checks passed

The 90 HIP tests require a ROCm RDNA4 device and were skipped on this host. PR #3's existing R9700 validation establishes the upstream hardware baseline; the new split-plan cases are pending a RDNA4 runner.