Repository navigation
Conversation
Grouped MXFP8 dequantization stored its per-tensor TMA descriptor storage in file-global `static __device__` arrays (`g_tensor_maps_input` / `g_tensor_maps_output`). Like the quantize path, that storage is shared by every kernel launch across all CUDA streams, so concurrent `group_dequantize` calls on different streams race on the descriptors: one stream's `update_tma_descriptors` overwrites tensor maps another stream is still consuming, producing incorrect output. Allocate per-invocation `CUtensorMap` arrays with `cudaMallocAsync` and pass them to both the descriptor-update and consumer kernels, then release them with `cudaFreeAsync`. The storage stays in global memory (required for TMA to consume the descriptors) but is disjoint across concurrent invocations; the stream-ordered allocate/free keeps the descriptors alive until the kernels that use them finish. This is the same fix as the quantize path in NVIDIA#3774, applied to the sibling dequantize path, which had the identical file-global buffer. Signed-off-by: zhaoye <772971548@qq.com>
github-actions
Bot
added
the
community-contribution
for more information, see https://pre-commit.ci
Split the grouped MXFP8 dequantize race fix into its own PR (NVIDIA#3775) so the quantize and dequantize fixes can be reviewed independently. This reverts the dequantize change from this branch; the quantize fix (grouped_tma.cuh + group_quantize_mxfp8.cuh) and its regression test remain here. Signed-off-by: zhaoye <772971548@qq.com>
 [Medium impact] The change affects grouped MXFP8 dequantization calls that use multiple CUDA streams. The race fix appears safe to merge, with non-blocking cleanup and test improvements. FindingsSummaryReplaces shared MXFP8 dequantize descriptors with arrays owned by each call.
DiagramsequenceDiagram
participant L as Launcher
participant S as CUDA stream
participant M as Private descriptor arrays
L->>S: Allocate input and output arrays
S->>M: Write descriptors
S->>M: Read descriptors during dequantize
S->>M: Free arrays after kernel completion
Reviews (1) · Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." · Reviewed by Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| NVTE_CHECK_CUDA( | ||
| cudaMallocAsync(reinterpret_cast<void **>(&tensor_maps_input), tmap_bytes, stream)); | ||
| NVTE_CHECK_CUDA(cudaMallocAsync(reinterpret_cast<void **>(&tensor_maps_output), | ||
| tmap_bytes, stream)); |
There was a problem hiding this comment.
If the first cudaMallocAsync succeeds but the second runs out of memory, NVTE_CHECK_CUDA throws before either cudaFreeAsync runs. The raw tensor_maps_input pointer is then lost, leaving its device memory allocated. Use scope-based cleanup or one allocation for both arrays so a failed call releases what it already allocated.
| if (!is_single_tensor) { | ||
| const size_t tmap_bytes = MAX_SUPPORTED_TENSOR_DESCRIPTORS * sizeof(CUtensorMap); | ||
| NVTE_CHECK_CUDA( | ||
| cudaMallocAsync(reinterpret_cast<void **>(&tensor_maps_input), tmap_bytes, stream)); | ||
| NVTE_CHECK_CUDA(cudaMallocAsync(reinterpret_cast<void **>(&tensor_maps_output), | ||
| tmap_bytes, stream)); |
There was a problem hiding this comment.
The new allocation path has no regression test in this PR. The existing test_group_dequantize_cudagraph_capturable uses equal shapes, which selects is_single_tensor and skips these allocations entirely. Add a test with different last dimensions that checks independent outputs on two concurrent CUDA streams. Also capture and replay that shape case to cover the new allocation path.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
|
This has the same issue as #3774. |
Same file-global TMA-descriptor race as #3774, but in the grouped MXFP8 dequantize path. Allocate per-invocation CUtensorMap arrays with cudaMallocAsync and pass them to both the descriptor-update and consumer kernels, releasing with cudaFreeAsync. Independent of #3774 (no shared files), so the two PRs do not conflict. Grouped-quantize counterpart: #3774.