From bb760a9f35680d81bb7c09e8e9bce1f28e5529ce Mon Sep 17 00:00:00 2001 From: PixelFlat <329975279+PixelFlat@users.noreply.github.com> Date: Sun, 25 Oct 2026 20:57:50 +0800 Subject: [PATCH] bugfix: preserve infinity masks in Ascend TileLang builds. --- .../apply_token_bitmask_wrapper_test.cpp | 34 +++++++++++++++++++ .../targets/ascend/aot/apply_token_bitmask.py | 5 ++- 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/tests/core/kernels/npu/tilelang/apply_token_bitmask_wrapper_test.cpp b/tests/core/kernels/npu/tilelang/apply_token_bitmask_wrapper_test.cpp index ee0f77c980..38a387045d 100644 --- a/tests/core/kernels/npu/tilelang/apply_token_bitmask_wrapper_test.cpp +++ b/tests/core/kernels/npu/tilelang/apply_token_bitmask_wrapper_test.cpp @@ -194,6 +194,40 @@ TEST_F(ApplyTokenBitmaskWrapperTest, MatchesReferenceAcrossDtypesAndPatterns) { } } +TEST_F(ApplyTokenBitmaskWrapperTest, + PreservesAllowedLogitsAndMasksWithInfinity) { + const torch::Device device("npu:0"); + const std::vector dtypes = { + torch::kFloat16, torch::kBFloat16, torch::kFloat32}; + for (const torch::ScalarType dtype : dtypes) { + SCOPED_TRACE(static_cast(dtype)); + const auto options = torch::TensorOptions().dtype(dtype).device(device); + torch::Tensor input = torch::arange(64, options).view({2, 32}) - 32; + input.select(/*dim=*/1, /*index=*/0) + .fill_(-std::numeric_limits::infinity()); + const torch::Tensor bitmask = make_mask(2, 32, "boundaries").to(device); + torch::Tensor actual = input.clone(); + + ASSERT_TRUE(can_apply_token_bitmask_inplace(actual, bitmask)); + apply_token_bitmask_inplace(actual, bitmask); + ASSERT_EQ(aclrtSynchronizeStream( + c10_npu::getCurrentNPUStream(device.index()).stream()), + ACL_SUCCESS); + + const torch::Tensor result = actual.to(torch::kCPU); + const torch::Tensor original = input.to(torch::kCPU); + EXPECT_TRUE(torch::equal(result.select(/*dim=*/1, /*index=*/0), + original.select(/*dim=*/1, /*index=*/0))); + EXPECT_TRUE(torch::equal(result.select(/*dim=*/1, /*index=*/31), + original.select(/*dim=*/1, /*index=*/31))); + EXPECT_TRUE( + torch::isneginf(result.slice(/*dim=*/1, /*start=*/1, /*end=*/31)) + .all() + .item()); + EXPECT_FALSE(torch::isnan(result).any().item()); + } +} + TEST_F(ApplyTokenBitmaskWrapperTest, RejectsInputsThatRequireFallback) { const auto device = torch::Device("npu:0"); const auto options = diff --git a/xllm/compiler/tilelang/targets/ascend/aot/apply_token_bitmask.py b/xllm/compiler/tilelang/targets/ascend/aot/apply_token_bitmask.py index 3ad1ca016c..9e1af34a68 100644 --- a/xllm/compiler/tilelang/targets/ascend/aot/apply_token_bitmask.py +++ b/xllm/compiler/tilelang/targets/ascend/aot/apply_token_bitmask.py @@ -52,4 +52,7 @@ def generate_source(dtype: str) -> str: ) with tilelang.tvm.transform.PassContext(opt_level=3, config=DEFAULT_ASCEND_PASS_CONFIGS): kernel = tilelang.engine.lower(tilelang_kernel) - return kernel.kernel_source + # TileLang can emit CUDA's infinity macro for Ascend float constants. + # Keep the mask at -inf without requiring CUDA headers or build flags. + compat_header = "#ifndef CUDART_INF\n#define CUDART_INF (__builtin_inff())\n#endif\n" + return compat_header + kernel.kernel_source