214 lines
10 KiB
Markdown
214 lines
10 KiB
Markdown
# Current Bug: CuTeDSLMoERunner — Status & Debug History
|
||
|
||
## Current Status (May 17, 2026 13:30 UTC)
|
||
|
||
**vLLM container runs, cudagraph capture succeeds, but model output is garbage (BOS token repeat).**
|
||
|
||
- ✅ `layertest.py` — 0.988 cosine
|
||
- ✅ `cudagraph_test.py` — capture + replay works
|
||
- ✅ Container builds, loads weights, warmup gs computed (no L2 gs=0)
|
||
- ✅ Cudagraph capture completes (51 sizes, ~15 min)
|
||
- ✅ Server accepts requests, generates tokens
|
||
- ❌ Model output is `<｜begin▁of▁sentence｜>` token repeated — garbage logits
|
||
|
||
**Current theory:** Scale assembly layout mismatch between the fixed 128-row-per-expert approach and what the GEMM actually expects. The latest fix pads slot_hidden to `num_experts * 128` rows and passes `padded_expert_offsets=[0, 128, 256, ...]` to the GEMM. Build is in progress on B200 to test.
|
||
|
||
---
|
||
|
||
## Bugs Found & Fixed
|
||
|
||
### Bug 1: Scale Assembly — Global Swizzle vs Per-Expert Swizzle
|
||
|
||
**Symptom:** GEMM produced all zeros even with correct global_scale.
|
||
|
||
**Root cause:** The original `_assemble_scales_cudagraph_safe` called `pad_and_swizzle_single()` on the ENTIRE padded buffer (all experts concatenated). But the kernel expects each expert's 128-row block to be swizzled independently (matching `assemble_scales_2d_side` which pads+swizzles each expert separately before concatenation).
|
||
|
||
**Fix:** Two-phase approach:
|
||
1. Scatter x_sf rows into 128-aligned positions in a padded buffer (GPU-only, no CPU sync)
|
||
2. Per-expert: copy 128 rows from padded buffer, `pad_and_swizzle_single()` each expert's block independently, then concatenate
|
||
|
||
### Bug 2: `searchsorted(right=False)` — Wrong Expert Assignment
|
||
|
||
**Symptom:** Scale data in wrong positions after scatter.
|
||
|
||
**Root cause:** `torch.searchsorted([4, 8, 8], 4, right=False)` returns 0, assigning row 4 (expert 1's first token) to expert 0.
|
||
|
||
**Fix:** Changed to `right=True`:
|
||
```python
|
||
expert_assign = torch.searchsorted(expert_offsets[1:], row_indices, right=True)
|
||
```
|
||
|
||
### Bug 3: CuTeDSL `cute.compile` GPU Memory Corruption — CRITICAL
|
||
|
||
**Symptom:** `_token_indices` was all zeros, making every token map to token 0.
|
||
|
||
**Root cause:** CuTeDSL's `cute.compile` (JIT compilation) corrupts GPU memory. Tensors allocated on GPU before or during JIT compilation get zeroed.
|
||
|
||
**Fix:** Allocate `_token_indices` with `_fill_token_indices()` which builds on CPU and copies to GPU. Added `_needs_token_refill` flag to handle GEMM JIT corruption on first call.
|
||
|
||
### Bug 4: `expert_offsets` With Leading 0
|
||
|
||
**Symptom:** GEMM produced wrong output with correct scale data.
|
||
|
||
**Root cause:** The runner passed `expert_offsets[:num_experts + 1]` = `[0, 4, 8, 8]` (4 elements with leading 0) but the kernel expects `[4, 8, 8]` (cumulative sum without leading 0).
|
||
|
||
**Fix:** Pass `expert_offsets[1:num_experts + 1]` to the GEMM.
|
||
|
||
### Bug 5: Checkpoint `input_scale` Is Wrong for Activation Global Scale
|
||
|
||
**Symptom:** Block scales all saturate at float8 max (448), producing garbage quantization.
|
||
|
||
**Root cause:** The checkpoint's `input_scale` (~0.000286) is a calibration value computed from a different input magnitude (amax ≈ 0.77) than what runtime produces (amax ≈ 8.17). Too-small gs → block scale overflow → garbage.
|
||
|
||
**Fix:** `compute_activation_global_scales()` warmup method that runs `quantize_to_nvfp4` (dynamic gs with `.max()`) before cudagraph capture.
|
||
|
||
### Bug 6: L1 and L2 Need Separate Activation Global Scales
|
||
|
||
**Symptom:** L2 output was garbage even with correct L1 gs.
|
||
|
||
**Root cause:** After SiLU(gate)*up, the activation has amax ~286. The L1 gs is 30x too small for L2.
|
||
|
||
**Fix:** `compute_activation_global_scales()` computes L2 gs from the actual L1 output (after SiLU*up).
|
||
|
||
### Bug 7: L1 and L2 Need Separate Padded Scale Buffers
|
||
|
||
**Symptom:** IndexError when quantizing L2 activation — K_sf differs between L1 (448) and L2 (192).
|
||
|
||
**Fix:** Separate `_padded_x_sf_buf_l1` and `_padded_x_sf_buf_l2`, plus separate `_per_expert_scale_bufs_l1` and `_per_expert_scale_bufs_l2`.
|
||
|
||
### Bug 8: Global→Local Expert ID Mismatch — CUDA_ERROR_ASSERT
|
||
|
||
**Symptom:** `IndexKernel.cu:111` assertion failed, cascading into CUDA_ERROR_ASSERT (710) across all workers.
|
||
|
||
**Root cause:** With EP=8, `topk_ids` contains global expert IDs (0-255), but `CuTeDSLMoERunner` treated them as local IDs (0-31/48).
|
||
|
||
**Fix:** Added `experts_start_idx` param; in `run()`, remap global→local and mask non-local tokens.
|
||
|
||
### Bug 8b: `.cpu()` Sync Breaking Cudagraph Compatibility
|
||
|
||
**Fix:** Moved `_token_indices` to GPU, added `_fill_token_indices()` (CPU→GPU copy), `_needs_token_refill` for GEMM JIT.
|
||
|
||
### Bug 9: `padded_x_sf` Buffer Too Small — Index Out of Bounds
|
||
|
||
**Symptom:** `IndexKernel.cu:111` OOB in scale assembly scatter. `dst_rows` exceeded buffer size.
|
||
|
||
**Root cause:** Buffer was sized for `num_experts * 128` rows, but scatter positions were computed from actual token distribution (not fixed 128 per expert). With 8192 tokens and top_k=6, dst_rows could exceed 6144.
|
||
|
||
**Fix (attempted):** Sized buffer for `max_num_tokens * top_k` rows. Later reverted to `num_experts * 128` with fixed 128-row-per-expert scatter layout.
|
||
|
||
### Bug 10: Wrong `top_k` and `max_num_tokens` Defaults
|
||
|
||
**Symptom:** `_token_indices` max=6143 instead of 8191 (built with top_k=8, actual top_k=6).
|
||
|
||
**Root cause:** `CuTeDSLMoERunner` defaulted to `max_num_tokens=8192, top_k=8`, but vLLM uses top_k=6. `deepseek_v4.py` didn't pass these values.
|
||
|
||
**Fix:** Pass `max_num_tokens` and `top_k` from `deepseek_v4.py` to the runner constructor.
|
||
|
||
### Bug 11: Full-Buffer Swizzle Produced Wrong GEMM Input
|
||
|
||
**Symptom:** L2 gs=0.0 on EP5/EP7 during warmup. Model produced BOS token.
|
||
|
||
**Root cause:** Applied the Blackwell 32_4_4 swizzle to the entire padded buffer at once, but the GEMM expects per-expert swizzled blocks. The combined swizzle layout doesn't match `expert_offsets` indexing.
|
||
|
||
**Fix (in progress):** Reverted to per-expert swizzle with fixed 128-row slots.
|
||
|
||
### Bug 12: `torch.full()` During Cudagraph Capture
|
||
|
||
**Symptom:** `cudaErrorStreamCaptureUnsupported` on all 8 workers during cudagraph capture.
|
||
|
||
**Root cause:** `torch.full()` in `run()` allocates a new tensor during stream capture, which CUDA doesn't allow.
|
||
|
||
**Fix:** Pre-allocated `_l1_gsa_buf` and `_l2_gsa_buf`, use `.fill_()` instead of `torch.full()`. Also pre-allocated `_output_buf`, `_row_indices_buf`.
|
||
|
||
### Bug 13: Warmup Passed Global Expert IDs Instead of Local
|
||
|
||
**Symptom:** L2 gs=0.0 on EP5/EP7 (all ranks except EP0).
|
||
|
||
**Root cause:** `_warmup_activation_global_scales()` passed global IDs (e.g. 336+) to `compute_activation_global_scales()`, which matches against `expert_id_range` (0..47). No tokens matched → zero L1 GEMM output → L2 gs=0.
|
||
|
||
**Fix:** Pass local expert IDs (0..num_experts-1) in warmup.
|
||
|
||
### Bug 14 (CURRENT): GEMM Scale Layout Mismatch — 128-Row Fixed vs Variable
|
||
|
||
**Symptom:** Model generates BOS token repeatedly. Tokens are produced but logits are garbage.
|
||
|
||
**Root cause:** Scale assembly places data at fixed `e*128` offsets (128 rows per expert). But the GEMM reads `scale_a` according to `expert_offsets` (real token counts, e.g. expert 0 = 500 rows). When expert 0 has 500 tokens, GEMM reads `scale_a[0:500]` but only rows 0-127 have valid scale data. Rows 128-499 are zeros → GEMM produces zeros for those tokens → garbage output.
|
||
|
||
**Fix (in progress):** Pad `slot_hidden` to `num_experts * 128` rows (128 per expert) and pass `padded_expert_offsets=[0, 128, 256, ...]` to the GEMM. The GEMM processes exactly 128 tokens per expert. Padding tokens' output is discarded by scatter_add. Pre-allocated `_padded_hidden_buf`, `_padded_activated_buf`, `_padded_expert_offsets_buf`.
|
||
|
||
---
|
||
|
||
## vLLM Integration Status
|
||
|
||
| Component | Status | Notes |
|
||
|-----------|--------|-------|
|
||
| Weight loading | ✅ | Direct NVFP4 path, no BF16 round-trip |
|
||
| Weight stacking | ✅ | `make_b_k_major` + `assemble_scales_3d_side` |
|
||
| Global→local ID remap | ✅ | `experts_start_idx`, mask non-local tokens |
|
||
| Warmup gs computation | ✅ | Per-layer, local expert IDs, L1+L2 gs |
|
||
| Scale assembly | ⚠️ | 128-row fixed layout, pending GEMM alignment fix |
|
||
| Cudagraph capture | ✅ | No dynamic allocations, no CPU syncs |
|
||
| Model output | ❌ | Garbage (BOS repeat) — scale/GEMM layout mismatch |
|
||
|
||
---
|
||
|
||
## Test Files
|
||
|
||
| File | Purpose |
|
||
|------|---------|
|
||
| `tests/layertest.py` | Reference: moe_pipeline with dynamic gs, 3 experts, layer 0. Must pass (≥0.98 cosine). |
|
||
| `tests/cudagraph_test.py` | CuTeDSLMoERunner cudagraph capture + replay. Must pass. |
|
||
| `tests/test_warmup_gs.py` | Warmup gs computation, safety margin sweep. |
|
||
| `tests/test_runner_vs_pipeline.py` | Compare runner.run() vs moe_pipeline. |
|
||
| `tests/test_scale_assembly.py` | Compare cudagraph-safe vs reference scale assembly. |
|
||
|
||
**Run order after any code change:**
|
||
1. `python3 tests/layertest.py` — must pass
|
||
2. `python3 tests/cudagraph_test.py` — must pass
|
||
|
||
---
|
||
|
||
## Key Architecture: CuTeDSL NVFP4 MoE
|
||
|
||
### Data Flow
|
||
```
|
||
hidden_states (BF16) ──→ global→local remap ──→ sort by expert
|
||
│
|
||
├── L1 (gate+up)
|
||
│ quantize_activation_nvfp4 → x_fp4, x_sf
|
||
│ _assemble_scales_cudagraph_safe → scale_a (swizzled)
|
||
│ run_nvfp4_grouped_gemm → l1_out (BF16)
|
||
│
|
||
├── SiLU(gate) * up → activated
|
||
│
|
||
├── L2 (down)
|
||
│ quantize_activation_nvfp4 → l2_x_fp4, l2_x_sf
|
||
│ _assemble_scales_cudagraph_safe → scale_a (swizzled)
|
||
│ run_nvfp4_grouped_gemm → l2_out (BF16)
|
||
│
|
||
└── scatter_add → y (BF16)
|
||
```
|
||
|
||
### Cudagraph Constraints
|
||
- No `.item()`, `.cpu()`, `.tolist()` — zero CPU-GPU syncs
|
||
- No `torch.zeros/ones/full/empty/arange` during capture — pre-allocate everything
|
||
- No dynamic shapes — `num_tokens` equals the captured budget
|
||
- Per-expert Python loops are OK (fixed `num_experts`, unrolled at capture time)
|
||
- `pad_and_swizzle_single` is OK on pre-padded 128×4-aligned buffers (no internal allocation)
|
||
|
||
### EP Configuration (DeepSeek-V4-Pro on 8×B200)
|
||
- 256 total experts, top_k=6
|
||
- EP=8 → 32 local experts per rank (in practice 48 based on logs)
|
||
- `experts_start_idx` = rank * 32 (0, 32, 64, ..., 224)
|
||
- `max_num_tokens` from `scheduler_config.max_num_batched_tokens`
|
||
|
||
---
|
||
|
||
## Repo Info
|
||
|
||
- **Kernel:** `sweetapi.com/biondizzle/nvfp4-megamoe-kernel` (master)
|
||
- **Local:** `~/dev/nvfp4-megamoe-kernel/`
|
||
- **B200:** `/root/nvfp4-megamoe-kernel/`
|
||
- **Model:** `/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4` (read-only)
|
||
- **Never edit on B200 directly** — edit locally → commit → push → pull on B200
|