Refactor tensor class in C++ unit tests

[timmoon10]

Description

The tensor wrapper in the C++ unit tests has become unwieldy, with complicated interactions between recipes and memory management. This has recently resulted in bugs where we accidently didn't allocate a required buffer (#2943). This PR disentangles the memory management from the recipe logic by adding a simple RAII class to manage GPU and CPU buffers. I've also added more explicit checks, e.g. when we assume a tensor is a single FP32.

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

• Add class to manage GPU buffer, CPU buffer, and memory transfers between them.
• Remove memory management logic from tensor class in C++ tests.
• Add checks to accessors that make implicit assumptions on buffer size and dtype.

Checklist:

• I have read and followed the contributing guidelines
• 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

[greptile-apps]

Greptile Summary

This PR replaces ad-hoc cudaMalloc/cudaFree calls and scattered CPU mirror pointers in the C++ test Tensor class with a compact inner Buffer RAII class that owns matched CPU and GPU allocations and exposes typed accessors. Recipe-specific allocation logic is reorganised into a switch in the constructor, removing a previously unreachable duplicate MXFP8 block and the exception-based scaling-mode probe.

• Tensor::Buffer: new inner class allocating both CPU and GPU buffers at construction, with to_cpu()/from_cpu() helpers and a GPUDeleter functor.
• Tensor member migration: old raw pointer fields replaced by optional<Buffer> and shared_ptr<Buffer>; ~Tensor simplified to = default.
• Operator test fixes: all 13 affected test files now cache ref_scale before launching the GPU kernel to avoid post-kernel reads of the scale buffer.

Confidence Score: 5/5

Safe to merge — pure test-infrastructure refactor with no production code affected and no correctness regressions found.

All allocations and transfers are covered by the new RAII Buffer class; NVTE_CHECK guards make failures loud. The duplicate MXFP8 block removal is safe because get_scales still handles NVTE_MXFP8_1D_SCALING in its first branch. The only findings are minor redundant cudaMemcpy calls and a removed dtype guard, neither of which causes incorrect test results.

No files require special attention beyond the minor observations on the shared scale_inv double-sync and missing type guards on rowwise_cpu_scale_inv_ptr.

Important Files Changed

Filename	Overview
tests/cpp/test_common.h	Adds inner Buffer RAII class; migrates Tensor fields to optional<Buffer> / shared_ptr<Buffer>; accessors strengthened with NVTE_CHECK guards.
tests/cpp/test_common.cu	Implements Buffer and Tensor memory methods; removes duplicate MXFP8 block from get_scales; replaces exception-based scaling-mode probe with explicit switch.
tests/cpp/operator/test_cast.cu	Captures ref_scale before the kernel runs and guards .scale() call with isFp8Type.
tests/cpp_distributed/test_comm_gemm.cu	Guards set_scale/set_scale_inv calls behind isFp8Type check.
tests/cpp/operator/test_cast_current_scaling.cu	Initialises ref_scale = 1.0 to fix prior UB where the variable was passed uninitialised in the non-FP8 path.
tests/cpp/operator/test_normalization.cu	Replaces post-kernel z.scale() call with pre-captured ref_scale.

Class Diagram

%%{init: {'theme': 'neutral'}}%%
classDiagram
    class Tensor {
        -TensorWrapper tensor_
        -optional~Buffer~ data_rowwise_
        -optional~Buffer~ data_columnwise_
        -shared_ptr~Buffer~ scale_inv_rowwise_
        -shared_ptr~Buffer~ scale_inv_columnwise_
        -optional~Buffer~ amax_rowwise_
        -optional~Buffer~ amax_columnwise_
        -optional~Buffer~ scale_
        +to_cpu()
        +from_cpu()
        +fill_uniform_rowwise_scale_inv()
        +fill_uniform_columnwise_scale_inv()
    }
    class Buffer {
        -unique_ptr cpu_buffer_
        -unique_ptr gpu_buffer_
        -size_t bytes_
        +to_cpu()
        +from_cpu()
        +cpu_buffer~T~() T*
        +gpu_buffer~T~() T*
    }
    class GPUDeleter {
        +operator()(void*)
    }
    Tensor *-- Buffer : owns
    Buffer *-- GPUDeleter : uses for cudaFree

Loading

_{Reviews (6): Last reviewed commit: "Typo" | Re-trigger Greptile}

[timmoon10]

This is weird, but it approximates the previous behavior.

[timmoon10]

/te-ci core L1

[Oleg-Goncharov]

LGTM, this looks much cleaner now, but the cast+transpose current scaling tests are failing with a segmentation fault.

[timmoon10]

/te-ci core L1

[timmoon10]

We handle MXFP8 earlier in this function, so this if-statement was redundant.