biondizzle
cf2b7ab7ec
The dense router now uses NVFP4 GEMM via Nvfp4Linear for the gate projection when NVFP4 scales are available in the checkpoint. This replaces the BF16 cuBLAS GEMM with Blackwell SM100 tensor-core NVFP4 acceleration. Changes: - dsv4/layers/router.py: add gate_lin (Nvfp4Linear) alongside W_gate fallback. New load_nvfp4_gate() method. - dsv4/kernels/router/dense_router_decode.py: add dense_router_dispatch_nvfp4() using Nvfp4Linear + activation_topk - dsv4/kernels/router/__init__.py: export new function - single_shot_inference.py: load NVFP4 gate weights when available, fall back to BF16 when not
26 lines
996 B
Python
26 lines
996 B
Python
"""DSV4 Router kernels — dispatch and CUDA kernel wrappers.
|
|
|
|
Exports:
|
|
dense_router_dispatch: BF16 GEMM + fused activation + top-k (fallback)
|
|
dense_router_dispatch_nvfp4: NVFP4 GEMM + fused activation + top-k (production)
|
|
hash_router_dispatch: Hash routing via precomputed LUT gather
|
|
"""
|
|
|
|
from dsv4.kernels.router.dense_router_decode import dense_router_dispatch, dense_router_dispatch_nvfp4
|
|
|
|
|
|
def hash_router_dispatch(
|
|
token_ids, # [N] int32
|
|
hash_lut, # [vocab_size, k] int32
|
|
top_k, # k=6
|
|
out_weights, # [N, k] float32, pre-allocated
|
|
out_ids, # [N, k] int32, pre-allocated
|
|
):
|
|
"""Hash router dispatch: gather expert IDs from precomputed LUT.
|
|
|
|
Wraps the hash_router CUDA kernel (dsv4/kernels/cuda/hash_router.cu).
|
|
One kernel launch, no intermediate buffers, no CPU-GPU sync.
|
|
"""
|
|
from dsv4.kernels.cuda._hash_router import run_hash_router
|
|
return run_hash_router(token_ids, hash_lut, top_k, out_weights, out_ids)
|