Skip to content

HANDOFF.md

Latest commit

 

History

History
108 lines (86 loc) · 4.68 KB

HANDOFF.md

File metadata and controls

108 lines (86 loc) · 4.68 KB
Raw

Megakernel Reproduction Handoff

Goal

Reproduce the luce-megakernel benchmark (Qwen3.5-0.8B fused CUDA kernel) on this RTX 3090 machine natively on Windows. Compare results to the reported 413 tok/s decode, 37,800 tok/s prefill.

Repo

  • Source: https://github.com/PixelML/luce-megakernel (forked from Luce-Org)
  • Clone location on this machine: C:\Users\ADMIN\Desktop\WIP\luce-megakernel-win
  • Original: https://github.com/Luce-Org/luce-megakernel

What's Already Done on This Machine

  1. CUDA 12.3 installed (nvcc in PATH)
  2. MSVC Build Tools 2022 at C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools
    • vcvarsall: "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\VC\Auxiliary\Build\vcvarsall.bat" amd64
  3. Python 3.9 at C:\Program Files (x86)\Microsoft Visual Studio\Shared\Python39_64
  4. Venv at C:\Users\ADMIN\Desktop\WIP\luce-megakernel\.winvenv with:
    • PyTorch 2.5.1+cu121
    • transformers (from git, v4.57.6)
    • accelerate
    • safetensors
  5. The repo was cloned to luce-megakernel-win but build was NOT completed there yet.

Steps to Build & Run

:: 1. Activate MSVC
call "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\VC\Auxiliary\Build\vcvarsall.bat" amd64

:: 2. Activate venv (reuse existing one or create new)
C:\Users\ADMIN\Desktop\WIP\luce-megakernel\.winvenv\Scripts\activate.bat

:: 3. Set required env var
set DISTUTILS_USE_SDK=1

:: 4. cd to repo
cd C:\Users\ADMIN\Desktop\WIP\luce-megakernel-win

:: 5. Build CUDA extension
python setup.py build_ext --inplace

:: 6. Run benchmark
python bench_pp_tg.py

Known Issues & Fixes Needed

1. transformers doesn't recognize qwen3_5 model type

The Qwen3.5-0.8B model uses model_type: "qwen3_5" which isn't in any released transformers version yet.

Fix: Add this to the TOP of model.py (after import torch):

try:
    from transformers.models.auto.configuration_auto import CONFIG_MAPPING
    if "qwen3_5" not in CONFIG_MAPPING:
        CONFIG_MAPPING.register("qwen3_5", CONFIG_MAPPING["qwen3_next"])
except Exception:
    pass

BUT this alias causes a crash in GenerationConfig.from_model_config because the config structures differ.

Better fix: Bypass AutoModelForCausalLM entirely. Replace the load_weights() function body to load safetensors directly:

from safetensors.torch import load_file
import glob, os

tokenizer = AutoTokenizer.from_pretrained(model_name)
cache_dir = os.path.expanduser("~/.cache/huggingface/hub/models--" + model_name.replace("/", "--"))
snap = os.path.join(cache_dir, "snapshots", os.listdir(os.path.join(cache_dir, "snapshots"))[0])
raw = {}
for sf in glob.glob(os.path.join(snap, "*.safetensors")):
    raw.update(load_file(sf, device="cuda"))
# Safetensors keys have "model.language_model." prefix, model.py expects "model."
state = {k.replace("model.language_model.", "model."): v.to(torch.bfloat16).contiguous() for k, v in raw.items()}
del raw

Delete the del model line too since there's no model object.

2. Decode kernel deadlocks (WSL only, may not affect native Windows)

On WSL2, the decode kernel's cooperative grid sync deadlocked because <<<>>> launch doesn't guarantee all 82 blocks run simultaneously. Fixed by using cudaLaunchCooperativeKernel in kernel.cu.

On native Windows this MAY work without the fix since the CUDA driver is native (not virtualized). Try without the fix first. If decode hangs, apply the cooperative launch patch from the WSL repo at C:\Users\ADMIN\Desktop\WIP\luce-megakernel\kernel.cu.

3. WSL Results (for comparison)

With cooperative launch fix on WSL2:

  • pp302: 2,338 tok/s
  • tg128: 62.2 tok/s
  • Correctness: FAIL (prefill+decode diverged from reference)

These are much worse than reported (37,800 / 413). WSL GPU-PV overhead + NTFS filesystem explain most of it.

Architecture Summary

  • kernel.cu (960 lines): Single persistent CUDA kernel for decode. All 24 layers (18 DeltaNet + 6 Full Attention) in one dispatch. 82 blocks x 512 threads. Uses cooperative grid sync between layers.
  • prefill.cu (462 lines): Prefill uses cuBLAS bf16 GEMM + standalone kernels.
  • model.py: Weight loading from HuggingFace, packs layer pointers into device blob.
  • bench_pp_tg.py: Correctness test + pp512/tg128 benchmarks.
  • setup.py: Build config targeting sm_86 (RTX 3090).

Model

  • Qwen/Qwen3.5-0.8B (bf16, ~1.5GB VRAM)
  • 24 layers: 18 DeltaNet (linear attention with recurrence) + 6 Full Attention
  • Layer pattern: [DN,DN,DN,FA, DN,DN,DN,FA, ...] x6

Expected Results (from repo README)

Method pp512 (tok/s) tg128 (tok/s)
Megakernel 37,800 413
llama.cpp BF16 11,247 267
PyTorch HF 7,578 108