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