[KDA-Pilot] Add LTX2 QKNorm split-RoPE CUDA fast path - #29708
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a CUDA fast path for LTX2 Q/K RMSNorm + split RoPE, including its Python FFI bindings, integration into the LTX2 model forward pass, benchmarks, and unit tests. The review feedback suggests several improvements to enhance robustness and code cleanliness: using int64_t for tensor dimensions in the CUDA kernel to prevent potential integer overflows, removing redundant explicit type casts in both CUDA and Python wrappers, simplifying parameter checks for RMSNorm, and aligning the reference implementation in the unit tests with the benchmark code.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| __device__ inline float compute_rstd( | ||
| const __nv_bfloat16* __restrict__ xrow, | ||
| int hidden_size, | ||
| float eps, | ||
| int tid, | ||
| int lane, | ||
| int warp_id, | ||
| float* warp_sum, | ||
| float* s_rstd) { | ||
| float local = 0.f; | ||
| const int n_vec = hidden_size >> 2; | ||
| for (int i = tid; i < n_vec; i += kThreads) { |
There was a problem hiding this comment.
To prevent potential integer overflows with large tensors, it's safer to use int64_t for dimensions passed from Python, as PyTorch tensor dimensions are 64-bit. hidden_size and related variables in this function should be int64_t.
__device__ inline float compute_rstd(
const __nv_bfloat16* __restrict__ xrow,
int64_t hidden_size,
float eps,
int tid,
int lane,
int warp_id,
float* warp_sum,
float* s_rstd) {
float local = 0.f;
const int64_t n_vec = hidden_size >> 2;
for (int64_t i = tid; i < n_vec; i += kThreads) {
| int seq_len, | ||
| int num_heads, | ||
| int head_dim, | ||
| int64_t stride_cos_b, | ||
| int64_t stride_cos_h, | ||
| int64_t stride_cos_t, | ||
| int64_t stride_sin_b, | ||
| int64_t stride_sin_h, | ||
| int64_t stride_sin_t) { | ||
| const int row = blockIdx.x; | ||
| const int batch = row / seq_len; | ||
| const int token = row - batch * seq_len; | ||
| const int hidden_size = num_heads * head_dim; | ||
| const int half_dim = head_dim >> 1; | ||
| const auto* __restrict__ xrow = x + static_cast<int64_t>(row) * hidden_size; | ||
| auto* __restrict__ outrow = out + static_cast<int64_t>(row) * hidden_size; | ||
| const int tid = threadIdx.x + threadIdx.y * 32; | ||
| const int lane = threadIdx.x; | ||
| const int warp_id = threadIdx.y; | ||
|
|
||
| __shared__ float warp_sum[4]; | ||
| __shared__ float s_rstd; | ||
| const float rstd = compute_rstd(xrow, hidden_size, eps, tid, lane, warp_id, warp_sum, &s_rstd); | ||
|
|
||
| const int num_pairs = num_heads * half_dim; | ||
| for (int pair = tid; pair < num_pairs; pair += kThreads) { |
There was a problem hiding this comment.
To prevent potential integer overflows, dimensions passed from Python should be handled as int64_t. This includes the kernel parameters, derived dimension variables, and loop counters.
int64_t seq_len,
int64_t num_heads,
int64_t head_dim,
int64_t stride_cos_b,
int64_t stride_cos_h,
int64_t stride_cos_t,
int64_t stride_sin_b,
int64_t stride_sin_h,
int64_t stride_sin_t) {
const int row = blockIdx.x;
const int batch = row / seq_len;
const int token = row - batch * seq_len;
const int64_t hidden_size = static_cast<int64_t>(num_heads) * head_dim;
const int64_t half_dim = head_dim >> 1;
const auto* __restrict__ xrow = x + static_cast<int64_t>(row) * hidden_size;
auto* __restrict__ outrow = out + static_cast<int64_t>(row) * hidden_size;
const int tid = threadIdx.x + threadIdx.y * 32;
const int lane = threadIdx.x;
const int warp_id = threadIdx.y;
__shared__ float warp_sum[4];
__shared__ float s_rstd;
const float rstd = compute_rstd(xrow, hidden_size, eps, tid, lane, warp_id, warp_sum, &s_rstd);
const int64_t num_pairs = num_heads * half_dim;
for (int64_t pair = tid; pair < num_pairs; pair += kThreads) {
| static_cast<int>(seq_len), | ||
| static_cast<int>(num_heads), | ||
| static_cast<int>(head_dim), |
There was a problem hiding this comment.
With the kernel arguments updated to int64_t, these static_cast<int> are no longer necessary and should be removed to pass the full 64-bit values.
seq_len,
num_heads,
head_dim,
| int(q.shape[0] * q.shape[1]), | ||
| int(q.shape[1]), | ||
| int(k.shape[0] * k.shape[1]), | ||
| int(k.shape[1]), | ||
| int(num_heads), | ||
| int(head_dim), |
There was a problem hiding this comment.
Corresponding to the change to int64_t in the CUDA kernel, these explicit casts to int should be removed. The TVM FFI will handle passing Python integers as int64_t to the C++ backend.
| int(q.shape[0] * q.shape[1]), | |
| int(q.shape[1]), | |
| int(k.shape[0] * k.shape[1]), | |
| int(k.shape[1]), | |
| int(num_heads), | |
| int(head_dim), | |
| q.shape[0] * q.shape[1], | |
| q.shape[1], | |
| k.shape[0] * k.shape[1], | |
| k.shape[1], | |
| num_heads, | |
| head_dim, |
| or not isinstance(q_norm, nn.RMSNorm) | ||
| or not isinstance(k_norm, nn.RMSNorm) | ||
| or q_norm.weight is None | ||
| or k_norm.weight is None | ||
| or q_norm.eps is None | ||
| or k_norm.eps is None | ||
| or float(q_norm.eps) != float(eps) | ||
| or float(k_norm.eps) != float(eps) |
There was a problem hiding this comment.
The checks for weight and eps being None appear to be redundant. torch.nn.RMSNorm (used when tp_size==1) is created with elementwise_affine=True (the default), so weight will be a Parameter. eps is also a required float argument. The custom LTX2TPRMSNormAcrossHeads also ensures these attributes are not None. Removing these unnecessary checks would improve code clarity.
| or not isinstance(q_norm, nn.RMSNorm) | |
| or not isinstance(k_norm, nn.RMSNorm) | |
| or q_norm.weight is None | |
| or k_norm.weight is None | |
| or q_norm.eps is None | |
| or k_norm.eps is None | |
| or float(q_norm.eps) != float(eps) | |
| or float(k_norm.eps) != float(eps) | |
| or not isinstance(q_norm, nn.RMSNorm) | |
| or not isinstance(k_norm, nn.RMSNorm) | |
| or float(q_norm.eps) != float(eps) | |
| or float(k_norm.eps) != float(eps) |
| out = split_x * cos_u | ||
| first_out = out[..., :1, :] | ||
| second_out = out[..., 1:, :] | ||
| first_out.addcmul_(-sin_u, second_x) | ||
| second_out.addcmul_(sin_u, first_x) | ||
| out = out.reshape(*out.shape[:-2], last) |
There was a problem hiding this comment.
This part of the reference implementation can be simplified by applying addcmul_ directly on slices of the out tensor, which would make it more concise and consistent with the implementation in the benchmark file.
| out = split_x * cos_u | |
| first_out = out[..., :1, :] | |
| second_out = out[..., 1:, :] | |
| first_out.addcmul_(-sin_u, second_x) | |
| second_out.addcmul_(sin_u, first_x) | |
| out = out.reshape(*out.shape[:-2], last) | |
| out = split_x * cos_u | |
| out[..., :1, :].addcmul_(-sin_u, second_x) | |
| out[..., 1:, :].addcmul_(sin_u, first_x) | |
| out = out.reshape(*out.shape[:-2], last) |
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
900251f to
8b5ee20
Compare
7399c56 to
06a3d4c
Compare
|
1 participant
Powered by KDA-Pilot, we propose a native CUDA JIT fast path for the LTX-2.3 Q/K RMSNorm + split-RoPE attention preprocessing pattern.
Motivation
Add a native CUDA JIT fast path for the LTX-2.3 Q/K RMSNorm + split-RoPE attention preprocessing pattern from KDA-Pilot task
b200_ltx2_qknorm_split_rope__bitwise.The hot pattern is equivalent to:
for BF16 contiguous
[B, S, H]Q/K tensors and 4D split-RoPE cos/sin tensors with production LTX-2.3 layouts. The KDA final candidate preserves bitwise equality to this BF16 attention-input contract while avoiding the separate eager RMSNorm/RoPE materialization path.Modifications
diffusion_ltx2_qknorm_split_ropelightweight JIT CUDA custom op.torch.compile(fullgraph=True)sees an opaque custom op instead of tracing JIT/module loading.head_dim in {64, 128}, BF16 inputs/weights/cos/sin, real non-contiguous cos/sin strides, and independent Q/K sequence lengths.LTX2Attentionto try this CUDA path by default for TP=1,torch.nn.RMSNorm, 4D split-RoPE inputs, then fall back to the existing RMSNorm + RoPE implementation if unsupported or if JIT load/launch fails once.torch.compile(fullgraph=True)custom-op coverage.Accuracy Tests
Result: local syntax/format/lint/diff checks passed.
B200 unit test:
Result:
5 passed.Speed Tests and Profiling
KDA-Pilot task
b200_ltx2_qknorm_split_rope__bitwise, final k22 run on an idle B200:torch.equal/ zero tolerance for both Q and K outputs.head_dim in {64, 128}, video/audio/cross rows, sequence lengths from 126 to 32640.fallback_count == 0across the production grid.Integrated benchmark script CI-small check on B200:
B200 LTX-2.3 HQ end-to-end A/B with the same k22 fused path, no
torch.compile:The fused run recorded
hit=16384, fallback=0for this path.B200 Fused-vs-Unfused Accuracy
The fused CUDA path was validated on B200 against the original unfused PyTorch implementation for the exact Q/K attention-preprocessing contract it replaces:
For the supported LTX-2.3 production shapes, the CUDA output is bitwise-identical to the unfused PyTorch output before attention consumes Q/K, so this optimization does not change model accuracy relative to the original non-fused path.
head_dim in {64, 128}, video/audio/cross rows, sequence lengths 126 to 32640torch.equalfor both Q and K on every rowtorch.compile(fullgraph=True)custom-op coverage5 passedhit=16384,fallback=0Accuracy conclusion: on B200, the optimization produces the same Q/K tensors as the original non-fused PyTorch path for all covered production inputs; precision is unchanged.
Result Image Comparison
B200 LTX-2.3 HQ A/B output sample (
seed=42,1920x1088, 121 frames). The figure compares the original unfused path against the fused CUDA QKNorm + split-RoPE path at frames 0, 60, and 120.Video-level decode comparison:
SSIM All=1.000000,PSNR=inf.Checklist
torch.compilecompatibility.Review and Merge Process
This is a focused KDA-Pilot kernel integration PR. The expected useful effect is the faster LTX-2.3 Q/K normalization + split-RoPE preprocessing group; whole-model speedup is visible but remains bounded because attention, MLP, and decode still dominate full video generation runtime.
CI States
Latest PR Test (Base): ⏳ Run #28491339822
Latest PR Test (Extra): ⏳ Run #28491339740