Add NVFP4 linear runner + attention projection test
- CuTeDSLNvfp4Linear: generic single-GEMM runner for any NVFP4 projection - test_attention.py: tests q_a_proj, q_b_proj, kv_proj, o_b_proj vs BF16 - Same pad+swizzle pattern as shared expert, but no SiLU/fusion
This commit is contained in:
158
cutedsl/nvfp4_linear.py
Normal file
158
cutedsl/nvfp4_linear.py
Normal file
@@ -0,0 +1,158 @@
|
||||
"""CuTeDSL NVFP4 Linear (single GEMM)
|
||||
|
||||
Generic NVFP4 GEMM runner for attention projections and any single
|
||||
linear layer. Uses ScaledGroupedGemmKernel with num_groups=1.
|
||||
|
||||
CUDA-graph-compatible: all buffers pre-allocated, no CPU-GPU syncs.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from cutedsl.bridge import (
|
||||
quantize_activation_nvfp4,
|
||||
quantize_to_nvfp4,
|
||||
make_b_k_major,
|
||||
assemble_scales_3d_side,
|
||||
run_nvfp4_grouped_gemm,
|
||||
)
|
||||
from cutedsl.kernel.moe.torch_scaled_grouped_mm import (
|
||||
ceil_div as cutedsl_ceil_div,
|
||||
pad_and_swizzle_single,
|
||||
)
|
||||
|
||||
|
||||
class CuTeDSLNvfp4Linear:
|
||||
"""Single NVFP4 GEMM using CuTeDSL (num_groups=1).
|
||||
|
||||
Handles any (K, N) weight matrix in NVFP4 format.
|
||||
Simple: quantize activation → GEMM → BF16 output.
|
||||
No SiLU, no fusion, no routing.
|
||||
|
||||
CUDA-graph-compatible: all buffers pre-allocated, no CPU-GPU syncs.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_features: int,
|
||||
out_features: int,
|
||||
max_num_tokens: int = 8192,
|
||||
device: str = "cuda",
|
||||
):
|
||||
self.in_features = in_features
|
||||
self.out_features = out_features
|
||||
self.max_num_tokens = max_num_tokens
|
||||
self.device = device
|
||||
|
||||
# Weights (set after construction, then call finalize_weights)
|
||||
self.fp4 = None # list of 1 tensor
|
||||
self.sf = None # list of 1 tensor
|
||||
self.gs = None # list of 1 float
|
||||
|
||||
# Processed weights
|
||||
self._mat_b = None
|
||||
self._scale_b = None
|
||||
self._gsb = None
|
||||
|
||||
# Activation global scale
|
||||
self._activation_global_scale = 1.0 / (6.0 * 448.0)
|
||||
|
||||
# Pre-allocated buffers
|
||||
self._padded_x_fp4_buf = None
|
||||
self._expert_offsets_buf = None
|
||||
self._gsa_buf = None
|
||||
self._buffers_allocated = False
|
||||
|
||||
import os
|
||||
print(f"[CLAWMINE] Nvfp4Linear init: in={in_features} out={out_features} "
|
||||
f"max_tokens={max_num_tokens} pid={os.getpid()}", flush=True)
|
||||
|
||||
def finalize_weights(self):
|
||||
"""Process weights for CuTeDSL GEMM."""
|
||||
self._mat_b = make_b_k_major(torch.stack(self.fp4)) # (1, K_packed, N_packed)
|
||||
self._scale_b = assemble_scales_3d_side(self.sf)
|
||||
self._gsb = torch.tensor(self.gs, dtype=torch.float32, device=self.device)
|
||||
|
||||
# Free raw weights
|
||||
self.fp4 = None
|
||||
self.sf = None
|
||||
self.gs = None
|
||||
|
||||
def _allocate_buffers(self):
|
||||
"""Pre-allocate buffers at max size for cudagraph compatibility."""
|
||||
max_rows = cutedsl_ceil_div(self.max_num_tokens, 128) * 128
|
||||
|
||||
self._padded_x_fp4_buf = torch.zeros(
|
||||
max_rows, self.in_features // 2, dtype=torch.uint8, device=self.device
|
||||
).view(torch.float4_e2m1fn_x2)
|
||||
|
||||
self._expert_offsets_buf = torch.zeros(1, dtype=torch.int32, device=self.device)
|
||||
self._gsa_buf = torch.zeros(1, dtype=torch.float32, device=self.device)
|
||||
self._buffers_allocated = True
|
||||
|
||||
def _ensure_initialized(self):
|
||||
if self._mat_b is None:
|
||||
self.finalize_weights()
|
||||
if not self._buffers_allocated:
|
||||
self._allocate_buffers()
|
||||
|
||||
def _assemble_scales_single_group(self, x_sf):
|
||||
"""Assemble 2D-side activation scales for num_groups=1."""
|
||||
num_rows, num_cols = x_sf.shape
|
||||
padded_rows = cutedsl_ceil_div(num_rows, 128) * 128
|
||||
padded_cols = cutedsl_ceil_div(num_cols, 4) * 4
|
||||
|
||||
buf = torch.zeros(padded_rows, padded_cols, dtype=torch.float16, device=x_sf.device).to(torch.float8_e4m3fn)
|
||||
buf[:num_rows, :num_cols] = x_sf
|
||||
swizzled_flat = pad_and_swizzle_single(buf)
|
||||
return swizzled_flat.reshape(padded_rows, padded_cols)
|
||||
|
||||
def compute_activation_global_scale(self, hidden_states_sample):
|
||||
"""Compute activation global scale from a warmup forward."""
|
||||
self._ensure_initialized()
|
||||
with torch.no_grad():
|
||||
_, _, gs = quantize_to_nvfp4(hidden_states_sample)
|
||||
self._activation_global_scale = gs
|
||||
print(f" Nvfp4Linear warmup: gs={gs:.10f}", flush=True)
|
||||
|
||||
def run(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
"""Forward: BF16 input → NVFP4 GEMM → BF16 output."""
|
||||
self._ensure_initialized()
|
||||
|
||||
num_tokens = hidden_states.shape[0]
|
||||
padded_rows = cutedsl_ceil_div(num_tokens, 128) * 128
|
||||
|
||||
# Quantize activation
|
||||
x_fp4, x_sf = quantize_activation_nvfp4(
|
||||
hidden_states, self._activation_global_scale
|
||||
)
|
||||
|
||||
# Scatter x_fp4 into padded buffer
|
||||
padded_x_fp4 = self._padded_x_fp4_buf
|
||||
padded_x_fp4.view(torch.uint8).zero_()
|
||||
padded_x_fp4.view(torch.uint8)[:num_tokens] = x_fp4.view(torch.uint8)
|
||||
|
||||
# Assemble A-side scales
|
||||
scale_a = self._assemble_scales_single_group(x_sf)
|
||||
|
||||
# Expert offsets: [padded_rows] for 1 group
|
||||
expert_offsets = self._expert_offsets_buf
|
||||
expert_offsets.fill_(padded_rows)
|
||||
|
||||
# Global scales
|
||||
gsa = self._gsa_buf.fill_(self._activation_global_scale)
|
||||
|
||||
# Run GEMM
|
||||
out = run_nvfp4_grouped_gemm(
|
||||
mat_a=padded_x_fp4,
|
||||
mat_b=self._mat_b,
|
||||
scale_a=scale_a,
|
||||
scale_b=self._scale_b,
|
||||
expert_offsets=expert_offsets,
|
||||
global_scale_a=gsa,
|
||||
global_scale_b=self._gsb,
|
||||
)
|
||||
|
||||
return out[:num_tokens]
|
||||
|
||||
def __call__(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
return self.run(hidden_states)
|
||||
173
tests/test_attention.py
Normal file
173
tests/test_attention.py
Normal file
@@ -0,0 +1,173 @@
|
||||
"""Standalone test: Attention projections using CuTeDSL NVFP4 linear runner.
|
||||
|
||||
Tests q_a_proj, q_b_proj, kv_proj, o_b_proj against BF16 reference.
|
||||
o_a_proj is BF16 (not NVFP4) — not tested here.
|
||||
|
||||
Usage: python3 test_attention.py
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import sys, os, json
|
||||
from safetensors import safe_open
|
||||
|
||||
MODEL_PATH = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4"
|
||||
DEVICE = "cuda:0"
|
||||
LAYER_IDX = 0
|
||||
HIDDEN_SIZE = 7168
|
||||
NUM_TOKENS = 4
|
||||
|
||||
E2M1_LUT = torch.tensor([0., 0.5, 1., 1.5, 2., 3., 4., 6., -0., -0.5, -1., -1.5, -2., -3., -4., -6.],
|
||||
dtype=torch.float32)
|
||||
|
||||
_cache = {}
|
||||
|
||||
def load_tensor(key, wm, model_dir):
|
||||
if key in _cache:
|
||||
return _cache[key]
|
||||
shard_path = os.path.join(model_dir, wm[key])
|
||||
with safe_open(shard_path, framework="pt") as f:
|
||||
t = f.get_tensor(key)
|
||||
_cache[key] = t
|
||||
return t
|
||||
|
||||
|
||||
def dequant_nvfp4(packed_uint8, scale_e4m3, global_scale):
|
||||
"""Dequantize NVFP4 weight to BF16 for reference."""
|
||||
device = packed_uint8.device
|
||||
lut = E2M1_LUT.to(device)
|
||||
lower = lut[(packed_uint8 & 0x0F).long()]
|
||||
upper = lut[((packed_uint8 >> 4) & 0x0F).long()]
|
||||
out_features = packed_uint8.shape[0]
|
||||
in_features = packed_uint8.shape[1] * 2
|
||||
unpacked = torch.empty(out_features, in_features, dtype=torch.float32, device=device)
|
||||
unpacked[:, 0::2] = lower
|
||||
unpacked[:, 1::2] = upper
|
||||
block_scale = scale_e4m3.float()
|
||||
block_expanded = block_scale.repeat_interleave(16, dim=1)[:out_features, :in_features]
|
||||
return (unpacked * block_expanded * global_scale).to(torch.bfloat16)
|
||||
|
||||
|
||||
def test_projection(name, weight, weight_sf, weight_gs, hidden_states, in_features, out_features):
|
||||
"""Test a single NVFP4 projection."""
|
||||
sys.path.insert(0, "/root/nvfp4-megamoe-kernel")
|
||||
from cutedsl.nvfp4_linear import CuTeDSLNvfp4Linear
|
||||
|
||||
# Convert weight to CuTeDSL format: (out, in_packed) uint8 → (in_packed, out) float4
|
||||
fp4 = [weight.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous()]
|
||||
sf = [weight_sf.permute(1, 0).contiguous()]
|
||||
gs = [weight_gs]
|
||||
|
||||
runner = CuTeDSLNvfp4Linear(
|
||||
in_features=in_features,
|
||||
out_features=out_features,
|
||||
max_num_tokens=8192,
|
||||
device=DEVICE,
|
||||
)
|
||||
runner.fp4 = fp4
|
||||
runner.sf = sf
|
||||
runner.gs = gs
|
||||
runner.finalize_weights()
|
||||
|
||||
# Warmup
|
||||
runner._ensure_initialized()
|
||||
runner.compute_activation_global_scale(hidden_states)
|
||||
|
||||
# Run CuTeDSL
|
||||
with torch.no_grad():
|
||||
output = runner.run(hidden_states)
|
||||
|
||||
# BF16 reference
|
||||
bf16_w = dequant_nvfp4(weight, weight_sf, weight_gs)
|
||||
with torch.no_grad():
|
||||
ref = hidden_states @ bf16_w.T
|
||||
|
||||
# Compare
|
||||
cos = F.cosine_similarity(ref.flatten().unsqueeze(0),
|
||||
output.flatten().unsqueeze(0)).item()
|
||||
mse = (ref - output).pow(2).mean().item()
|
||||
status = "✅" if cos >= 0.98 else "❌"
|
||||
print(f" {name}: cosine={cos:.6f} MSE={mse:.6e} amax_ref={ref.amax():.4f} amax_out={output.amax():.4f} {status}")
|
||||
return cos
|
||||
|
||||
|
||||
def main():
|
||||
torch.cuda.set_device(0)
|
||||
torch.manual_seed(42)
|
||||
|
||||
with open(os.path.join(MODEL_PATH, "model.safetensors.index.json")) as f:
|
||||
wm = json.load(f)["weight_map"]
|
||||
P = lambda key: load_tensor(key, wm, MODEL_PATH).to(DEVICE)
|
||||
|
||||
prefix = f"model.layers.{LAYER_IDX}.self_attn"
|
||||
|
||||
print("=== Attention Projection Tests (CuTeDSL NVFP4 Linear) ===\n")
|
||||
|
||||
# Load weights and determine dimensions from shapes
|
||||
projs = {
|
||||
"q_a_proj": {"key": f"{prefix}.q_a_proj"},
|
||||
"q_b_proj": {"key": f"{prefix}.q_b_proj"},
|
||||
"kv_proj": {"key": f"{prefix}.kv_proj"},
|
||||
"o_b_proj": {"key": f"{prefix}.o_b_proj"},
|
||||
}
|
||||
|
||||
for name, info in projs.items():
|
||||
key = info["key"]
|
||||
w = P(f"{key}.weight")
|
||||
sf = P(f"{key}.weight_scale")
|
||||
gs = P(f"{key}.weight_scale_2").item()
|
||||
out_features = w.shape[0]
|
||||
in_features = w.shape[1] * 2 # unpacked
|
||||
info["weight"] = w
|
||||
info["sf"] = sf
|
||||
info["gs"] = gs
|
||||
info["in_features"] = in_features
|
||||
info["out_features"] = out_features
|
||||
print(f" {name}: weight={w.shape} → in={in_features} out={out_features} gs={gs:.8f}")
|
||||
|
||||
print()
|
||||
|
||||
# Test each projection
|
||||
# q_a_proj: input is hidden_states (HIDDEN_SIZE=7168)
|
||||
hidden = torch.randn(NUM_TOKENS, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE) * 2.0
|
||||
|
||||
cos_qa = test_projection("q_a_proj", projs["q_a_proj"]["weight"],
|
||||
projs["q_a_proj"]["sf"], projs["q_a_proj"]["gs"],
|
||||
hidden, projs["q_a_proj"]["in_features"], projs["q_a_proj"]["out_features"])
|
||||
|
||||
# q_b_proj: input is q_a output (1536 features)
|
||||
q_a_out_features = projs["q_a_proj"]["out_features"]
|
||||
q_a_out = torch.randn(NUM_TOKENS, q_a_out_features, dtype=torch.bfloat16, device=DEVICE) * 2.0
|
||||
cos_qb = test_projection("q_b_proj", projs["q_b_proj"]["weight"],
|
||||
projs["q_b_proj"]["sf"], projs["q_b_proj"]["gs"],
|
||||
q_a_out, projs["q_b_proj"]["in_features"], projs["q_b_proj"]["out_features"])
|
||||
|
||||
# kv_proj: input is hidden_states (7168)
|
||||
cos_kv = test_projection("kv_proj", projs["kv_proj"]["weight"],
|
||||
projs["kv_proj"]["sf"], projs["kv_proj"]["gs"],
|
||||
hidden, projs["kv_proj"]["in_features"], projs["kv_proj"]["out_features"])
|
||||
|
||||
# o_b_proj: input is o_a output (16384 features after attention)
|
||||
o_b_in_features = projs["o_b_proj"]["in_features"]
|
||||
o_b_input = torch.randn(NUM_TOKENS, o_b_in_features, dtype=torch.bfloat16, device=DEVICE) * 2.0
|
||||
cos_ob = test_projection("o_b_proj", projs["o_b_proj"]["weight"],
|
||||
projs["o_b_proj"]["sf"], projs["o_b_proj"]["gs"],
|
||||
o_b_input, projs["o_b_proj"]["in_features"], projs["o_b_proj"]["out_features"])
|
||||
|
||||
print(f"\n=== SUMMARY ===")
|
||||
results = {"q_a_proj": cos_qa, "q_b_proj": cos_qb, "kv_proj": cos_kv, "o_b_proj": cos_ob}
|
||||
all_pass = True
|
||||
for name, cos in results.items():
|
||||
status = "✅" if cos >= 0.98 else "❌"
|
||||
if cos < 0.98:
|
||||
all_pass = False
|
||||
print(f" {name}: cosine={cos:.6f} {status}")
|
||||
|
||||
if all_pass:
|
||||
print("\n✅ ALL PASS")
|
||||
else:
|
||||
print("\n❌ SOME FAILED")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user