From 8737fd57c0d80a937134e060eb4280b43f8bfdf9 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Tue, 12 May 2026 14:53:42 +0000 Subject: [PATCH] remove crap --- .env | 1 - build_push.sh | 1 - patches/test_nvfp4_mega_moe.py | 104 +++++++++++++++++++++++++++++++++ 3 files changed, 104 insertions(+), 2 deletions(-) create mode 100644 patches/test_nvfp4_mega_moe.py diff --git a/.env b/.env index 2d1100d..58df547 100644 --- a/.env +++ b/.env @@ -6,7 +6,6 @@ PASSWORD=6)Jr)B@dcX[mN?dx # B200 Paths DOCKER_COMPOSE=/root/nvidia-meeting/docker-compose.yml WEIGHTS=/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4 -PATCHES=/root/nvidia-meeting/patches REPO=/root/nvidia-meeting/deepseek-v4-quant # Docker diff --git a/build_push.sh b/build_push.sh index e331a4b..7b66491 100644 --- a/build_push.sh +++ b/build_push.sh @@ -43,7 +43,6 @@ services: - "8000:8000" volumes: - /root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4:/model - - /root/nvidia-meeting/patches:/patches environment: - VLLM_USE_FLASHINFER_MOE_FP4=1 deploy: diff --git a/patches/test_nvfp4_mega_moe.py b/patches/test_nvfp4_mega_moe.py new file mode 100644 index 0000000..0c733bd --- /dev/null +++ b/patches/test_nvfp4_mega_moe.py @@ -0,0 +1,104 @@ +"""Minimal test for fp8_nvfp4_mega_moe kernel with synthetic data. + +Run inside the vllm container: + python3 /patches/test_nvfp4_mega_moe.py +""" +import torch +import torch.distributed as dist +import os +import sys + +def test_nvfp4_mega_moe(): + # Small but aligned dimensions + num_experts = 2 + num_tokens = 8 # must be multiple of alignment (8 for block_m=8) + top_k = 2 + hidden = 256 # must be multiple of 128 and 64 + intermediate_hidden = 512 # must be multiple of 128 and 64 + + device = "cuda" + torch.cuda.set_device(0) + + # Single-rank process group for SymmBuffer + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29501") + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + if not dist.is_initialized(): + dist.init_process_group("nccl") + group = dist.new_group() + + from deep_gemm.mega import ( + fp8_nvfp4_mega_moe, + get_symm_buffer_for_nvfp4_mega_moe, + transform_nvfp4_weights_for_mega_moe, + ) + + # --- Weights: random NVFP4 --- + w13_weight = torch.randint(0, 256, (num_experts, 2 * intermediate_hidden, hidden // 2), + dtype=torch.uint8, device=device).view(torch.int8) + w13_weight_scale = torch.randn(num_experts, 2 * intermediate_hidden, hidden // 16, + device=device).abs().clamp(0.1, 10.0).to(torch.float8_e4m3fn) + w13_weight_scale_2 = torch.ones(num_experts, device=device) # global scale = 1 for simplicity + w2_weight = torch.randint(0, 256, (num_experts, hidden, intermediate_hidden // 2), + dtype=torch.uint8, device=device).view(torch.int8) + w2_weight_scale = torch.randn(num_experts, hidden, intermediate_hidden // 16, + device=device).abs().clamp(0.1, 10.0).to(torch.float8_e4m3fn) + w2_weight_scale_2 = torch.ones(num_experts, device=device) + + print("Transforming weights...") + l1_weights, l2_weights = transform_nvfp4_weights_for_mega_moe( + (w13_weight, w13_weight_scale), + (w2_weight, w2_weight_scale), + l1_weight_scale_2=w13_weight_scale_2, + l2_weight_scale_2=w2_weight_scale_2, + ) + for name, t in [("l1_w", l1_weights[0]), ("l1_w_sf", l1_weights[1]), + ("l2_w", l2_weights[0]), ("l2_w_sf", l2_weights[1])]: + print(f" {name}: dtype={t.dtype} shape={tuple(t.shape)} strides={t.stride()} contig={t.is_contiguous()}") + + # --- Symm buffer --- + print("Creating symm buffer...") + symm_buffer = get_symm_buffer_for_nvfp4_mega_moe( + group, num_experts, num_tokens, top_k, hidden, intermediate_hidden) + for name, t in [("x", symm_buffer.x), ("x_sf", symm_buffer.x_sf), + ("l1_acts", symm_buffer.l1_acts), ("l1_acts_sf", symm_buffer.l1_acts_sf), + ("l2_acts", symm_buffer.l2_acts), ("l2_acts_sf", symm_buffer.l2_acts_sf)]: + print(f" symm_{name}: dtype={t.dtype} shape={tuple(t.shape)} strides={t.stride()}") + + # --- Stage inputs (BF16 hidden_states → FP4 packed + UE4M3 scales) --- + print("Staging inputs...") + hidden_states = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=device) * 0.5 + topk_weights = torch.softmax(torch.randn(num_tokens, top_k, device=device), dim=-1) + topk_ids = torch.randint(0, num_experts, (num_tokens, top_k), device=device) + + # Import the staging kernel directly (can't import from vllm model without full init) + import triton + import triton.language as tl + # Just manually pack a few tokens into FP4 for testing + # For now, zero-fill the symm buffer (the kernel should still launch without crash + # even with zero data — we just want to verify the kernel runs at all) + symm_buffer.x.zero_() + symm_buffer.x_sf.zero_() + # Write topk data directly + symm_buffer.topk_idx[:num_tokens].zero_() + for i in range(num_tokens): + for j in range(top_k): + symm_buffer.topk_idx[i, j] = topk_ids[i, j].item() + symm_buffer.topk_weights[i, j] = topk_weights[i, j].item() + torch.cuda.synchronize() + print("Buffer populated (zeros for activations, real topk)") + + # --- Run kernel --- + y = torch.zeros(num_tokens, hidden, dtype=torch.bfloat16, device=device) + print("Calling fp8_nvfp4_mega_moe...") + try: + fp8_nvfp4_mega_moe(y, l1_weights, l2_weights, symm_buffer) + torch.cuda.synchronize() + print(f"SUCCESS! y stats: min={y.min().item():.4f} max={y.max().item():.4f} mean={y.mean().item():.4f} nonzero={torch.count_nonzero(y).item()}") + except Exception as e: + print(f"FAILED: {e}") + raise + +if __name__ == "__main__": + test_nvfp4_mega_moe()