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
62 lines
2.3 KiB
Python
62 lines
2.3 KiB
Python
"""DSV4 Dense Router — NVFP4 GEMM + sqrt(softplus) + bias + top-k.
|
|
|
|
Production path: NVFP4 GEMM via Nvfp4Linear (Blackwell tensor cores)
|
|
followed by the fused activation_topk CUDA kernel for sqrt(softplus) +
|
|
bias + top-k + renorm.
|
|
|
|
BF16 cuBLAS fallback: When NVFP4 scales are not available in the
|
|
checkpoint, dense_router_dispatch uses torch.nn.functional.linear
|
|
(cuBLAS, SM100 tensor cores) instead.
|
|
|
|
The CuTeDSL fused GEMM+epilogue kernel (dense_router_decode_kernel.py)
|
|
is the future single-kernel path but is not yet production-ready.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
from typing import Tuple, Optional
|
|
import torch
|
|
|
|
|
|
def dense_router_dispatch(
|
|
hidden_states: torch.Tensor, # [N, hidden_size] BF16
|
|
W_gate: torch.Tensor, # [hidden_size, num_experts] BF16
|
|
e_bias: torch.Tensor, # [num_experts] FP32
|
|
routed_scaling_factor: float,
|
|
top_k: int,
|
|
out_weights: torch.Tensor, # [N, top_k] FP32, pre-allocated
|
|
out_ids: torch.Tensor, # [N, top_k] int32, pre-allocated
|
|
):
|
|
"""Dispatch the dense router (BF16 cuBLAS fallback).
|
|
|
|
BF16 GEMM via torch.nn.functional.linear (cuBLAS, SM100 tensor cores),
|
|
then fused activation + top-k via the CUDA kernel.
|
|
"""
|
|
logits = torch.nn.functional.linear(hidden_states.float(), W_gate.T.float())
|
|
from dsv4.kernels.router._activation_topk import run_fused_activation_topk
|
|
run_fused_activation_topk(
|
|
logits, e_bias, routed_scaling_factor, top_k,
|
|
out_weights, out_ids,
|
|
)
|
|
|
|
|
|
def dense_router_dispatch_nvfp4(
|
|
hidden_states: torch.Tensor, # [N, hidden_size] BF16
|
|
gate_lin, # Nvfp4Linear instance
|
|
e_bias: torch.Tensor, # [num_experts] FP32
|
|
routed_scaling_factor: float,
|
|
top_k: int,
|
|
out_weights: torch.Tensor, # [N, top_k] FP32, pre-allocated
|
|
out_ids: torch.Tensor, # [N, top_k] int32, pre-allocated
|
|
):
|
|
"""Dispatch the dense router (NVFP4 production GEMM).
|
|
|
|
NVFP4 GEMM via Nvfp4Linear (Blackwell SM100 tensor cores),
|
|
then fused activation + top-k via the CUDA kernel.
|
|
"""
|
|
logits = gate_lin(hidden_states).float() # (N, E) FP32
|
|
from dsv4.kernels.router._activation_topk import run_fused_activation_topk
|
|
run_fused_activation_topk(
|
|
logits, e_bias, routed_scaling_factor, top_k,
|
|
out_weights, out_ids,
|
|
)
|