[BUG][Ascend] Launcher workspace size overflows in 32-bit multiplication before uint64_t assignment

[YaooXu]

Summary

The Ascend generated C++ launcher computes the workspace byte count using 32-bit unsigned multiplication, then assigns the already-overflowed result to uint64_t. A launch requiring 4.5 GiB can therefore allocate only 0.5 GiB, leading to out-of-bounds device workspace accesses.

This is still present at the currently checked Ascend branch head, triton_v3.5.x @ 1aaab9566f011f3050da0f8c69cd7a1d773a9436:

third_party/ascend/backend/driver.py, lines 858–862

uint32_t blockNum4Workspace = gridX * gridY * gridZ;
// ... pre_launch hook ...
uint64_t totalWorkSpaceSize = {workspace_size} * blockNum4Workspace;

For example, {workspace_size} becomes the decimal literal 147456. Its multiplication with uint32_t blockNum4Workspace takes place in 32 bits. The uint64_t destination does not widen the multiplication operands.

Minimal reproducer — no NPU required

Save as workspace_overflow.cpp:

#include <cassert>
#include <cstdint>
#include <iostream>

int main() {
    uint32_t blocks = 32768;
    uint64_t bad4 = 131072 * blocks;
    uint64_t good4 = static_cast<uint64_t>(131072) * blocks;
    uint64_t bad45 = 147456 * blocks;
    uint64_t good45 = static_cast<uint64_t>(147456) * blocks;
    uint64_t bad6 = 196608 * blocks;
    uint64_t good6 = static_cast<uint64_t>(196608) * blocks;
    std::cout << "4 GiB: " << bad4 << " vs " << good4 << '\n';
    std::cout << "4.5 GiB: " << bad45 << " vs " << good45 << '\n';
    std::cout << "6 GiB: " << bad6 << " vs " << good6 << '\n';
    assert(bad4 == 0 && good4 == 4294967296ULL);
    assert(bad45 == 536870912 && good45 == 4831838208ULL);
    assert(bad6 == 2147483648ULL && good6 == 6442450944ULL);
    blocks = 16384;
    uint64_t small_bad = 147456 * blocks;
    uint64_t small_good = static_cast<uint64_t>(147456) * blocks;
    assert(small_bad == small_good && small_good == 2415919104ULL);
}

Run:

g++ -std=c++17 -Wall -Wextra -Werror -O2 workspace_overflow.cpp -o workspace_overflow
./workspace_overflow

Observed output (buggy vs expected bytes):

4 GiB: 0 vs 4294967296
4.5 GiB: 536870912 vs 4831838208
6 GiB: 2147483648 vs 6442450944

I reran this reproducer, plus 22 workspace-product boundary cases and an exact-4-GiB case, while preparing this report. This is a host arithmetic reproduction, not a GPU/NPU allocation test.

Prior device reproduction and patch validation

Device evidence was obtained on an installed FlagTree build, not on a newly built copy of the current branch head:

• Ascend 910B4-1; driver 25.2.0.
• FlagTree 0.6.0+ascend.gitc286cba6; imported Triton 3.5.1.
• PyTorch 2.9.0+cpu; torch_npu 2.9.0.post2; Python 3.11.

A plain-Triton attention kernel compiled with 147456 bytes of workspace per logical block. At 32768 blocks, the generated launcher allocated 536870912 bytes, instead of 4831838208 bytes. The launch failed with an illegal GM-address/access-timeout error. Reducing KV length, while retaining the same grid and workspace requirement, still failed; a longer-compute case with 16384 blocks and a non-overflowing 2.25-GiB workspace succeeded. This distinguishes the allocation-size defect from simply exceeding a compute-duration limit.

After applying the one-line cast below to the installed launcher generator, independent fresh processes successfully ran the original 4.5-GiB case and a 6-GiB case. We checked the generated C++ expression, allocator peak memory, finite outputs, and sampled reference correctness. Both cases were repeated twice. This is targeted launcher validation, not a claim of full attention correctness or unrestricted large-grid support.

The current upstream source still has the same defective expression; the new host reproducer above confirms its arithmetic. No new NPU test of current upstream HEAD is claimed.

Suggested fix

Widen an operand before multiplying:

-  uint64_t totalWorkSpaceSize = {workspace_size} * blockNum4Workspace;
+  uint64_t totalWorkSpaceSize = static_cast<uint64_t>({workspace_size}) * blockNum4Workspace;

Please add a launcher-generation/arithmetic regression covering below 4 GiB, exactly 4 GiB, 4.5 GiB, and 6 GiB, without requiring multi-GiB device allocation in every CI run. Also retain normal allocation-failure handling for genuinely insufficient device memory.

This change fixes the workspace-byte multiplication only; it does not claim to fix separate grid-product overflow, launch-dimension limits, or all other large-grid failures. A separate larger-grid diagnostic stalled after the local patch and remains unexplained, so it is not counted as validation of this fix.

[ZQ-Struggle]

The issue is in the Ascend launcher’s workspace-size calculation: multiplication is performed using 32-bit unsigned int, so overflow occurs before the result is assigned to uint64_t.
For example, 147456 × 32768 should require 4.5 GiB but evaluates to only 512 MiB, causing workspace under-allocation and potentially illegal memory accesses or timeouts. The fix is to cast an operand to uint64_t before multiplication so that the calculation uses 64-bit arithmetic.

[ZQ-Struggle]

I've opened #1195 to address this issue, targeting triton_v3.5.x. The fix widens the workspace-size operand to uint64_t before multiplication and adds regression coverage for the 4 GiB boundary and larger workspace sizes.