From d7dcdf3efd2c4e9b15d3f0a8040ccc83a9b1f493 Mon Sep 17 00:00:00 2001 From: Kaixi Date: Wed, 8 Jul 2026 02:40:59 +0200 Subject: [PATCH] [DSV4] perf: Make FP8 quant output tensor contiguous (#27926) Co-authored-by: liqichao Co-authored-by: chenbong --- .../fp8_wo_a_group_major_quant.cuh | 169 ++++++++++++++ python/sglang/jit_kernel/dsv4/__init__.py | 2 + python/sglang/jit_kernel/dsv4/fp8_wo_a.py | 93 ++++++++ python/sglang/srt/models/deepseek_v4.py | 10 +- .../jit/deepseek_v4/test_fp8_wo_a.py | 212 ++++++++++++++++++ 5 files changed, 479 insertions(+), 7 deletions(-) create mode 100644 python/sglang/jit_kernel/csrc/deepseek_v4/fp8_wo_a_group_major_quant.cuh create mode 100644 python/sglang/jit_kernel/dsv4/fp8_wo_a.py create mode 100644 test/registered/jit/deepseek_v4/test_fp8_wo_a.py diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/fp8_wo_a_group_major_quant.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/fp8_wo_a_group_major_quant.cuh new file mode 100644 index 000000000..aec79faee --- /dev/null +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/fp8_wo_a_group_major_quant.cuh @@ -0,0 +1,169 @@ +// DeepSeek-V4 wo_a activation quantization for DeepGEMM fp8_einsum. +// +// This is intentionally narrower than the generic per_token_group_quant_8bit_v2 +// kernel: input is a [T, G, D] view with contiguous hidden groups, output_q is +// contiguous [T, G, D], group_size is fixed to 128, scales are fp32 UE8M0 +// power-of-two values, and output_s is a logical [T, G, D/128] view backed by +// group-major [G, T, D/128] storage. +// +// The generic kernel cannot read the strided DSV4 view while producing +// contiguous [T, G, D] codes and group-major scales without an extra full-tensor +// copy. +#include // TensorMatcher, SymbolicSize/Device +#include // RuntimeCheck + +#include // fp8 aliases, PDL helpers +#include // warp::reduce_max + +#include // UE8M0 and FP8 helpers + +#include // tvm::ffi::TensorView + +#include +#include + +namespace { + +using deepseek_v4::fp8::cast_to_ue8m0; +using deepseek_v4::fp8::inv_scale_ue8m0; +using deepseek_v4::fp8::pack_fp8; + +constexpr float LOCAL_ABSMAX_ABS = 1e-10f; +constexpr uint32_t GROUP_SIZE = 128; +constexpr uint32_t THREADS_PER_GROUP = 8; +constexpr uint32_t SUBWARPS_PER_BLOCK = 16; +constexpr uint32_t INPUT_VEC_NUM_BYTES = 32; +constexpr uint32_t INPUT_INT4_SIZE = INPUT_VEC_NUM_BYTES / sizeof(int4); + +template +SGL_DEVICE float GroupReduceMax(float val) { + static_assert( + (THREADS_PER_SUBWARP & (THREADS_PER_SUBWARP - 1)) == 0 && THREADS_PER_SUBWARP <= 16 && THREADS_PER_SUBWARP >= 1, + "THREADS_PER_SUBWARP must be 1, 2, 4, 8, or 16"); + // Tail subwarps can be inactive at the bounds check, so reduce with only the + // current subgroup's lanes rather than a full-warp mask. + constexpr device::warp::mask_t kSub = (device::warp::mask_t{1} << THREADS_PER_SUBWARP) - 1; + const device::warp::mask_t mask = kSub << (THREADS_PER_SUBWARP * ((threadIdx.x % 32) / THREADS_PER_SUBWARP)); + return device::warp::reduce_max(val, mask); +} + +template +__global__ void fp8_wo_a_group_major_quant_ue8m0_kernel( + const T* __restrict__ input, + fp8_e4m3_t* __restrict__ output_q, + float* __restrict__ output_s, + int64_t total_scale_groups, + int64_t num_tokens, + int hidden_dim_groups, + int num_outer_groups, + int64_t input_stride_t) { + device::PDLWaitPrimary(); + + const int64_t subwarp_id = threadIdx.x / THREADS_PER_GROUP; + const int lane_id = threadIdx.x % THREADS_PER_GROUP; + const int64_t group_id = static_cast(blockIdx.x) * SUBWARPS_PER_BLOCK + subwarp_id; + if (group_id < total_scale_groups) { + const int hidden_group = group_id % hidden_dim_groups; + const int64_t token_outer = group_id / hidden_dim_groups; + const int outer_idx = token_outer % num_outer_groups; + const int64_t token_idx = token_outer / num_outer_groups; + + constexpr uint32_t INPUT_VEC_SIZE = INPUT_VEC_NUM_BYTES / sizeof(T); + static_assert(INPUT_VEC_SIZE * THREADS_PER_GROUP == GROUP_SIZE); + + const int64_t input_group_start_offset = + token_idx * input_stride_t + outer_idx * GROUP_SIZE * hidden_dim_groups + hidden_group * GROUP_SIZE; + const int64_t output_group_start_offset = group_id * GROUP_SIZE; + + int4 input_int4[INPUT_INT4_SIZE]; + T* input_vec = reinterpret_cast(input_int4); + +#pragma unroll + for (uint32_t j = 0; j < INPUT_INT4_SIZE; ++j) { + input_int4[j] = reinterpret_cast(input + input_group_start_offset + lane_id * INPUT_VEC_SIZE)[j]; + } + + float local_absmax = LOCAL_ABSMAX_ABS; +#pragma unroll + for (uint32_t j = 0; j < INPUT_VEC_SIZE; ++j) { + const float val = static_cast(input_vec[j]); + local_absmax = fmaxf(local_absmax, fabsf(val)); + } + + local_absmax = GroupReduceMax(local_absmax); + + constexpr float kFp8MaxInv = 1.0f / kFP8E4M3Max; + const int32_t scale_ue8m0 = cast_to_ue8m0(local_absmax * kFp8MaxInv); + const float y_scale = inv_scale_ue8m0(scale_ue8m0); + const float y_scale_inv = __uint_as_float(static_cast(scale_ue8m0) << 23); + + int4 output_buf; + auto* output_buf_ptr = reinterpret_cast(&output_buf); +#pragma unroll + for (uint32_t j = 0; j < INPUT_VEC_SIZE; j += 2) { + output_buf_ptr[j / 2] = + pack_fp8(static_cast(input_vec[j]) * y_scale, static_cast(input_vec[j + 1]) * y_scale); + } + + *reinterpret_cast(output_q + output_group_start_offset + lane_id * INPUT_VEC_SIZE) = output_buf; + + if (lane_id == 0) { + output_s[(outer_idx * num_tokens + token_idx) * hidden_dim_groups + hidden_group] = y_scale_inv; + } + } + + device::PDLTriggerSecondary(); +} + +template +struct FP8WoAGroupMajorQuantUE8M0Kernel { + static void run(tvm::ffi::TensorView input, tvm::ffi::TensorView output_q, tvm::ffi::TensorView output_s) { + using namespace host; + + auto device = SymbolicDevice{}; + device.set_options(); + auto TSize = SymbolicSize{"num_tokens"}; + auto GSize = SymbolicSize{"num_outer_groups"}; + auto DSize = SymbolicSize{"hidden_dim"}; + auto SSize = SymbolicSize{"hidden_dim_groups"}; + + TensorMatcher({TSize, GSize, DSize}).with_strides({-1, DSize, 1}).with_dtype().with_device(device).verify(input); + TensorMatcher({TSize, GSize, DSize}).with_dtype().with_device(device).verify(output_q); + TensorMatcher({GSize, TSize, SSize}).with_dtype().with_device(device).verify(output_s); + + const auto num_tokens = TSize.unwrap(); + const auto num_outer_groups = GSize.unwrap(); + const auto hidden_dim = DSize.unwrap(); + const auto hidden_dim_groups = SSize.unwrap(); + const auto input_stride_t = input.stride(0); + constexpr int64_t kInputAlignElements = sizeof(int4) / sizeof(T); + + RuntimeCheck(hidden_dim % GROUP_SIZE == 0, "hidden_dim must be divisible by 128"); + RuntimeCheck(hidden_dim_groups == hidden_dim / GROUP_SIZE, "output_s hidden dim mismatch"); + RuntimeCheck( + reinterpret_cast(input.data_ptr()) % sizeof(int4) == 0, + "input base pointer must be 16-byte aligned"); + RuntimeCheck( + num_tokens <= 1 || input_stride_t % kInputAlignElements == 0, + "input token stride must preserve 16-byte vector-load alignment"); + + const int64_t total_scale_groups = num_tokens * num_outer_groups * hidden_dim_groups; + if (total_scale_groups == 0) return; + + const auto grid = dim3((total_scale_groups + SUBWARPS_PER_BLOCK - 1) / SUBWARPS_PER_BLOCK); + const auto block = dim3(SUBWARPS_PER_BLOCK * THREADS_PER_GROUP); + host::LaunchKernel(grid, block, device.unwrap()) + .enable_pdl(kUsePDL)( + fp8_wo_a_group_major_quant_ue8m0_kernel, + static_cast(input.data_ptr()), + static_cast(output_q.data_ptr()), + static_cast(output_s.data_ptr()), + total_scale_groups, + static_cast(num_tokens), + static_cast(hidden_dim_groups), + static_cast(num_outer_groups), + static_cast(input_stride_t)); + } +}; + +} // namespace diff --git a/python/sglang/jit_kernel/dsv4/__init__.py b/python/sglang/jit_kernel/dsv4/__init__.py index 21a9c9b28..a5b4d48b2 100644 --- a/python/sglang/jit_kernel/dsv4/__init__.py +++ b/python/sglang/jit_kernel/dsv4/__init__.py @@ -19,6 +19,7 @@ from .elementwise import ( fused_q_norm_rope, fused_rope_inplace, ) +from .fp8_wo_a import sglang_per_token_group_quant_fp8_dsv4_wo_a from .gemm import linear_bf16_fp32 from .moe import ( hash_topk, @@ -45,6 +46,7 @@ __all__ = [ "fused_q_indexer_rope_hadamard_fp4_quant", "fused_q_indexer_rope_hadamard_quant", "fused_k_norm_rope_flashmla", + "sglang_per_token_group_quant_fp8_dsv4_wo_a", "make_name", "linear_bf16_fp32", "get_paged_mqa_logits_metadata", diff --git a/python/sglang/jit_kernel/dsv4/fp8_wo_a.py b/python/sglang/jit_kernel/dsv4/fp8_wo_a.py new file mode 100644 index 000000000..907900bdb --- /dev/null +++ b/python/sglang/jit_kernel/dsv4/fp8_wo_a.py @@ -0,0 +1,93 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Tuple + +import torch + +from sglang.jit_kernel.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) +from sglang.kernel_api_logging import debug_kernel_api +from sglang.srt.utils.custom_op import register_custom_op + +from .utils import make_name + +if TYPE_CHECKING: + from tvm_ffi.module import Module + +_GROUP_SIZE = 128 + + +@cache_once +def _jit_module(in_dtype: torch.dtype, use_pdl: bool) -> Module: + args = make_cpp_args(in_dtype, use_pdl) + return load_jit( + make_name("fp8_wo_a_group_major_quant_ue8m0"), + *args, + cuda_files=["deepseek_v4/fp8_wo_a_group_major_quant.cuh"], + cuda_wrappers=[ + ( + "fp8_wo_a_group_major_quant_ue8m0", + f"FP8WoAGroupMajorQuantUE8M0Kernel<{args}>::run", + ) + ], + # Match the AOT/JIT v2 quant path's fast-math build so FP8 rounding stays + # bit-identical for the DSV4 wo_a replacement. + extra_cuda_cflags=["--use_fast_math"], + ) + + +@register_custom_op( + op_name="fp8_wo_a_group_major_quant_ue8m0", + mutates_args=["output_q", "output_s"], +) +def _fp8_wo_a_group_major_quant_ue8m0_custom_op( + input: torch.Tensor, + output_q: torch.Tensor, + output_s: torch.Tensor, +) -> None: + """Opaque custom-op boundary for the DeepSeek-V4 wo_a quant JIT kernel.""" + assert input.dtype in (torch.bfloat16, torch.float16) + + module = _jit_module(input.dtype, is_arch_support_pdl()) + module.fp8_wo_a_group_major_quant_ue8m0(input, output_q, output_s) + + +@debug_kernel_api +def fp8_wo_a_group_major_quant_ue8m0( + input: torch.Tensor, + output_q: torch.Tensor, + output_s: torch.Tensor, +) -> None: + _fp8_wo_a_group_major_quant_ue8m0_custom_op(input, output_q, output_s) + + +def sglang_per_token_group_quant_fp8_dsv4_wo_a( + x: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Quantize DSV4 wo_a activations for DeepGEMM fp8_einsum. + + The input is a [T, G, D] bf16/fp16 tensor whose hidden dimension is + contiguous. The output codes are contiguous [T, G, D] fp8 values. The scale + tensor is returned as logical [T, G, D/128] fp32 UE8M0 values backed by + contiguous [G, T, D/128] storage, so each group/head [T, S] panel is + contiguous for the DeepGEMM recipe=(1, 1, 128) consumer. Group size is fixed + to 128 and the absmax floor is fixed to 1e-10. + """ + num_tokens, num_groups, hidden = x.shape + hidden_groups = hidden // _GROUP_SIZE + x_q = torch.empty(x.shape, device=x.device, dtype=torch.float8_e4m3fn) + x_s_storage = torch.empty( + (num_groups, num_tokens, hidden_groups), + device=x.device, + dtype=torch.float32, + ) + + if x.numel() > 0: + fp8_wo_a_group_major_quant_ue8m0(x, x_q, x_s_storage) + + # DeepGEMM fp8_einsum consumes each group/head [T, S] scale panel contiguously. + return x_q, x_s_storage.transpose(0, 1) diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index f470bcd60..33a26896d 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -24,6 +24,7 @@ from sglang.jit_kernel.dsv4 import ( fused_norm_rope_inplace, fused_q_norm_rope, fused_rope_inplace, + sglang_per_token_group_quant_fp8_dsv4_wo_a, ) from sglang.srt.compilation.compilation_config import register_split_op from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config @@ -76,7 +77,6 @@ from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.moe import get_moe_a2a_backend, should_use_dp_reduce_scatterv from sglang.srt.layers.moe.fused_moe_triton import FusedMoE -from sglang.srt.layers.quantization.fp8_kernel import sglang_per_token_group_quant_fp8 from sglang.srt.layers.rotary_embedding import get_rope_wrapper from sglang.srt.layers.utils import PPMissingLayer, get_layer_id from sglang.srt.layers.utils.cp_utils import ( @@ -1090,15 +1090,11 @@ class MQALayer(nn.Module): T, G, D = o.shape R = self.o_lora_rank - o_fp8, o_s = sglang_per_token_group_quant_fp8( - o.reshape(T * G, D).contiguous(), - group_size=128, - scale_ue8m0=True, - ) + o_fp8, o_s = sglang_per_token_group_quant_fp8_dsv4_wo_a(o) output = torch.empty(T, G, R, device=o.device, dtype=torch.bfloat16) deep_gemm.fp8_einsum( "bhr,hdr->bhd", - (o_fp8.view(T, G, D), o_s.view(T, G, -1)), + (o_fp8, o_s), (self.wo_a.weight.view(G, R, D), self.wo_a.weight_scale_inv.data), output, recipe=(1, 1, 128), diff --git a/test/registered/jit/deepseek_v4/test_fp8_wo_a.py b/test/registered/jit/deepseek_v4/test_fp8_wo_a.py new file mode 100644 index 000000000..96139aa3a --- /dev/null +++ b/test/registered/jit/deepseek_v4/test_fp8_wo_a.py @@ -0,0 +1,212 @@ +"""DeepSeek-V4 wo_a FP8 activation quant for DeepGEMM fp8_einsum. + +Covers the dedicated DSV4 wo_a quant helper: bit-exact FP8/scales against the +ordinary flat UE8M0 quant values, group-major scale storage, large / DSV4-shaped +token axes, and the DeepGEMM fp8_einsum consumer contract. +""" + +import unittest + +import torch + +import sglang.jit_kernel.dsv4.fp8_wo_a as fp8_wo_a_module +from sglang.jit_kernel.dsv4 import sglang_per_token_group_quant_fp8_dsv4_wo_a +from sglang.srt.layers.quantization.fp8_kernel import ( + fp8_dtype, + sglang_per_token_group_quant_fp8, +) +from sglang.srt.layers.quantization.fp8_utils import ( + block_quant_dequant, + quant_weight_ue8m0, + transform_scale_ue8m0, +) +from sglang.srt.utils import get_device_sm +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="4-gpu-b200") + +_GROUP_SIZE = 128 + + +class TestDeepSeekV4FP8WoA(CustomTestCase): + @classmethod + def setUpClass(cls): + if not torch.cuda.is_available(): + raise unittest.SkipTest("CUDA is not available") + if get_device_sm() < 100: + raise unittest.SkipTest("Test requires CUDA SM 100 or higher") + + try: + import deep_gemm + except ImportError as exc: + raise unittest.SkipTest("deep_gemm is required") from exc + + cls.deep_gemm = deep_gemm + + def _flat_reference(self, o): + T, G, D = o.shape + q_ref, s_ref = sglang_per_token_group_quant_fp8( + o.contiguous().view(T * G, D), + _GROUP_SIZE, + scale_ue8m0=True, + ) + return q_ref.view(T, G, D), s_ref.view(T, G, D // _GROUP_SIZE) + + def _strided_tgd(self, T, G, D, dtype, device): + storage = ( + torch.randn(T, G + 1, D, device=device, dtype=torch.float32) * 0.25 + ).to(dtype) + o = storage[:, 1:, :] + self.assertFalse(o.is_contiguous()) + self.assertEqual(o.stride(-1), 1) + return o + + def _assert_matches_flat_reference(self, o, o_fp8, o_s): + T, G, D = o.shape + q_ref, s_ref = self._flat_reference(o) + torch.cuda.synchronize() + + self.assertEqual(o_fp8.shape, (T, G, D)) + self.assertEqual(o_fp8.dtype, fp8_dtype) + self.assertEqual(o_s.shape, (T, G, D // _GROUP_SIZE)) + self.assertEqual(o_s.dtype, torch.float32) + self.assertEqual(o_s.stride(), (D // _GROUP_SIZE, T * (D // _GROUP_SIZE), 1)) + self.assertTrue(o_s[:, 0, :].is_contiguous()) + self.assertTrue( + torch.equal(o_fp8.view(torch.int8), q_ref.view(torch.int8)), + "fp8 codes differ", + ) + self.assertTrue(torch.equal(o_s, s_ref), "scales differ") + + def test_dsv4_wo_a_quant_matches_flat_reference(self): + torch.manual_seed(1) + torch.cuda.manual_seed_all(1) + + device = torch.device("cuda") + for dtype, T, G, D in [ + (torch.bfloat16, 9, 5, 384), + (torch.float16, 7, 3, 512), + ]: + with self.subTest(dtype=dtype, T=T, G=G, D=D): + o = ( + torch.randn(T, G, D, device=device, dtype=torch.float32) * 0.25 + ).to(dtype) + o_fp8, o_s = sglang_per_token_group_quant_fp8_dsv4_wo_a(o) + self._assert_matches_flat_reference(o, o_fp8, o_s) + + o = self._strided_tgd(T, G, D, dtype, device) + o_fp8, o_s = sglang_per_token_group_quant_fp8_dsv4_wo_a(o) + self._assert_matches_flat_reference(o, o_fp8, o_s) + + def test_dsv4_wo_a_quant_empty_token_dimension(self): + o = torch.empty(0, 3, 256, device="cuda", dtype=torch.bfloat16) + o_fp8, o_s = sglang_per_token_group_quant_fp8_dsv4_wo_a(o) + + self.assertEqual(o_fp8.shape, o.shape) + self.assertEqual(o_fp8.dtype, fp8_dtype) + self.assertEqual(o_s.shape, (0, 3, 2)) + self.assertEqual(o_s.dtype, torch.float32) + self.assertEqual(o_s.stride(), (2, 2, 1)) + + def test_dsv4_wo_a_quant_large_token_dimension(self): + torch.manual_seed(2) + torch.cuda.manual_seed_all(2) + + cases = [ + (10_001, 1, 128, False), + (10_001, 8, 4096, False), + (300_000, 1, 128, True), + ] + + device = torch.device("cuda") + for T, G, D, strided in cases: + with self.subTest(T=T, G=G, D=D, strided=strided): + if strided: + o = self._strided_tgd(T, G, D, torch.bfloat16, device) + else: + o = ( + torch.randn(T, G, D, device=device, dtype=torch.float32) * 0.25 + ).to(torch.bfloat16) + o_fp8, o_s = sglang_per_token_group_quant_fp8_dsv4_wo_a(o) + self._assert_matches_flat_reference(o, o_fp8, o_s) + + def test_dsv4_wo_a_quant_uses_dedicated_jit(self): + torch.manual_seed(3) + torch.cuda.manual_seed_all(3) + + original_jit_module = fp8_wo_a_module._jit_module + jit_module_calls = 0 + + def wrapped_jit_module(*args, **kwargs): + nonlocal jit_module_calls + jit_module_calls += 1 + return original_jit_module(*args, **kwargs) + + fp8_wo_a_module._jit_module = wrapped_jit_module + try: + o = self._strided_tgd(3, 2, 256, torch.bfloat16, "cuda") + sglang_per_token_group_quant_fp8_dsv4_wo_a(o) + torch.cuda.synchronize() + self.assertGreater(jit_module_calls, 0) + finally: + fp8_wo_a_module._jit_module = original_jit_module + + def test_fp8_wo_a_einsum_uses_group_major_activation_scales(self): + torch.manual_seed(0) + torch.cuda.manual_seed_all(0) + + cases = [ + (5, 2, 256, 256), + (7, 8, 4096, 1024), + ] + device = torch.device("cuda") + for T, G, D, R in cases: + with self.subTest(T=T, G=G, D=D, R=R): + token_scale = torch.linspace(0.5, 1.5, T, device=device).view(T, 1, 1) + group_scale = torch.pow( + 2.0, (torch.arange(G, device=device).float() % 5) - 2.0 + ).view(1, G, 1) + o = ( + torch.randn(T, G, D, device=device, dtype=torch.float32) + * token_scale + * group_scale + * 0.2 + ).to(torch.bfloat16) + weight = ( + torch.randn(G, R, D, device=device, dtype=torch.float32) * 0.2 + ).to(torch.bfloat16) + + weight_fp8, weight_s_raw = quant_weight_ue8m0( + weight, weight_block_size=[_GROUP_SIZE, _GROUP_SIZE] + ) + weight_s = transform_scale_ue8m0(weight_s_raw, mn=R) + + q_dsv4, s_dsv4 = sglang_per_token_group_quant_fp8_dsv4_wo_a(o) + out = torch.empty(T, G, R, device=device, dtype=torch.bfloat16) + self.deep_gemm.fp8_einsum( + "bhr,hdr->bhd", + (q_dsv4, s_dsv4), + (weight_fp8, weight_s), + out, + recipe=(1, 1, _GROUP_SIZE), + ) + torch.cuda.synchronize() + + o_dequant = q_dsv4.float().view( + T, G, D // _GROUP_SIZE, _GROUP_SIZE + ) * s_dsv4.unsqueeze(-1) + weight_dequant = block_quant_dequant( + weight_fp8, + weight_s_raw, + block_size=[_GROUP_SIZE, _GROUP_SIZE], + dtype=torch.float32, + ) + ref = torch.einsum( + "tgd,grd->tgr", o_dequant.view(T, G, D), weight_dequant.float() + ).to(torch.bfloat16) + torch.testing.assert_close(out, ref, atol=1e-1, rtol=2e-2) + + +if __name__ == "__main__": + unittest.main()