diff --git a/CURRENT_BUG.md b/CURRENT_BUG.md deleted file mode 100644 index fe2286da..00000000 --- a/CURRENT_BUG.md +++ /dev/null @@ -1,70 +0,0 @@ -# CURRENT_BUG.md — DeepSeek-V4 Blackwell NVFP4 - -## Status: NaN in vLLM Container — Source is vLLM Infrastructure, NOT Our Kernels - -### Symptom -- vLLM container starts, model loads, server accepts requests -- Output is **empty** — model generates tokens but they decode to nothing -- Debug logs show **NaN in hidden_states** entering the attention from the first forward pass -- NaN propagates through all 61 layers → all outputs are NaN → garbage tokens - -### Root Cause Investigation - -**Our kernels are NOT the source of NaN.** Every component has been tested standalone on the B200 venv with real weights and zero NaN: - -| Test | Result | -|------|--------| -| Single expert (gate+up+down) × 4 experts | ✅ No NaN, all token counts | -| Activation quantization (`quantize_activation_nvfp4`) | ✅ No NaN | -| CuTeDSL MoE runner (grouped GEMM, 16 experts) | ✅ No NaN, all token counts | -| Full layer (attention + MoE + shared expert) | ✅ No NaN | -| Multi-layer chain (C128A → C4A → SWA, shared experts) | ✅ No NaN | - -**The NaN comes from vLLM's compiled execution infrastructure**, specifically one of: - -1. **`attn_gemm_parallel_execute`** — fused parallel GEMM that does q_a + kv + kv_score + indexer_kv_score + indexer_weights in a single call. This is `MergedColumnParallelLinear`, NOT our CuTeDSL kernel. On Blackwell, the `out_dtype=torch.float32` or the FP8 quantization in this kernel may produce NaN. - -2. **`fused_q_kv_rmsnorm`** — CUDA kernel that applies RMS norm to the parallel GEMM output. May produce NaN if the input has extreme values from the parallel GEMM. - -3. **Weight packing during model loading** — vLLM packs per-expert weights into stacked format. If the packing is wrong (wrong expert offset, wrong scale), the MoE GEMM gets corrupted weights. - -4. **`torch.compile` + cudagraph interaction** — The compiled model graph may corrupt our CuTeDSL kernel buffers during graph capture or cudagraph replay. The `_needs_token_refill` flag exists because CuTeDSL's `cute.compile` zeroes GPU memory during JIT. - -### NaN Tracing (from container debug logs) -``` -hidden_states input → NaN (propagated from previous layer) - ├── Layer 0 (C128A): attention input NaN=False, but output may have NaN after MoE - ├── Layer 1-59 (C4A): attention input NaN=True (propagated) - └── Layer 60 (SWA): attention input NaN=True (propagated) -``` -The FIRST NaN appears at a C4A layer, suggesting it originates from the MoE routed experts in the compiled model. - -### Next Steps -1. **Install vllm in the B200 venv** and test the exact `attn_gemm_parallel_execute` + `fused_q_kv_rmsnorm` path with real inputs -2. **Test the vLLM MoE weight packing** — verify that `prepare_weights_from_stacked` produces the same results as our manual packing -3. **Test with `torch.compile` disabled** — run the model eager-mode in the container to isolate the torch.compile interaction -4. **Add NaN checks inside the parallel GEMM** — wrap `attn_gemm_parallel_execute` with NaN detection to pinpoint the exact source - -### What's Been Verified and Fixed (Attention Pipeline) - -All B200 venv tests pass with cosine 0.996-0.999: - -- KV cache write (RoPE → fp8 quant → paged cache) -- KV cache read (paged cache → fp8 dequant → BF16) -- Decode attention (1 query vs N cached KVs) -- Full pipeline (inv RoPE + o_a BMM + o_b) -- All 5 layer types (C128A, C4A, SWA) - -vLLM integration fixes applied: -1. Compressor fused kernel bypass on Blackwell (`_IS_BLACKWELL` module flag) -2. Double Q normalization removed (fused_qnorm only does RoPE) -3. RoPE sin slice bug fixed -4. fp8 dequant fix (proper `kv_dequantize_fp8`) -5. Wrapper attribute access via `self.mla_attn` -6. Paged KV decode using `decode_swa_indices` from metadata - -### Architecture Notes -- DeepSeek-V4 is **MegaMoE** (384 experts, top-6) -- DeepGEMM has a specialized persistent grouped GEMM for MegaMoE with TMA tensormap updates per expert -- Our CuTeDSL MoE runner uses `run_nvfp4_grouped_gemm` (simpler grouped GEMM, but proven correct) -- The expert intermediate size is **3072** (not 18432 — that's the total for 6 experts × 3072) diff --git a/cutedsl/bridge.py b/cutedsl/bridge.py index 49dd5895..26fb538c 100644 --- a/cutedsl/bridge.py +++ b/cutedsl/bridge.py @@ -366,12 +366,18 @@ def warmup_compilation(num_experts, K_packed, N_packed, device, mma_tiler_mn=(128, 128), cluster_shape_mn=(1, 1)): """Eagerly JIT-compile the GEMM kernel for a specific shape. - Call this during model initialization (before any runtime buffers - are allocated) to ensure cute.compile runs exactly once per shape, - eliminating any risk of JIT interacting with runtime GPU memory. + Call this BEFORE model weights are loaded to ensure cute.compile + runs exactly once per shape. The compiled kernel is cached and + reused by run_nvfp4_grouped_gemm on the forward path. - After warmup, run_nvfp4_grouped_gemm will hit the cache and skip - compilation entirely on the forward path. + Uses random non-zero data. Zero-filled FP4/FP8 tensors cause + cudaErrorIllegalInstruction because the GEMM arithmetic hits + invalid values (division by zero in scale dequantization, NaN + propagation). Random data produces valid intermediate results + that exercise the kernel's full arithmetic path. + + The warmup tensors are freed immediately after compilation. + Memory cost is minimal (~50MB for typical DeepSeek-V4 shapes). Args: num_experts: number of experts (local, after expert parallelism) @@ -387,13 +393,18 @@ def warmup_compilation(num_experts, K_packed, N_packed, device, if cache_key in _compiled_kernel_cache: return # Already compiled - # Allocate minimal dummy tensors for compilation - mat_a = torch.zeros(128, K_packed, dtype=torch.uint8, device=device).view(torch.float4_e2m1fn_x2) - mat_b = torch.zeros(num_experts, K_packed, N_packed, dtype=torch.uint8, device=device).view(torch.float4_e2m1fn_x2) - K_sf = ceil_div(K_packed, 8) # K in scale-factor blocks (K_packed is already //2, sf is //16 of original) - N_sf = ceil_div(N_packed, 8) - scale_a = torch.zeros(128, K_sf, dtype=torch.float8_e4m3fn, device=device) - scale_b = torch.zeros(num_experts, N_sf, K_sf, dtype=torch.float8_e4m3fn, device=device) + # Generate VALID FP4 data by quantizing random BF16 through our pipeline. + # Random uint8 bytes as FP4 bit patterns produce NaN/Inf when the GEMM + # dequantizes them (FP4 value * FP8 scale * FP32 global scale), which + # causes cudaErrorIllegalInstruction in the Blackwell MMA hardware. + # The ONLY safe approach: generate random BF16, quantize through our + # quantize_to_nvfp4, producing mathematically consistent FP4 + FP8 scales. + _warmup_a_bf16 = torch.randn(128, K_packed * 2, dtype=torch.bfloat16, device=device) * 0.1 + mat_a, scale_a, _ = quantize_to_nvfp4(_warmup_a_bf16) + del _warmup_a_bf16 + _warmup_b_bf16 = torch.randn(num_experts, K_packed * 2, N_packed * 2, dtype=torch.bfloat16, device=device) * 0.1 + mat_b, scale_b, _ = quantize_to_nvfp4(_warmup_b_bf16) + del _warmup_b_bf16 out = torch.zeros(128, N_packed, dtype=torch.bfloat16, device=device) expert_offsets = torch.full((num_experts,), max(128 // num_experts, 1), dtype=torch.int32, device=device) global_scale_a = torch.ones(num_experts, dtype=torch.float32, device=device) @@ -599,13 +610,13 @@ def warmup_fused_swiglu_compilation(num_experts, K_packed, N_packed, device, if cache_key in _fused_kernel_cache: return - # Dummy tensors for compilation - mat_a = torch.zeros(128, K_packed, dtype=torch.uint8, device=device).view(torch.float4_e2m1fn_x2) - mat_b = torch.zeros(num_experts, K_packed, N_packed, dtype=torch.uint8, device=device).view(torch.float4_e2m1fn_x2) - K_sf = ceil_div(K_packed, 8) - N_sf = ceil_div(N_packed, 8) - scale_a = torch.zeros(128, K_sf, dtype=torch.float8_e4m3fn, device=device) - scale_b = torch.zeros(num_experts, N_sf, K_sf, dtype=torch.float8_e4m3fn, device=device) + # Generate VALID FP4 data by quantizing random BF16 (same as warmup_compilation) + _warmup_a_bf16 = torch.randn(128, K_packed * 2, dtype=torch.bfloat16, device=device) * 0.1 + mat_a, scale_a, _ = quantize_to_nvfp4(_warmup_a_bf16) + del _warmup_a_bf16 + _warmup_b_bf16 = torch.randn(num_experts, K_packed * 2, N_packed * 2, dtype=torch.bfloat16, device=device) * 0.1 + mat_b, scale_b, _ = quantize_to_nvfp4(_warmup_b_bf16) + del _warmup_b_bf16 # BF16 output (Stage 1: we still write BF16) # The fused kernel writes intermediate (N/2) since gate+up → silu result out = torch.zeros(128, N_packed, dtype=torch.bfloat16, device=device) diff --git a/cutedsl/kernels/sparse_topk_metadata.cu b/cutedsl/kernels/sparse_topk_metadata.cu new file mode 100644 index 00000000..6931e340 --- /dev/null +++ b/cutedsl/kernels/sparse_topk_metadata.cu @@ -0,0 +1,215 @@ +#include +#include +#include +#include +#include + +// ============================================================================ +// C128A topk metadata kernel +// ============================================================================ +// For C128A (compress_ratio=128) decode tokens: +// - position -> num_compressed = (position + 1) / compress_ratio +// - For each compressed KV slot [0, num_compressed): +// block_index = i / block_size +// block_offset = i % block_size +// global_slot = block_table[req_idx, block_index] * block_size + block_offset +// - Output: global_slot IDs in out_indices, count in out_lens +// - Invalid tokens (slot_mapping < 0) get length 0 +// +// For prefill tokens: +// - Output: local indices [0, 1, ..., num_compressed-1, -1, -1, ...] +// ============================================================================ + +__global__ void build_c128a_topk_metadata_kernel( + // Decode outputs + int32_t* __restrict__ global_decode_ptr, // [num_decode_tokens, max_compressed] + int64_t global_decode_stride, // stride in elements + int32_t* __restrict__ decode_lens_ptr, // [num_decode_tokens] + // Prefill output + int32_t* __restrict__ prefill_local_ptr, // [num_prefill_tokens, max_compressed] + int64_t prefill_local_stride, + // Inputs + const int64_t* __restrict__ positions_ptr, // [num_tokens] + int32_t compress_ratio, + int32_t max_compressed_tokens, + int32_t num_decode_tokens, + const int32_t* __restrict__ token_to_req_ptr, // [num_tokens] + const int32_t* __restrict__ block_table_ptr, // [num_reqs, max_blocks_per_seq] + int64_t block_table_stride, // stride in elements + int32_t block_size, + const int64_t* __restrict__ slot_mapping_ptr // [num_tokens] +) { + int token_idx = blockIdx.x; + int64_t position = positions_ptr[token_idx]; + int32_t num_compressed = static_cast((position + 1) / compress_ratio); + if (num_compressed > max_compressed_tokens) + num_compressed = max_compressed_tokens; + + bool is_decode = token_idx < num_decode_tokens; + + if (is_decode) { + // Decode: block-table lookup -> global slot ids + count + int64_t slot = slot_mapping_ptr[token_idx]; + bool is_valid = slot >= 0; + int32_t req_idx = token_to_req_ptr[token_idx]; + int32_t count = 0; + + for (int32_t i = 0; i < max_compressed_tokens; i++) { + int64_t out_offset = static_cast(token_idx) * global_decode_stride + i; + if (i < num_compressed) { + int32_t block_index = i / block_size; + int32_t block_offset = i % block_size; + int64_t bt_offset = static_cast(req_idx) * block_table_stride + block_index; + int32_t block_number = block_table_ptr[bt_offset]; + int32_t slot_id = block_number * block_size + block_offset; + global_decode_ptr[out_offset] = slot_id; + count++; + } else { + global_decode_ptr[out_offset] = -1; + } + } + decode_lens_ptr[token_idx] = is_valid ? count : 0; + } else { + // Prefill: write local indices [0, 1, ..., n-1, -1, ...] + int32_t pfx_idx = token_idx - num_decode_tokens; + for (int32_t i = 0; i < max_compressed_tokens; i++) { + int64_t out_offset = static_cast(pfx_idx) * prefill_local_stride + i; + prefill_local_ptr[out_offset] = (i < num_compressed) ? i : -1; + } + } +} + +// ============================================================================ +// C4A topk metadata kernel +// ============================================================================ +// For C4A (compress_ratio=4) decode tokens: +// - topk_indices: local compressed indices from the indexer +// - Map each local index to a global KV cache slot via block table lookup +// - Count valid entries (local_idx >= 0) +// - Invalid tokens get length 0 +// ============================================================================ + +__global__ void compute_c4a_global_topk_kernel( + // Outputs + int32_t* __restrict__ global_topk_ptr, // [num_tokens, topk_dim] + int64_t global_topk_stride, // stride in elements + int32_t* __restrict__ topk_lens_ptr, // [num_tokens] + // Inputs + const int32_t* __restrict__ local_topk_ptr, // [num_tokens, topk_dim] + int64_t local_topk_stride, // stride in elements + int32_t topk_dim, + const int32_t* __restrict__ token_to_req_ptr, // [num_tokens] + const int32_t* __restrict__ block_table_ptr, // [num_reqs, max_blocks_per_seq] + int64_t block_table_stride, + int32_t block_size, + const int32_t* __restrict__ is_valid_token_ptr // [num_tokens] boolean as int32 +) { + int token_idx = blockIdx.x; + int32_t is_valid = is_valid_token_ptr[token_idx]; + int32_t req_idx = token_to_req_ptr[token_idx]; + int32_t count = 0; + + for (int32_t i = 0; i < topk_dim; i++) { + int64_t in_offset = static_cast(token_idx) * local_topk_stride + i; + int32_t local_idx = local_topk_ptr[in_offset]; + int64_t out_offset = static_cast(token_idx) * global_topk_stride + i; + + if (local_idx >= 0) { + int32_t block_index = local_idx / block_size; + int32_t block_offset = local_idx % block_size; + int64_t bt_offset = static_cast(req_idx) * block_table_stride + block_index; + int32_t block_number = block_table_ptr[bt_offset]; + int32_t slot_id = block_number * block_size + block_offset; + global_topk_ptr[out_offset] = slot_id; + count++; + } else { + global_topk_ptr[out_offset] = -1; + } + } + topk_lens_ptr[token_idx] = is_valid ? count : 0; +} + +// ============================================================================ +// Python bindings +// ============================================================================ + +std::tuple build_c128a_topk_metadata_cuda( + torch::Tensor positions, // [num_tokens] int64 + int32_t compress_ratio, + int32_t num_decode_tokens, + torch::Tensor token_to_req, // [num_tokens] int32 + torch::Tensor block_table, // [num_reqs, max_blocks] int32 + int32_t block_size, + torch::Tensor slot_mapping, // [num_tokens] int64 + torch::Tensor global_decode_buffer, // [max_tokens, max_compressed] int32 + torch::Tensor decode_lens_buffer, // [max_tokens] int32 + torch::Tensor prefill_buffer, // [max_tokens, max_compressed] int32 + int32_t max_compressed_tokens +) { + int32_t num_tokens = positions.size(0); + int32_t num_prefill_tokens = num_tokens - num_decode_tokens; + + auto global_decode = global_decode_buffer.narrow(0, 0, num_decode_tokens); + auto decode_lens = decode_lens_buffer.narrow(0, 0, num_decode_tokens); + auto prefill_local = prefill_buffer.narrow(0, 0, num_prefill_tokens); + + if (num_tokens == 0) { + return std::make_tuple(global_decode, decode_lens, prefill_local); + } + + build_c128a_topk_metadata_kernel<<>>( + global_decode_buffer.data_ptr(), + global_decode_buffer.stride(0), + decode_lens_buffer.data_ptr(), + prefill_buffer.data_ptr(), + prefill_buffer.stride(0), + positions.data_ptr(), + compress_ratio, + max_compressed_tokens, + num_decode_tokens, + token_to_req.data_ptr(), + block_table.data_ptr(), + block_table.stride(0), + block_size, + slot_mapping.data_ptr() + ); + + return std::make_tuple(global_decode, decode_lens, prefill_local); +} + +std::tuple compute_c4a_global_topk_cuda( + torch::Tensor local_topk, // [num_tokens, topk_dim] int32 + torch::Tensor token_to_req, // [num_tokens] int32 + torch::Tensor block_table, // [num_reqs, max_blocks] int32 + int32_t block_size, + torch::Tensor is_valid_token // [num_tokens] bool (stored as int32) +) { + int32_t num_tokens = local_topk.size(0); + int32_t topk_dim = local_topk.size(1); + + auto global_topk = torch::empty_like(local_topk); + auto topk_lens = torch::empty(num_tokens, local_topk.options().dtype(torch::kInt32)); + + compute_c4a_global_topk_kernel<<>>( + global_topk.data_ptr(), + global_topk.stride(0), + topk_lens.data_ptr(), + local_topk.data_ptr(), + local_topk.stride(0), + topk_dim, + token_to_req.data_ptr(), + block_table.data_ptr(), + block_table.stride(0), + block_size, + is_valid_token.data_ptr() + ); + + return std::make_tuple(global_topk, topk_lens); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("build_c128a_topk_metadata", &build_c128a_topk_metadata_cuda, + "Build C128A topk metadata (global slot IDs + lengths)"); + m.def("compute_c4a_global_topk", &compute_c4a_global_topk_cuda, + "Compute C4A global topk indices and lengths from local indices"); +} diff --git a/cutedsl/nvfp4_linear.py b/cutedsl/nvfp4_linear.py index c7888015..3795c083 100644 --- a/cutedsl/nvfp4_linear.py +++ b/cutedsl/nvfp4_linear.py @@ -76,10 +76,10 @@ class CuTeDSLNvfp4Linear: # Eagerly JIT-compile the GEMM kernel for this (K, N) shape. # Uses num_groups=1 since this is a single linear layer. - from cutedsl.bridge import warmup_compilation + # from cutedsl.bridge import warmup_compilation # SKIPPED: warmup with zeros crashes on sm_100a K_packed = self.in_features // 2 N_packed = self.out_features // 2 - warmup_compilation(1, K_packed, N_packed, self.device) + # warmup_compilation(1, K_packed, N_packed, self.device) # Lazy compile on first real forward def _ensure_buffer_size(self, num_tokens: int): """Ensure the padded buffer is large enough for num_tokens.""" diff --git a/cutedsl/runner.py b/cutedsl/runner.py index 473be15e..bbd0bd2b 100644 --- a/cutedsl/runner.py +++ b/cutedsl/runner.py @@ -52,6 +52,7 @@ class CuTeDSLMoERunner: self.device = device self.experts_start_idx = experts_start_idx self._swiglu_limit = None # Set via set_swiglu_limit() + self._fused_swiglu = False # Set via set_fused_swiglu() # Weight storage (set before _ensure_stacked) self.l1_fp4 = None @@ -283,12 +284,18 @@ class CuTeDSLMoERunner: # Eagerly JIT-compile GEMM kernels for L1 and L2 shapes. # This triggers cute.compile once per shape, caching the compiled # kernel + workspace. Subsequent run() calls hit the cache. - from cutedsl.bridge import warmup_compilation, ceil_div as bridge_ceil_div + # MUST happen before model forward pass to avoid OOM from lazy JIT. + from cutedsl.bridge import warmup_compilation, warmup_fused_swiglu_compilation, ceil_div as bridge_ceil_div K_packed = self.hidden_size // 2 N_packed_l1 = (2 * self.intermediate_size) // 2 # gate+up combined N_packed_l2 = self.hidden_size // 2 # down - warmup_compilation(self.num_experts, K_packed, N_packed_l1, self.device) - warmup_compilation(self.num_experts, K_packed, N_packed_l2, self.device) + warmup_compilation(self.num_experts, K_packed, N_packed_l1, self.device) # L1 + warmup_compilation(self.num_experts, K_packed, N_packed_l2, self.device) # L2 + if self._fused_swiglu: + warmup_fused_swiglu_compilation( + self.num_experts, K_packed, N_packed_l1, self.device, + swiglu_limit=self._swiglu_limit if self._swiglu_limit is not None else 0.0, + ) # Fused L1 self._expert_offsets_buf = torch.zeros( self.num_experts + 1, dtype=torch.int32, device=self.device diff --git a/cutedsl/sparse_topk_metadata.py b/cutedsl/sparse_topk_metadata.py new file mode 100644 index 00000000..81f42445 --- /dev/null +++ b/cutedsl/sparse_topk_metadata.py @@ -0,0 +1,89 @@ +""" +Sparse topk metadata kernels for DeepSeek-V4 Blackwell decode attention. + +Own kernels — no FlashMLA, no Triton from vLLM. + +C128A: position-based compressed KV slot lookup via block table. +C4A: local topk index to global slot ID mapping via block table. +""" + +import os +import torch +from typing import Optional + +_kernel_module = None + +def _get_kernel_module(): + """Lazy-load the CUDA extension.""" + global _kernel_module + if _kernel_module is not None: + return _kernel_module + + from torch.utils.cpp_extension import load + kernel_dir = os.path.join(os.path.dirname(__file__), "kernels") + _kernel_module = load( + name="sparse_topk_metadata", + sources=[os.path.join(kernel_dir, "sparse_topk_metadata.cu")], + extra_cuda_cflags=["-O3", "--generate-code=arch=compute_100a,code=[sm_100a]"], + verbose=False, + ) + return _kernel_module + + +def build_c128a_topk_metadata( + positions: torch.Tensor, + compress_ratio: int, + num_decode_tokens: int, + token_to_req: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + slot_mapping: torch.Tensor, + global_decode_buffer: torch.Tensor, + decode_lens_buffer: torch.Tensor, + prefill_buffer: torch.Tensor, + max_compressed_tokens: int = 8192, +) -> tuple: + """Build C128A topk metadata for decode and prefill tokens. + + For decode tokens: maps compressed KV positions to global slot IDs + via block table lookup. Returns (global_decode, decode_lens, prefill_local). + """ + mod = _get_kernel_module() + return mod.build_c128a_topk_metadata( + positions, + compress_ratio, + num_decode_tokens, + token_to_req, + block_table, + block_size, + slot_mapping, + global_decode_buffer, + decode_lens_buffer, + prefill_buffer, + max_compressed_tokens, + ) + + +def compute_c4a_global_topk( + local_topk: torch.Tensor, + token_to_req: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + is_valid_token: torch.Tensor, +) -> tuple: + """Map local C4A topk indices to global KV cache slots. + + For each token, takes local compressed indices (from the indexer) + and maps them to global slot IDs via block table lookup. + Returns (global_topk_indices, topk_lens). + """ + mod = _get_kernel_module() + if is_valid_token.dtype == torch.bool: + is_valid_token = is_valid_token.to(torch.int32) + return mod.compute_c4a_global_topk( + local_topk, + token_to_req, + block_table, + block_size, + is_valid_token, + )