Hacktoberfest 2026: những issue maintainer đã đánh dấu cho tháng Mười, đang mở và phù hợp người mới. Xem issue Hacktoberfest

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

Đang mở
  • Có gì không ổn?
#3,630 0 bình luận 0 reaction 0 người được giao Xem trên GitHub

Maintainer thường phản hồi trong vòng 2 ngày

@kainzhong đang làm issue này rồi.

Từ ngày 5/10/2026.

  • #3632 của @kainzhong — đang mở

Đánh giá

Độ khó
4/5
Thời gian dự kiến
3-5 ngày
Mức phù hợp với người mới
25/100
Loại issue
Lỗi
Độ rõ ràng
Đặc tả rõ ràng
Mức độ hoạt động
Đình trệ
Công nghệ
cpp, python, pytorch
Lĩnh vực
machine-learning, performance

Hướng nghiên cứu

Start by reading transformer_engine/common/cast/core/grouped_tma.cuh, especially the global g_tensor_maps storage described in the issue, and review linked pull request #3632 to see what work is already underway. Run the supplied two-stream reproduction on a B200 and compare its outputs with the serial references; done means overlapping quantizations produce the same outputs as serial execution.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

bug

Describe the bug

In transformer_engine/common/cast/core/grouped_tma.cuh, TensorMapStorage is a global static variable

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;

which would be not concurrency safe if multiple kernels are launched and their execution overlaps

Steps/Code to reproduce bug

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")

Expected behavior

This should output the same result as when they are quantized in serial, but now it's 90/100 iterations produced wrong output

Environment overview (please complete the following information)

Not related

Environment details

Not related

Device details

  • B200

Additional context

N/A

Ngôn ngữ chính
Python
Star
3.6k
Fork
862
Merge trung bình
4 ngày 1 giờ
Pull request đã merge (30 ngày)
57

Chuẩn bị môi trường

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của NVIDIA/TransformerEngine

Tất cả issue của NVIDIA/TransformerEngine

Issue tương tự

Thêm issue về Python

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.