Repository navigation
refactor: make data type mappings and shared CUDA headers device-aware - #38
Merged
Merged
Conversation
…6`, and `DispatchFunc`
Collaborator
Author
|
There was a problem hiding this comment.
Pull request overview
This PR refactors the dtype/type-dispatch layer to be device-aware so each backend can provide its own float16/bfloat16 C++ types, and removes backend-specific CUDA symbols from shared src/cuda/ headers to avoid include-order fragility across GPU platforms.
Changes:
- Make
TypeMap/IsFP16/IsBFloat16andDispatchFuncdevice-parameterized, with per-backenddevice_.hspecializations. - Refactor shared CUDA-like kernels (add/swiglu/rms_norm/causal_softmax) to use device-aware
TypeMapType<kDev,...>andCaster<kDev>, and move block-size selection behindBackend::GetOptimalBlockSize(). - Change BLAS handle management to a static singleton handle.
Reviewed changes
Copilot reviewed 43 out of 44 changed files in this pull request and generated 21 comments.
Show a summary per file
| File | Description |
|---|---|
| src/tensor.cc | Updates dtype dispatch to pass Device::Type::kCpu. |
| src/pybind11_utils.h | Builds Torch device-name map from ActiveDevices list. |
| src/dispatcher.h | Adds Device::Type parameterization for DataType dispatch adapters and list-alias dispatch. |
| src/data_type.h | Makes TypeMap device-aware; updates IsFP16/IsBFloat16 predicates accordingly. |
| src/common/constexpr_map.h | Minor formatting update. |
| src/cpu/device_.h | New CPU TypeMap specializations for FP16/BF16. |
| src/cpu/caster_.h | Includes CPU device mapping and updates FP16/BF16 predicates to be device-aware. |
| src/cpu/add/add.h | Updates dtype dispatch and FP16/BF16 predicate usage for CPU. |
| src/cpu/swiglu/swiglu.h | Updates dtype dispatch and FP16/BF16 predicate usage for CPU. |
| src/cpu/rms_norm/rms_norm.h | Updates dtype dispatch for CPU. |
| src/cpu/causal_softmax/causal_softmax.h | Updates dtype dispatch for CPU. |
| src/cpu/gemm/gemm.h | Updates dtype dispatch for CPU. |
| src/cuda/caster_.h | Introduces CudaCasterImpl<kDev> and removes platform header selection from shared CUDA caster. |
| src/cuda/kernel_commons.h | Removes platform-specific includes and device-property helpers from shared CUDA header. |
| src/cuda/add/kernel.h | Uses Backend::GetOptimalBlockSize() and device-aware TypeMapType. |
| src/cuda/add/kernel.cuh | Makes the add kernel/device predicates device-aware. |
| src/cuda/swiglu/kernel.h | Uses Backend::GetOptimalBlockSize() and device-aware TypeMapType; passes kDev into kernel. |
| src/cuda/swiglu/kernel.cuh | Reworks swiglu to use IsFP16<kDev,T> / IsBFloat16<kDev,T> and Caster<kDev>. |
| src/cuda/rms_norm/kernel.h | Uses Backend::GetOptimalBlockSize() and device-aware TypeMapType; passes kDev into kernel; changes dtype checks to assert. |
| src/cuda/rms_norm/kernel.cuh | Makes casts device-aware via Caster<kDev>. |
| src/cuda/causal_softmax/kernel.h | Uses Backend::GetOptimalBlockSize() and device-aware TypeMapType; passes kDev into kernel; changes dtype checks to assert. |
| src/cuda/causal_softmax/kernel.cuh | Makes casts device-aware via Caster<kDev>. |
| src/cuda/gemm/blas.h | Replaces per-instance BLAS handle with a static singleton handle. |
| src/nvidia/device_.h | Adds NVIDIA TypeMap specializations, CUDA prop cache, QueryMaxThreadsPerBlock(), and Caster specialization. |
| src/nvidia/add/kernel.h | Adds kDeviceType and GetOptimalBlockSize(), and adjusts includes. |
| src/nvidia/swiglu/kernel.h | Adds kDeviceType and GetOptimalBlockSize(), and adjusts includes. |
| src/nvidia/rms_norm/kernel.h | Adds kDeviceType and GetOptimalBlockSize(), and adjusts includes. |
| src/nvidia/causal_softmax/kernel.h | Adds kDeviceType and GetOptimalBlockSize(), and adjusts includes. |
| src/iluvatar/device_.h | Adds Iluvatar TypeMap specializations, CUDA prop cache, QueryMaxThreadsPerBlock(), and Caster specialization. |
| src/iluvatar/add/kernel.h | Adds kDeviceType and GetOptimalBlockSize(), and adjusts includes. |
| src/iluvatar/swiglu/kernel.h | Adds kDeviceType and GetOptimalBlockSize(), and adjusts includes. |
| src/iluvatar/rms_norm/kernel.h | Adds kDeviceType and GetOptimalBlockSize(), and adjusts includes. |
| src/iluvatar/causal_softmax/kernel.h | Adds kDeviceType and GetOptimalBlockSize(), and adjusts includes. |
| src/metax/device_.h | Adds MetaX TypeMap specializations, QueryMaxThreadsPerBlock(), and Caster specialization. |
| src/metax/add/kernel.h | Adds kDeviceType and GetOptimalBlockSize(), and adjusts includes. |
| src/metax/swiglu/kernel.h | Adds kDeviceType and GetOptimalBlockSize(), and adjusts includes. |
| src/metax/rms_norm/kernel.h | Adds kDeviceType and GetOptimalBlockSize(), and adjusts includes. |
| src/metax/causal_softmax/kernel.h | Adds kDeviceType and GetOptimalBlockSize(), and adjusts includes. |
| src/moore/device_.h | Adds Moore TypeMap specializations, QueryMaxThreadsPerBlock(), and Caster specialization. |
| src/moore/add/kernel.h | Adds kDeviceType and GetOptimalBlockSize(), and adjusts includes. |
| src/moore/swiglu/kernel.h | Adds kDeviceType and GetOptimalBlockSize(), and adjusts includes. |
| src/cambricon/device_.h | New Cambricon TypeMap specializations for FP16/BF16. |
| src/cambricon/rms_norm/rms_norm.h | Includes Cambricon device mapping and updates dispatch to pass device type. |
| src/cambricon/rms_norm/kernel.mlu | Formatting and pointer-cast cleanup; keeps explicit template instantiations. |
Comments suppressed due to low confidence (1)
src/cuda/kernel_commons.h:13
kernel_commons.hnow usessize_t/ptrdiff_tin the public API but doesn't include<cstddef>. Relying on transitive includes makes this header non-self-contained and can break builds depending on include order. Add the direct standard includes needed for the types it uses.
#include "caster.h"
namespace infini::ops {
using AllCudaBlockSizes = List<128, 256, 512, 1024, 2048>;
__forceinline__ __device__ __host__ size_t
IndexToOffset(size_t flat_index, size_t ndim, const size_t* shape,
const ptrdiff_t* strides) {
size_t res = 0;
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
… `clang-format off` and `clang-format on`
Collaborator
Author
|
Moore |
Collaborator
|
Ziminli
requested changes
Apr 1, 2026
Ziminli
requested changes
Apr 1, 2026
Collaborator
Author
|
英伟达上的错误是因为 CI 运行时选择的卡 OOM 了，我换了张卡单独跑了一下： |
Ziminli
requested changes
Apr 1, 2026
Collaborator
Author
|
|
Ziminli
approved these changes
Apr 1, 2026
10 of 19 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
TypeMap,IsFP16,IsBFloat16, andDispatchFuncaccept aDevice::Typetemplate parameter so each platform can specialize float16/bfloat16 to its own C++ type (half,__half,__mt_bfloat16, etc.)QueryMaxThreadsPerBlock,half,DevicePropertyCache) from sharedcuda/headers; each platform'sdevice_.hnow provides its ownTypeMapspecializations,Casterdefinition, andBackend::GetOptimalBlockSize()cuda/swiglukernels to useIsFP16<kDev, T>/IsBFloat16<kDev, T>predicates andCaster<kDev>instead of direct CUDA type names, making them include-order-independentBlasby replacing per-instanceblasHandle_twith a static singletonMotivation
The shared
cuda/headers previously referenced platform-specific types likehalfandhalf2directly. This created implicit include-order dependencies —clang-format'sIncludeBlocks: Regroupwould silently break compilation by reordering includes. Every file that touched these headers needed// clang-format off/onguards.Additionally,
TypeMap<dtype>mappedkFloat16/kBFloat16to a single C++ type, ignoring that each platform uses a different underlying type.