From 101e6af45f45e02c8bd83465805394fa2c90a56b Mon Sep 17 00:00:00 2001 From: tangchengxiang <2064027004@qq.com> Date: Sat, 19 Sep 2026 13:08:06 +0000 Subject: [PATCH 1/2] issue/1565 fix(nvidia): complete runtime support for Qwen MTP Map E4M3 and BOOL through the existing ATen adaptor and preserve the caller's CUDA device across NCCL communicator destruction. Reuse the existing paged Prefill warp kernel for NVIDIA head size 256, without changing other vendors' default dispatch. Extend existing multi-page/long-context coverage and add finite FP8/mask cast checks. Validation: fresh SM86 build, 88 paged Prefill cases, 2 cast tests, and TP2 communicator teardown from both caller devices. Closes #1565 --- include/infinicore/adaptor/aten_adaptor.hpp | 4 ++ src/infiniccl/cuda/infiniccl_cuda.cu | 8 +++- .../cuda/kernel_v2.cuh | 2 +- .../nvidia/paged_attention_prefill_nvidia.cu | 15 +++++++ test/infinicore/ops/fp8_cast.py | 44 +++++++++++++++++++ .../infinicore/ops/paged_attention_prefill.py | 6 ++- 6 files changed, 76 insertions(+), 3 deletions(-) create mode 100644 test/infinicore/ops/fp8_cast.py diff --git a/include/infinicore/adaptor/aten_adaptor.hpp b/include/infinicore/adaptor/aten_adaptor.hpp index 0c9eed3ec..b71c0d63a 100644 --- a/include/infinicore/adaptor/aten_adaptor.hpp +++ b/include/infinicore/adaptor/aten_adaptor.hpp @@ -45,6 +45,8 @@ inline at::ScalarType to_at_dtype(DataType dtype) { return at::kHalf; case DataType::BF16: return at::kBFloat16; + case DataType::F8: + return at::kFloat8_e4m3fn; case DataType::I8: return at::kChar; case DataType::U8: @@ -53,6 +55,8 @@ inline at::ScalarType to_at_dtype(DataType dtype) { return at::kInt; case DataType::I64: return at::kLong; + case DataType::BOOL: + return at::kBool; default: throw std::runtime_error("Unsupported dtype for ATen"); } diff --git a/src/infiniccl/cuda/infiniccl_cuda.cu b/src/infiniccl/cuda/infiniccl_cuda.cu index 9b305afd0..303516418 100644 --- a/src/infiniccl/cuda/infiniccl_cuda.cu +++ b/src/infiniccl/cuda/infiniccl_cuda.cu @@ -119,7 +119,13 @@ infiniStatus_t commInitRank( } infiniStatus_t commDestroy(infinicclComm_t comm) { - CHECK_NCCL(ncclCommDestroy(getNcclComm(comm))); + // NCCL teardown may activate the communicator's device. Preserve the + // caller's device so its existing stream and runtime context stay valid. + int previous_device; + CHECK_INTERNAL(cudaGetDevice(&previous_device), cudaSuccess); + const auto status = ncclCommDestroy(getNcclComm(comm)); + CHECK_INTERNAL(cudaSetDevice(previous_device), cudaSuccess); + CHECK_NCCL(status); delete comm; return INFINI_STATUS_SUCCESS; } diff --git a/src/infiniop/ops/paged_attention_prefill/cuda/kernel_v2.cuh b/src/infiniop/ops/paged_attention_prefill/cuda/kernel_v2.cuh index bf58a167d..9033c9018 100644 --- a/src/infiniop/ops/paged_attention_prefill/cuda/kernel_v2.cuh +++ b/src/infiniop/ops/paged_attention_prefill/cuda/kernel_v2.cuh @@ -280,7 +280,7 @@ __global__ void PagedAttentionPrefillWarpGlobalKernel( ptrdiff_t o_head_stride) { constexpr int kWarpSize = 32; - static_assert(HEAD_SIZE == 64 || HEAD_SIZE == 128 || HEAD_SIZE == 192, "Only head_size 64/128/192 supported in v0.4."); + static_assert(HEAD_SIZE == 64 || HEAD_SIZE == 128 || HEAD_SIZE == 192 || HEAD_SIZE == 256, "Unsupported head_size."); static_assert(HEAD_SIZE % kWarpSize == 0, "HEAD_SIZE must be divisible by 32."); constexpr int DIMS_PER_THREAD = HEAD_SIZE / kWarpSize; diff --git a/src/infiniop/ops/paged_attention_prefill/nvidia/paged_attention_prefill_nvidia.cu b/src/infiniop/ops/paged_attention_prefill/nvidia/paged_attention_prefill_nvidia.cu index c3183525b..7fd686938 100644 --- a/src/infiniop/ops/paged_attention_prefill/nvidia/paged_attention_prefill_nvidia.cu +++ b/src/infiniop/ops/paged_attention_prefill/nvidia/paged_attention_prefill_nvidia.cu @@ -27,7 +27,11 @@ inline const char *default_prefill_kernel(const PagedAttentionPrefillInfo &info) return "ref"; } if (info.head_size == 256) { +#if defined(ENABLE_NVIDIA_API) + return "warp"; +#else return "ref"; +#endif } // Iluvatar/Hygon: use warp for the non-MLA shapes where it is the stable path. #if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) @@ -1020,6 +1024,17 @@ infiniStatus_t launch_prefill_warp( v_batch_stride, v_row_stride, v_head_stride, o_stride, o_head_stride); return INFINI_STATUS_SUCCESS; + case 256: + op::paged_attention_prefill::cuda::PagedAttentionPrefillWarpGlobalKernel + <<>>( + out, q, k_cache, v_cache, block_tables, total_kv_lens, cu_seqlens_q, alibi_slopes, + num_heads, num_seqs, num_kv_heads, total_q_tokens, scale, max_num_blocks_per_seq, + page_block_size, block_table_batch_stride, + q_stride, q_head_stride, + k_batch_stride, k_row_stride, k_head_stride, + v_batch_stride, v_row_stride, v_head_stride, + o_stride, o_head_stride); + return INFINI_STATUS_SUCCESS; default: return INFINI_STATUS_BAD_TENSOR_SHAPE; } diff --git a/test/infinicore/ops/fp8_cast.py b/test/infinicore/ops/fp8_cast.py new file mode 100644 index 000000000..feeab59ec --- /dev/null +++ b/test/infinicore/ops/fp8_cast.py @@ -0,0 +1,44 @@ +"""Check FP8 weights and boolean acceptance masks across the ATen cast bridge.""" + +import torch +from infinicore.lib import _infinicore + +import infinicore + + +def test_fp8_cast(): + bits = torch.arange(256, dtype=torch.int16) + bits = bits[(bits != 127) & (bits != 255)].to(torch.uint8) + source = bits.view(torch.float8_e4m3fn).cuda() + for dtype in (torch.float32, torch.bfloat16): + output = torch.empty(source.shape, dtype=dtype, device="cuda") + _infinicore.cast_( + infinicore.from_torch(output)._underlying, + infinicore.from_torch(source)._underlying, + ) + infinicore.sync_device() + torch.testing.assert_close(output, source.to(dtype), rtol=0, atol=0) + + +def test_bool_cast(): + candidates = torch.tensor([3, 5, 7, 9], device="cuda") + expected = torch.tensor([4, 5, 7, 8], device="cuda") + source = infinicore.equal( + infinicore.from_torch(candidates), infinicore.from_torch(expected) + ) + for dtype in (torch.float32, torch.int64): + output = torch.empty(source.shape, dtype=dtype, device="cuda") + _infinicore.cast_( + infinicore.from_torch(output)._underlying, + source._underlying, + ) + infinicore.sync_device() + torch.testing.assert_close( + output, (candidates == expected).to(dtype), rtol=0, atol=0 + ) + + +if __name__ == "__main__": + test_fp8_cast() + test_bool_cast() + print("Finite E4M3 and boolean acceptance mask casts passed") diff --git a/test/infinicore/ops/paged_attention_prefill.py b/test/infinicore/ops/paged_attention_prefill.py index 99036ec4f..a1e109fc1 100644 --- a/test/infinicore/ops/paged_attention_prefill.py +++ b/test/infinicore/ops/paged_attention_prefill.py @@ -26,6 +26,8 @@ (1, 24, 4, 256, 8, 8, 1), (1, 12, 2, 256, 8, 8, 1), (1, 6, 1, 256, 8, 8, 1), + (2, 12, 2, 256, 64, 128, 2), + (1, 24, 4, 256, 64, 1024, 1), # New DeepSeek MLA wrapper case: verifies prefill supports q/k head # size 576 with value head size 512. (1, 16, 1, 576, 8, 8, 1, 512), @@ -94,7 +96,9 @@ def parse_test_cases(): value_size, ) = case scale = head_size**-0.5 - num_blocks = 8192 + num_blocks = num_seqs * ( + (max_step_len * num_rounds + block_size - 1) // block_size + ) manager = SimpleCacheManager(num_blocks, block_size) kv_lens = torch.zeros(num_seqs, dtype=torch.int32) From a3ac4df44fc902fe223b441e0255631cc6b73a71 Mon Sep 17 00:00:00 2001 From: tangchengxiang <2064027004@qq.com> Date: Sun, 20 Sep 2026 11:05:06 +0000 Subject: [PATCH 2/2] issue/1565 fix(graph): include recurrent replay prerequisites Consolidate the allocator ownership, reduction/scalar-power recording and MetaX concatenation fixes from #1560 into the runtime support PR. Preserve the original implementations and focused regressions without adding a separate compiler or Prefill graph path. --- src/infinicore-test/memory_test.cc | 22 ++++++ .../allocators/pinnable_block_allocator.cc | 13 +++- src/infinicore/ops/cat/cat.cc | 4 +- src/infinicore/ops/float_power/float_power.cc | 22 ++++++ src/infinicore/ops/sum/sum_infiniop.cc | 67 +++++++++---------- src/infiniop/ops/sum/nvidia/sum_nvidia.cu | 3 +- test/infinicore/graph/test_cat.py | 48 +++++++++++++ .../infinicore/graph/test_sum_scalar_power.py | 45 +++++++++++++ 8 files changed, 185 insertions(+), 39 deletions(-) create mode 100644 test/infinicore/graph/test_cat.py create mode 100644 test/infinicore/graph/test_sum_scalar_power.py diff --git a/src/infinicore-test/memory_test.cc b/src/infinicore-test/memory_test.cc index 2029e0a81..59f753ba2 100644 --- a/src/infinicore-test/memory_test.cc +++ b/src/infinicore-test/memory_test.cc @@ -1,4 +1,5 @@ #include "memory_test.h" +#include "../infinicore/context/allocators/pinnable_block_allocator.hpp" #include #include #include @@ -78,6 +79,27 @@ TestResult BasicMemoryTest::run() { } spdlog::debug("BasicMemoryTest: Pinned memory test completed"); + PinnableBlockAllocator allocator(current_device); + auto reinstantiated = allocator.allocate(1024); + allocator.deallocate(reinstantiated); + allocator.mark_in_use_(reinstantiated, true); + allocator.trim(); + allocator.deallocate(reinstantiated); + allocator.trim(); + + auto captured = allocator.allocate(1024); + allocator.deallocate(captured); + allocator.set_pin_mode(true); + if (allocator.allocate(1024) != captured) { + return false; + } + allocator.deallocate(captured); + allocator.set_pin_mode(false); + allocator.trim(); + // A graph may reinstantiate storage after its temporary owner expires. + allocator.mark_in_use_(captured, true); + allocator.deallocate(captured); + return true; } catch (const std::exception &e) { std::cerr << "BasicMemoryTest failed with exception: " << e.what() << std::endl; diff --git a/src/infinicore/context/allocators/pinnable_block_allocator.cc b/src/infinicore/context/allocators/pinnable_block_allocator.cc index 32e5c5e9b..a4f2f7273 100644 --- a/src/infinicore/context/allocators/pinnable_block_allocator.cc +++ b/src/infinicore/context/allocators/pinnable_block_allocator.cc @@ -64,6 +64,7 @@ std::byte *PinnableBlockAllocator::allocate(size_t size) { cls.free_blocks.pop_back(); block->in_use = true; block->use_count = 1; + block->frozen = block->frozen || pinned_mode_; return reinterpret_cast(block->ptr); } } @@ -151,6 +152,16 @@ size_t PinnableBlockAllocator::mark_in_use_(void *ptr, bool in_use) { auto block = it->second; if (in_use) { + if (!block->in_use) { + // Reinstantiated graph storage is no longer available for trimming. + for (auto &cls : size_classes_) { + if (block->size == cls.block_size) { + auto &free = cls.free_blocks; + free.erase(std::remove(free.begin(), free.end(), block), free.end()); + break; + } + } + } block->in_use = true; ++block->use_count; } else if (block->use_count > 0) { @@ -166,7 +177,7 @@ void PinnableBlockAllocator::trim() { // Free non-frozen size-class blocks for (auto &cls : size_classes_) { for (auto it = cls.free_blocks.begin(); it != cls.free_blocks.end();) { - if (!(*it)->frozen) { + if (!(*it)->frozen && !(*it)->in_use) { INFINICORE_CHECK_ERROR(infinirtFree((*it)->ptr)); all_blocks_.erase((*it)->ptr); it = cls.free_blocks.erase(it); diff --git a/src/infinicore/ops/cat/cat.cc b/src/infinicore/ops/cat/cat.cc index 217ac7ec9..c48e48718 100644 --- a/src/infinicore/ops/cat/cat.cc +++ b/src/infinicore/ops/cat/cat.cc @@ -12,6 +12,7 @@ bool use_slice_copy_cat(Device::Type device_type, int dim, int ndim) { // correct through copy_from, but their performance impact is unverified. return dim == ndim - 1 && (device_type == Device::Type::NVIDIA + || device_type == Device::Type::METAX || device_type == Device::Type::HYGON || device_type == Device::Type::ILUVATAR || device_type == Device::Type::ALI); @@ -172,7 +173,8 @@ void cat_(Tensor out, std::vector tensors, int dim) { // index. Concatenating MLA tensors on the last dimension can therefore // enqueue hundreds of tiny D2D copy calls per layer. A strided output // slice is semantically identical and copy_from lowers it to a single - // rearrange kernel per input tensor on CUDA-like backends. + // rearrange kernel per input tensor on CUDA-like backends. Unlike + // direct memcpy calls, these copies are also recorded for graph replay. size_t offset = 0; for (auto &tensor : tensors) { if (tensor->ndim() == 1) { diff --git a/src/infinicore/ops/float_power/float_power.cc b/src/infinicore/ops/float_power/float_power.cc index c3bf5003b..a8e17b514 100644 --- a/src/infinicore/ops/float_power/float_power.cc +++ b/src/infinicore/ops/float_power/float_power.cc @@ -1,7 +1,25 @@ #include "infinicore/ops/float_power.hpp" +#include "infinicore/graph/graph.hpp" #include "infinicore/tensor.hpp" namespace infinicore::op { +namespace { + +class RecordedScalarPower final : public graph::GraphOperator { +public: + RecordedScalarPower(Tensor output, Tensor input, double exponent) + : output_(output), input_(input), exponent_(exponent) {} + + void run() const override { + FloatPower::dispatcher_scalar().lookup(input_->device().getType())(output_, input_, exponent_); + } + +private: + graph::GraphTensor output_, input_; + double exponent_; +}; + +} // namespace // ======================================================================= // 1. Dispatcher 单例 @@ -22,6 +40,10 @@ common::OpDispatcher &FloatPower::dispatcher_tensor() // ======================================================================= void FloatPower::execute(Tensor output, Tensor input, double exponent) { + if (context::isGraphRecording()) { + context::addGraphOperator(std::make_shared(output, input, exponent)); + return; + } dispatcher_scalar() .lookup(context::getDevice().getType())(output, input, exponent); } diff --git a/src/infinicore/ops/sum/sum_infiniop.cc b/src/infinicore/ops/sum/sum_infiniop.cc index 700ab04fa..c1f960959 100644 --- a/src/infinicore/ops/sum/sum_infiniop.cc +++ b/src/infinicore/ops/sum/sum_infiniop.cc @@ -1,47 +1,44 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" +#include "../infiniop_impl.hpp" +#include "infinicore/graph/graph.hpp" #include "infinicore/ops/sum.hpp" -#include namespace infinicore::op::sum_impl::infiniop { -thread_local common::OpCache caches( - 100, // capacity - [](infiniopSumDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroySumDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input, std::vector dim, bool keepdim) { - size_t seed = hash_combine(output, input, dim.size(), keepdim); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Sum, 100); - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopSumDescriptor_t desc = nullptr; +class RecordedSum final : public graph::GraphOperator { +public: + RecordedSum(Tensor output, Tensor input, std::vector dim, bool keepdim) + : output_(output), input_(input), dim_(std::move(dim)), keepdim_(keepdim) { + size_t seed = hash_combine(output, input, dim_.size(), keepdim_); + for (auto axis : dim_) { + hash_combine(seed, axis); + } + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( + Descriptor, descriptor, Sum, seed, + output->desc(), input->desc(), dim_.data(), dim_.size(), keepdim_); + INFINIOP_WORKSPACE_TENSOR(workspace, Sum, descriptor); + descriptor_ = std::move(descriptor); + workspace_ = std::make_unique(workspace); + } - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateSumDescriptor( - context::getInfiniopHandle(output->device()), &desc, - output->desc(), input->desc(), dim.data(), dim.size(), keepdim)); - cache.put(seed, desc); - } else { - desc = *desc_opt; + void run() const override { + auto output = output_; + INFINICORE_CHECK_ERROR(infiniopSum( + descriptor_->desc, (*workspace_)->data(), (*workspace_)->numel(), + output->data(), input_->data(), dim_.data(), dim_.size(), keepdim_, context::getStream())); } - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetSumWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); +private: + graph::GraphTensor output_, input_; + mutable std::vector dim_; + bool keepdim_; + std::shared_ptr descriptor_; + std::unique_ptr workspace_; +}; - INFINICORE_CHECK_ERROR(infiniopSum( - desc, workspace->data(), workspace_size, - output->data(), input->data(), dim.data(), dim.size(), keepdim, context::getStream())); +void calculate(Tensor output, Tensor input, std::vector dim, bool keepdim) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(RecordedSum, output, input, std::move(dim), keepdim); } static bool registered = []() { diff --git a/src/infiniop/ops/sum/nvidia/sum_nvidia.cu b/src/infiniop/ops/sum/nvidia/sum_nvidia.cu index 9f165d271..224a57f6c 100644 --- a/src/infiniop/ops/sum/nvidia/sum_nvidia.cu +++ b/src/infiniop/ops/sum/nvidia/sum_nvidia.cu @@ -59,8 +59,7 @@ infiniStatus_t launchKernel( CHECK_CUDA(cudaMemcpyAsync(output_strides_cuda, info.output_strides.data(), output_ndim * sizeof(ptrdiff_t), cudaMemcpyHostToDevice, stream)); if (info.reduce_num == input_size) { - T zero = static_cast(0.0f); - CHECK_CUDA(cudaMemcpyAsync(output, &zero, sizeof(T), cudaMemcpyHostToDevice, stream)); + CHECK_CUDA(cudaMemsetAsync(output, 0, sizeof(T), stream)); size_t grid_size = (input_size + BLOCK_SIZE - 1) / BLOCK_SIZE; sumAllKernel<<>>( output, input, input_size, input_ndim, permuted_input_shape_cuda, permuted_input_strides_cuda); diff --git a/test/infinicore/graph/test_cat.py b/test/infinicore/graph/test_cat.py new file mode 100644 index 000000000..1448755ac --- /dev/null +++ b/test/infinicore/graph/test_cat.py @@ -0,0 +1,48 @@ +"""Last-axis concatenation must replay its producers and strided copies.""" + +import pytest +import torch + +import infinicore + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="A GPU is required.") +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("batch", [1, 4]) +def test_last_axis_cat_replays_changed_inputs(dtype, batch): + device = infinicore.device("cuda", 0) + infinicore.set_device(device) + source = ( + torch.linspace(-1, 1, batch * 129, device="cuda") + .reshape(1, batch, 129) + .to(dtype) + ) + # Match projection views with a larger row stride than the sliced width. + inputs = [source[..., :64], source[..., 64:96], source[..., 96:]] + tensors = [ + infinicore.strided_from_blob( + tensor.data_ptr(), + list(tensor.shape), + list(tensor.stride()), + dtype=infinicore.utils.to_infinicore_dtype(dtype), + device=device, + ) + for tensor in inputs + ] + output = torch.empty_like(source) + target = infinicore.from_torch(output) + torch.cuda.synchronize() + + infinicore.start_graph_recording(device) + projected = infinicore.mul(tensors[0], tensors[0]) + joined = infinicore.cat([projected, *tensors[1:]], dim=-1) + infinicore.mul(joined, joined, out=target) + graph = infinicore.stop_graph_recording() + + for scale in (1.0, -0.5, 2.0): + source.mul_(scale) + torch.cuda.synchronize() + graph.run() + infinicore.sync_stream() + expected = torch.cat([inputs[0].square(), *inputs[1:]], dim=-1).square() + torch.testing.assert_close(output, expected) diff --git a/test/infinicore/graph/test_sum_scalar_power.py b/test/infinicore/graph/test_sum_scalar_power.py new file mode 100644 index 000000000..0cd9ed442 --- /dev/null +++ b/test/infinicore/graph/test_sum_scalar_power.py @@ -0,0 +1,45 @@ +"""Exercise the reduction and reciprocal norm used by tensor-parallel models.""" + +import pytest +import torch + +import infinicore + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="A CUDA device is required.") +def test_sum_descriptor_cache_distinguishes_axes(): + source = torch.arange(16, device="cuda", dtype=torch.float32).reshape(4, 4) + tensor = infinicore.from_torch(source) + for axis in (0, 1, 0): + output = torch.empty(4, device="cuda") + torch.cuda.synchronize() + infinicore.sum(tensor, dim=axis, out=infinicore.from_torch(output)) + infinicore.sync_device() + torch.testing.assert_close(output, source.sum(axis)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="A CUDA device is required.") +@pytest.mark.parametrize("keepdim", [False, True]) +@pytest.mark.parametrize("batch", [1, 4]) +def test_sum_scalar_power_replay(keepdim, batch): + device = infinicore.device("cuda", 0) + infinicore.set_device(device) + source = torch.linspace(0.25, 2.0, batch * 768, device="cuda").reshape(batch, 768) + output = torch.empty((batch, 1) if keepdim else (batch,), device="cuda") + input_tensor = infinicore.from_torch(source) + output_tensor = infinicore.from_torch(output) + torch.cuda.synchronize() + + infinicore.start_graph_recording(device) + squared = infinicore.mul(input_tensor, input_tensor) + reduced = infinicore.sum(squared, dim=1, keepdim=keepdim) + infinicore.float_power(reduced, -0.5, out=output_tensor) + graph = infinicore.stop_graph_recording() + + for scale in (1.0, 0.5, 3.0): + source.mul_(scale) + torch.cuda.synchronize() + graph.run() + infinicore.sync_stream() + expected = source.square().sum(dim=1, keepdim=keepdim).rsqrt() + torch.testing.assert_close(output, expected, atol=1e-6, rtol=1e-5)