diff --git a/fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.cu b/fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.cu index 0023774b5..6d54bf22c 100644 --- a/fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.cu +++ b/fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.cu @@ -4,6 +4,7 @@ #include #include #include +#include #include "common/common.hpp" #include "norm/layernorm.hpp" @@ -14,10 +15,6 @@ auto layer_norm( std::optional const B, std::optional Output ) { - using ElementIn = float; - using ElementOut = float; - using ElementWeight = float; - int64_t const m = Input.size(0); int64_t const n = Input.size(1); torch::Device const input_device = Input.device(); @@ -26,31 +23,70 @@ auto layer_norm( Output.emplace( torch::empty( {m, n}, - torch::TensorOptions().device(input_device).dtype(torch::kFloat32) + torch::TensorOptions().device(input_device).dtype(Input.scalar_type()) ) ); } + TORCH_CHECK(Output.value().scalar_type() == Input.scalar_type(), + "Output dtype must match Input dtype. Got Output=", + Output.value().scalar_type(), ", Input=", Input.scalar_type()); + if (W.has_value()) { + TORCH_CHECK(W.value().scalar_type() == Input.scalar_type(), + "W dtype must match Input dtype. Got W=", + W.value().scalar_type(), ", Input=", Input.scalar_type()); + } + if (B.has_value()) { + TORCH_CHECK(B.value().scalar_type() == Input.scalar_type(), + "B dtype must match Input dtype. Got B=", + B.value().scalar_type(), ", Input=", Input.scalar_type()); + } void *Iptr = Input.data_ptr(); void *Wptr = W.has_value() ? W.value().data_ptr() : nullptr; void *Bptr = B.has_value() ? B.value().data_ptr() : nullptr; void *Optr = Output.value().data_ptr(); - BOOL_SWITCH(B.has_value(), BIAS, [&]{ - BOOL_SWITCH(W.has_value(), AFFINE, [&]{ - CONFIG_SWITCH(n, [&]{ - layernorm< - ElementIn, ElementOut, ElementWeight, - AFFINE, BIAS, - MAX_HIDDEN_SIZE, NUM_THR_PER_CTA> ( - Iptr, Wptr, Bptr, - Optr, eps, m, n, - at::cuda::getCurrentCUDAStream().stream() - ); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + if (Input.scalar_type() == at::kHalf) { + using ElementIn = cutlass::half_t; + using ElementOut = cutlass::half_t; + using ElementWeight = cutlass::half_t; + BOOL_SWITCH(B.has_value(), BIAS, [&]{ + BOOL_SWITCH(W.has_value(), AFFINE, [&]{ + CONFIG_SWITCH(n, [&]{ + layernorm( + Iptr, Wptr, Bptr, Optr, eps, m, n, stream); + }); }); }); - }); + } else if (Input.scalar_type() == at::kBFloat16) { + using ElementIn = cutlass::bfloat16_t; + using ElementOut = cutlass::bfloat16_t; + using ElementWeight = cutlass::bfloat16_t; + BOOL_SWITCH(B.has_value(), BIAS, [&]{ + BOOL_SWITCH(W.has_value(), AFFINE, [&]{ + CONFIG_SWITCH(n, [&]{ + layernorm( + Iptr, Wptr, Bptr, Optr, eps, m, n, stream); + }); + }); + }); + } else if (Input.scalar_type() == at::kFloat) { + using ElementIn = float; + using ElementOut = float; + using ElementWeight = float; + BOOL_SWITCH(B.has_value(), BIAS, [&]{ + BOOL_SWITCH(W.has_value(), AFFINE, [&]{ + CONFIG_SWITCH(n, [&]{ + layernorm( + Iptr, Wptr, Bptr, Optr, eps, m, n, stream); + }); + }); + }); + } else { + TORCH_CHECK(false, "Unsupported dtype for layer_norm_cuda: ", Input.scalar_type()); + } diff --git a/fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.hpp b/fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.hpp index b129e0229..733e27550 100644 --- a/fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.hpp +++ b/fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.hpp @@ -68,9 +68,22 @@ public: // mean reduction float u = _reduce_sum(x, shared_data) / params.n; + // IMPORTANT: + // Loader pads out-of-range lanes with 0. That is OK for the sum, but after + // subtracting mean, those padded lanes become -u and would incorrectly + // contribute to the variance. Mask them back to 0 before variance reduction. + // We launch exactly NumThrPerCta threads for a 1xMaxHiddenSize tile, + // so each thread is responsible for a contiguous chunk in N. + int thr_n_offset = tidx * NumElementPerThread; CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < NumElementPerThread; ++i) - x[i] -= u; + for (int i = 0; i < NumElementPerThread; ++i) { + int idx = thr_n_offset + i; + if (idx < params.n) { + x[i] -= u; + } else { + x[i] = 0.f; + } + } __syncthreads(); // var reduction diff --git a/fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.cu b/fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.cu index dd1c4d8a8..17b64e890 100644 --- a/fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.cu +++ b/fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.cu @@ -4,6 +4,7 @@ #include #include #include +#include #include #include "common/common.hpp" @@ -16,10 +17,6 @@ auto rms_norm( std::optional& Output ) { - using ElementIn = float; - using ElementOut = float; - using ElementWeight = float; - int64_t const m = Input.size(0); int64_t const n = Input.size(1); torch::Device const input_device = Input.device(); @@ -28,27 +25,51 @@ auto rms_norm( Output.emplace( torch::empty( {m, n}, - torch::TensorOptions().device(input_device).dtype(torch::kFloat32) + torch::TensorOptions().device(input_device).dtype(Input.scalar_type()) ) ); } + TORCH_CHECK(Output.value().scalar_type() == Input.scalar_type(), + "Output dtype must match Input dtype. Got Output=", + Output.value().scalar_type(), ", Input=", Input.scalar_type()); + if (Weight.has_value()) { + TORCH_CHECK(Weight.value().scalar_type() == Input.scalar_type(), + "Weight dtype must match Input dtype. Got Weight=", + Weight.value().scalar_type(), ", Input=", Input.scalar_type()); + } + void *Iptr = Input.data_ptr(); void *Wptr = Weight.has_value() ? Weight.value().data_ptr() : nullptr; void *Optr = Output.value().data_ptr(); - - CONFIG_SWITCH(n, [&]{ - rmsnorm< - ElementIn, ElementOut, ElementWeight, - MAX_HIDDEN_SIZE, NUM_THR_PER_CTA - > ( - Iptr, Wptr, - Optr, - eps, m, n, - at::cuda::getCurrentCUDAStream().stream() - ); -}); + if (Input.scalar_type() == at::kHalf) { + using ElementIn = cutlass::half_t; + using ElementOut = cutlass::half_t; + using ElementWeight = cutlass::half_t; + CONFIG_SWITCH(n, [&]{ + rmsnorm( + Iptr, Wptr, Optr, eps, m, n, at::cuda::getCurrentCUDAStream().stream()); + }); + } else if (Input.scalar_type() == at::kBFloat16) { + using ElementIn = cutlass::bfloat16_t; + using ElementOut = cutlass::bfloat16_t; + using ElementWeight = cutlass::bfloat16_t; + CONFIG_SWITCH(n, [&]{ + rmsnorm( + Iptr, Wptr, Optr, eps, m, n, at::cuda::getCurrentCUDAStream().stream()); + }); + } else if (Input.scalar_type() == at::kFloat) { + using ElementIn = float; + using ElementOut = float; + using ElementWeight = float; + CONFIG_SWITCH(n, [&]{ + rmsnorm( + Iptr, Wptr, Optr, eps, m, n, at::cuda::getCurrentCUDAStream().stream()); + }); + } else { + TORCH_CHECK(false, "Unsupported dtype for rms_norm_cuda: ", Input.scalar_type()); + } return Output; diff --git a/fastvideo-kernel/pyproject.toml b/fastvideo-kernel/pyproject.toml index 22ec90bac..d593c6f38 100644 --- a/fastvideo-kernel/pyproject.toml +++ b/fastvideo-kernel/pyproject.toml @@ -9,7 +9,7 @@ build-backend = "scikit_build_core.build" [project] name = "fastvideo-kernel" -version = "0.2.3" +version = "0.2.4" description = "Unified CUDA kernels for FastVideo" readme = "README.md" requires-python = ">=3.10" diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py index 0ed935dbd..7206034fd 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os from typing import Tuple import torch @@ -23,6 +24,12 @@ def _is_sm90() -> bool: return major == 9 and minor == 0 +def _force_triton() -> bool: + # Force Triton even on SM90 and even if the compiled extension is available. + # Useful for CI / debugging / parity testing. + return os.environ.get("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", "0") == "1" + + def _map_to_index_torch(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """ Pure-torch (no triton) conversion: @@ -281,7 +288,7 @@ def block_sparse_attn( - Otherwise: uses Triton implementation (requires q/k/v to have same padded length today). """ block_sparse_fwd, block_sparse_bwd = _get_sm90_ops() - if _is_sm90() and (block_sparse_fwd is not None) and (block_sparse_bwd is not None): + if (not _force_triton()) and _is_sm90() and (block_sparse_fwd is not None) and (block_sparse_bwd is not None): return block_sparse_attn_sm90(q, k, v, block_map, variable_block_sizes) # Triton path: generally assumes q/k/v share the same padded length if q.shape[2] != k.shape[2] or q.shape[2] != v.shape[2]: diff --git a/fastvideo-kernel/python/fastvideo_kernel/ops.py b/fastvideo-kernel/python/fastvideo_kernel/ops.py index f3910c9b2..7173351c0 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/ops.py +++ b/fastvideo-kernel/python/fastvideo_kernel/ops.py @@ -46,14 +46,22 @@ def sliding_tile_attention( flag = shape_map[seq_shape] for head_idx, (t, h, w) in enumerate(window_size): + # Per-head slices are not contiguous in the batch dimension when batch>1 + # (they keep the original head-stride). The TK kernel assumes contiguous + # [B, H, S, D] layout, so we materialize a contiguous [B,1,S,D] view. + q_h = q[:, head_idx:head_idx + 1].contiguous() + k_h = k[:, head_idx:head_idx + 1].contiguous() + v_h = v[:, head_idx:head_idx + 1].contiguous() + o_h = torch.empty_like(q_h) sta_fwd( - q[:, head_idx:head_idx + 1], k[:, head_idx:head_idx + 1], - v[:, head_idx:head_idx + 1], output[:, head_idx:head_idx + 1], + q_h, k_h, + v_h, o_h, t, h, w, text_length, False, has_text, flag ) + output[:, head_idx:head_idx + 1] = o_h if has_text: - sta_fwd(q, k, v, output, 3, 3, 3, text_length, True, True, flag) + sta_fwd(q.contiguous(), k.contiguous(), v.contiguous(), output, 3, 3, 3, text_length, True, True, flag) return output[:, :, :seq_length] diff --git a/fastvideo-kernel/python/fastvideo_kernel/version.py b/fastvideo-kernel/python/fastvideo_kernel/version.py index d31c31eae..788da1fb3 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/version.py +++ b/fastvideo-kernel/python/fastvideo_kernel/version.py @@ -1 +1 @@ -__version__ = "0.2.3" +__version__ = "0.2.4" diff --git a/fastvideo/tests/training/VSA/test_training_loss_VSA.py b/fastvideo/tests/training/VSA/test_training_loss_VSA.py index 0a3e6743d..6b69ecd24 100644 --- a/fastvideo/tests/training/VSA/test_training_loss_VSA.py +++ b/fastvideo/tests/training/VSA/test_training_loss_VSA.py @@ -6,6 +6,11 @@ import json from huggingface_hub import snapshot_download import torch +# Ensure backend selection happens during import-time initialization +os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN" +# Force VSA to use Triton implementation even on H100 / when CUDA extension is available +# os.environ["FASTVIDEO_KERNEL_VSA_FORCE_TRITON"] = "1" + # Import the training pipeline sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent)) from fastvideo.training.wan_training_pipeline import main @@ -19,8 +24,6 @@ h200_reference_wandb_summary_file = "fastvideo/tests/training/VSA/h200_reference NUM_NODES = "1" NUM_GPUS_PER_NODE = "2" -os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN" - def run_worker(): """Worker function that will be run on each GPU""" # Create and populate args @@ -118,10 +121,10 @@ def test_distributed_training(): wandb_summary = json.load(open(summary_file)) fields_and_thresholds = { - 'avg_step_time': 1.0, - 'grad_norm': 0.1, - 'step_time': 1.0, - 'train_loss': 0.02 + 'avg_step_time': 3, + 'grad_norm': 0.2, + 'step_time': 2.5, + 'train_loss': 0.04 } failures = [] diff --git a/pyproject.toml b/pyproject.toml index 0f0ad8609..d46492542 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -63,7 +63,7 @@ dependencies = [ "remote-pdb", # Kernel & Packaging - "fastvideo-kernel==0.2.3", + "fastvideo-kernel==0.2.4", "wheel", # Training Dependencies diff --git a/pyproject_other.toml b/pyproject_other.toml index 76de271f6..746501a6d 100644 --- a/pyproject_other.toml +++ b/pyproject_other.toml @@ -63,7 +63,7 @@ dependencies = [ "remote-pdb", # Kernel & Packaging - "fastvideo-kernel==0.2.3", + "fastvideo-kernel==0.2.4", "wheel", # Training Dependencies