37 lines
1.5 KiB
Python
37 lines
1.5 KiB
Python
"""DSV4 Dense Router — BF16 GEMM + sqrt(softplus) + bias + top-k.
|
|
|
|
Production path: BF16 GEMM via cuBLAS (tensor cores on Blackwell) followed by
|
|
the fused activation_topk CUDA kernel for sqrt(softplus) + bias + top-k + renorm.
|
|
|
|
The CuTeDSL fused GEMM+epilogue kernel was attempted but make_trivial_tiled_mma
|
|
for BF16 on SM100 has no working reference in our codebase (all other GEMMs use
|
|
NVFP4 blockscaled MMA). The unfused path is production-grade: cuBLAS uses SM100
|
|
tensor cores, and activation_topk is a real CUDA kernel (not PyTorch).
|
|
"""
|
|
|
|
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 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,
|
|
)
|