Skip to content

Phase 4A2 — rocWMMA Width-64 Production Inference for RDNA4 / gfx1201

Choose a tag to compare

View all tags
@Painter3000 Painter3000 released this 26 Jul 04:45
phase4a2-width64-production-inference-gfx1201-pass
cd1330a

Phase 4A2 — rocWMMA Width-64 Production Inference for RDNA4 / gfx1201

This release introduces and qualifies the first opt-in rocWMMA production inference backend in tiny-rdna4-nn: RocWMMAWidth64MLP.

The backend is designed for AMD RDNA4 / gfx1201 and was developed and validated on an AMD Radeon AI PRO R9700 with ROCm 7.2.

Highlights

  • New fused rocWMMA inference backend for AMD RDNA4 / gfx1201
  • Public tinycudann.Network integration
  • Explicit opt-in build switch
  • Default-off and fail-closed behavior
  • FP16 parameters, input, and output
  • Fixed 64 -> 64 -> 64 -> 64 topology
  • Two ReLU hidden layers
  • Caller-provided HIP stream support
  • Non-multiple public batch sizes handled through internal padding
  • Production kernel bound to the generated gfx1201 code object
  • Full runtime lifecycle and reproducibility validation

Qualified scope

RocWMMAWidth64MLP is deliberately narrow:

Architecture:       gfx1201 / Wave32
Input width:        64
Hidden width:       64
Hidden layers:      2
Output width:       64
Activation:         ReLU
Output activation:  None
Precision:          FP16
Execution:          inference only

The qualified parameter buffer contains exactly:

12,480 FP16 values
24,960 bytes

The backend is not selected automatically and is not used as a silent fallback for FullyFusedMLP or another network type.

Build

The backend is disabled by default. Enable it explicitly when building the PyTorch extension:

export ROCM_PATH=/opt/rocm
export PYTORCH_ROCM_ARCH=gfx1201
export TCNN_ENABLE_ROCWMMA_WIDTH64_MLP=1
export TCNN_HALF_PRECISION=1

cd bindings/torch
python -m pip install --no-build-isolation .

Without TCNN_ENABLE_ROCWMMA_WIDTH64_MLP=1, the backend is not compiled and the factory rejects RocWMMAWidth64MLP.

Minimal usage

import torch
import tinycudann as tcnn

model = tcnn.Network(
    n_input_dims=64,
    n_output_dims=64,
    network_config={
        "otype": "RocWMMAWidth64MLP",
        "activation": "ReLU",
        "output_activation": "None",
        "n_neurons": 64,
        "n_hidden_layers": 2,
    },
).to("cuda").eval()

x = torch.randn(257, 64, device="cuda", dtype=torch.float16)

with torch.inference_mode():
    y = model(x)

assert y.shape == (257, 64)
assert y.dtype == torch.float16
assert torch.isfinite(y).all()

Runtime validation

The production path passed:

  • correctness comparison against a CPU FP32 reference
  • public batches including 1, 16, 17, 255, 256, and 257
  • a 20-case runtime batch matrix
  • internal padding boundaries from 256 through 1024
  • 64 bitwise-identical repeated launches
  • tile-prefix invariance
  • parameter hot swap A -> B -> A
  • two model instances on two non-default streams
  • replay and stream isolation checks
  • construction of the existing portable and hipBLASLt factories
  • fail-closed rejection of training-forward execution
  • two fresh-process runtime replays with byte-identical JSON results

Production code-object audit

The linked production object was audited through its embedded gfx1201 AMDGPU code object.

Recorded kernel facts:

Group segment / LDS:     2048 bytes
Private segment:            0 bytes
VGPR field:                92
SGPR field:                74

MFMA/WMMA instructions:    12
ds_load_b128:               8
ds_store_b128:              2
ds_bpermute_b32:          192
Block barriers:             6
Scratch instructions:       0

The audit also verified:

  • exactly one executable production-kernel symbol
  • correct classification of compiler-generated metadata symbols
  • exact DS mnemonic inventory
  • zero private-segment allocation
  • zero scratch instructions
  • reproducible extraction of the gfx1201 code object
  • byte identity of the production object, linked extension, extracted code object, source, mapping header, network bridge, and runtime replay results across the final audit correction

Important limitations

This release does not claim:

  • backward support for RocWMMAWidth64MLP
  • training support for RocWMMAWidth64MLP
  • automatic backend selection
  • compatibility with arbitrary MLP widths or depths
  • compatibility with AMD architectures other than the qualified gfx1201 target
  • replacement of upstream NVIDIA FullyFusedMLP
  • an occupancy result derived from register counts
  • a performance advantage over the existing hipBLASLt backends
  • CUDA RTC or JIT fusion support

No performance claim is made for the rocWMMA backend in this release. Performance qualification is intentionally deferred to a later phase.

The existing Phase 3B1 performance figures continue to apply only to the qualified hipBLASLt FP16 backend.

Validated environment

Operating system:  Ubuntu 24.04
GPU:               AMD Radeon AI PRO R9700
Architecture:      gfx1201
ROCm:              7.2
Python:            3.12
PyTorch:           2.13.0 + ROCm 7.2

Release anchor

Tag:
phase4a2-width64-production-inference-gfx1201-pass

Commit:
cd1330a21452f7e2edab9e676567b7a040f922bc

Project status

The existing qualified backends remain available:

  • PortableMLP — FP32 forward, backward, and training
  • HipBLASLtMLP — FP32 forward, backward, and training
  • HipBLASLtMLPFP16 — FP16 forward, backward, and training
  • RocWMMAWidth64MLP — opt-in FP16 fused inference for the exact qualified Width-64 topology

This is an independent community port of NVlabs/tiny-cuda-nn and is not affiliated with or endorsed by NVIDIA or AMD.

Assets 2
Loading