From 6ce6a47be9969c748111f827dc81fb7be31f7e09 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 18 May 2026 20:14:03 +0000 Subject: [PATCH] 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 --- cutedsl/nvfp4_linear.py | 158 ++++++++++++++++++++++++++++++++++++ tests/test_attention.py | 173 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 331 insertions(+) create mode 100644 cutedsl/nvfp4_linear.py create mode 100644 tests/test_attention.py diff --git a/cutedsl/nvfp4_linear.py b/cutedsl/nvfp4_linear.py new file mode 100644 index 00000000..b9d8c7cc --- /dev/null +++ b/cutedsl/nvfp4_linear.py @@ -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) diff --git a/tests/test_attention.py b/tests/test_attention.py new file mode 100644 index 00000000..841f0950 --- /dev/null +++ b/tests/test_attention.py @@ -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()