test: add CuTeDSL NVFP4 GEMM test using reference ScaledGroupedGemmKernel
Tests the NVIDIA reference kernel with our quantization pipeline: 1. Quantize BF16 → NVFP4 (our stage_activation logic) 2. Pad and swizzle scale factors (to_blocked) 3. Run ScaledGroupedGemmKernel (2Dx3D scenario) 4. Compare against BF16 matmul reference Also adds cutedsl/moe.py module for the future pipeline integration.
This commit is contained in:
14
src/nvfp4_megamoe_kernel/cutedsl/__init__.py
Normal file
14
src/nvfp4_megamoe_kernel/cutedsl/__init__.py
Normal file
@@ -0,0 +1,14 @@
|
||||
"""
|
||||
NVFP4 MoE kernel using NVIDIA's CuTeDSL ScaledGroupedGemmKernel.
|
||||
|
||||
This replaces the broken C++ CUTLASS kernel. The CuTeDSL kernel handles:
|
||||
- NVFP4 (Float4E2M1FN + Float8E4M3FN, sf_vec_size=16) natively
|
||||
- Block-scaled SF layouts (no manual remap needed)
|
||||
- Full Blackwell pipeline (TMA → MMA → Epilogue overlap)
|
||||
- Per-expert global scales for NVFP4
|
||||
|
||||
We just need to:
|
||||
1. Quantize activations to FP4 (stage_activation)
|
||||
2. Call the kernel with the right tensor layout
|
||||
3. Apply MoE routing (gate/up fusion, SiLU, scatter)
|
||||
"""
|
||||
171
src/nvfp4_megamoe_kernel/cutedsl/moe.py
Normal file
171
src/nvfp4_megamoe_kernel/cutedsl/moe.py
Normal file
@@ -0,0 +1,171 @@
|
||||
"""
|
||||
NVFP4 MoE pipeline using CuTeDSL ScaledGroupedGemmKernel.
|
||||
|
||||
Replaces the broken C++ CUTLASS path. Uses NVIDIA's official MoE scaled
|
||||
grouped GEMM kernel from the CUTLASS CuTeDSL examples.
|
||||
|
||||
Usage:
|
||||
from nvfp4_megamoe_kernel.cutedsl.moe import nvfp4_mega_moe_full
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
import torch
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
import cutlass.torch as cutlass_torch
|
||||
import cutlass.utils as utils
|
||||
import cutlass.utils.blockscaled_layout as blockscaled_utils
|
||||
|
||||
# Add the CuTeDSL examples to the path so we can import the kernel
|
||||
_CUTLASS_ROOT = os.environ.get("CUTLASS_ROOT", "/root/cutlass")
|
||||
_CUTEDSL_EXAMPLES = os.path.join(_CUTLASS_ROOT, "examples/python/CuTeDSL")
|
||||
if _CUTEDSL_EXAMPLES not in sys.path:
|
||||
sys.path.insert(0, _CUTEDSL_EXAMPLES)
|
||||
|
||||
from cute.blackwell.kernel.moe.torch_scaled_grouped_mm import ScaledGroupedGemmKernel
|
||||
|
||||
from nvfp4_megamoe_kernel.nvfp4_mega_moe import (
|
||||
stage_activation,
|
||||
_quantize_to_e2m1,
|
||||
)
|
||||
|
||||
# ── Module-level compiled kernel cache ──
|
||||
_compiled_l1_kernel = None
|
||||
_compiled_l2_kernel = None
|
||||
_l1_kernel_config = None
|
||||
_l2_kernel_config = None
|
||||
|
||||
|
||||
def _get_torch_dtype(cutlass_dtype):
|
||||
"""Convert CUTLASS dtype to PyTorch dtype."""
|
||||
mapping = {
|
||||
cutlass.Float4E2M1FN: torch.float4_e2m1fn_x2,
|
||||
cutlass.Float8E4M3FN: torch.float8_e4m3fn,
|
||||
cutlass.Float8E8M0FNU: torch.float8_e8m0fnu,
|
||||
cutlass.BFloat16: torch.bfloat16,
|
||||
cutlass.Float16: torch.float16,
|
||||
cutlass.Float32: torch.float32,
|
||||
}
|
||||
return mapping.get(cutlass_dtype)
|
||||
|
||||
|
||||
def _torch_tensor_to_cute(torch_tensor: torch.Tensor) -> cute.Tensor:
|
||||
"""Convert a PyTorch GPU tensor to a CuTe tensor with dynamic layout."""
|
||||
cute_tensor = cutlass_torch.from_dlpack(torch_tensor)
|
||||
leading_dim = cutlass_torch.get_leading_dim(torch_tensor)
|
||||
cute_tensor = cute_tensor.mark_layout_dynamic(leading_dim=leading_dim)
|
||||
return cute_tensor
|
||||
|
||||
|
||||
def _compile_kernel_once(kernel, sample_tensors, global_scale_a=None, global_scale_b=None):
|
||||
"""Compile the CuTeDSL kernel on first call, cache the result."""
|
||||
import cuda.bindings.driver as cuda
|
||||
|
||||
a_cute, b_cute, sfa_cute, sfb_cute, c_cute, offs_cute, ws_cute = sample_tensors
|
||||
|
||||
gsa_cute = _torch_tensor_to_cute(global_scale_a) if global_scale_a is not None else None
|
||||
gsb_cute = _torch_tensor_to_cute(global_scale_b) if global_scale_b is not None else None
|
||||
|
||||
cluster_size = kernel.cluster_shape_mn[0] * kernel.cluster_shape_mn[1]
|
||||
max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size)
|
||||
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
|
||||
|
||||
compiled = cute.compile(
|
||||
kernel,
|
||||
a_cute, b_cute, sfa_cute, sfb_cute, c_cute, offs_cute, ws_cute,
|
||||
max_active_clusters, stream,
|
||||
global_scale_a=gsa_cute,
|
||||
global_scale_b=gsb_cute,
|
||||
)
|
||||
return compiled
|
||||
|
||||
|
||||
def run_scaled_grouped_gemm(
|
||||
mat_a: torch.Tensor, # (tokens_sum, K_packed) float4_e2m1fn_x2 — row-major (K-major for CuTe)
|
||||
mat_b: torch.Tensor, # (experts, K_packed, N) float4_e2m1fn_x2 — K-major
|
||||
scale_a: torch.Tensor, # (tokens_sum, K_sf) float8_e4m3fn — row-major
|
||||
scale_b: torch.Tensor, # (experts, K_sf, N) float8_e4m3fn — K-major after transpose
|
||||
expert_offsets: torch.Tensor, # (experts,) int32 — cumulative token offsets
|
||||
global_scale_a: torch.Tensor = None, # (experts,) float32 — NVFP4 per-expert activation scale
|
||||
global_scale_b: torch.Tensor = None, # (experts,) float32 — NVFP4 per-expert weight scale
|
||||
mma_tiler_mn: tuple = (128, 128),
|
||||
cluster_shape_mn: tuple = (1, 1),
|
||||
) -> torch.Tensor:
|
||||
"""Run the CuTeDSL NVFP4 scaled grouped GEMM.
|
||||
|
||||
2Dx3D scenario: A(tokens, K) x B(experts, K, N) -> C(tokens, N)
|
||||
|
||||
Args:
|
||||
mat_a: Activation tensor (tokens_sum, K_packed) in FP4
|
||||
mat_b: Weight tensor (experts, K_packed, N) in FP4
|
||||
scale_a: Activation block scales (tokens_sum, K_sf) in E4M3
|
||||
scale_b: Weight block scales (experts, K_sf, N) in E4M3
|
||||
expert_offsets: Cumulative token end offsets per expert
|
||||
global_scale_a: Per-expert float32 activation global scale (NVFP4)
|
||||
global_scale_b: Per-expert float32 weight global scale (NVFP4)
|
||||
|
||||
Returns:
|
||||
Output tensor (tokens_sum, N) in BF16
|
||||
"""
|
||||
global _compiled_l1_kernel, _l1_kernel_config
|
||||
|
||||
tokens_sum = mat_a.shape[0]
|
||||
k_packed = mat_a.shape[1]
|
||||
num_experts = mat_b.shape[0]
|
||||
n_dim = mat_b.shape[2]
|
||||
k_dim = k_packed * 2 # 2 FP4 values per byte
|
||||
|
||||
# Output tensor
|
||||
out = torch.zeros(tokens_sum, n_dim, dtype=torch.bfloat16, device=mat_a.device)
|
||||
|
||||
# Create kernel config
|
||||
kernel = ScaledGroupedGemmKernel(
|
||||
scenario="2Dx3D",
|
||||
sf_vec_size=16,
|
||||
accumulate_on_output=False,
|
||||
separate_tensormap_init=True,
|
||||
consistent_token_padding=False,
|
||||
mma_tiler_mnk=(*mma_tiler_mn, 256),
|
||||
cluster_shape_mnk=(*cluster_shape_mn, 1),
|
||||
)
|
||||
|
||||
# Convert to CuTe tensors
|
||||
a_cute = _torch_tensor_to_cute(mat_a)
|
||||
b_cute = _torch_tensor_to_cute(mat_b)
|
||||
sfa_cute = _torch_tensor_to_cute(scale_a)
|
||||
sfb_cute = _torch_tensor_to_cute(scale_b)
|
||||
c_cute = _torch_tensor_to_cute(out)
|
||||
offs_cute = _torch_tensor_to_cute(expert_offsets)
|
||||
|
||||
# Workspace
|
||||
workspace_size = kernel.get_workspace_size(num_experts)
|
||||
workspace = torch.full((workspace_size,), 255, dtype=torch.uint8, device=mat_a.device)
|
||||
ws_cute = _torch_tensor_to_cute(workspace)
|
||||
|
||||
gsa_cute = _torch_tensor_to_cute(global_scale_a) if global_scale_a is not None else None
|
||||
gsb_cute = _torch_tensor_to_cute(global_scale_b) if global_scale_b is not None else None
|
||||
|
||||
import cuda.bindings.driver as cuda
|
||||
cluster_size = kernel.cluster_shape_mn[0] * kernel.cluster_shape_mn[1]
|
||||
max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size)
|
||||
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
|
||||
|
||||
# Compile and run
|
||||
compiled = cute.compile(
|
||||
kernel,
|
||||
a_cute, b_cute, sfa_cute, sfb_cute, c_cute, offs_cute, ws_cute,
|
||||
max_active_clusters, stream,
|
||||
global_scale_a=gsa_cute,
|
||||
global_scale_b=gsb_cute,
|
||||
)
|
||||
|
||||
compiled(
|
||||
a_cute, b_cute, sfa_cute, sfb_cute, c_cute, offs_cute, ws_cute,
|
||||
stream,
|
||||
global_scale_a=gsa_cute,
|
||||
global_scale_b=gsb_cute,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
return out
|
||||
306
tests/test_cutedsl.py
Normal file
306
tests/test_cutedsl.py
Normal file
@@ -0,0 +1,306 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
CuTeDSL NVFP4 GEMM test — verify the reference kernel works with our data.
|
||||
|
||||
Uses NVIDIA's ScaledGroupedGemmKernel from the CUTLASS CuTeDSL examples
|
||||
with NVFP4 (Float4E2M1FN + Float8E4M3FN, sf_vec_size=16).
|
||||
|
||||
This tests a single GEMM: A(tokens, K) @ B(experts, K, N) = C(tokens, N)
|
||||
with proper scale factor padding/swizzling.
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
import math
|
||||
import torch
|
||||
|
||||
# Add CuTeDSL examples to path
|
||||
CUTLASS_ROOT = os.environ.get("CUTLASS_ROOT", "/root/cutlass")
|
||||
CUTEDSL_EXAMPLES = os.path.join(CUTLASS_ROOT, "examples/python/CuTeDSL")
|
||||
sys.path.insert(0, CUTEDSL_EXAMPLES)
|
||||
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
import cutlass.torch as cutlass_torch
|
||||
import cutlass.utils as utils
|
||||
import cutlass.utils.blockscaled_layout as blockscaled_utils
|
||||
|
||||
from cute.blackwell.kernel.moe.torch_scaled_grouped_mm import (
|
||||
ScaledGroupedGemmKernel,
|
||||
pad_and_swizzle_single,
|
||||
assemble_raw_scales_2d3d_2d_side,
|
||||
assemble_raw_scales_2d3d_3d_side,
|
||||
cat_byte_reinterpretable_tensors,
|
||||
stack_byte_reinterpretable_tensors,
|
||||
offs_to_group_sizes,
|
||||
)
|
||||
|
||||
# ── Helpers ────────────────────────────────────────────────────────────
|
||||
|
||||
def ceil_div(a, b):
|
||||
return (a + b - 1) // b
|
||||
|
||||
def round_up(a, b):
|
||||
return ceil_div(a, b) * b
|
||||
|
||||
|
||||
def quantize_bf16_to_nvfp4(x_bf16, block_size=16):
|
||||
"""Quantize BF16 tensor to NVFP4 (E2M1 + E4M3 block scales + global scale).
|
||||
|
||||
Returns (x_fp4, block_scales, global_scale) where:
|
||||
x_fp4: torch.float4_e2m1fn_x2 with same shape (packed along last dim)
|
||||
block_scales: float8_e4m3fn with shape (..., ceil_div(last_dim, block_size))
|
||||
global_scale: float32 scalar
|
||||
"""
|
||||
x_f32 = x_bf16.float()
|
||||
amax = x_f32.abs().max().clamp(min=1e-8).float()
|
||||
global_scale = amax / (6.0 * 448.0)
|
||||
|
||||
x_norm = x_f32 / global_scale
|
||||
|
||||
# Per-block amax for block scales
|
||||
last_dim = x_norm.shape[-1]
|
||||
n_blocks = ceil_div(last_dim, block_size)
|
||||
|
||||
# Pad last dim to multiple of block_size
|
||||
if last_dim % block_size != 0:
|
||||
pad_size = n_blocks * block_size - last_dim
|
||||
x_norm = torch.nn.functional.pad(x_norm, (0, pad_size))
|
||||
|
||||
x_reshaped = x_norm.reshape(*x_norm.shape[:-1], n_blocks, block_size)
|
||||
block_amax = x_reshaped.abs().amax(dim=-1).clamp(min=1e-8)
|
||||
block_scale = (block_amax / 6.0).to(torch.float8_e4m3fn)
|
||||
|
||||
# Quantize to E2M1
|
||||
E2M1_MAGNITUDES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
|
||||
# For each value, find nearest E2M1 magnitude
|
||||
x_blocks = x_reshaped # (..., n_blocks, block_size)
|
||||
block_sf_expanded = block_scale.float().unsqueeze(-1) # (..., n_blocks, 1)
|
||||
x_scaled = x_blocks / block_sf_expanded.clamp(min=1e-8) # normalize by block scale
|
||||
|
||||
# Nearest E2M1
|
||||
magnitudes = torch.tensor(E2M1_MAGNITUDES, dtype=torch.float32, device=x_bf16.device)
|
||||
signs = torch.sign(x_scaled)
|
||||
abs_scaled = x_scaled.abs().unsqueeze(-1) # (..., block_size, 1)
|
||||
distances = (abs_scaled - magnitudes).abs() # (..., block_size, 8)
|
||||
indices = distances.argmin(dim=-1) # (..., block_size)
|
||||
|
||||
# Sign: positive = 0-7, negative = 8-15
|
||||
nibbles = torch.where(signs < 0, indices + 8, indices).to(torch.uint8)
|
||||
|
||||
# Pack pairs: byte = (odd_nibble << 4) | even_nibble
|
||||
even = nibbles[..., ::2]
|
||||
odd = nibbles[..., 1::2]
|
||||
packed = (odd << 4) | even
|
||||
|
||||
# Reshape back to original shape (with packed last dim)
|
||||
orig_shape = list(x_bf16.shape)
|
||||
orig_shape[-1] = ceil_div(orig_shape[-1], 2)
|
||||
x_fp4 = packed.view(torch.float4_e2m1fn_x2).reshape(orig_shape)
|
||||
|
||||
# Reshape block scales
|
||||
sf_shape = list(x_bf16.shape[:-1]) + [n_blocks]
|
||||
block_scale = block_scale.reshape(sf_shape)
|
||||
|
||||
return x_fp4, block_scale, global_scale
|
||||
|
||||
|
||||
def dequantize_nvfp4(x_fp4, block_scales, global_scale):
|
||||
"""Dequantize NVFP4 back to BF16 for reference comparison."""
|
||||
E2M1_LUT = torch.tensor([
|
||||
0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
|
||||
-0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0,
|
||||
], dtype=torch.float32, device=x_fp4.device)
|
||||
|
||||
raw = x_fp4.view(torch.uint8)
|
||||
lo = E2M1_LUT[(raw & 0x0F).long()]
|
||||
hi = E2M1_LUT[((raw >> 4) & 0x0F).long()]
|
||||
|
||||
unpacked = torch.empty(*raw.shape[:-1], raw.shape[-1] * 2, dtype=torch.float32, device=x_fp4.device)
|
||||
unpacked[..., ::2] = lo
|
||||
unpacked[..., 1::2] = hi
|
||||
|
||||
# Expand block scales
|
||||
n_blocks = block_scales.shape[-1]
|
||||
block_size = (unpacked.shape[-1]) // n_blocks
|
||||
block_sf = block_scales.float().unsqueeze(-1).expand(*block_scales.shape, block_size)
|
||||
block_sf = block_sf.reshape(*unpacked.shape)
|
||||
|
||||
return (unpacked * block_sf * global_scale).to(torch.bfloat16)
|
||||
|
||||
|
||||
# ── Main Test ──────────────────────────────────────────────────────────
|
||||
|
||||
def main():
|
||||
torch.manual_seed(42)
|
||||
device = "cuda"
|
||||
|
||||
# Problem sizes
|
||||
num_experts = 2
|
||||
tokens_per_expert = 64
|
||||
hidden = 256 # K dimension
|
||||
intermediate = 128 # N dimension
|
||||
sf_vec_size = 16
|
||||
block_size = 16
|
||||
|
||||
tokens_sum = num_experts * tokens_per_expert
|
||||
|
||||
print(f"Test: {num_experts} experts, {tokens_per_expert} tokens each, K={hidden}, N={intermediate}")
|
||||
|
||||
# ── Create BF16 reference data ──
|
||||
x_bf16 = torch.randn(tokens_sum, hidden, dtype=torch.bfloat16, device=device) * 2.0
|
||||
w_bf16 = torch.randn(num_experts, hidden, intermediate, dtype=torch.bfloat16, device=device) * 0.5
|
||||
|
||||
# BF16 reference: for each expert, matmul its tokens with its weight
|
||||
ref_out = torch.zeros(tokens_sum, intermediate, dtype=torch.bfloat16, device=device)
|
||||
for e in range(num_experts):
|
||||
start = e * tokens_per_expert
|
||||
end = (e + 1) * tokens_per_expert
|
||||
ref_out[start:end] = x_bf16[start:end] @ w_bf16[e]
|
||||
|
||||
print(f"BF16 ref: amax={ref_out.abs().max():.4f} mean={ref_out.float().mean():.6f}")
|
||||
|
||||
# ── Quantize to NVFP4 ──
|
||||
x_fp4, x_sf, x_gs = quantize_bf16_to_nvfp4(x_bf16)
|
||||
w_fp4_list, w_sf_list, w_gs_list = [], [], []
|
||||
for e in range(num_experts):
|
||||
w_fp4, w_sf, w_gs = quantize_bf16_to_nvfp4(w_bf16[e])
|
||||
w_fp4_list.append(w_fp4)
|
||||
w_sf_list.append(w_sf)
|
||||
w_gs_list.append(w_gs)
|
||||
|
||||
# Verify quantization roundtrip
|
||||
x_deq = dequantize_nvfp4(x_fp4, x_sf, x_gs)
|
||||
cos_quant = torch.nn.functional.cosine_similarity(
|
||||
x_bf16.flatten().unsqueeze(0).float(),
|
||||
x_deq.flatten().unsqueeze(0).float(),
|
||||
).item()
|
||||
print(f"Quantization roundtrip cosine: {cos_quant:.6f}")
|
||||
|
||||
# ── Prepare CuTeDSL kernel inputs ──
|
||||
# The kernel expects:
|
||||
# mat_a: (tokens_sum, K_packed) float4_e2m1fn_x2
|
||||
# mat_b: (experts, K_packed, N_packed) float4_e2m1fn_x2 — K-major
|
||||
# scale_a: assembled 2D side (padded + swizzled)
|
||||
# scale_b: assembled 3D side (padded + swizzled per expert)
|
||||
# offs: (experts,) int32 cumulative token offsets
|
||||
# global_scale_a: (experts,) float32
|
||||
# global_scale_b: (experts,) float32
|
||||
|
||||
# Expert offsets (cumulative sum of tokens per expert)
|
||||
offs = torch.tensor([tokens_per_expert * (e + 1) for e in range(num_experts)],
|
||||
dtype=torch.int32, device=device)
|
||||
|
||||
# Assemble scale_a (2D side: concatenate per-expert, pad to 128, swizzle)
|
||||
raw_scale_a = [x_sf[e*tokens_per_expert:(e+1)*tokens_per_expert] for e in range(num_experts)]
|
||||
scale_a = assemble_raw_scales_2d3d_2d_side(raw_scale_a)
|
||||
|
||||
# Assemble scale_b (3D side: per-expert, pad and swizzle each)
|
||||
scale_b = assemble_raw_scales_2d3d_3d_side(w_sf_list)
|
||||
|
||||
# Global scales
|
||||
global_scale_a = torch.tensor([x_gs] * num_experts, dtype=torch.float32, device=device)
|
||||
global_scale_b = torch.tensor([w_gs_list[e] for e in range(num_experts)], dtype=torch.float32, device=device)
|
||||
|
||||
# mat_a is already (tokens_sum, K_packed)
|
||||
mat_a = x_fp4
|
||||
|
||||
# mat_b needs to be (experts, K_packed, N_packed) — K-major
|
||||
mat_b = torch.stack(w_fp4_list) # (experts, K_packed, N_packed)
|
||||
|
||||
print(f"\nKernel inputs:")
|
||||
print(f" mat_a: {mat_a.shape} {mat_a.dtype}")
|
||||
print(f" mat_b: {mat_b.shape} {mat_b.dtype}")
|
||||
print(f" scale_a: {scale_a.shape} {scale_a.dtype}")
|
||||
print(f" scale_b: {scale_b.shape} {scale_b.dtype}")
|
||||
print(f" offs: {offs.tolist()}")
|
||||
print(f" global_scale_a: {global_scale_a.tolist()}")
|
||||
print(f" global_scale_b: {[f'{v:.6e}' for v in global_scale_b.tolist()]}")
|
||||
|
||||
# ── Run CuTeDSL kernel ──
|
||||
print("\nCompiling and running CuTeDSL kernel (first run takes ~1 min to compile)...")
|
||||
|
||||
out = torch.zeros(tokens_sum, intermediate, dtype=torch.bfloat16, device=device)
|
||||
|
||||
kernel = ScaledGroupedGemmKernel(
|
||||
scenario="2Dx3D",
|
||||
sf_vec_size=sf_vec_size,
|
||||
accumulate_on_output=False,
|
||||
separate_tensormap_init=True,
|
||||
consistent_token_padding=False,
|
||||
mma_tiler_mnk=(128, 128, 256),
|
||||
cluster_shape_mnk=(1, 1, 1),
|
||||
)
|
||||
|
||||
# Convert to CuTe tensors
|
||||
a_cute = cutlass_torch.from_dlpack(mat_a)
|
||||
a_cute = a_cute.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(mat_a))
|
||||
|
||||
b_cute = cutlass_torch.from_dlpack(mat_b)
|
||||
b_cute = b_cute.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(mat_b))
|
||||
|
||||
sfa_cute = cutlass_torch.from_dlpack(scale_a)
|
||||
sfa_cute = sfa_cute.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(scale_a))
|
||||
|
||||
sfb_cute = cutlass_torch.from_dlpack(scale_b)
|
||||
sfb_cute = sfb_cute.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(scale_b))
|
||||
|
||||
c_cute = cutlass_torch.from_dlpack(out)
|
||||
c_cute = c_cute.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(out))
|
||||
|
||||
offs_cute = cutlass_torch.from_dlpack(offs)
|
||||
offs_cute = offs_cute.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(offs))
|
||||
|
||||
workspace_size = kernel.get_workspace_size(num_experts)
|
||||
workspace = torch.full((workspace_size,), 255, dtype=torch.uint8, device=device)
|
||||
ws_cute = cutlass_torch.from_dlpack(workspace)
|
||||
ws_cute = ws_cute.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(workspace))
|
||||
|
||||
gsa_cute = cutlass_torch.from_dlpack(global_scale_a)
|
||||
gsa_cute = gsa_cute.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(global_scale_a))
|
||||
|
||||
gsb_cute = cutlass_torch.from_dlpack(global_scale_b)
|
||||
gsb_cute = gsb_cute.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(global_scale_b))
|
||||
|
||||
import cuda.bindings.driver as cuda
|
||||
cluster_size = 1
|
||||
max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size)
|
||||
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
|
||||
|
||||
compiled = cute.compile(
|
||||
kernel,
|
||||
a_cute, b_cute, sfa_cute, sfb_cute, c_cute, offs_cute, ws_cute,
|
||||
max_active_clusters, stream,
|
||||
global_scale_a=gsa_cute,
|
||||
global_scale_b=gsb_cute,
|
||||
)
|
||||
|
||||
compiled(
|
||||
a_cute, b_cute, sfa_cute, sfb_cute, c_cute, offs_cute, ws_cute,
|
||||
stream,
|
||||
global_scale_a=gsa_cute,
|
||||
global_scale_b=gsb_cute,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# ── Compare ──
|
||||
cosine = torch.nn.functional.cosine_similarity(
|
||||
out.flatten().unsqueeze(0).float(),
|
||||
ref_out.flatten().unsqueeze(0).float(),
|
||||
).item()
|
||||
mse = (out.float() - ref_out.float()).pow(2).mean().item()
|
||||
|
||||
print(f"\n{'='*70}")
|
||||
print(f" RESULT: cosine={cosine:.6f} MSE={mse:.6e}")
|
||||
print(f"{'='*70}")
|
||||
|
||||
if cosine > 0.99:
|
||||
print(f" ✅ PASS: CuTeDSL kernel matches BF16 reference")
|
||||
elif cosine > 0.95:
|
||||
print(f" ⚠️ Close but not perfect — quantization loss?")
|
||||
else:
|
||||
print(f" ❌ FAIL: kernel output doesn't match reference")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user