[BUG] Grouped MXFP8 quantization is not concurrency safe with multiple streams #3630

struct alignas(128) TensorMapStorage {
  alignas(128) CUtensorMap input[MAX_SUPPORTED_TENSOR_DESCRIPTORS];
  alignas(128) CUtensorMap act_input[MAX_SUPPORTED_TENSOR_DESCRIPTORS];
  alignas(128) CUtensorMap output_rowwise[MAX_SUPPORTED_TENSOR_DESCRIPTORS];
  alignas(128) CUtensorMap output_colwise[MAX_SUPPORTED_TENSOR_DESCRIPTORS];
  size_t rows[MAX_SUPPORTED_TENSOR_DESCRIPTORS];
  size_t cols[MAX_SUPPORTED_TENSOR_DESCRIPTORS];
  size_t offsets[MAX_SUPPORTED_TENSOR_DESCRIPTORS];
};

// Internal linkage avoids device-link ODR issues when this header is included by multiple .cu TUs.
static __device__ TensorMapStorage g_tensor_maps;

import torch
import transformer_engine.pytorch as te
import transformer_engine_torch as tex

NUM_GROUPS, ROWS, COLS_PER_GROUP = 8, 8192, 2048


def group_quantize(x):
    quantizer = te.MXFP8Quantizer(te.DType.kFloat8E4M3, rowwise=True, columnwise=False)
    last_dims = torch.full((NUM_GROUPS,), COLS_PER_GROUP, dtype=torch.int64, device="cuda")
    out = tex.group_quantize(x, quantizer, NUM_GROUPS, None, last_dims)
    return out.rowwise_data, out.scale_inv


x_a = torch.randn(ROWS, NUM_GROUPS * COLS_PER_GROUP, dtype=torch.bfloat16, device="cuda")
x_b = torch.randn_like(x_a) * 1000
ref_a = [t.clone() for t in group_quantize(x_a)]
ref_b = [t.clone() for t in group_quantize(x_b)]

s1, s2 = torch.cuda.Stream(), torch.cuda.Stream()
failures, iters = 0, 100
for _ in range(iters):
    torch.cuda.synchronize()
    with torch.cuda.stream(s1):
        out_a = group_quantize(x_a)
    with torch.cuda.stream(s2):
        out_b = group_quantize(x_b)
    torch.cuda.synchronize()

    if any(not torch.equal(o, r) for o, r in zip(out_a + out_b, ref_a + ref_b)):
        failures += 1
    # Poison freed outputs: the allocator reuses these blocks next iteration, and stale correct
    # bytes would hide tiles that were never written.
    for t in out_a + out_b:
        t.view(torch.uint8).fill_(0xA5)

print(f"{failures}/{iters} iterations produced wrong output")