diff --git a/src/h3_blackwell_runtime/csrc/nvfp4_scale.cpp b/src/h3_blackwell_runtime/csrc/nvfp4_scale.cpp index 9291ed5..ecabdc8 100644 --- a/src/h3_blackwell_runtime/csrc/nvfp4_scale.cpp +++ b/src/h3_blackwell_runtime/csrc/nvfp4_scale.cpp @@ -5,7 +5,12 @@ torch::Tensor nvfp4_activation_scale_cuda(torch::Tensor input, double divisor); torch::Tensor nvfp4_activation_scale_into_cuda(torch::Tensor input, double divisor, torch::Tensor partials, torch::Tensor output, int64_t blocks, int64_t threads); std::vector quantize_nvfp4_bf16_cuda(torch::Tensor input, torch::Tensor scale, bool pad_16x, int64_t threads); +void quantize_nvfp4_bf16_into_cuda(torch::Tensor input, torch::Tensor scale, torch::Tensor qdata, torch::Tensor block_scale, int64_t threads, bool hi_first); +std::vector quantize_nvfp4_modulated_bf16_cuda(torch::Tensor input, torch::Tensor shift, torch::Tensor scale, torch::Tensor row_index, double divisor, bool pad_16x, int64_t blocks, int64_t scale_threads, int64_t pack_threads); +std::vector quantize_nvfp4_swiglu_bf16_cuda(torch::Tensor input, double divisor, bool pad_16x, int64_t blocks, int64_t scale_threads, int64_t pack_threads); std::vector qkv_to_bshd_cuda(torch::Tensor qkv, int64_t heads, int64_t head_dim); +std::vector sage2_prepare_qk_cuda(torch::Tensor q, torch::Tensor k, torch::Tensor rotation, torch::Tensor q_weight, torch::Tensor k_weight, double epsilon, bool materialize_q); +std::vector sage2_prepare_v_cuda(torch::Tensor v, double scale_max); torch::Tensor nvfp4_activation_scale(torch::Tensor input, double divisor) { TORCH_CHECK(input.is_cuda(), "nvfp4_activation_scale expects a CUDA tensor"); @@ -42,6 +47,50 @@ std::vector quantize_nvfp4_bf16(torch::Tensor input, torch::Tenso return quantize_nvfp4_bf16_cuda(input, scale, pad_16x, threads); } +void quantize_nvfp4_bf16_into(torch::Tensor input, torch::Tensor scale, torch::Tensor qdata, torch::Tensor block_scale, int64_t threads, bool hi_first) { + TORCH_CHECK(input.is_cuda() && scale.is_cuda() && qdata.is_cuda() && block_scale.is_cuda(), "quantize_nvfp4_bf16_into expects CUDA tensors"); + TORCH_CHECK(input.device() == scale.device() && input.device() == qdata.device() && input.device() == block_scale.device(), "quantize_nvfp4_bf16_into tensors must share a device"); + TORCH_CHECK(input.is_contiguous() && qdata.is_contiguous() && block_scale.is_contiguous(), "quantize_nvfp4_bf16_into tensors must be contiguous"); + TORCH_CHECK(input.dim() == 2 && qdata.dim() == 2 && block_scale.dim() == 2, "quantize_nvfp4_bf16_into expects 2D input and outputs"); + TORCH_CHECK(input.scalar_type() == torch::kBFloat16, "quantize_nvfp4_bf16_into expects bfloat16 input"); + TORCH_CHECK(scale.scalar_type() == torch::kFloat32 && scale.numel() == 1, "quantize_nvfp4_bf16_into expects a scalar float32 scale"); + TORCH_CHECK(qdata.scalar_type() == torch::kUInt8, "quantize_nvfp4_bf16_into expects uint8 qdata"); + TORCH_CHECK(block_scale.scalar_type() == torch::kFloat8_e4m3fn, "quantize_nvfp4_bf16_into expects float8_e4m3fn block scales"); + TORCH_CHECK(qdata.size(0) >= input.size(0) && qdata.size(1) >= input.size(1) / 2, "quantize_nvfp4_bf16_into qdata output is too small"); + const int64_t required_scale_rows = ((qdata.size(0) + 127) / 128) * 128; + const int64_t required_scale_cols = (((qdata.size(1) / 8) + 3) / 4) * 4; + TORCH_CHECK(block_scale.size(0) >= required_scale_rows && block_scale.size(1) >= required_scale_cols, "quantize_nvfp4_bf16_into block-scale output is too small"); + TORCH_CHECK(threads == 64 || threads == 128 || threads == 256 || threads == 512, "quantize_nvfp4_bf16_into threads must be 64, 128, 256, or 512"); + quantize_nvfp4_bf16_into_cuda(input, scale, qdata, block_scale, threads, hi_first); +} + +std::vector quantize_nvfp4_modulated_bf16(torch::Tensor input, torch::Tensor shift, torch::Tensor scale, torch::Tensor row_index, double divisor, bool pad_16x, int64_t blocks, int64_t scale_threads, int64_t pack_threads) { + TORCH_CHECK(input.is_cuda() && shift.is_cuda() && scale.is_cuda() && row_index.is_cuda(), "quantize_nvfp4_modulated_bf16 expects CUDA tensors"); + TORCH_CHECK(input.device() == shift.device() && input.device() == scale.device() && input.device() == row_index.device(), "quantize_nvfp4_modulated_bf16 tensors must share a device"); + TORCH_CHECK(input.is_contiguous() && shift.is_contiguous() && scale.is_contiguous() && row_index.is_contiguous(), "quantize_nvfp4_modulated_bf16 tensors must be contiguous"); + TORCH_CHECK(input.dim() == 2 && shift.dim() == 2 && scale.dim() == 2, "quantize_nvfp4_modulated_bf16 expects 2D input and tables"); + TORCH_CHECK(input.scalar_type() == torch::kBFloat16, "quantize_nvfp4_modulated_bf16 expects bfloat16 input"); + TORCH_CHECK(shift.scalar_type() == scale.scalar_type() && (shift.scalar_type() == torch::kBFloat16 || shift.scalar_type() == torch::kFloat32), "quantize_nvfp4_modulated_bf16 tables must share bfloat16 or float32 dtype"); + TORCH_CHECK(shift.sizes() == scale.sizes() && shift.size(1) == input.size(1), "quantize_nvfp4_modulated_bf16 table shapes must match the input width"); + TORCH_CHECK(row_index.scalar_type() == torch::kInt32 && row_index.dim() == 1 && row_index.numel() == input.size(0), "quantize_nvfp4_modulated_bf16 row_index must be int32 with one entry per row"); + TORCH_CHECK(divisor > 0.0, "quantize_nvfp4_modulated_bf16 divisor must be positive"); + TORCH_CHECK(blocks > 0, "quantize_nvfp4_modulated_bf16 blocks must be positive"); + TORCH_CHECK(scale_threads == 128 || scale_threads == 256 || scale_threads == 512, "quantize_nvfp4_modulated_bf16 scale_threads must be 128, 256, or 512"); + TORCH_CHECK(pack_threads == 64 || pack_threads == 128 || pack_threads == 256 || pack_threads == 512, "quantize_nvfp4_modulated_bf16 pack_threads must be 64, 128, 256, or 512"); + return quantize_nvfp4_modulated_bf16_cuda(input, shift, scale, row_index, divisor, pad_16x, blocks, scale_threads, pack_threads); +} + +std::vector quantize_nvfp4_swiglu_bf16(torch::Tensor input, double divisor, bool pad_16x, int64_t blocks, int64_t scale_threads, int64_t pack_threads) { + TORCH_CHECK(input.is_cuda() && input.is_contiguous(), "quantize_nvfp4_swiglu_bf16 expects contiguous CUDA input"); + TORCH_CHECK(input.dim() == 2 && input.scalar_type() == torch::kBFloat16, "quantize_nvfp4_swiglu_bf16 expects 2D bfloat16 input"); + TORCH_CHECK(input.size(1) % 2 == 0, "quantize_nvfp4_swiglu_bf16 input width must be even"); + TORCH_CHECK(divisor > 0.0 && blocks > 0, "quantize_nvfp4_swiglu_bf16 divisor and blocks must be positive"); + TORCH_CHECK(scale_threads == 128 || scale_threads == 256 || scale_threads == 512, "quantize_nvfp4_swiglu_bf16 scale_threads must be 128, 256, or 512"); + TORCH_CHECK(pack_threads == 64 || pack_threads == 128 || pack_threads == 256 || pack_threads == 512, "quantize_nvfp4_swiglu_bf16 pack_threads must be 64, 128, 256, or 512"); + return quantize_nvfp4_swiglu_bf16_cuda(input, divisor, pad_16x, blocks, scale_threads, pack_threads); +} + + std::vector qkv_to_bshd(torch::Tensor qkv, int64_t heads, int64_t head_dim) { TORCH_CHECK(qkv.is_cuda(), "qkv_to_bshd expects a CUDA tensor"); TORCH_CHECK(qkv.is_contiguous(), "qkv_to_bshd expects contiguous input"); @@ -51,9 +100,38 @@ std::vector qkv_to_bshd(torch::Tensor qkv, int64_t heads, int64_t return qkv_to_bshd_cuda(qkv, heads, head_dim); } +std::vector sage2_prepare_qk(torch::Tensor q, torch::Tensor k, torch::Tensor rotation, torch::Tensor q_weight, torch::Tensor k_weight, double epsilon, bool materialize_q) { + TORCH_CHECK(q.is_cuda() && k.is_cuda() && rotation.is_cuda() && q_weight.is_cuda() && k_weight.is_cuda(), "sage2_prepare_qk expects CUDA tensors"); + TORCH_CHECK(q.device() == k.device() && q.device() == rotation.device() && q.device() == q_weight.device() && q.device() == k_weight.device(), "sage2_prepare_qk tensors must share a device"); + TORCH_CHECK(q.scalar_type() == torch::kBFloat16 && k.scalar_type() == torch::kBFloat16, "sage2_prepare_qk expects BF16 Q and K"); + TORCH_CHECK(rotation.scalar_type() == torch::kBFloat16 && q_weight.scalar_type() == torch::kBFloat16 && k_weight.scalar_type() == torch::kBFloat16, "sage2_prepare_qk expects BF16 rotation and weights"); + TORCH_CHECK(q.dim() == 4 && k.dim() == 4 && q.sizes() == k.sizes(), "sage2_prepare_qk expects matching 4D Q and K"); + TORCH_CHECK(q.size(0) == 1 && q.size(3) == 128, "sage2_prepare_qk currently requires batch 1 and head dimension 128"); + TORCH_CHECK(q.stride(3) == 1 && k.stride(3) == 1, "sage2_prepare_qk requires contiguous head dimensions"); + TORCH_CHECK(rotation.dim() == 6 && rotation.size(0) == 1 && rotation.size(1) == q.size(1) && rotation.size(2) == 1 && rotation.size(3) > 0 && rotation.size(3) <= 64 && rotation.size(4) == 2 && rotation.size(5) == 2, "sage2_prepare_qk expects rotation shape [1, sequence, 1, pairs, 2, 2] with pairs <= 64; got rotation ", rotation.sizes(), " for Q ", q.sizes()); + TORCH_CHECK(q_weight.is_contiguous() && k_weight.is_contiguous() && q_weight.numel() == 128 && k_weight.numel() == 128, "sage2_prepare_qk expects contiguous 128-element weights"); + TORCH_CHECK(epsilon > 0.0, "sage2_prepare_qk epsilon must be positive"); + return sage2_prepare_qk_cuda(q, k, rotation, q_weight, k_weight, epsilon, materialize_q); +} + +std::vector sage2_prepare_v(torch::Tensor v, double scale_max) { + TORCH_CHECK(v.is_cuda(), "sage2_prepare_v expects a CUDA tensor"); + TORCH_CHECK(v.scalar_type() == torch::kBFloat16, "sage2_prepare_v expects BF16 V"); + TORCH_CHECK(v.dim() == 4, "sage2_prepare_v expects a 4D NHD tensor"); + TORCH_CHECK(v.size(0) == 1 && v.size(3) == 128, "sage2_prepare_v currently requires batch 1 and head dimension 128"); + TORCH_CHECK(v.stride(3) == 1, "sage2_prepare_v requires a contiguous head dimension"); + TORCH_CHECK(scale_max > 0.0, "sage2_prepare_v scale_max must be positive"); + return sage2_prepare_v_cuda(v, scale_max); +} + PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("nvfp4_activation_scale", &nvfp4_activation_scale, "Vortex NVFP4 activation scale"); m.def("nvfp4_activation_scale_into", &nvfp4_activation_scale_into, "Vortex NVFP4 activation scale with caller workspace"); m.def("quantize_nvfp4_bf16", &quantize_nvfp4_bf16, "Vortex BF16 to TensorCore NVFP4 quantizer"); + m.def("quantize_nvfp4_bf16_into", &quantize_nvfp4_bf16_into, "Vortex BF16 to TensorCore NVFP4 quantizer with caller outputs", py::arg("input"), py::arg("scale"), py::arg("qdata"), py::arg("block_scale"), py::arg("threads"), py::arg("hi_first") = true); + m.def("quantize_nvfp4_modulated_bf16", &quantize_nvfp4_modulated_bf16, "Vortex fused H3 modulation and TensorCore NVFP4 quantizer"); + m.def("quantize_nvfp4_swiglu_bf16", &quantize_nvfp4_swiglu_bf16, "Vortex fused H3 SwiGLU and TensorCore NVFP4 quantizer"); m.def("qkv_to_bshd", &qkv_to_bshd, "Fused H3 QKV split to BSHD tensors"); + m.def("sage2_prepare_qk", &sage2_prepare_qk, "Exact fused H3 Q/K RMSNorm, split-half RoPE, and Sage2 Q quantization", py::arg("q"), py::arg("k"), py::arg("rotation"), py::arg("q_weight"), py::arg("k_weight"), py::arg("epsilon"), py::arg("materialize_q") = false); + m.def("sage2_prepare_v", &sage2_prepare_v, "Direct Sage2 NHD V scale, FP8 quantization, and layout preparation", py::arg("v"), py::arg("scale_max") = 2.25); } diff --git a/src/h3_blackwell_runtime/csrc/nvfp4_scale.cu b/src/h3_blackwell_runtime/csrc/nvfp4_scale.cu index f6b85a5..3c323bf 100644 --- a/src/h3_blackwell_runtime/csrc/nvfp4_scale.cu +++ b/src/h3_blackwell_runtime/csrc/nvfp4_scale.cu @@ -3,6 +3,7 @@ #include #include +#include #include #include @@ -54,6 +55,63 @@ __inline__ __device__ float bf16_bits_to_float(uint16_t bits) { return __uint_as_float(static_cast(bits) << 16); } +__inline__ __device__ float round_bf16(float value) { + return __bfloat162float(__float2bfloat16_rn(value)); +} + +template +__inline__ __device__ float modulated_bf16_value( + const uint16_t* __restrict__ input, + const table_t* __restrict__ shift, + const table_t* __restrict__ scale, + const int32_t* __restrict__ row_index, + int64_t row, + int64_t col, + int64_t cols) { + const int64_t table_offset = static_cast(row_index[row]) * cols + col; + const float x = bf16_bits_to_float(input[row * cols + col]); + const float scale_bf16 = round_bf16(static_cast(scale[table_offset])); + const float shift_bf16 = round_bf16(static_cast(shift[table_offset])); + const float factor = round_bf16(1.0f + scale_bf16); + const float multiplied = round_bf16(x * factor); + return round_bf16(multiplied + shift_bf16); +} + +__inline__ __device__ float swiglu_bf16_value( + const uint16_t* __restrict__ input, + int64_t row, + int64_t col, + int64_t output_cols) { + const int64_t row_offset = row * output_cols * 2; + const float gate = bf16_bits_to_float(input[row_offset + col]); + const float up = bf16_bits_to_float(input[row_offset + output_cols + col]); + const float silu = round_bf16(gate / (1.0f + expf(-gate))); + return round_bf16(silu * up); +} + +__inline__ __device__ uint32_t warp_reduce_max_u32(uint32_t value) { + for (int offset = 16; offset > 0; offset >>= 1) { + value = max(value, __shfl_down_sync(0xffffffff, value, offset)); + } + return value; +} + +__inline__ __device__ uint32_t block_reduce_max_u32(uint32_t value) { + __shared__ uint32_t warp_values[32]; + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + value = warp_reduce_max_u32(value); + if (lane == 0) { + warp_values[warp] = value; + } + __syncthreads(); + value = threadIdx.x < ((blockDim.x + 31) >> 5) ? warp_values[lane] : 0; + if (warp == 0) { + value = warp_reduce_max_u32(value); + } + return value; +} + __inline__ __device__ uint8_t encode_fp4_e2m1(float value) { const bool negative = signbit(value); float abs_value = fabsf(value); @@ -87,6 +145,12 @@ __inline__ __device__ float decode_fp8_e4m3(uint8_t value) { return static_cast(encoded); } +__inline__ __device__ float reciprocal_approx_ftz(float value) { + float result; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(result) : "f"(value)); + return result; +} + __global__ void quantize_nvfp4_bf16_kernel( const uint16_t* __restrict__ input, const float* __restrict__ scale, @@ -97,7 +161,8 @@ __global__ void quantize_nvfp4_bf16_kernel( int64_t q_rows, int64_t q_cols, int64_t scale_rows, - int64_t scale_cols) { + int64_t scale_cols, + bool hi_first) { const int64_t row = blockIdx.x; const float tensor_scale = scale[0]; for (int64_t block_col = threadIdx.x; block_col < scale_cols; block_col += blockDim.x) { @@ -113,43 +178,197 @@ __global__ void quantize_nvfp4_bf16_kernel( values[i] = value; local_max = fmaxf(local_max, fabsf(value)); } - uint8_t scale_byte = 0; - float block_scale_value = 0.0f; - if (local_max > 0.0f && tensor_scale > 0.0f) { - scale_byte = encode_fp8_e4m3(local_max / (tensor_scale * 6.0f)); - block_scale_value = decode_fp8_e4m3(scale_byte); - } + const float raw_block_scale = (local_max / 6.0f) / tensor_scale; + const uint8_t scale_byte = encode_fp8_e4m3(fminf(raw_block_scale, 448.0f)); + const float block_scale_value = decode_fp8_e4m3(scale_byte); if (row < scale_rows) { - int64_t scale_row = row; - int64_t scale_col = block_col; - if (scale_cols == 336 || scale_cols == 896) { - const int64_t row_in_tile = row % 128; - const int64_t group = block_col / 4; - const int64_t pair = group / 2; - const int64_t phase = group % 2; - const int64_t col_stride = scale_cols == 336 ? 16 : 128; - const int64_t phase_offset = scale_cols == 336 ? 176 : 512; - int64_t swizzled_col_base = pair * col_stride + phase * phase_offset; - int64_t swizzled_col = (swizzled_col_base % scale_cols) + (block_col % 4) + (row_in_tile / 32) * 4 + (row_in_tile % 32) * 16; - scale_row = (row / 128) * 128 + pair * (scale_cols == 336 ? 3 : 1) + (scale_cols == 336 ? phase : 0) + swizzled_col_base / scale_cols + swizzled_col / scale_cols; - scale_col = swizzled_col % scale_cols; - if (scale_col >= scale_cols) { - scale_col -= scale_cols; - scale_row += 1; - } - } else if (scale_cols >= 32) { - const int64_t row_in_tile = row % 128; - scale_row = (row / 128) * 128 + (block_col / 4) * 16 + ((row_in_tile % 32) / 2); - scale_col = (block_col % 4) + (row_in_tile / 32) * 4 + (row_in_tile % 2) * 16; + const int64_t row_in_tile = row % 128; + const int64_t tile = ((row / 128) * (scale_cols / 4)) + (block_col / 4); + const int64_t within_tile = + ((row_in_tile % 32) / 2) * 32 + + (block_col % 4) + + (row_in_tile / 32) * 4 + + (row_in_tile % 2) * 16; + block_scale[tile * 512 + within_tile] = scale_byte; + } + if (row < q_rows && block_col * 8 < q_cols) { + #pragma unroll + for (int pair = 0; pair < 8; ++pair) { + const float encode_scale = fminf( + reciprocal_approx_ftz(tensor_scale * block_scale_value), + 3.402823466e+38f); + const uint8_t even = encode_fp4_e2m1(values[pair * 2] * encode_scale); + const uint8_t odd = encode_fp4_e2m1(values[pair * 2 + 1] * encode_scale); + qdata[row * q_cols + block_col * 8 + pair] = hi_first + ? static_cast((even << 4) | odd) + : static_cast(even | (odd << 4)); } - block_scale[scale_row * scale_cols + scale_col] = scale_byte; } + } +} + +template +__global__ void partial_absmax_modulated_bf16_kernel( + const uint16_t* __restrict__ input, + const table_t* __restrict__ shift, + const table_t* __restrict__ scale, + const int32_t* __restrict__ row_index, + float* __restrict__ partials, + int64_t rows, + int64_t cols) { + const int64_t numel = rows * cols; + const int64_t stride = static_cast(blockDim.x) * gridDim.x; + uint32_t local_bits = 0; + for (int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + index < numel; + index += stride) { + const int64_t row = index / cols; + const int64_t col = index - row * cols; + const float value = modulated_bf16_value(input, shift, scale, row_index, row, col, cols); + local_bits = max(local_bits, __float_as_uint(value) >> 16 & 0x7fffu); + } + const uint32_t block_bits = block_reduce_max_u32(local_bits); + if (threadIdx.x == 0) { + partials[blockIdx.x] = bf16_abs_bits_to_float(block_bits); + } +} + +__global__ void final_scale_bf16_compat_kernel( + const float* __restrict__ partials, + float* __restrict__ output, + int64_t count, + float divisor) { + float local_max = 0.0f; + for (int64_t index = threadIdx.x; index < count; index += blockDim.x) { + local_max = fmaxf(local_max, partials[index]); + } + const float block_max = block_reduce_max(local_max); + if (threadIdx.x == 0) { + output[0] = round_bf16(block_max / divisor); + } +} + +template +__global__ void quantize_nvfp4_modulated_bf16_kernel( + const uint16_t* __restrict__ input, + const table_t* __restrict__ shift, + const table_t* __restrict__ adaln_scale, + const int32_t* __restrict__ row_index, + const float* __restrict__ tensor_scale_ptr, + uint8_t* __restrict__ qdata, + uint8_t* __restrict__ block_scale, + int64_t rows, + int64_t cols, + int64_t q_rows, + int64_t q_cols, + int64_t scale_rows, + int64_t scale_cols) { + const int64_t row = blockIdx.x; + const float tensor_scale = tensor_scale_ptr[0]; + for (int64_t block_col = threadIdx.x; block_col < scale_cols; block_col += blockDim.x) { + float local_max = 0.0f; + float values[16]; + #pragma unroll + for (int i = 0; i < 16; ++i) { + const int64_t col = block_col * 16 + i; + float value = 0.0f; + if (row < rows && col < cols) { + value = modulated_bf16_value(input, shift, adaln_scale, row_index, row, col, cols); + } + values[i] = value; + local_max = fmaxf(local_max, fabsf(value)); + } + const float raw_block_scale = (local_max / 6.0f) / tensor_scale; + const uint8_t scale_byte = encode_fp8_e4m3(fminf(raw_block_scale, 448.0f)); + const float block_scale_value = decode_fp8_e4m3(scale_byte); + const int64_t row_in_tile = row % 128; + const int64_t tile = ((row / 128) * (scale_cols / 4)) + (block_col / 4); + const int64_t within_tile = + ((row_in_tile % 32) / 2) * 32 + + (block_col % 4) + + (row_in_tile / 32) * 4 + + (row_in_tile % 2) * 16; + block_scale[tile * 512 + within_tile] = scale_byte; if (row < q_rows && block_col * 8 < q_cols) { + const float encode_scale = fminf( + reciprocal_approx_ftz(tensor_scale * block_scale_value), + 3.402823466e+38f); #pragma unroll for (int pair = 0; pair < 8; ++pair) { - const float denom = tensor_scale * block_scale_value; - const uint8_t even = denom > 0.0f ? encode_fp4_e2m1(values[pair * 2] / denom) : 0; - const uint8_t odd = denom > 0.0f ? encode_fp4_e2m1(values[pair * 2 + 1] / denom) : 0; + const uint8_t even = encode_fp4_e2m1(values[pair * 2] * encode_scale); + const uint8_t odd = encode_fp4_e2m1(values[pair * 2 + 1] * encode_scale); + qdata[row * q_cols + block_col * 8 + pair] = static_cast((even << 4) | odd); + } + } + } +} + +__global__ void partial_absmax_swiglu_bf16_kernel( + const uint16_t* __restrict__ input, + float* __restrict__ partials, + int64_t rows, + int64_t output_cols) { + const int64_t numel = rows * output_cols; + const int64_t stride = static_cast(blockDim.x) * gridDim.x; + uint32_t local_bits = 0; + for (int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + index < numel; + index += stride) { + const int64_t row = index / output_cols; + const int64_t col = index - row * output_cols; + const float value = swiglu_bf16_value(input, row, col, output_cols); + local_bits = max(local_bits, __float_as_uint(value) >> 16 & 0x7fffu); + } + const uint32_t block_bits = block_reduce_max_u32(local_bits); + if (threadIdx.x == 0) { + partials[blockIdx.x] = bf16_abs_bits_to_float(block_bits); + } +} + +__global__ void quantize_nvfp4_swiglu_bf16_kernel( + const uint16_t* __restrict__ input, + const float* __restrict__ tensor_scale_ptr, + uint8_t* __restrict__ qdata, + uint8_t* __restrict__ block_scale, + int64_t rows, + int64_t output_cols, + int64_t q_rows, + int64_t q_cols, + int64_t scale_cols) { + const int64_t row = blockIdx.x; + const float tensor_scale = tensor_scale_ptr[0]; + for (int64_t block_col = threadIdx.x; block_col < scale_cols; block_col += blockDim.x) { + float local_max = 0.0f; + float values[16]; + #pragma unroll + for (int i = 0; i < 16; ++i) { + const int64_t col = block_col * 16 + i; + float value = 0.0f; + if (row < rows && col < output_cols) { + value = swiglu_bf16_value(input, row, col, output_cols); + } + values[i] = value; + local_max = fmaxf(local_max, fabsf(value)); + } + const float raw_block_scale = (local_max / 6.0f) / tensor_scale; + const uint8_t scale_byte = encode_fp8_e4m3(fminf(raw_block_scale, 448.0f)); + const float block_scale_value = decode_fp8_e4m3(scale_byte); + const int64_t row_in_tile = row % 128; + const int64_t tile = ((row / 128) * (scale_cols / 4)) + (block_col / 4); + const int64_t within_tile = + ((row_in_tile % 32) / 2) * 32 + + (block_col % 4) + + (row_in_tile / 32) * 4 + + (row_in_tile % 2) * 16; + block_scale[tile * 512 + within_tile] = scale_byte; + if (row < q_rows && block_col * 8 < q_cols) { + const float encode_scale = fminf( + reciprocal_approx_ftz(tensor_scale * block_scale_value), + 3.402823466e+38f); + #pragma unroll + for (int pair = 0; pair < 8; ++pair) { + const uint8_t even = encode_fp4_e2m1(values[pair * 2] * encode_scale); + const uint8_t odd = encode_fp4_e2m1(values[pair * 2 + 1] * encode_scale); qdata[row * q_cols + block_col * 8 + pair] = static_cast((even << 4) | odd); } } @@ -277,6 +496,298 @@ __global__ void qkv_to_bshd_vec16_kernel( } } +__device__ __forceinline__ float sage2_warp_sum(float value) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffffu, value, offset); + } + return __shfl_sync(0xffffffffu, value, 0); +} + +__device__ __forceinline__ float sage2_warp_max(float value) { + #pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + value = fmaxf(value, __shfl_xor_sync(0xffffffffu, value, mask)); + } + return value; +} + +__device__ __forceinline__ int8_t sage2_int8_rn(float value) { + uint32_t result; + asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(result) : "f"(value)); + return static_cast(result); +} + +__device__ __forceinline__ nv_bfloat16 sage2_rotate_value( + const nv_bfloat16* __restrict__ normalized, + const nv_bfloat16* __restrict__ rotation, + int dimension, + int rotation_pairs, + int64_t rotation_pair_stride, + int64_t rotation_row_stride, + int64_t rotation_component_stride) { + if (dimension >= rotation_pairs * 2) { + return normalized[dimension]; + } + const int pair = dimension < rotation_pairs ? dimension : dimension - rotation_pairs; + const float x0 = __bfloat162float(normalized[pair]); + const float x1 = __bfloat162float(normalized[pair + rotation_pairs]); + const nv_bfloat16* pair_rotation = rotation + static_cast(pair) * rotation_pair_stride; + const int row = dimension < rotation_pairs ? 0 : 1; + const float f0 = __bfloat162float(pair_rotation[static_cast(row) * rotation_row_stride]); + const float f1 = __bfloat162float(pair_rotation[static_cast(row) * rotation_row_stride + rotation_component_stride]); + return __float2bfloat16_rn(f0 * x0 + f1 * x1); +} + +__global__ __launch_bounds__(1024) void sage2_prepare_qk_bf16_kernel( + const nv_bfloat16* __restrict__ q, + nv_bfloat16* __restrict__ k, + const nv_bfloat16* __restrict__ rotation, + const nv_bfloat16* __restrict__ q_weight, + const nv_bfloat16* __restrict__ k_weight, + int8_t* __restrict__ q_int8, + float* __restrict__ q_scale, + nv_bfloat16* __restrict__ q_prepared, + int64_t sequence, + int64_t heads, + int64_t q_s0, + int64_t q_s1, + int64_t q_s2, + int64_t k_s0, + int64_t k_s1, + int64_t k_s2, + int64_t f_s0, + int64_t f_s1, + int64_t f_s2, + int64_t f_s3, + int64_t f_s4, + int64_t f_s5, + int rotation_pairs, + float epsilon) { + constexpr int kHeadDim = 128; + constexpr int kTokensPerBlock = 32; + __shared__ nv_bfloat16 normalized[kTokensPerBlock][kHeadDim]; + __shared__ float warp_maxima[32]; + __shared__ float block_amax; + + const int lane = threadIdx.x & 31; + const int token_slot = threadIdx.x >> 5; + const int64_t token = static_cast(blockIdx.x) * kTokensPerBlock + token_slot; + const int64_t head = blockIdx.y; + const int64_t batch = blockIdx.z; + const bool valid = token < sequence; + const int64_t q_base = batch * q_s0 + token * q_s1 + head * q_s2; + const int64_t k_base = batch * k_s0 + token * k_s1 + head * k_s2; + const int64_t out_base = (batch * sequence * heads + token * heads + head) * kHeadDim; + const int64_t freq_base = batch * f_s0 + token * f_s1; + + float q_raw[4]; + float q_sum = 0.0f; + #pragma unroll + for (int item = 0; item < 4; ++item) { + const int dimension = lane + item * 32; + const float value = valid ? __bfloat162float(q[q_base + dimension]) : 0.0f; + q_raw[item] = value; + q_sum = fmaf(value, value, q_sum); + } + const float q_rrms = rsqrtf(sage2_warp_sum(q_sum) / static_cast(kHeadDim) + epsilon); + #pragma unroll + for (int item = 0; item < 4; ++item) { + const int dimension = lane + item * 32; + normalized[token_slot][dimension] = __float2bfloat16_rn( + q_raw[item] * q_rrms * __bfloat162float(q_weight[dimension])); + } + __syncthreads(); + + const nv_bfloat16* token_rotation = rotation + freq_base; + nv_bfloat16 q_values[4]; + float local_amax = 0.0000001f; + #pragma unroll + for (int item = 0; item < 4; ++item) { + const int dimension = lane + item * 32; + const nv_bfloat16 value = valid + ? sage2_rotate_value(normalized[token_slot], token_rotation, dimension, rotation_pairs, f_s3, f_s4, f_s5) + : __float2bfloat16_rn(0.0f); + q_values[item] = value; + local_amax = fmaxf(local_amax, fabsf(__bfloat162float(value))); + if (valid && q_prepared != nullptr) { + q_prepared[out_base + dimension] = value; + } + } + + const float warp_amax = sage2_warp_max(local_amax); + if (lane == 0) { + warp_maxima[token_slot] = warp_amax; + } + __syncthreads(); + float block_value = threadIdx.x < 32 ? warp_maxima[lane] : -1e20f; + block_value = sage2_warp_max(block_value); + if (threadIdx.x == 0) { + block_amax = block_value; + q_scale[(batch * heads + head) * gridDim.x + blockIdx.x] = + block_value * reciprocal_approx_ftz(127.0f); + } + __syncthreads(); + const float reciprocal_scale = 127.0f * reciprocal_approx_ftz(block_amax); + if (valid) { + #pragma unroll + for (int item = 0; item < 4; ++item) { + const int dimension = lane + item * 32; + q_int8[out_base + dimension] = sage2_int8_rn(__bfloat162float(q_values[item]) * reciprocal_scale); + } + } + __syncthreads(); + + float k_raw[4]; + float k_sum = 0.0f; + #pragma unroll + for (int item = 0; item < 4; ++item) { + const int dimension = lane + item * 32; + const float value = valid ? __bfloat162float(k[k_base + dimension]) : 0.0f; + k_raw[item] = value; + k_sum = fmaf(value, value, k_sum); + } + const float k_rrms = rsqrtf(sage2_warp_sum(k_sum) / static_cast(kHeadDim) + epsilon); + #pragma unroll + for (int item = 0; item < 4; ++item) { + const int dimension = lane + item * 32; + normalized[token_slot][dimension] = __float2bfloat16_rn( + k_raw[item] * k_rrms * __bfloat162float(k_weight[dimension])); + } + __syncthreads(); + if (valid) { + #pragma unroll + for (int item = 0; item < 4; ++item) { + const int dimension = lane + item * 32; + k[k_base + dimension] = sage2_rotate_value( + normalized[token_slot], token_rotation, dimension, rotation_pairs, f_s3, f_s4, f_s5); + } + } +} + +__global__ void sage2_v_partial_absmax_bf16_kernel( + const nv_bfloat16* __restrict__ v, + float* __restrict__ partials, + int64_t sequence, + int64_t heads, + int64_t tiles, + int64_t v_s0, + int64_t v_s1, + int64_t v_s2) { + constexpr int kHeadDim = 128; + constexpr int kTileRows = 64; + const int dimension = threadIdx.x; + const int64_t tile = blockIdx.x; + const int64_t head = blockIdx.y; + const int64_t batch = blockIdx.z; + const int64_t first_token = tile * kTileRows; + float local_max = 0.0f; + + if (dimension < kHeadDim) { + #pragma unroll + for (int row = 0; row < kTileRows; ++row) { + const int64_t token = first_token + row; + if (token < sequence) { + const int64_t offset = batch * v_s0 + token * v_s1 + head * v_s2 + dimension; + local_max = fmaxf(local_max, fabsf(__bfloat162float(v[offset]))); + } + } + partials[((batch * heads + head) * kHeadDim + dimension) * tiles + tile] = local_max; + } +} + +__global__ void sage2_v_finalize_scale_kernel( + const float* __restrict__ partials, + float* __restrict__ scale, + float* __restrict__ reciprocal_scale, + int64_t heads, + int64_t tiles, + float scale_max) { + constexpr int kHeadDim = 128; + const int64_t head = blockIdx.x; + const int64_t batch = blockIdx.y; + const int64_t dimension = blockIdx.z; + const int64_t partial_base = ((batch * heads + head) * kHeadDim + dimension) * tiles; + float local_max = 0.0f; + for (int64_t tile = threadIdx.x; tile < tiles; tile += blockDim.x) { + local_max = fmaxf(local_max, partials[partial_base + tile]); + } + const float channel_max = block_reduce_max(local_max); + if (threadIdx.x == 0) { + scale[(batch * heads + head) * kHeadDim + dimension] = + channel_max * reciprocal_approx_ftz(scale_max); + reciprocal_scale[(batch * heads + head) * kHeadDim + dimension] = + channel_max == 0.0f ? 0.0f : scale_max * reciprocal_approx_ftz(channel_max); + } +} + +__device__ __forceinline__ uint32_t sage2_e4m3x4(float x0, float x1, float x2, float x3) { + uint32_t packed; + asm volatile( + "{\n" + ".reg .b16 lo;\n" + ".reg .b16 hi;\n" + "cvt.rn.satfinite.e4m3x2.f32 lo, %2, %1;\n" + "cvt.rn.satfinite.e4m3x2.f32 hi, %4, %3;\n" + "mov.b32 %0, {lo, hi};\n" + "}\n" + : "=r"(packed) : "f"(x0), "f"(x1), "f"(x2), "f"(x3)); + return packed; +} + +__global__ __launch_bounds__(1024) void sage2_v_quantize_permute_bf16_kernel( + const nv_bfloat16* __restrict__ v, + const float* __restrict__ reciprocal_scale, + uint8_t* __restrict__ output, + int64_t sequence, + int64_t heads, + int64_t padded_sequence, + int64_t v_s0, + int64_t v_s1, + int64_t v_s2, + float scale_max) { + constexpr int kHeadDim = 128; + constexpr int kTileRows = 64; + __shared__ nv_bfloat16 tile[kTileRows][kHeadDim]; + + const int lane_group = threadIdx.x & 15; + const int source_row = threadIdx.x >> 4; + const int dimension_base = lane_group * 8; + const int64_t tile_index = blockIdx.x; + const int64_t head = blockIdx.y; + const int64_t batch = blockIdx.z; + const int64_t token = tile_index * kTileRows + source_row; + const int row_mod = source_row & 15; + const int permuted_row = (source_row & ~15) + (row_mod / 8) * 2 + ((row_mod / 2) & 3) * 4 + (row_mod & 1); + + #pragma unroll + for (int item = 0; item < 8; ++item) { + const int dimension = dimension_base + item; + tile[permuted_row][dimension] = token < sequence + ? v[batch * v_s0 + token * v_s1 + head * v_s2 + dimension] + : __float2bfloat16_rn(0.0f); + } + __syncthreads(); + + const int output_dimension = threadIdx.x >> 3; + const int output_group = threadIdx.x & 7; + const int output_row = output_group * 8; + const float channel_reciprocal_scale = + reciprocal_scale[(batch * heads + head) * kHeadDim + output_dimension]; + float values[8]; + #pragma unroll + for (int item = 0; item < 8; ++item) { + values[item] = __bfloat162float(tile[output_row + item][output_dimension]) * channel_reciprocal_scale; + } + const uint2 packed = make_uint2( + sage2_e4m3x4(values[0], values[1], values[2], values[3]), + sage2_e4m3x4(values[4], values[5], values[6], values[7])); + const int64_t output_base = + ((batch * kHeadDim + output_dimension) * heads + head) * padded_sequence + + tile_index * kTileRows + output_row; + *reinterpret_cast(output + output_base) = packed; +} + } // namespace torch::Tensor nvfp4_activation_scale_cuda(torch::Tensor input, double divisor) { @@ -327,6 +838,8 @@ torch::Tensor nvfp4_activation_scale_into_cuda(torch::Tensor input, double divis return output; } +void quantize_nvfp4_bf16_into_cuda(torch::Tensor input, torch::Tensor scale, torch::Tensor qdata, torch::Tensor block_scale, int64_t threads, bool hi_first); + std::vector quantize_nvfp4_bf16_cuda(torch::Tensor input, torch::Tensor scale, bool pad_16x, int64_t threads) { c10::cuda::CUDAGuard device_guard(input.device()); const int64_t rows = input.size(0); @@ -336,8 +849,21 @@ std::vector quantize_nvfp4_bf16_cuda(torch::Tensor input, torch:: const int64_t q_cols = pad_16x ? roundup(cols, 16) / 2 : cols / 2; const int64_t scale_rows = roundup(q_rows, 128); const int64_t scale_cols = roundup(q_cols / 8, 4); + TORCH_CHECK(scale_cols % 4 == 0, "NVFP4 block-scale width must be divisible by four"); auto qdata = torch::empty({q_rows, q_cols}, input.options().dtype(torch::kUInt8)); auto block_scale = torch::zeros({scale_rows, scale_cols}, input.options().dtype(torch::kUInt8)).view(torch::kFloat8_e4m3fn); + quantize_nvfp4_bf16_into_cuda(input, scale, qdata, block_scale, threads, true); + return {qdata, block_scale}; +} + +void quantize_nvfp4_bf16_into_cuda(torch::Tensor input, torch::Tensor scale, torch::Tensor qdata, torch::Tensor block_scale, int64_t threads, bool hi_first) { + c10::cuda::CUDAGuard device_guard(input.device()); + const int64_t rows = input.size(0); + const int64_t cols = input.size(1); + const int64_t q_rows = qdata.size(0); + const int64_t q_cols = qdata.size(1); + const int64_t scale_rows = block_scale.size(0); + const int64_t scale_cols = block_scale.size(1); auto stream = at::cuda::getCurrentCUDAStream(); quantize_nvfp4_bf16_kernel<<(threads), 0, stream>>>( reinterpret_cast(input.data_ptr()), @@ -349,11 +875,110 @@ std::vector quantize_nvfp4_bf16_cuda(torch::Tensor input, torch:: q_rows, q_cols, scale_rows, + scale_cols, + hi_first); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +std::vector quantize_nvfp4_modulated_bf16_cuda( + torch::Tensor input, + torch::Tensor shift, + torch::Tensor adaln_scale, + torch::Tensor row_index, + double divisor, + bool pad_16x, + int64_t blocks, + int64_t scale_threads, + int64_t pack_threads) { + c10::cuda::CUDAGuard device_guard(input.device()); + const int64_t rows = input.size(0); + const int64_t cols = input.size(1); + const int64_t q_rows = pad_16x ? roundup(rows, 16) : rows; + const int64_t q_cols = pad_16x ? roundup(cols, 16) / 2 : cols / 2; + const int64_t scale_rows = roundup(q_rows, 128); + const int64_t scale_cols = roundup(q_cols / 8, 4); + auto partials = torch::empty({blocks}, input.options().dtype(torch::kFloat32)); + auto tensor_scale = torch::empty({}, input.options().dtype(torch::kFloat32)); + auto qdata = torch::empty({q_rows, q_cols}, input.options().dtype(torch::kUInt8)); + auto block_scale = torch::zeros({scale_rows, scale_cols}, input.options().dtype(torch::kUInt8)).view(torch::kFloat8_e4m3fn); + auto stream = at::cuda::getCurrentCUDAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND(at::ScalarType::BFloat16, shift.scalar_type(), "vortex_nvfp4_modulated_bf16", [&] { + partial_absmax_modulated_bf16_kernel<<(blocks), static_cast(scale_threads), 0, stream>>>( + reinterpret_cast(input.data_ptr()), + shift.data_ptr(), + adaln_scale.data_ptr(), + row_index.data_ptr(), + partials.data_ptr(), + rows, + cols); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + final_scale_bf16_compat_kernel<<<1, static_cast(scale_threads), 0, stream>>>( + partials.data_ptr(), tensor_scale.data_ptr(), blocks, static_cast(divisor)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + quantize_nvfp4_modulated_bf16_kernel<<(scale_rows), static_cast(pack_threads), 0, stream>>>( + reinterpret_cast(input.data_ptr()), + shift.data_ptr(), + adaln_scale.data_ptr(), + row_index.data_ptr(), + tensor_scale.data_ptr(), + qdata.data_ptr(), + reinterpret_cast(block_scale.data_ptr()), + rows, + cols, + q_rows, + q_cols, + scale_rows, + scale_cols); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); + return {tensor_scale, qdata, block_scale}; +} + +std::vector quantize_nvfp4_swiglu_bf16_cuda( + torch::Tensor input, + double divisor, + bool pad_16x, + int64_t blocks, + int64_t scale_threads, + int64_t pack_threads) { + c10::cuda::CUDAGuard device_guard(input.device()); + const int64_t rows = input.size(0); + const int64_t output_cols = input.size(1) / 2; + const int64_t q_rows = pad_16x ? roundup(rows, 16) : rows; + const int64_t q_cols = pad_16x ? roundup(output_cols, 16) / 2 : output_cols / 2; + const int64_t scale_rows = roundup(q_rows, 128); + const int64_t scale_cols = roundup(q_cols / 8, 4); + auto partials = torch::empty({blocks}, input.options().dtype(torch::kFloat32)); + auto tensor_scale = torch::empty({}, input.options().dtype(torch::kFloat32)); + auto qdata = torch::empty({q_rows, q_cols}, input.options().dtype(torch::kUInt8)); + auto block_scale = torch::zeros({scale_rows, scale_cols}, input.options().dtype(torch::kUInt8)).view(torch::kFloat8_e4m3fn); + auto stream = at::cuda::getCurrentCUDAStream(); + + partial_absmax_swiglu_bf16_kernel<<(blocks), static_cast(scale_threads), 0, stream>>>( + reinterpret_cast(input.data_ptr()), + partials.data_ptr(), + rows, + output_cols); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + final_scale_bf16_compat_kernel<<<1, static_cast(scale_threads), 0, stream>>>( + partials.data_ptr(), tensor_scale.data_ptr(), blocks, static_cast(divisor)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + quantize_nvfp4_swiglu_bf16_kernel<<(scale_rows), static_cast(pack_threads), 0, stream>>>( + reinterpret_cast(input.data_ptr()), + tensor_scale.data_ptr(), + qdata.data_ptr(), + reinterpret_cast(block_scale.data_ptr()), + rows, + output_cols, + q_rows, + q_cols, scale_cols); C10_CUDA_KERNEL_LAUNCH_CHECK(); - return {qdata, block_scale}; + return {tensor_scale, qdata, block_scale}; } + std::vector qkv_to_bshd_cuda(torch::Tensor qkv, int64_t heads, int64_t head_dim) { c10::cuda::CUDAGuard device_guard(qkv.device()); const int64_t sequence = qkv.size(0); @@ -396,3 +1021,97 @@ std::vector qkv_to_bshd_cuda(torch::Tensor qkv, int64_t heads, in C10_CUDA_KERNEL_LAUNCH_CHECK(); return {q, k, v}; } + +std::vector sage2_prepare_qk_cuda( + torch::Tensor q, + torch::Tensor k, + torch::Tensor rotation, + torch::Tensor q_weight, + torch::Tensor k_weight, + double epsilon, + bool materialize_q) { + c10::cuda::CUDAGuard device_guard(q.device()); + const int64_t batch = q.size(0); + const int64_t sequence = q.size(1); + const int64_t heads = q.size(2); + const int64_t groups = ((sequence + 127) / 128) * 4; + auto q_int8 = torch::empty(q.sizes(), q.options().dtype(torch::kInt8)); + auto q_scale = torch::empty({batch, heads, groups}, q.options().dtype(torch::kFloat32)); + auto q_prepared = materialize_q + ? torch::empty(q.sizes(), q.options()) + : torch::empty({0}, q.options()); + auto stream = at::cuda::getCurrentCUDAStream(); + const dim3 grid(static_cast(groups), static_cast(heads), static_cast(batch)); + sage2_prepare_qk_bf16_kernel<<>>( + reinterpret_cast(q.data_ptr()), + reinterpret_cast(k.data_ptr()), + reinterpret_cast(rotation.data_ptr()), + reinterpret_cast(q_weight.data_ptr()), + reinterpret_cast(k_weight.data_ptr()), + q_int8.data_ptr(), + q_scale.data_ptr(), + materialize_q ? reinterpret_cast(q_prepared.data_ptr()) : nullptr, + sequence, + heads, + q.stride(0), q.stride(1), q.stride(2), + k.stride(0), k.stride(1), k.stride(2), + rotation.stride(0), rotation.stride(1), rotation.stride(2), + rotation.stride(3), rotation.stride(4), rotation.stride(5), + static_cast(rotation.size(3)), + static_cast(epsilon)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {q_int8, q_scale, q_prepared}; +} + +std::vector sage2_prepare_v_cuda(torch::Tensor v, double scale_max) { + c10::cuda::CUDAGuard device_guard(v.device()); + constexpr int64_t kHeadDim = 128; + constexpr int64_t kTileRows = 64; + const int64_t batch = v.size(0); + const int64_t sequence = v.size(1); + const int64_t heads = v.size(2); + const int64_t tiles = (sequence + kTileRows - 1) / kTileRows; + const int64_t padded_sequence = tiles * kTileRows; + auto partials = torch::empty({batch, heads, kHeadDim, tiles}, v.options().dtype(torch::kFloat32)); + auto v_scale = torch::empty({batch, heads, kHeadDim}, v.options().dtype(torch::kFloat32)); + auto reciprocal_scale = torch::empty_like(v_scale); + auto v_fp8 = torch::empty( + {batch, kHeadDim, heads, padded_sequence}, + v.options().dtype(torch::kFloat8_e4m3fn)); + auto stream = at::cuda::getCurrentCUDAStream(); + const dim3 tile_grid( + static_cast(tiles), + static_cast(heads), + static_cast(batch)); + sage2_v_partial_absmax_bf16_kernel<<>>( + reinterpret_cast(v.data_ptr()), + partials.data_ptr(), + sequence, + heads, + tiles, + v.stride(0), v.stride(1), v.stride(2)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + const dim3 scale_grid( + static_cast(heads), + static_cast(batch), + static_cast(kHeadDim)); + sage2_v_finalize_scale_kernel<<>>( + partials.data_ptr(), + v_scale.data_ptr(), + reciprocal_scale.data_ptr(), + heads, + tiles, + static_cast(scale_max)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + sage2_v_quantize_permute_bf16_kernel<<>>( + reinterpret_cast(v.data_ptr()), + reciprocal_scale.data_ptr(), + reinterpret_cast(v_fp8.data_ptr()), + sequence, + heads, + padded_sequence, + v.stride(0), v.stride(1), v.stride(2), + static_cast(scale_max)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {v_fp8, v_scale}; +} diff --git a/src/h3_blackwell_runtime/nvfp4_quant.py b/src/h3_blackwell_runtime/nvfp4_quant.py index 97ea031..b964944 100644 --- a/src/h3_blackwell_runtime/nvfp4_quant.py +++ b/src/h3_blackwell_runtime/nvfp4_quant.py @@ -214,3 +214,131 @@ def vortex_native_quantize_nvfp4( ), ), ) + + +def vortex_native_quantize_nvfp4_into( + tensor: torch.Tensor, + scale: torch.Tensor, + qdata: torch.Tensor, + block_scale: torch.Tensor, + *, + hi_first: bool = True, +) -> None: + """Pack BF16 activation into caller-owned NVFP4 ring buffers.""" + if tensor.dim() != 2 or tensor.dtype != torch.bfloat16: + raise ValueError("vortex_native_quantize_nvfp4_into expects a 2D BF16 tensor") + if not tensor.is_cuda or not tensor.is_contiguous(): + raise ValueError("vortex_native_quantize_nvfp4_into expects contiguous CUDA input") + _vortex_scale_extension().quantize_nvfp4_bf16_into( + tensor, + torch.as_tensor(scale, device=tensor.device, dtype=torch.float32), + qdata, + block_scale, + _env_int("H3_NVFP4_PACK_THREADS", 256), + hi_first, + ) + + +def vortex_native_quantize_modulated_nvfp4( + tensor: torch.Tensor, + shift: torch.Tensor, + scale: torch.Tensor, + row_index: torch.Tensor, +): + """Produce exact Comfy-layout NVFP4 bytes without materializing modulation.""" + if tensor.dim() != 2 or tensor.dtype != torch.bfloat16: + raise ValueError("vortex_native_quantize_modulated_nvfp4 expects a 2D BF16 tensor") + if not tensor.is_cuda or not tensor.is_contiguous(): + raise ValueError("vortex_native_quantize_modulated_nvfp4 expects contiguous CUDA input") + if shift.shape != scale.shape or shift.dim() != 2 or shift.shape[1] != tensor.shape[1]: + raise ValueError("shift and scale must have matching [table_rows, hidden] shapes") + if shift.device != tensor.device or scale.device != tensor.device: + raise ValueError("modulation tables must be on the input device") + if shift.dtype != scale.dtype or shift.dtype not in {torch.bfloat16, torch.float32}: + raise ValueError("modulation tables must share BF16 or FP32 dtype") + shift = shift.contiguous() + scale = scale.contiguous() + if row_index.shape != (tensor.shape[0],) or row_index.dtype != torch.int32: + raise ValueError("row_index must be int32 with one entry per input row") + if row_index.device != tensor.device or not row_index.is_contiguous(): + raise ValueError("row_index must be contiguous on the input device") + + from comfy_kitchen.float_utils import F4_E2M1_MAX, F8_E4M3_MAX + from comfy_kitchen.tensor import TensorCoreNVFP4Layout + + blocks, scale_threads = _vortex_scale_geometry(tensor.numel()) + return _vortex_scale_extension().quantize_nvfp4_modulated_bf16( + tensor, + shift, + scale, + row_index, + float(F8_E4M3_MAX * F4_E2M1_MAX), + TensorCoreNVFP4Layout.get_padded_shape(tuple(tensor.shape)) != tuple(tensor.shape), + blocks, + scale_threads, + _env_int("H3_NVFP4_PACK_THREADS", 256), + ) + + +def vortex_quantize_modulated_nvfp4( + tensor: torch.Tensor, + shift: torch.Tensor, + scale: torch.Tensor, + row_index: torch.Tensor, +): + """Wrap fused modulation output as a Comfy TensorCoreNVFP4Layout tensor.""" + from comfy_kitchen.tensor import QuantizedTensor, TensorCoreNVFP4Layout + + tensor_scale, qdata, block_scale = vortex_native_quantize_modulated_nvfp4( + tensor, shift, scale, row_index, + ) + return QuantizedTensor( + qdata, + "TensorCoreNVFP4Layout", + TensorCoreNVFP4Layout.Params( + scale=tensor_scale, + orig_dtype=tensor.dtype, + orig_shape=tuple(tensor.shape), + block_scale=block_scale, + ), + ) + + +def vortex_native_quantize_swiglu_nvfp4(tensor: torch.Tensor): + """Produce exact Comfy-layout NVFP4 bytes without materializing SwiGLU.""" + if tensor.dim() != 2 or tensor.dtype != torch.bfloat16 or tensor.shape[1] % 2: + raise ValueError("vortex_native_quantize_swiglu_nvfp4 expects even-width 2D BF16 input") + if not tensor.is_cuda or not tensor.is_contiguous(): + raise ValueError("vortex_native_quantize_swiglu_nvfp4 expects contiguous CUDA input") + + from comfy_kitchen.float_utils import F4_E2M1_MAX, F8_E4M3_MAX + from comfy_kitchen.tensor import TensorCoreNVFP4Layout + + output_shape = (tensor.shape[0], tensor.shape[1] // 2) + blocks, scale_threads = _vortex_scale_geometry(tensor.numel() // 2) + return _vortex_scale_extension().quantize_nvfp4_swiglu_bf16( + tensor, + float(F8_E4M3_MAX * F4_E2M1_MAX), + TensorCoreNVFP4Layout.get_padded_shape(output_shape) != output_shape, + blocks, + scale_threads, + _env_int("H3_NVFP4_PACK_THREADS", 256), + ) + + +def vortex_quantize_swiglu_nvfp4(tensor: torch.Tensor): + """Wrap fused SwiGLU output as a Comfy TensorCoreNVFP4Layout tensor.""" + from comfy_kitchen.tensor import QuantizedTensor, TensorCoreNVFP4Layout + + tensor_scale, qdata, block_scale = vortex_native_quantize_swiglu_nvfp4(tensor) + output_shape = (tensor.shape[0], tensor.shape[1] // 2) + return QuantizedTensor( + qdata, + "TensorCoreNVFP4Layout", + TensorCoreNVFP4Layout.Params( + scale=tensor_scale, + orig_dtype=tensor.dtype, + orig_shape=output_shape, + block_scale=block_scale, + ), + )