NVFP4-1.1 integration: GPU-only quantize kernel + MoE pipeline wiring
- Add quantize_nvfp4.cu: BF16→FP4 GPU kernel (no CPU sync, warp shuffle amax) - Add quantize_nvfp4_gpu() bridge in ops/quantize.py - Fix deinterleave_quantize kernel path (dsv4/ops/kernels → dsv4/kernels/cuda) - Wire GPU quantize into Nvfp4MoE._run_impl(): - L1 input: quantize_nvfp4_gpu (replaces quantize_activation_nvfp4) - Fused SwiGLU L2: deinterleave_quantize_nvfp4_cuda (single kernel) - Non-fused L2: quantize_nvfp4_gpu - Add test_nvfp4_gpu_quantize.py for both kernels
This commit is contained in:
116
dsv4/kernels/cuda/quantize_nvfp4.cu
Normal file
116
dsv4/kernels/cuda/quantize_nvfp4.cu
Normal file
@@ -0,0 +1,116 @@
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp8.h>
|
||||
#include <cuda_fp8.hpp>
|
||||
#include <ATen/ATen.h>
|
||||
#include <torch/extension.h>
|
||||
#include <cstdint>
|
||||
|
||||
// BF16 → NVFP4 quantization kernel (no deinterleave, GPU-only).
|
||||
// Reads BF16 from GMEM, quantizes to NVFP4 (FP4 data + FP8 E4M3 scales),
|
||||
// writes both to GMEM. No CPU-GPU syncs.
|
||||
//
|
||||
// This replaces quantize_activation_nvfp4() which uses .amax() (CPU sync).
|
||||
// Global scale is passed in as a pre-computed scalar.
|
||||
//
|
||||
// Grid: (N / 16, M, 1) — each CTA processes one 16-element block in one row.
|
||||
// Block: 16 threads (1 thread per element in the 16-element microblock).
|
||||
|
||||
__device__ __forceinline__ int half_step_to_e2m1(int hs) {
|
||||
// Matches Python step_to_idx LUT:
|
||||
// 0→0, 1→1, 2→2, 3→3, 4→4, 5→4, 6→5, 7→5, 8→6, 9→6, 10→6, 11→7, 12→7
|
||||
if (hs <= 4) return hs;
|
||||
if (hs <= 5) return 4;
|
||||
if (hs <= 7) return 5;
|
||||
if (hs <= 10) return 6;
|
||||
return 7;
|
||||
}
|
||||
|
||||
__global__ void quantize_nvfp4_kernel(
|
||||
const __nv_bfloat16* __restrict__ input,
|
||||
int M, int N,
|
||||
float global_scale,
|
||||
uint8_t* __restrict__ out_fp4,
|
||||
uint8_t* __restrict__ out_sf
|
||||
) {
|
||||
int m = blockIdx.y;
|
||||
int n_block = blockIdx.x;
|
||||
if (m >= M || n_block * 16 >= N) return;
|
||||
|
||||
int lane = threadIdx.x; // 0..15
|
||||
|
||||
// Step 1: Read 1 BF16 element per thread, normalize by global_scale
|
||||
int col = n_block * 16 + lane;
|
||||
float val = 0.0f;
|
||||
if (col < N) {
|
||||
val = __bfloat162float(input[m * N + col]);
|
||||
val = val / global_scale;
|
||||
}
|
||||
|
||||
// Step 2: Warp-level amax reduction (16 threads = half-warp)
|
||||
// Use warp shuffle for the reduction
|
||||
float abs_val = fabsf(val);
|
||||
for (int offset = 8; offset > 0; offset >>= 1) {
|
||||
abs_val = fmaxf(abs_val, __shfl_down_sync(0xFFFF, abs_val, offset));
|
||||
}
|
||||
float block_amax = abs_val; // Same value in all 16 lanes after reduction
|
||||
|
||||
// Step 3: Compute FP8 E4M3 block scale = amax / 6.0
|
||||
float bsf = block_amax / 6.0f;
|
||||
if (block_amax < 6.0f * 0.001953125f) { // FP8 E4M3 underflow threshold
|
||||
bsf = 0;
|
||||
val = 0;
|
||||
}
|
||||
__nv_fp8_e4m3 bsf8_obj(bsf);
|
||||
float bs = (float)bsf8_obj;
|
||||
uint8_t bsf8 = *(uint8_t*)&bsf8_obj;
|
||||
|
||||
// Step 4: Quantize each value to FP4 E2M1
|
||||
uint8_t nibble = 0;
|
||||
if (bs >= 1e-8f) {
|
||||
float s = val / bs;
|
||||
int hs = __float2int_rn(fminf(fabsf(s), 6.0f) * 2.0f);
|
||||
if (hs > 12) hs = 12;
|
||||
int idx = half_step_to_e2m1(hs);
|
||||
if (s < 0) idx += 8;
|
||||
nibble = idx;
|
||||
}
|
||||
|
||||
// Step 5: Pack pairs of FP4 nibbles into bytes
|
||||
// Even lanes write the packed byte: (odd_nibble << 4) | even_nibble
|
||||
if (lane % 2 == 0) {
|
||||
uint8_t odd_nibble = __shfl_down_sync(0xFFFF, nibble, 1);
|
||||
uint8_t packed = (odd_nibble << 4) | nibble;
|
||||
int byte_idx = m * (N / 2) + n_block * 8 + (lane / 2);
|
||||
out_fp4[byte_idx] = packed;
|
||||
}
|
||||
|
||||
// Step 6: Write FP8 block scale (1 thread per block)
|
||||
if (lane == 0) {
|
||||
out_sf[m * (N / 16) + n_block] = bsf8;
|
||||
}
|
||||
}
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor> quantize_nvfp4_cuda(
|
||||
torch::Tensor input_bf16, double global_scale
|
||||
) {
|
||||
int M = input_bf16.size(0);
|
||||
int N = input_bf16.size(1);
|
||||
TORCH_CHECK(N % 16 == 0, "N must be a multiple of 16 for NVFP4 quantization");
|
||||
auto opts = input_bf16.options();
|
||||
auto out_fp4 = torch::zeros({M, N / 2}, opts.dtype(torch::kUInt8));
|
||||
auto out_sf = torch::zeros({M, N / 16}, opts.dtype(torch::kUInt8));
|
||||
int nb = N / 16;
|
||||
dim3 grid(nb, M);
|
||||
dim3 block(16);
|
||||
quantize_nvfp4_kernel<<<grid, block>>>(
|
||||
reinterpret_cast<const __nv_bfloat16*>(input_bf16.data_ptr<at::BFloat16>()),
|
||||
M, N, (float)global_scale,
|
||||
out_fp4.data_ptr<uint8_t>(), out_sf.data_ptr<uint8_t>()
|
||||
);
|
||||
return {out_fp4.view(torch::kFloat4_e2m1fn_x2), out_sf.view(torch::kFloat8_e4m3fn)};
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("quantize_nvfp4", &quantize_nvfp4_cuda);
|
||||
}
|
||||
@@ -19,6 +19,8 @@ from dsv4.ops.quantize import (
|
||||
quantize_activation_nvfp4,
|
||||
quantize_weight_to_nvfp4,
|
||||
quantize_to_nvfp4,
|
||||
quantize_nvfp4_gpu,
|
||||
deinterleave_quantize_nvfp4_cuda,
|
||||
)
|
||||
from dsv4.ops.layouts import (
|
||||
make_b_k_major,
|
||||
@@ -563,14 +565,10 @@ class Nvfp4MoE:
|
||||
padded_dst = padded_expert_offsets[expert_assign] + local_row
|
||||
|
||||
# === L1: gate + up ===
|
||||
# Quantize slot_hidden (sorted tokens), NOT padded_hidden.
|
||||
# padded_hidden is padded with zeros; quantizing it produces
|
||||
# x_sf rows at padded positions, but x_sf[:num_slots] would
|
||||
# only get scales for the first num_slots PADDED rows (expert 0),
|
||||
# not the scattered token positions. Quantizing slot_hidden
|
||||
# gives x_sf with num_slots rows (one per token), which the
|
||||
# scale assembly correctly scatters into padded layout.
|
||||
slot_x_fp4, slot_x_sf = quantize_activation_nvfp4(
|
||||
# Quantize slot_hidden using GPU-only kernel (no CPU-GPU sync).
|
||||
# slot_hidden is the sorted tokens (not padded). The GPU kernel
|
||||
# replaces quantize_activation_nvfp4 which uses .amax() (CPU sync).
|
||||
slot_x_fp4, slot_x_sf = quantize_nvfp4_gpu(
|
||||
slot_hidden, self._l1_activation_global_scale
|
||||
)
|
||||
# Scatter x_fp4 into padded layout for the GEMM
|
||||
@@ -596,9 +594,13 @@ class Nvfp4MoE:
|
||||
swiglu_limit=self._swiglu_limit if self._swiglu_limit is not None else 0.0,
|
||||
)
|
||||
l1_out_real = l1_out[padded_dst]
|
||||
# De-interleave: odd 8-col groups = silu(gate)*up (the SwiGLU result)
|
||||
l1_deil = deinterleave_l1_weights(l1_out_real.unsqueeze(0).contiguous())[0]
|
||||
activated = l1_deil[:, self.intermediate_size:]
|
||||
# De-interleave + quantize to FP4 in one GPU kernel.
|
||||
# l1_out_real has interleaved [silu(gate)*8, swiglu*8, ...].
|
||||
# The CUDA kernel extracts odd 8-col groups (SwiGLU result)
|
||||
# and quantizes to NVFP4. No CPU sync, no Python deinterleave.
|
||||
slot_l2_x_fp4, slot_l2_x_sf = deinterleave_quantize_nvfp4_cuda(
|
||||
l1_out_real, self.intermediate_size, self._l2_activation_global_scale
|
||||
)
|
||||
else:
|
||||
# === Non-fused L1 GEMM + PyTorch SiLU(gate)*up ===
|
||||
l1_out = run_nvfp4_grouped_gemm(
|
||||
@@ -618,10 +620,12 @@ class Nvfp4MoE:
|
||||
activated = gate_silu * up
|
||||
|
||||
# === L2: down ===
|
||||
# Quantize activated (per-token), scatter into padded FP4 buffer
|
||||
slot_l2_x_fp4, slot_l2_x_sf = quantize_activation_nvfp4(
|
||||
activated, self._l2_activation_global_scale
|
||||
)
|
||||
# Quantize activated (per-token) using GPU-only kernel, scatter into padded FP4 buffer.
|
||||
# For fused_swiglu path, slot_l2_x_fp4/sf already set by deinterleave_quantize_nvfp4_cuda.
|
||||
if not self._fused_swiglu:
|
||||
slot_l2_x_fp4, slot_l2_x_sf = quantize_nvfp4_gpu(
|
||||
activated, self._l2_activation_global_scale
|
||||
)
|
||||
padded_activated_fp4 = self._shared_bufs['activated_fp4']
|
||||
padded_activated_fp4.view(torch.uint8).zero_()
|
||||
padded_activated_fp4.view(torch.uint8)[padded_dst] = slot_l2_x_fp4.view(torch.uint8)
|
||||
|
||||
@@ -244,7 +244,8 @@ def deinterleave_quantize_nvfp4_cuda(fused_bf16, intermediate, global_scale, gra
|
||||
"""
|
||||
from torch.utils.cpp_extension import load
|
||||
import os
|
||||
kernel_dir = os.path.join(os.path.dirname(__file__), "kernels")
|
||||
# dsv4/ops/quantize.py → dsv4/kernels/cuda/
|
||||
kernel_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "kernels", "cuda")
|
||||
mod = load(
|
||||
name="deinterleave_quantize_nvfp4",
|
||||
sources=[os.path.join(kernel_dir, "deinterleave_quantize.cu")],
|
||||
@@ -252,3 +253,30 @@ def deinterleave_quantize_nvfp4_cuda(fused_bf16, intermediate, global_scale, gra
|
||||
verbose=False,
|
||||
)
|
||||
return mod.deinterleave_quantize_nvfp4(fused_bf16, intermediate, granularity, global_scale)
|
||||
|
||||
|
||||
def quantize_nvfp4_gpu(x_bf16, global_scale):
|
||||
"""Quantize BF16 tensor to NVFP4 using a custom CUDA kernel (GPU-only, no CPU sync).
|
||||
|
||||
Replaces quantize_activation_nvfp4() which uses .amax() (CPU sync).
|
||||
The global_scale must be pre-computed (from warmup or known value).
|
||||
|
||||
Args:
|
||||
x_bf16: (M, N) BF16 tensor. N must be a multiple of 16.
|
||||
global_scale: float32 scalar (pre-computed, NOT from .max())
|
||||
|
||||
Returns:
|
||||
x_fp4: (M, N//2) float4_e2m1fn_x2
|
||||
x_sf: (M, N//16) float8_e4m3fn
|
||||
"""
|
||||
from torch.utils.cpp_extension import load
|
||||
import os
|
||||
# dsv4/ops/quantize.py → dsv4/kernels/cuda/
|
||||
kernel_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "kernels", "cuda")
|
||||
mod = load(
|
||||
name="quantize_nvfp4",
|
||||
sources=[os.path.join(kernel_dir, "quantize_nvfp4.cu")],
|
||||
extra_cuda_cflags=["-gencode=arch=compute_100a,code=sm_100a"],
|
||||
verbose=False,
|
||||
)
|
||||
return mod.quantize_nvfp4(x_bf16, global_scale)
|
||||
|
||||
193
tests/unit/test_nvfp4_gpu_quantize.py
Normal file
193
tests/unit/test_nvfp4_gpu_quantize.py
Normal file
@@ -0,0 +1,193 @@
|
||||
"""
|
||||
NVFP4 GPU Quantize Kernel Test.
|
||||
|
||||
Tests:
|
||||
1. quantize_nvfp4_gpu: BF16 → FP4 (no deinterleave, GPU-only, no CPU sync)
|
||||
2. deinterleave_quantize_nvfp4_cuda: interleaved BF16 → FP4 (deinterleave + quantize)
|
||||
3. Integration: Both kernels match the Python quantize_activation_nvfp4 reference
|
||||
|
||||
Run: ~/.openclaw/workspace/fire_b200_test tests/unit/test_nvfp4_gpu_quantize.py
|
||||
"""
|
||||
import torch
|
||||
import math
|
||||
import sys, os
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
|
||||
|
||||
from dsv4.ops.quantize import (
|
||||
quantize_activation_nvfp4,
|
||||
quantize_nvfp4_gpu,
|
||||
deinterleave_quantize_nvfp4_cuda,
|
||||
SF_VEC_SIZE,
|
||||
)
|
||||
from dsv4.ops.layouts import interleave_l1_weights
|
||||
|
||||
|
||||
def test_quantize_nvfp4_gpu_basic():
|
||||
"""Test GPU-only quantize kernel against Python reference."""
|
||||
print("\n=== Test 1: quantize_nvfp4_gpu basic correctness ===")
|
||||
torch.manual_seed(42)
|
||||
M, N = 128, 512
|
||||
x = torch.randn(M, N, dtype=torch.bfloat16, device='cuda')
|
||||
global_scale = 1.0
|
||||
|
||||
# Python reference
|
||||
ref_fp4, ref_sf = quantize_activation_nvfp4(x, global_scale)
|
||||
|
||||
# GPU kernel
|
||||
gpu_fp4, gpu_sf = quantize_nvfp4_gpu(x, global_scale)
|
||||
|
||||
# Compare shapes and dtypes
|
||||
assert gpu_fp4.shape == ref_fp4.shape, f"FP4 shape mismatch: {gpu_fp4.shape} vs {ref_fp4.shape}"
|
||||
assert gpu_sf.shape == ref_sf.shape, f"SF shape mismatch: {gpu_sf.shape} vs {ref_sf.shape}"
|
||||
assert gpu_fp4.dtype == torch.float4_e2m1fn_x2, f"FP4 dtype: {gpu_fp4.dtype}"
|
||||
assert gpu_sf.dtype == torch.float8_e4m3fn, f"SF dtype: {gpu_sf.dtype}"
|
||||
|
||||
# Byte-exact comparison of FP4 data
|
||||
ref_bytes = ref_fp4.view(torch.uint8)
|
||||
gpu_bytes = gpu_fp4.view(torch.uint8)
|
||||
fp4_match = (ref_bytes == gpu_bytes).float().mean().item()
|
||||
print(f" FP4 byte match: {fp4_match*100:.1f}%")
|
||||
|
||||
# SF comparison
|
||||
ref_sf_bytes = ref_sf.view(torch.uint8)
|
||||
gpu_sf_bytes = gpu_sf.view(torch.uint8)
|
||||
sf_match = (ref_sf_bytes == gpu_sf_bytes).float().mean().item()
|
||||
print(f" SF byte match: {sf_match*100:.1f}%")
|
||||
|
||||
# Round-trip cosine similarity
|
||||
ref_deq = _dequantize_nvfp4(ref_fp4, ref_sf, global_scale, N)
|
||||
gpu_deq = _dequantize_nvfp4(gpu_fp4, gpu_sf, global_scale, N)
|
||||
|
||||
cos_ref = torch.nn.functional.cosine_similarity(
|
||||
x.flatten().float().unsqueeze(0), ref_deq.flatten().float().unsqueeze(0)
|
||||
).item()
|
||||
cos_gpu = torch.nn.functional.cosine_similarity(
|
||||
x.flatten().float().unsqueeze(0), gpu_deq.flatten().float().unsqueeze(0)
|
||||
).item()
|
||||
cos_cross = torch.nn.functional.cosine_similarity(
|
||||
ref_deq.flatten().float().unsqueeze(0), gpu_deq.flatten().float().unsqueeze(0)
|
||||
).item()
|
||||
|
||||
print(f" Python round-trip cos: {cos_ref:.6f}")
|
||||
print(f" GPU round-trip cos: {cos_gpu:.6f}")
|
||||
print(f" Python vs GPU cos: {cos_cross:.6f}")
|
||||
|
||||
assert cos_gpu >= 0.95, f"GPU round-trip cosine too low: {cos_gpu}"
|
||||
assert cos_cross >= 0.99, f"Python vs GPU cosine too low: {cos_cross}"
|
||||
print(f" ✅ PASS")
|
||||
|
||||
|
||||
def test_quantize_nvfp4_gpu_larger():
|
||||
"""Test GPU quantize with larger dimensions (MoE intermediate size)."""
|
||||
print("\n=== Test 2: quantize_nvfp4_gpu larger shape ===")
|
||||
torch.manual_seed(42)
|
||||
M, N = 64, 4096
|
||||
x = torch.randn(M, N, dtype=torch.bfloat16, device='cuda')
|
||||
global_scale = 1.0 / (6.0 * 448.0)
|
||||
|
||||
ref_fp4, ref_sf = quantize_activation_nvfp4(x, global_scale)
|
||||
gpu_fp4, gpu_sf = quantize_nvfp4_gpu(x, global_scale)
|
||||
|
||||
ref_deq = _dequantize_nvfp4(ref_fp4, ref_sf, global_scale, N)
|
||||
gpu_deq = _dequantize_nvfp4(gpu_fp4, gpu_sf, global_scale, N)
|
||||
|
||||
cos_ref = torch.nn.functional.cosine_similarity(
|
||||
x.flatten().float().unsqueeze(0), ref_deq.flatten().float().unsqueeze(0)
|
||||
).item()
|
||||
cos_gpu = torch.nn.functional.cosine_similarity(
|
||||
x.flatten().float().unsqueeze(0), gpu_deq.flatten().float().unsqueeze(0)
|
||||
).item()
|
||||
|
||||
print(f" Python round-trip cos: {cos_ref:.6f}")
|
||||
print(f" GPU round-trip cos: {cos_gpu:.6f}")
|
||||
assert cos_gpu >= 0.95, f"GPU round-trip cosine too low: {cos_gpu}"
|
||||
print(f" ✅ PASS")
|
||||
|
||||
|
||||
def test_quantize_nvfp4_gpu_no_cpu_sync():
|
||||
"""Verify quantize_nvfp4_gpu has no CPU-GPU sync points."""
|
||||
print("\n=== Test 3: No CPU-GPU sync verification ===")
|
||||
torch.manual_seed(42)
|
||||
M, N = 32, 512
|
||||
x = torch.randn(M, N, dtype=torch.bfloat16, device='cuda')
|
||||
global_scale = 0.001
|
||||
|
||||
# This should NOT trigger a CUDA synchronization
|
||||
# If it does, cudagraph capture would fail
|
||||
gpu_fp4, gpu_sf = quantize_nvfp4_gpu(x, global_scale)
|
||||
|
||||
# Force a sync to check results
|
||||
result = gpu_fp4.view(torch.uint8).sum().item()
|
||||
print(f" FP4 sum (verification): {result}")
|
||||
print(f" ✅ PASS (no crash = no CPU sync in kernel path)")
|
||||
|
||||
|
||||
def test_deinterleave_quantize_correctness():
|
||||
"""Test deinterleave + quantize CUDA kernel against Python reference."""
|
||||
print("\n=== Test 4: deinterleave_quantize_nvfp4_cuda correctness ===")
|
||||
torch.manual_seed(42)
|
||||
M = 64
|
||||
intermediate = 512
|
||||
N = 2 * intermediate # fused output has gate + up interleaved
|
||||
|
||||
# Create SwiGLU-style interleaved data
|
||||
gate = torch.randn(M, intermediate, dtype=torch.bfloat16, device='cuda')
|
||||
up = torch.randn(M, intermediate, dtype=torch.bfloat16, device='cuda')
|
||||
gate_silu = torch.nn.functional.silu(gate)
|
||||
swiglu_result = gate_silu * up
|
||||
|
||||
# Create interleaved layout: [silu(gate)*8, up*8, ...]
|
||||
fused = torch.cat([gate_silu, up], dim=-1).unsqueeze(0) # (1, M, N)
|
||||
fused = interleave_l1_weights(fused)[0] # (M, N) interleaved
|
||||
|
||||
global_scale = 1.0
|
||||
|
||||
# Python reference: deinterleave then quantize
|
||||
ref_fp4, ref_sf = quantize_activation_nvfp4(swiglu_result, global_scale)
|
||||
|
||||
# CUDA kernel: deinterleave + quantize in one pass
|
||||
gpu_fp4, gpu_sf = deinterleave_quantize_nvfp4_cuda(fused, intermediate, global_scale)
|
||||
|
||||
# Compare round-trip
|
||||
ref_deq = _dequantize_nvfp4(ref_fp4, ref_sf, global_scale, intermediate)
|
||||
gpu_deq = _dequantize_nvfp4(gpu_fp4, gpu_sf, global_scale, intermediate)
|
||||
|
||||
cos_cross = torch.nn.functional.cosine_similarity(
|
||||
ref_deq.flatten().float().unsqueeze(0), gpu_deq.flatten().float().unsqueeze(0)
|
||||
).item()
|
||||
|
||||
print(f" Python vs CUDA kernel cos: {cos_cross:.6f}")
|
||||
assert cos_cross >= 0.99, f"Python vs CUDA kernel cosine too low: {cos_cross}"
|
||||
print(f" ✅ PASS")
|
||||
|
||||
|
||||
def test():
|
||||
print("=== NVFP4 GPU Quantize Kernel Tests ===")
|
||||
test_quantize_nvfp4_gpu_basic()
|
||||
test_quantize_nvfp4_gpu_larger()
|
||||
test_quantize_nvfp4_gpu_no_cpu_sync()
|
||||
test_deinterleave_quantize_correctness()
|
||||
print("\n=== ALL TESTS PASSED ===")
|
||||
|
||||
|
||||
def _dequantize_nvfp4(x_fp4, block_scale, global_scale, N):
|
||||
"""Dequantize NVFP4 back to BF16 for verification."""
|
||||
M = x_fp4.shape[0]
|
||||
block_size = SF_VEC_SIZE
|
||||
raw = x_fp4.view(torch.uint8)
|
||||
even = raw & 0x0F
|
||||
odd = (raw >> 4) & 0x0F
|
||||
indices = torch.stack([even, odd], dim=-1).reshape(M, N)
|
||||
signs = (indices >= 8).float() * -2 + 1
|
||||
mag = indices % 8
|
||||
idx_to_half_step = torch.tensor([0, 2, 4, 6, 8, 10, 12, 14],
|
||||
dtype=torch.float32, device='cuda')
|
||||
half_steps = idx_to_half_step[mag.long()]
|
||||
x_deq_fp32 = signs * half_steps / 2.0
|
||||
block_scale_exp = block_scale.repeat_interleave(block_size, dim=-1).float()
|
||||
x_deq_fp32 = x_deq_fp32 * block_scale_exp * global_scale
|
||||
return x_deq_fp32.to(torch.bfloat16)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test()
|
||||
Reference in New Issue
Block a user