179 lines
7.8 KiB
Markdown
179 lines
7.8 KiB
Markdown
# Current Bug: CuTeDSLMoERunner — Status & Debug History
|
||
|
||
## Current Status (May 17, 2026 16:52 UTC)
|
||
|
||
**vLLM runs, cudagraph capture succeeds, but model output is empty/invisible tokens (garbage logits). Going back to layer tests to debug GEMM output quality.**
|
||
|
||
- ✅ `layertest.py` — 0.988 cosine (with dynamic gs reference)
|
||
- ✅ `cudagraph_test.py` — capture + replay works
|
||
- ✅ Container builds, loads weights, warmup gs computed (no L2 gs=0)
|
||
- ✅ With `--gpu_memory_utilization=0.9` and `max_cudagraph_capture_size=8`, container starts and serves
|
||
- ❌ Model output is empty content (30 tokens of invisible/BS token) — MoE GEMM output is wrong
|
||
|
||
**Next step:** Debug WHY the runner produces 0.988 cosine in layertest but garbage in vLLM. Likely issues:
|
||
1. The warmup gs values (computed from random data) don't match real runtime activation magnitudes
|
||
2. The scale assembly layout is subtly wrong for 48 experts vs 3 experts in the test
|
||
3. The padded buffer scatter (clamped_local) is dropping real token data
|
||
|
||
**Current vLLM launch config:**
|
||
```
|
||
--gpu_memory_utilization=0.9
|
||
--compilation-config='{"cudagraph_mode": "FULL_DECODE_ONLY", "custom_ops": ["all"], "cudagraph_capture_sizes": [1, 2, 4, 8], "max_cudagraph_capture_size": 8}'
|
||
```
|
||
|
||
---
|
||
|
||
## Bugs Found & Fixed (1–21)
|
||
|
||
### Bug 1: Scale Assembly — Global vs Per-Expert Swizzle
|
||
**Fix:** Two-phase scatter + per-expert swizzle.
|
||
|
||
### Bug 2: `searchsorted(right=False)`
|
||
**Fix:** Changed to `right=True`.
|
||
|
||
### Bug 3: CuTeDSL `cute.compile` GPU Memory Corruption — CRITICAL
|
||
**Symptom:** `_token_indices` all zeros after JIT.
|
||
**Root cause:** `cute.compile` corrupts GPU memory.
|
||
**Fix:** `_fill_token_indices()` builds on CPU, copies to GPU. `_needs_token_refill` flag.
|
||
|
||
### Bug 4: `expert_offsets` With Leading 0
|
||
**Fix:** Pass `expert_offsets[1:]` to GEMM.
|
||
|
||
### Bug 5: Checkpoint `input_scale` Wrong for Runtime gs
|
||
**Root cause:** Calibration value, too-small gs → block scale overflow.
|
||
**Fix:** `compute_activation_global_scales()` warmup method.
|
||
|
||
### Bug 6: L1/L2 Need Separate gs
|
||
**Fix:** Compute L2 gs from L1 output after SiLU*up.
|
||
|
||
### Bug 7: L1/L2 Need Separate Scale Buffers
|
||
**Fix:** Separate `_padded_x_sf_buf_l1`/`_l2`, separate per-expert bufs.
|
||
|
||
### Bug 8: Global→Local Expert ID Mismatch — CUDA_ERROR_ASSERT
|
||
**Root cause:** `topk_ids` contains global IDs (0-255), runner treated as local.
|
||
**Fix:** `experts_start_idx`, remap global→local, mask non-local tokens.
|
||
|
||
### Bug 8b: `.cpu()` Sync Breaking Cudagraph
|
||
**Fix:** `_token_indices` on GPU, `_fill_token_indices()` CPU→GPU copy.
|
||
|
||
### Bug 9: `padded_x_sf` Buffer Too Small
|
||
**Fix:** Iterative — see Bugs 14, 16.
|
||
|
||
### Bug 10: Wrong `top_k`/`max_num_tokens` Defaults
|
||
**Fix:** Pass from `deepseek_v4.py`.
|
||
|
||
### Bug 11: Full-Buffer Swizzle Wrong for GEMM
|
||
**Fix:** Per-expert swizzle.
|
||
|
||
### Bug 12: `torch.full()` During Cudagraph Capture
|
||
**Symptom:** `cudaErrorStreamCaptureUnsupported`.
|
||
**Fix:** Pre-allocated buffers, `.fill_()` instead of `torch.full()`.
|
||
|
||
### Bug 13: Warmup Passed Global Expert IDs
|
||
**Symptom:** L2 gs=0.0 on EP5/EP7.
|
||
**Fix:** Pass local IDs.
|
||
|
||
### Bug 14: GEMM Scale Layout Mismatch — 128-Row Fixed vs Variable
|
||
**Symptom:** BOS token repeat (garbage logits).
|
||
**Root cause:** Scale assembly at `e*128` offsets, GEMM reads by real `expert_offsets`. Expert with 500 tokens → GEMM reads 500 scale rows but only 128 have data.
|
||
**Fix:** Fixed-layout padding: each expert gets `max_chunks * 128` rows. Pad `slot_hidden`. Pass `padded_expert_offsets` to GEMM. Extract via `l1_out[padded_dst]`.
|
||
|
||
### Bug 15: OOM — Per-Layer Padded Buffers (4.3 GB)
|
||
**Root cause:** `padded_hidden_buf` + `padded_activated_buf` at 72 MB × 60 layers.
|
||
**Fix:** Shared buffers (Bug 21).
|
||
|
||
### Bug 16: `padded_max_slots` Mismatch
|
||
**Root cause:** Sized for `max_tokens*top_k` but needed `num_experts*max_chunks*128`.
|
||
**Fix:** Size correctly.
|
||
|
||
### Bug 17: Shape Mismatch (49152 vs 3072)
|
||
**Root cause:** Cap `max_num_tokens` to 512 made buffers too small for 8192-token warmup.
|
||
**Fix:** Reverted cap, use shared buffers (Bug 21).
|
||
|
||
### Bug 18–20: Cudagraph Capture Failures
|
||
**Root cause:** Dynamic tensor allocation (`torch.zeros`), variable-trip loops, GPU scalars in Python control flow.
|
||
**Fix:** Pre-allocate everything, fixed loop counts, Python constants for offsets.
|
||
|
||
### Bug 21: OOM (correct fix) — Shared Padded Buffers
|
||
**Root cause:** Per-layer allocation of padded buffers.
|
||
**Fix:** Class-level shared buffers dict keyed by device. Layers execute sequentially → safe to share. Also shared `padded_x_sf_buf` and `output_buf`. Total ~150 MB instead of ~4.3 GB.
|
||
|
||
---
|
||
|
||
## Current Architecture: Fixed-Layout Padding
|
||
|
||
```
|
||
Each expert gets max_chunks * 128 rows at offset (e * max_chunks * 128).
|
||
|
||
Scatter: padded_dst = expert_assign * max_rows_per_expert + clamped_local_row
|
||
GEMM input: padded_hidden (total = num_experts * max_chunks * 128 rows)
|
||
GEMM offsets: [0, max_rows, 2*max_rows, ...] (fixed, pre-computed)
|
||
GEMM output: same total rows
|
||
Extract: l1_out[padded_dst] → only real token rows
|
||
|
||
Scale assembly:
|
||
Phase 1: Scatter x_sf into padded_x_sf at same fixed offsets
|
||
Phase 2: Per-expert, per-chunk swizzle (fixed loop: max_chunks iterations)
|
||
No dynamic tensor allocation, no GPU→CPU syncs
|
||
|
||
Shared buffers (class-level):
|
||
padded_hidden, padded_activated, padded_xsf_l1, padded_xsf_l2, output
|
||
~150 MB total (not per-layer)
|
||
```
|
||
|
||
### Cudagraph Constraints (All Resolved)
|
||
- No `.item()`, `.cpu()`, `.tolist()`
|
||
- No `torch.zeros/ones/full/empty/arange()` during capture
|
||
- No dynamic Python control flow from GPU values
|
||
- Per-expert Python loops OK (fixed `num_experts`)
|
||
- Shared buffers OK (layers sequential during capture and replay)
|
||
|
||
### EP Configuration (DeepSeek-V4-Pro on 8×B200)
|
||
- 256 total experts, top_k=6
|
||
- EP=8 → 48 local experts per rank
|
||
- `experts_start_idx` = rank × 32
|
||
- `max_num_tokens` = 8192
|
||
- `max_chunks_per_expert` = ceil(8192 × 6 / (48 × 128)) = 8
|
||
|
||
---
|
||
|
||
## Outstanding Issue: Garbage Model Output
|
||
|
||
**Symptom:** Model generates 30 tokens of empty/invisible content (BOS or thinking token). Not meaningful text.
|
||
|
||
**What works:** layertest gives 0.988 cosine with 3 experts, 8 tokens, top_k=8.
|
||
|
||
**What doesn't:** vLLM with 48 experts, variable tokens, top_k=6 produces garbage.
|
||
|
||
**Hypotheses to investigate:**
|
||
1. **Warmup gs from random data ≠ real activation magnitudes.** The warmup uses `torch.randn` (amax ~3) but real activations have amax ~8-10. The gs values would be wrong, causing quantization errors.
|
||
2. **Scale assembly with 48 experts × 8 chunks.** With max_chunks=8 and 48 experts, there are 384 swizzle blocks. The fixed-layout scatter with `clamped_local` may be dropping tokens that overflow the expert's max_rows section.
|
||
3. **`clamped_local = local_row.clamp(max=max_rows_per_expert - 1)`.** If an expert has more than `max_chunks*128` real tokens, overflow tokens all map to the same row, overwriting each other. This silently drops data.
|
||
4. **The `_needs_token_refill` path.** After GEMM JIT, `_token_indices` may get corrupted. The refill happens AFTER the first run, but the first run already used corrupted indices.
|
||
|
||
---
|
||
|
||
## Test Files
|
||
|
||
| File | Purpose |
|
||
|------|---------|
|
||
| `tests/layertest.py` | Reference vs runner, 3 experts. Must pass ≥0.98 cosine. |
|
||
| `tests/cudagraph_test.py` | Cudagraph capture + replay. Must pass. |
|
||
| `tests/test_runner_vs_pipeline.py` | Runner vs pipeline comparison. |
|
||
| `tests/test_scale_assembly.py` | Scale assembly comparison. |
|
||
| `tests/test_warmup_gs.py` | Warmup gs computation. |
|
||
|
||
**Run order after any code change:**
|
||
1. `python3 tests/layertest.py` — must pass
|
||
2. `python3 tests/cudagraph_test.py` — must pass
|
||
|
||
---
|
||
|
||
## 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
|