h3-blackwell-runtime/research/shared_cuda_experimental_worktree/patches/0001-shared-cuda-experiments.patch

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<torch::Tensor> 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<torch::Tensor> 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<torch::Tensor> 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<torch::Tensor> qkv_to_bshd_cuda(torch::Tensor qkv, int64_t heads, int64_t head_dim);
+std::vector<torch::Tensor> 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<torch::Tensor> 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<torch::Tensor> 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<torch::Tensor> 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<torch::Tensor> 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<torch::Tensor> 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<torch::Tensor> qkv_to_bshd(torch::Tensor qkv, int64_t heads, int64_t
   return qkv_to_bshd_cuda(qkv, heads, head_dim);
 }

+std::vector<torch::Tensor> 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<torch::Tensor> 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 <torch/extension.h>

 #include <cuda_fp8.h>
+#include <cuda_bf16.h>

 #include <cmath>
 #include <cstdint>
@@ -54,6 +55,63 @@ __inline__ __device__ float bf16_bits_to_float(uint16_t bits) {
   return __uint_as_float(static_cast<uint32_t>(bits) << 16);
 }

+__inline__ __device__ float round_bf16(float value) {
+  return __bfloat162float(__float2bfloat16_rn(value));
+}
+
+template <typename table_t>
+__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<int64_t>(row_index[row]) * cols + col;
+  const float x = bf16_bits_to_float(input[row * cols + col]);
+  const float scale_bf16 = round_bf16(static_cast<float>(scale[table_offset]));
+  const float shift_bf16 = round_bf16(static_cast<float>(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<float>(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<uint8_t>((even << 4) | odd)
+            : static_cast<uint8_t>(even | (odd << 4));
       }
-      block_scale[scale_row * scale_cols + scale_col] = scale_byte;
     }
+  }
+}
+
+template <typename table_t>
+__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<int64_t>(blockDim.x) * gridDim.x;
+  uint32_t local_bits = 0;
+  for (int64_t index = static_cast<int64_t>(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 <typename table_t>
+__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<uint8_t>((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<int64_t>(blockDim.x) * gridDim.x;
+  uint32_t local_bits = 0;
+  for (int64_t index = static_cast<int64_t>(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<uint8_t>((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<int8_t>(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<int64_t>(pair) * rotation_pair_stride;
+  const int row = dimension < rotation_pairs ? 0 : 1;
+  const float f0 = __bfloat162float(pair_rotation[static_cast<int64_t>(row) * rotation_row_stride]);
+  const float f1 = __bfloat162float(pair_rotation[static_cast<int64_t>(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<int64_t>(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<float>(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<float>(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<uint2*>(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<torch::Tensor> 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<torch::Tensor> 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<<<scale_rows, static_cast<int>(threads), 0, stream>>>(
       reinterpret_cast<const uint16_t*>(input.data_ptr()),
@@ -349,11 +875,110 @@ std::vector<torch::Tensor> 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<torch::Tensor> 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<scalar_t><<<static_cast<int>(blocks), static_cast<int>(scale_threads), 0, stream>>>(
+        reinterpret_cast<const uint16_t*>(input.data_ptr()),
+        shift.data_ptr<scalar_t>(),
+        adaln_scale.data_ptr<scalar_t>(),
+        row_index.data_ptr<int32_t>(),
+        partials.data_ptr<float>(),
+        rows,
+        cols);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+    final_scale_bf16_compat_kernel<<<1, static_cast<int>(scale_threads), 0, stream>>>(
+        partials.data_ptr<float>(), tensor_scale.data_ptr<float>(), blocks, static_cast<float>(divisor));
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+    quantize_nvfp4_modulated_bf16_kernel<scalar_t><<<static_cast<int>(scale_rows), static_cast<int>(pack_threads), 0, stream>>>(
+        reinterpret_cast<const uint16_t*>(input.data_ptr()),
+        shift.data_ptr<scalar_t>(),
+        adaln_scale.data_ptr<scalar_t>(),
+        row_index.data_ptr<int32_t>(),
+        tensor_scale.data_ptr<float>(),
+        qdata.data_ptr<uint8_t>(),
+        reinterpret_cast<uint8_t*>(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<torch::Tensor> 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<<<static_cast<int>(blocks), static_cast<int>(scale_threads), 0, stream>>>(
+      reinterpret_cast<const uint16_t*>(input.data_ptr()),
+      partials.data_ptr<float>(),
+      rows,
+      output_cols);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  final_scale_bf16_compat_kernel<<<1, static_cast<int>(scale_threads), 0, stream>>>(
+      partials.data_ptr<float>(), tensor_scale.data_ptr<float>(), blocks, static_cast<float>(divisor));
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  quantize_nvfp4_swiglu_bf16_kernel<<<static_cast<int>(scale_rows), static_cast<int>(pack_threads), 0, stream>>>(
+      reinterpret_cast<const uint16_t*>(input.data_ptr()),
+      tensor_scale.data_ptr<float>(),
+      qdata.data_ptr<uint8_t>(),
+      reinterpret_cast<uint8_t*>(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<torch::Tensor> 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<torch::Tensor> qkv_to_bshd_cuda(torch::Tensor qkv, int64_t heads, in
   C10_CUDA_KERNEL_LAUNCH_CHECK();
   return {q, k, v};
 }
+
+std::vector<torch::Tensor> 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<unsigned int>(groups), static_cast<unsigned int>(heads), static_cast<unsigned int>(batch));
+  sage2_prepare_qk_bf16_kernel<<<grid, 1024, 0, stream>>>(
+      reinterpret_cast<const nv_bfloat16*>(q.data_ptr()),
+      reinterpret_cast<nv_bfloat16*>(k.data_ptr()),
+      reinterpret_cast<const nv_bfloat16*>(rotation.data_ptr()),
+      reinterpret_cast<const nv_bfloat16*>(q_weight.data_ptr()),
+      reinterpret_cast<const nv_bfloat16*>(k_weight.data_ptr()),
+      q_int8.data_ptr<int8_t>(),
+      q_scale.data_ptr<float>(),
+      materialize_q ? reinterpret_cast<nv_bfloat16*>(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<int>(rotation.size(3)),
+      static_cast<float>(epsilon));
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return {q_int8, q_scale, q_prepared};
+}
+
+std::vector<torch::Tensor> 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<unsigned int>(tiles),
+      static_cast<unsigned int>(heads),
+      static_cast<unsigned int>(batch));
+  sage2_v_partial_absmax_bf16_kernel<<<tile_grid, kHeadDim, 0, stream>>>(
+      reinterpret_cast<const nv_bfloat16*>(v.data_ptr()),
+      partials.data_ptr<float>(),
+      sequence,
+      heads,
+      tiles,
+      v.stride(0), v.stride(1), v.stride(2));
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  const dim3 scale_grid(
+      static_cast<unsigned int>(heads),
+      static_cast<unsigned int>(batch),
+      static_cast<unsigned int>(kHeadDim));
+  sage2_v_finalize_scale_kernel<<<scale_grid, 256, 0, stream>>>(
+      partials.data_ptr<float>(),
+      v_scale.data_ptr<float>(),
+      reciprocal_scale.data_ptr<float>(),
+      heads,
+      tiles,
+      static_cast<float>(scale_max));
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  sage2_v_quantize_permute_bf16_kernel<<<tile_grid, 1024, 0, stream>>>(
+      reinterpret_cast<const nv_bfloat16*>(v.data_ptr()),
+      reciprocal_scale.data_ptr<float>(),
+      reinterpret_cast<uint8_t*>(v_fp8.data_ptr()),
+      sequence,
+      heads,
+      padded_sequence,
+      v.stride(0), v.stride(1), v.stride(2),
+      static_cast<float>(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,
+        ),
+    )