[7/n] Migrate pos_encoding and norm kernels to libtorch stable ABI (continued) #43209
+577
−419
There are no files selected for viewing
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
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
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -18,21 +18,20 @@ | |
| #include <cuda_runtime.h> | ||
| #include <type_traits> | ||
|
|
||
| #include "torch_utils.h" | ||
|
|
||
|
Contributor
Author
There was a problem hiding this comment. Merge conflicts location: deletion of |
||
| #include "../async_util.cuh" | ||
| #include "../cuda_compat.h" | ||
| #include "../type_convert.cuh" | ||
| #include "dispatch_utils.h" | ||
|
|
||
| #define CHECK_TYPE(x, st) \ | ||
| STD_TORCH_CHECK(x.scalar_type() == st, #x " dtype is ", x.scalar_type(), \ | ||
| ", while ", st, " is expected") | ||
| #define CHECK_TH_CUDA(x) \ | ||
| STD_TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor") | ||
| #define CHECK_CONTIGUOUS(x) \ | ||
| STD_TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") | ||
| #define CHECK_INPUT(x) \ | ||
| CHECK_TH_CUDA(x); \ | ||
| CHECK_CONTIGUOUS(x) | ||
|
|
@@ -589,8 +588,8 @@ void launchFusedQKNormRope(void* qkv, int const num_tokens, | |
| }); | ||
| break; | ||
| default: | ||
| STD_TORCH_CHECK( | ||
| false, "Unsupported head dimension for fusedQKNormRope: ", head_dim); | ||
| } | ||
| } | ||
|
|
||
|
|
@@ -604,10 +603,10 @@ void launchFusedQKNormRopeNTokenHeads( | |
| void const* k_weight, void const* cos_sin_cache, bool const interleave, | ||
| int64_t const* position_ids, int const token_heads_per_warp, | ||
| cudaStream_t stream) { | ||
|
Contributor
Author
There was a problem hiding this comment. This block wasn't here before so I had to update TORCH_CHECK. |
||
| STD_TORCH_CHECK(token_heads_per_warp == 1 || token_heads_per_warp == 2 || | ||
| token_heads_per_warp == 4 || token_heads_per_warp == 8, | ||
| "token_heads_per_warp must be 1, 2, 4, or 8, got ", | ||
| token_heads_per_warp); | ||
|
|
||
| // token_heads_per_warp == 1: delegate to the 1-head baseline kernel. | ||
| if (token_heads_per_warp == 1) { | ||
|
|
@@ -691,7 +690,7 @@ void launchFusedQKNormRopeNTokenHeads( | |
| }); \ | ||
| break; \ | ||
| default: \ | ||
|
Contributor
Author
There was a problem hiding this comment. This block wasn't here before so I had to update TORCH_CHECK. |
||
| STD_TORCH_CHECK(false, "Unsupported head dimension: ", head_dim); \ | ||
| } \ | ||
| } while (0) | ||
|
|
||
|
|
@@ -708,60 +707,64 @@ void launchFusedQKNormRopeNTokenHeads( | |
| } // namespace tensorrt_llm::kernels | ||
|
|
||
| void fused_qk_norm_rope( | ||
|
Contributor
Author
There was a problem hiding this comment. Merge conflicts location: addition of the "forced_token_heads_per_warp" |
||
| torch::stable::Tensor& | ||
| qkv, // Combined QKV tensor [num_tokens, | ||
| // (num_heads_q+num_heads_k+num_heads_v)*head_dim] | ||
| int64_t num_heads_q, // Number of query heads | ||
| int64_t num_heads_k, // Number of key heads | ||
| int64_t num_heads_v, // Number of value heads | ||
| int64_t head_dim, // Dimension per head | ||
| double eps, // Epsilon for RMS normalization | ||
| torch::stable::Tensor& q_weight, // RMSNorm weights for query [head_dim] | ||
| torch::stable::Tensor& k_weight, // RMSNorm weights for key [head_dim] | ||
| torch::stable::Tensor& cos_sin_cache, // Cos/sin cache [max_position, | ||
| // head_dim] | ||
| bool is_neox, // Whether RoPE is applied in Neox style | ||
| torch::stable::Tensor& position_ids, // Position IDs for RoPE [num_tokens] | ||
| int64_t forced_token_heads_per_warp // -1 = auto-select, >0 = forced value | ||
| ) { | ||
| // Input validation | ||
| CHECK_INPUT(qkv); | ||
| CHECK_INPUT(position_ids); | ||
| CHECK_INPUT(q_weight); | ||
| CHECK_INPUT(k_weight); | ||
| CHECK_INPUT(cos_sin_cache); | ||
| CHECK_TYPE(position_ids, torch::headeronly::ScalarType::Long); | ||
|
|
||
| STD_TORCH_CHECK(qkv.dim() == 2, | ||
| "QKV tensor must be 2D: [num_tokens, " | ||
| "(num_heads_q+num_heads_k+num_heads_v)*head_dim]"); | ||
| STD_TORCH_CHECK(position_ids.dim() == 1, | ||
| "Position IDs must be 1D: [num_tokens]"); | ||
| STD_TORCH_CHECK(q_weight.dim() == 1, "Query weights must be 1D: [head_dim]"); | ||
| STD_TORCH_CHECK(k_weight.dim() == 1, "Key weights must be 1D: [head_dim]"); | ||
| STD_TORCH_CHECK(cos_sin_cache.dim() == 2, | ||
| "Cos/sin cache must be 2D: [max_position, head_dim]"); | ||
| STD_TORCH_CHECK(q_weight.size(0) == head_dim, | ||
| "Query weights size must match head dimension"); | ||
| STD_TORCH_CHECK(k_weight.size(0) == head_dim, | ||
| "Key weights size must match head dimension"); | ||
|
|
||
| STD_TORCH_CHECK(cos_sin_cache.size(1) % 2 == 0, "rotary_dim must be even"); | ||
| STD_TORCH_CHECK(cos_sin_cache.size(1) <= head_dim, | ||
| "rotary_dim must be less than or equal to head_dim"); | ||
|
|
||
| STD_TORCH_CHECK(qkv.scalar_type() == q_weight.scalar_type() && | ||
| qkv.scalar_type() == k_weight.scalar_type(), | ||
| "qkv, q_weight and k_weight must have the same dtype"); | ||
|
|
||
| int64_t num_tokens = qkv.size(0); | ||
| STD_TORCH_CHECK(position_ids.size(0) == num_tokens, | ||
| "Number of tokens in position_ids must match QKV"); | ||
|
|
||
| int64_t total_heads = num_heads_q + num_heads_k + num_heads_v; | ||
| STD_TORCH_CHECK( | ||
| qkv.size(1) == total_heads * head_dim, | ||
| "QKV tensor size must match total number of heads and head dimension"); | ||
|
|
||
| const torch::stable::accelerator::DeviceGuard device_guard( | ||
| qkv.get_device_index()); | ||
| auto stream = get_current_cuda_stream(qkv.get_device_index()); | ||
|
|
||
|
Contributor
Author
There was a problem hiding this comment. Merge conflicts location: addition of device_id and stream. |
||
| // Select token_heads_per_warp: forced value if >0, else auto-select. | ||
| // Auto thresholds are calibrated on SM 9.0 (H100). On other architectures, | ||
|
|
@@ -771,8 +774,7 @@ void fused_qk_norm_rope( | |
| token_heads_per_warp = static_cast<int>(forced_token_heads_per_warp); | ||
| } else { | ||
| token_heads_per_warp = 1; | ||
| int sm_version = get_device_prop()->major * 10 + get_device_prop()->minor; | ||
| int64_t total_qk_units = num_tokens * (num_heads_q + num_heads_k); | ||
| if (sm_version == 90) { | ||
| if (head_dim >= 256) { | ||
|
|
@@ -795,21 +797,22 @@ void fused_qk_norm_rope( | |
| } | ||
| } | ||
|
|
||
| VLLM_STABLE_DISPATCH_HALF_TYPES( | ||
| qkv.scalar_type(), "fused_qk_norm_rope_kernel", [&] { | ||
| using qkv_scalar_t = scalar_t; | ||
| VLLM_STABLE_DISPATCH_FLOATING_TYPES( | ||
| cos_sin_cache.scalar_type(), "fused_qk_norm_rope_kernel", [&] { | ||
| using cache_scalar_t = scalar_t; | ||
| tensorrt_llm::kernels::launchFusedQKNormRopeNTokenHeads< | ||
| qkv_scalar_t, cache_scalar_t>( | ||
| qkv.data_ptr(), static_cast<int>(num_tokens), | ||
| static_cast<int>(num_heads_q), static_cast<int>(num_heads_k), | ||
| static_cast<int>(num_heads_v), static_cast<int>(head_dim), | ||
| static_cast<int>(cos_sin_cache.size(1)), | ||
| static_cast<float>(eps), q_weight.data_ptr(), | ||
| k_weight.data_ptr(), cos_sin_cache.data_ptr(), !is_neox, | ||
| reinterpret_cast<int64_t const*>(position_ids.data_ptr()), | ||
| token_heads_per_warp, stream); | ||
| }); | ||
| }); | ||
| } | ||
Oops, something went wrong.
Oops, something went wrong.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Outside scope of this PR, but technically this is exactly the same as #define VLLM_DISPATCH_GROUP_SIZE so we could use the stable version everywhere now