some optimizations
This commit is contained in:
344
NEXT_OPTIMIZATION.md
Normal file
344
NEXT_OPTIMIZATION.md
Normal file
@@ -0,0 +1,344 @@
|
||||
pre-remap the weight scale factors (SFB) once, and stop remapping/allocating them inside every GEMM.
|
||||
|
||||
This is lower hanging than rewriting stage_activation in Triton, because weights are static after load. Activation scales SFA still need per-call remap, but SFB does not.
|
||||
|
||||
Current hot path in cutlass_nvfp4_gemm.cu:
|
||||
```
|
||||
cutlass::device_memory::allocation<ElementSF> sfa_cutlass(sfa_size);
|
||||
cutlass::device_memory::allocation<ElementSF> sfb_cutlass(sfb_size);
|
||||
|
||||
cudaMemsetAsync(sfa_cutlass.get(), 0, sfa_size * sizeof(ElementSF), stream);
|
||||
cudaMemsetAsync(sfb_cutlass.get(), 0, sfb_size * sizeof(ElementSF), stream);
|
||||
|
||||
remap_sf_to_cutlass_kernel<<<...>>>(SFA_ptr, sfa_cutlass.get(), ...);
|
||||
remap_sf_to_cutlass_kernel<<<...>>>(SFB_ptr, sfb_cutlass.get(), ...);
|
||||
```
|
||||
|
||||
|
||||
|
||||
Target shape
|
||||
|
||||
Keep this per GEMM:
|
||||
```
|
||||
SFA row-major activation scales
|
||||
→ remap dynamically
|
||||
```
|
||||
|
||||
Change this:
|
||||
|
||||
```
|
||||
SFB weight scales
|
||||
→ already CUTLASS-remapped
|
||||
→ pass directly to GEMM
|
||||
```
|
||||
|
||||
So the GEMM run becomes:
|
||||
|
||||
```
|
||||
// still dynamic
|
||||
remap SFA
|
||||
|
||||
// no remap
|
||||
use prepacked_SFB_cutlass directly
|
||||
```
|
||||
|
||||
Step 1: add a prepack_sfb CUDA entrypoint
|
||||
|
||||
Add a new exported function next to cutlass_nvfp4_gemm_run.
|
||||
|
||||
Something like:
|
||||
```
|
||||
extern "C" int cutlass_nvfp4_prepack_sfb_run(
|
||||
const void* SFB_ptr,
|
||||
void* SFB_cutlass_ptr,
|
||||
int M, int N, int K,
|
||||
cudaStream_t stream
|
||||
) {
|
||||
using Sm1xxBlkScaledConfig =
|
||||
typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
|
||||
|
||||
LayoutSFB layout_SFB =
|
||||
Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(
|
||||
cute::make_shape(M, N, K, 1));
|
||||
|
||||
using ElementSF =
|
||||
typename Gemm::GemmKernel::CollectiveMainloop::ElementSF;
|
||||
|
||||
int sfb_size = cute::size(layout_SFB);
|
||||
int K_sf = K / InputSFVectorSize;
|
||||
|
||||
cudaMemsetAsync(
|
||||
static_cast<ElementSF*>(SFB_cutlass_ptr),
|
||||
0,
|
||||
sfb_size * sizeof(ElementSF),
|
||||
stream);
|
||||
|
||||
int block = 256;
|
||||
remap_sf_to_cutlass_kernel<<<(sfb_size + block - 1) / block, block, 0, stream>>>(
|
||||
static_cast<const ElementSF*>(SFB_ptr),
|
||||
static_cast<ElementSF*>(SFB_cutlass_ptr),
|
||||
layout_SFB,
|
||||
N,
|
||||
K_sf,
|
||||
true // SFB source is (K_sf, N)
|
||||
);
|
||||
|
||||
return 0;
|
||||
}
|
||||
```
|
||||
|
||||
You’ll also want a size query helper:
|
||||
```
|
||||
extern "C" int cutlass_nvfp4_sfb_size(
|
||||
int M, int N, int K,
|
||||
int* out_size
|
||||
) {
|
||||
using Sm1xxBlkScaledConfig =
|
||||
typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
|
||||
|
||||
LayoutSFB layout_SFB =
|
||||
Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(
|
||||
cute::make_shape(M, N, K, 1));
|
||||
|
||||
*out_size = cute::size(layout_SFB);
|
||||
return 0;
|
||||
}
|
||||
```
|
||||
|
||||
Important: pass the same M, N, K geometry you’ll use for the GEMM at first. Later you can test whether SFB layout is independent of M; I suspect it is effectively N/K-driven, but don’t assume that until you print sizes for several M.
|
||||
|
||||
Step 2: expose it in pytorch_binding.cpp
|
||||
|
||||
Add declarations:
|
||||
|
||||
```
|
||||
extern "C" int cutlass_nvfp4_sfb_size(
|
||||
int M, int N, int K,
|
||||
int* out_size
|
||||
);
|
||||
|
||||
extern "C" int cutlass_nvfp4_prepack_sfb_run(
|
||||
const void* SFB_ptr,
|
||||
void* SFB_cutlass_ptr,
|
||||
int M, int N, int K,
|
||||
cudaStream_t stream
|
||||
);
|
||||
```
|
||||
|
||||
Then add a Python-visible wrapper:
|
||||
```
|
||||
torch::Tensor prepack_sfb(
|
||||
torch::Tensor SFB,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K
|
||||
) {
|
||||
int size = 0;
|
||||
int rc = cutlass_nvfp4_sfb_size(
|
||||
static_cast<int>(M),
|
||||
static_cast<int>(N),
|
||||
static_cast<int>(K),
|
||||
&size
|
||||
);
|
||||
TORCH_CHECK(rc == 0, "sfb_size failed");
|
||||
|
||||
auto out = torch::empty(
|
||||
{size},
|
||||
torch::dtype(SFB.dtype()).device(SFB.device())
|
||||
);
|
||||
|
||||
auto stream = c10::cuda::getCurrentCUDAStream();
|
||||
|
||||
rc = cutlass_nvfp4_prepack_sfb_run(
|
||||
SFB.data_ptr(),
|
||||
out.data_ptr(),
|
||||
static_cast<int>(M),
|
||||
static_cast<int>(N),
|
||||
static_cast<int>(K),
|
||||
stream.stream()
|
||||
);
|
||||
|
||||
TORCH_CHECK(rc == 0, "prepack_sfb failed");
|
||||
|
||||
return out;
|
||||
}
|
||||
```
|
||||
|
||||
Register it:
|
||||
```
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("forward", &cutlass_nvfp4_gemm_forward,
|
||||
"CUTLASS NVFP4 block-scaled GEMM forward");
|
||||
|
||||
m.def("prepack_sfb", &prepack_sfb,
|
||||
"Pre-remap SFB weight scales into CUTLASS layout");
|
||||
}
|
||||
```
|
||||
|
||||
Step 3: add a GEMM path that accepts prepacked SFB
|
||||
Add a second C API entrypoint, or a boolean flag, for:
|
||||
```
|
||||
cutlass_nvfp4_gemm_run_prepacked_sfb(...)
|
||||
```
|
||||
|
||||
Inside it, delete this:
|
||||
```
|
||||
cutlass::device_memory::allocation<ElementSF> sfb_cutlass(sfb_size);
|
||||
cudaMemsetAsync(sfb_cutlass.get(), 0, ...);
|
||||
remap_sf_to_cutlass_kernel<<<...>>>(SFB_ptr, sfb_cutlass.get(), ...);
|
||||
```
|
||||
|
||||
and use the passed pointer directly:
|
||||
```
|
||||
static_cast<const ElementSF*>(SFB_cutlass_ptr), layout_SFB
|
||||
```
|
||||
|
||||
The relevant section becomes:
|
||||
|
||||
```
|
||||
cutlass::device_memory::allocation<ElementSF> sfa_cutlass(sfa_size);
|
||||
cudaMemsetAsync(sfa_cutlass.get(), 0, sfa_size * sizeof(ElementSF), stream);
|
||||
|
||||
remap_sf_to_cutlass_kernel<<<(sfa_size + block - 1) / block, block, 0, stream>>>(
|
||||
static_cast<const ElementSF*>(SFA_ptr),
|
||||
sfa_cutlass.get(),
|
||||
layout_SFA,
|
||||
M,
|
||||
K_sf,
|
||||
false);
|
||||
|
||||
typename Gemm::Arguments arguments {
|
||||
cutlass::gemm::GemmUniversalMode::kGemm,
|
||||
{M, N, K, 1},
|
||||
{
|
||||
static_cast<const ArrayElementA*>(A_ptr), stride_A,
|
||||
static_cast<const ArrayElementB*>(B_ptr), stride_B,
|
||||
sfa_cutlass.get(), layout_SFA,
|
||||
static_cast<const ElementSF*>(SFB_cutlass_ptr), layout_SFB
|
||||
},
|
||||
{
|
||||
{ alpha, beta },
|
||||
nullptr, stride_C,
|
||||
static_cast<typename Gemm::GemmKernel::CollectiveEpilogue::ElementD*>(D_ptr), stride_D
|
||||
}
|
||||
};
|
||||
```
|
||||
|
||||
Step 4: prepack in Python after weight transform
|
||||
|
||||
In weight_transform.py, you currently return:
|
||||
```
|
||||
return (l1_weight_out, l1_sf_out), (l2_weight_out, l2_sf_out)
|
||||
```
|
||||
|
||||
Do not do the prepack directly there unless extension import order is clean. Safer first pass: do it lazily in nvfp4_mega_moe_full() once.
|
||||
|
||||
Add helper:
|
||||
|
||||
```
|
||||
def _maybe_prepack_weight_sf(weights, weight_sf, N, K, tag):
|
||||
cache_attr = f"_prepacked_{tag}"
|
||||
|
||||
cached = getattr(_maybe_prepack_weight_sf, cache_attr, None)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
from nvfp4_megamoe_kernel.cutlass_nvfp4_gemm import _C
|
||||
|
||||
E = weight_sf.shape[0]
|
||||
packed = []
|
||||
|
||||
# Use M=128 initially to match the MMA tile boundary.
|
||||
# Later test if SFB size/layout is stable across M.
|
||||
M_for_layout = 128
|
||||
|
||||
for e in range(E):
|
||||
packed.append(
|
||||
_C.prepack_sfb(
|
||||
weight_sf[e],
|
||||
M_for_layout,
|
||||
N,
|
||||
K,
|
||||
)
|
||||
)
|
||||
|
||||
packed = torch.stack(packed, dim=0).contiguous()
|
||||
setattr(_maybe_prepack_weight_sf, cache_attr, packed)
|
||||
return packed
|
||||
```
|
||||
|
||||
Then before L1/L2 calls:
|
||||
|
||||
```
|
||||
l1_N = l1_w.shape[2]
|
||||
l1_K = l1_w.shape[1] * 2
|
||||
l1_sf_prepacked = _maybe_prepack_weight_sf(l1_w, l1_sf, l1_N, l1_K, "l1")
|
||||
|
||||
l2_N = l2_w.shape[2]
|
||||
l2_K = l2_w.shape[1] * 2
|
||||
l2_sf_prepacked = _maybe_prepack_weight_sf(l2_w, l2_sf, l2_N, l2_K, "l2")
|
||||
```
|
||||
|
||||
Then pass l1_sf_prepacked / l2_sf_prepacked into the new prepacked-SFB GEMM path.
|
||||
|
||||
Validation test
|
||||
|
||||
Before trusting it, compare old vs new on one expert:
|
||||
|
||||
```
|
||||
old = cutlass_nvfp4_blockscaled_gemm(
|
||||
expert_x,
|
||||
expert_x_sf,
|
||||
expert_w,
|
||||
expert_w_sf,
|
||||
M_expert,
|
||||
N,
|
||||
K,
|
||||
alpha=alpha,
|
||||
)
|
||||
|
||||
sfb_pre = _C.prepack_sfb(expert_w_sf, 128, N, K)
|
||||
|
||||
new = cutlass_nvfp4_blockscaled_gemm_prepacked_sfb(
|
||||
expert_x,
|
||||
expert_x_sf,
|
||||
expert_w,
|
||||
sfb_pre,
|
||||
M_expert,
|
||||
N,
|
||||
K,
|
||||
alpha=alpha,
|
||||
)
|
||||
|
||||
print("max diff", (old - new).abs().max())
|
||||
print("cos", torch.nn.functional.cosine_similarity(
|
||||
old.flatten().float(),
|
||||
new.flatten().float(),
|
||||
dim=0,
|
||||
))
|
||||
```
|
||||
|
||||
Expected:
|
||||
```
|
||||
max diff = 0 or tiny BF16-level difference
|
||||
cos ≈ 1.0
|
||||
```
|
||||
|
||||
If it fails only when M_expert != 128, then LayoutSFB is M-dependent. In that case, cache by an M bucket:
|
||||
```
|
||||
bucket_m = ((M_expert + 127) // 128) * 128
|
||||
cache_key = (tag, bucket_m, N, K)
|
||||
```
|
||||
|
||||
But test first. If SFB layout size and values are stable across M, keep one prepack per expert per layer.
|
||||
|
||||
Why this is worth doing now
|
||||
|
||||
This removes, per active expert GEMM:
|
||||
```
|
||||
1 cudaMalloc/free-ish allocation for SFB
|
||||
1 cudaMemsetAsync for SFB padding
|
||||
1 remap kernel launch for SFB
|
||||
```
|
||||
|
||||
Across L1 + L2 and many routed experts, that’s a very clean win without touching routing semantics, quantization, or the actual MMA tile.
|
||||
@@ -284,7 +284,7 @@ The CUTLASS extension builds inside the container during `pip install` of the nv
|
||||
|
||||
## Known Issues
|
||||
|
||||
1. **MoE dispatch is slow** — `cutlass_grouped_nvfp4_gemm` uses a Python loop over 48 experts with per-token scatter/gather. Needs a proper grouped GEMM or at least CUDA-side dispatch.
|
||||
1. ~~**MoE dispatch is slow**~~ — Fixed. Slot-based `index_add_` replaces the Python double loop over tokens×topk. Routing weights applied once at final scatter. Per-expert loop still exists (gather+GEMM) but scatter is vectorized.
|
||||
|
||||
2. **stage_activation is Python** — Re-quantization from L1 BF16 output to FP4 for L2 input runs in PyTorch. Should use the Triton staging kernel for speed and consistency with vLLM's built-in staging.
|
||||
|
||||
|
||||
@@ -213,6 +213,110 @@ int cutlass_nvfp4_gemm_run(
|
||||
return 0;
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// SFB prepack: pre-remap weight scale factors once at load time
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
extern "C" int cutlass_nvfp4_sfb_size(
|
||||
int M, int N, int K,
|
||||
int* out_size
|
||||
) {
|
||||
using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
|
||||
LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, 1));
|
||||
*out_size = cute::size(layout_SFB);
|
||||
return 0;
|
||||
}
|
||||
|
||||
extern "C" int cutlass_nvfp4_prepack_sfb_run(
|
||||
const void* SFB_ptr,
|
||||
void* SFB_cutlass_ptr,
|
||||
int M, int N, int K,
|
||||
cudaStream_t stream
|
||||
) {
|
||||
using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
|
||||
LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, 1));
|
||||
using ElementSF = typename Gemm::GemmKernel::CollectiveMainloop::ElementSF;
|
||||
|
||||
int sfb_size = cute::size(layout_SFB);
|
||||
int K_sf = K / InputSFVectorSize;
|
||||
|
||||
cudaMemsetAsync(static_cast<ElementSF*>(SFB_cutlass_ptr), 0, sfb_size * sizeof(ElementSF), stream);
|
||||
|
||||
int block = 256;
|
||||
remap_sf_to_cutlass_kernel<<<(sfb_size + block - 1) / block, block, 0, stream>>>(
|
||||
static_cast<const ElementSF*>(SFB_ptr),
|
||||
static_cast<ElementSF*>(SFB_cutlass_ptr),
|
||||
layout_SFB,
|
||||
N, K_sf, true // SFB source is (K_sf, N)
|
||||
);
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// GEMM with prepacked SFB — skips SFB allocation, memset, and remap
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
extern "C" int cutlass_nvfp4_gemm_run_prepacked_sfb(
|
||||
const void* A_ptr, const void* SFA_ptr,
|
||||
const void* B_ptr, const void* SFB_cutlass_ptr,
|
||||
void* D_ptr,
|
||||
int M, int N, int K,
|
||||
float alpha, float beta,
|
||||
cudaStream_t stream
|
||||
) {
|
||||
StrideA stride_A = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(M, K, 1));
|
||||
StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(N, K, 1));
|
||||
StrideC stride_C = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(M, N, 1));
|
||||
StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(M, N, 1));
|
||||
|
||||
using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
|
||||
LayoutSFA layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(M, N, K, 1));
|
||||
LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, 1));
|
||||
|
||||
using ArrayElementA = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementA;
|
||||
using ArrayElementB = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementB;
|
||||
using ElementSF = typename Gemm::GemmKernel::CollectiveMainloop::ElementSF;
|
||||
|
||||
int sfa_size = cute::size(layout_SFA);
|
||||
int K_sf = K / InputSFVectorSize;
|
||||
|
||||
// Only remap SFA (activation scales) — SFB is prepacked
|
||||
cutlass::device_memory::allocation<ElementSF> sfa_cutlass(sfa_size);
|
||||
cudaMemsetAsync(sfa_cutlass.get(), 0, sfa_size * sizeof(ElementSF), stream);
|
||||
|
||||
int block = 256;
|
||||
remap_sf_to_cutlass_kernel<<<(sfa_size + block - 1) / block, block, 0, stream>>>(
|
||||
static_cast<const ElementSF*>(SFA_ptr), sfa_cutlass.get(), layout_SFA, M, K_sf, false);
|
||||
|
||||
typename Gemm::Arguments arguments {
|
||||
cutlass::gemm::GemmUniversalMode::kGemm,
|
||||
{M, N, K, 1},
|
||||
{
|
||||
static_cast<const ArrayElementA*>(A_ptr), stride_A,
|
||||
static_cast<const ArrayElementB*>(B_ptr), stride_B,
|
||||
sfa_cutlass.get(), layout_SFA,
|
||||
static_cast<const ElementSF*>(SFB_cutlass_ptr), layout_SFB
|
||||
},
|
||||
{
|
||||
{ alpha, beta },
|
||||
nullptr, stride_C,
|
||||
static_cast<typename Gemm::GemmKernel::CollectiveEpilogue::ElementD*>(D_ptr), stride_D
|
||||
}
|
||||
};
|
||||
|
||||
Gemm gemm;
|
||||
CUTLASS_CHECK(gemm.can_implement(arguments));
|
||||
|
||||
size_t workspace_size = Gemm::get_workspace_size(arguments);
|
||||
cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
|
||||
|
||||
CUTLASS_CHECK(gemm.initialize(arguments, workspace.get(), stream));
|
||||
CUTLASS_CHECK(gemm.run(stream));
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
} // extern "C"
|
||||
|
||||
#endif
|
||||
|
||||
@@ -21,23 +21,46 @@ def cutlass_nvfp4_blockscaled_gemm(
|
||||
A_packed, # (M, K_half) int8 packed E2M1
|
||||
SFA, # scale factors for A (float8_e4m3fn)
|
||||
B_packed, # (K_half, N) int8 packed E2M1, column-major for CUTLASS
|
||||
SFB, # scale factors for B (sf_k, N) float8_e4m3fn, column-major for CUTLASS
|
||||
SFB, # scale factors for B — either (sf_k, N) float8_e4m3fn row-major, or prepacked CUTLASS layout
|
||||
M, N, K, # Problem dimensions (K in FP4 elements)
|
||||
alpha=1.0, # fp32 scalar applied in epilogue: D = alpha * A @ B + beta * C
|
||||
sfb_prepacked=False, # True if SFB is already in CUTLASS layout
|
||||
):
|
||||
"""Single NVFP4 block-scaled GEMM using CUTLASS."""
|
||||
"""Single NVFP4 block-scaled GEMM using CUTLASS.
|
||||
|
||||
If sfb_prepacked=True, SFB is assumed to be in CUTLASS interleaved layout
|
||||
(from prepack_sfb) and the SFB remap is skipped.
|
||||
"""
|
||||
if not _CUTLASS_AVAILABLE:
|
||||
raise RuntimeError("CUTLASS NVFP4 GEMM extension not available")
|
||||
return _C.forward(A_packed, SFA, B_packed, SFB, M, N, K, alpha)
|
||||
if sfb_prepacked:
|
||||
return _C.forward_prepacked_sfb(A_packed, SFA, B_packed, SFB, M, N, K, alpha)
|
||||
else:
|
||||
return _C.forward(A_packed, SFA, B_packed, SFB, M, N, K, alpha)
|
||||
|
||||
|
||||
def prepack_sfb(SFB, M, N, K):
|
||||
"""Pre-remap SFB weight scales into CUTLASS interleaved layout.
|
||||
|
||||
Call once after weight transform. Returns a tensor that can be passed
|
||||
to cutlass_nvfp4_blockscaled_gemm with sfb_prepacked=True.
|
||||
|
||||
M is used for layout sizing. Test with different M values to confirm
|
||||
SFB layout is M-independent; if so, any valid M works (e.g. 128).
|
||||
"""
|
||||
if not _CUTLASS_AVAILABLE:
|
||||
raise RuntimeError("CUTLASS NVFP4 GEMM extension not available")
|
||||
return _C.prepack_sfb(SFB, M, N, K)
|
||||
|
||||
|
||||
def cutlass_grouped_nvfp4_gemm(
|
||||
x_fp4, # (num_tokens, K_half) int8 packed E2M1
|
||||
x_sf, # (num_tokens, sf_k) float8_e4m3fn block scales
|
||||
weights, # (E_per_rank, K_half, N) int8 packed E2M1, column-major for CUTLASS
|
||||
weight_sf, # (E_per_rank, sf_k, N) float8_e4m3fn, column-major for CUTLASS
|
||||
weight_sf, # (E_per_rank, sf_k, N) float8_e4m3fn, column-major — or prepacked (E_per_rank, sfb_size) if sfb_prepacked=True
|
||||
topk_ids, # (num_tokens, NUM_TOPK) int32 — local expert IDs
|
||||
alpha=1.0, # fp32 scalar: D = alpha * A @ B (from stage_activation global scale)
|
||||
sfb_prepacked=False, # True if weight_sf is already prepacked into CUTLASS layout
|
||||
):
|
||||
"""Per-expert grouped GEMM for MoE dispatch using CUTLASS NVFP4.
|
||||
|
||||
@@ -56,23 +79,24 @@ def cutlass_grouped_nvfp4_gemm(
|
||||
num_topk = topk_ids.shape[1]
|
||||
|
||||
# Build slot mapping: which (token, topk) pairs land on local experts?
|
||||
local_mask = (topk_ids >= 0) & (topk_ids < num_experts) # (num_tokens, num_topk)
|
||||
slot_token, slot_k = local_mask.nonzero(as_tuple=True) # (num_slots,)
|
||||
slot_expert = topk_ids[slot_token, slot_k] # (num_slots,) local expert id
|
||||
local_mask = (topk_ids >= 0) & (topk_ids < num_experts)
|
||||
slot_token, slot_k = local_mask.nonzero(as_tuple=True)
|
||||
slot_expert = topk_ids[slot_token, slot_k]
|
||||
|
||||
num_slots = slot_token.shape[0]
|
||||
|
||||
if MEGA_MOE_DEBUG:
|
||||
print(f"[cutlass_grouped_gemm] tokens={num_tokens} K={K} N={N} "
|
||||
f"experts={num_experts} topk={num_topk} slots={num_slots}")
|
||||
f"experts={num_experts} topk={num_topk} slots={num_slots} "
|
||||
f"sfb_prepacked={sfb_prepacked}")
|
||||
|
||||
if num_slots == 0:
|
||||
slot_out = torch.empty(0, N, dtype=torch.bfloat16, device=x_fp4.device)
|
||||
return slot_out, slot_token
|
||||
|
||||
# Gather activations for all slots
|
||||
slot_x = x_fp4[slot_token] # (num_slots, K_half)
|
||||
slot_x_sf = x_sf[slot_token] # (num_slots, sf_k)
|
||||
slot_x = x_fp4[slot_token]
|
||||
slot_x_sf = x_sf[slot_token]
|
||||
|
||||
slot_out = torch.empty(num_slots, N, dtype=torch.bfloat16, device=x_fp4.device)
|
||||
|
||||
@@ -85,38 +109,26 @@ def cutlass_grouped_nvfp4_gemm(
|
||||
expert_x = slot_x[e_idx]
|
||||
expert_x_sf = slot_x_sf[e_idx]
|
||||
expert_w = weights[e]
|
||||
expert_w_sf = weight_sf[e]
|
||||
expert_w_sf = weight_sf[e] # prepacked or raw depending on flag
|
||||
M_expert = e_idx.shape[0]
|
||||
|
||||
if e < 3 and M_expert > 0:
|
||||
if MEGA_MOE_DEBUG and e < 3 and M_expert > 0:
|
||||
print(f"[GEMM-IN] expert={e} M={M_expert} N={N} K={K} "
|
||||
f"w shape={expert_w.shape} w_sf shape={expert_w_sf.shape} "
|
||||
f"w absmax={expert_w.view(torch.int8).abs().max().item()} "
|
||||
f"w_sf range=[{expert_w_sf.to(torch.float32).min().item():.4e}, "
|
||||
f"{expert_w_sf.to(torch.float32).max().item():.4e}] "
|
||||
f"w_sf nonzero_frac={(expert_w_sf.view(torch.uint8) != 0).float().mean().item():.4f}")
|
||||
f"w shape={expert_w.shape} sfb_prepacked={sfb_prepacked}")
|
||||
|
||||
expert_out = cutlass_nvfp4_blockscaled_gemm(
|
||||
expert_x, expert_x_sf,
|
||||
expert_w, expert_w_sf,
|
||||
M_expert, N, K,
|
||||
alpha=alpha,
|
||||
sfb_prepacked=sfb_prepacked,
|
||||
)
|
||||
|
||||
torch.cuda.current_stream().synchronize()
|
||||
|
||||
if torch.isnan(expert_out).any() or torch.isinf(expert_out).any():
|
||||
raise RuntimeError(
|
||||
f"expert {e} of {num_experts}: GEMM emitted NaN/Inf. "
|
||||
f"M={M_expert} N={N} K={K} | "
|
||||
f"x abs range [{expert_x.view(torch.int8).abs().max().item()}], "
|
||||
f"x_sf range [{expert_x_sf.to(torch.float32).min().item():.4e}, "
|
||||
f"{expert_x_sf.to(torch.float32).max().item():.4e}], "
|
||||
f"w_sf range [{expert_w_sf.to(torch.float32).min().item():.4e}, "
|
||||
f"{expert_w_sf.to(torch.float32).max().item():.4e}], "
|
||||
f"x_sf nan_frac={torch.isnan(expert_x_sf.to(torch.float32)).float().mean().item():.4f}, "
|
||||
f"w_sf nan_frac={torch.isnan(expert_w_sf.to(torch.float32)).float().mean().item():.4f}"
|
||||
)
|
||||
if MEGA_MOE_DEBUG:
|
||||
if torch.isnan(expert_out).any() or torch.isinf(expert_out).any():
|
||||
raise RuntimeError(
|
||||
f"expert {e} of {num_experts}: GEMM emitted NaN/Inf. "
|
||||
f"M={M_expert} N={N} K={K}")
|
||||
|
||||
slot_out[e_idx] = expert_out
|
||||
|
||||
|
||||
@@ -13,6 +13,27 @@ extern "C" int cutlass_nvfp4_gemm_run(
|
||||
cudaStream_t stream
|
||||
);
|
||||
|
||||
extern "C" int cutlass_nvfp4_gemm_run_prepacked_sfb(
|
||||
const void* A_ptr, const void* SFA_ptr,
|
||||
const void* B_ptr, const void* SFB_cutlass_ptr,
|
||||
void* D_ptr,
|
||||
int M, int N, int K,
|
||||
float alpha, float beta,
|
||||
cudaStream_t stream
|
||||
);
|
||||
|
||||
extern "C" int cutlass_nvfp4_sfb_size(
|
||||
int M, int N, int K,
|
||||
int* out_size
|
||||
);
|
||||
|
||||
extern "C" int cutlass_nvfp4_prepack_sfb_run(
|
||||
const void* SFB_ptr,
|
||||
void* SFB_cutlass_ptr,
|
||||
int M, int N, int K,
|
||||
cudaStream_t stream
|
||||
);
|
||||
|
||||
torch::Tensor cutlass_nvfp4_gemm_forward(
|
||||
torch::Tensor A_packed,
|
||||
torch::Tensor SFA,
|
||||
@@ -40,6 +61,71 @@ torch::Tensor cutlass_nvfp4_gemm_forward(
|
||||
return D;
|
||||
}
|
||||
|
||||
torch::Tensor cutlass_nvfp4_gemm_forward_prepacked_sfb(
|
||||
torch::Tensor A_packed,
|
||||
torch::Tensor SFA,
|
||||
torch::Tensor B_packed,
|
||||
torch::Tensor SFB_cutlass,
|
||||
int64_t M, int64_t N, int64_t K,
|
||||
double alpha = 1.0
|
||||
) {
|
||||
auto D = torch::empty({M, N}, torch::dtype(torch::kBFloat16).device(A_packed.device()));
|
||||
|
||||
auto stream = c10::cuda::getCurrentCUDAStream();
|
||||
cudaStream_t cuda_stream = stream.stream();
|
||||
|
||||
int rc = cutlass_nvfp4_gemm_run_prepacked_sfb(
|
||||
A_packed.data_ptr(), SFA.data_ptr(),
|
||||
B_packed.data_ptr(), SFB_cutlass.data_ptr(),
|
||||
D.data_ptr(),
|
||||
static_cast<int>(M), static_cast<int>(N), static_cast<int>(K),
|
||||
static_cast<float>(alpha), 0.0f,
|
||||
cuda_stream
|
||||
);
|
||||
|
||||
TORCH_CHECK(rc == 0, "CUTLASS NVFP4 GEMM (prepacked SFB) failed with error code ", rc);
|
||||
|
||||
return D;
|
||||
}
|
||||
|
||||
torch::Tensor prepack_sfb(
|
||||
torch::Tensor SFB,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K
|
||||
) {
|
||||
int size = 0;
|
||||
int rc = cutlass_nvfp4_sfb_size(
|
||||
static_cast<int>(M),
|
||||
static_cast<int>(N),
|
||||
static_cast<int>(K),
|
||||
&size
|
||||
);
|
||||
TORCH_CHECK(rc == 0, "sfb_size failed");
|
||||
|
||||
auto out = torch::empty(
|
||||
{size},
|
||||
torch::dtype(SFB.dtype()).device(SFB.device())
|
||||
);
|
||||
|
||||
auto stream = c10::cuda::getCurrentCUDAStream();
|
||||
|
||||
rc = cutlass_nvfp4_prepack_sfb_run(
|
||||
SFB.data_ptr(),
|
||||
out.data_ptr(),
|
||||
static_cast<int>(M),
|
||||
static_cast<int>(N),
|
||||
static_cast<int>(K),
|
||||
stream.stream()
|
||||
);
|
||||
|
||||
TORCH_CHECK(rc == 0, "prepack_sfb failed");
|
||||
|
||||
return out;
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("forward", &cutlass_nvfp4_gemm_forward, "CUTLASS NVFP4 block-scaled GEMM forward");
|
||||
m.def("forward_prepacked_sfb", &cutlass_nvfp4_gemm_forward_prepacked_sfb, "CUTLASS NVFP4 GEMM forward with prepacked SFB");
|
||||
m.def("prepack_sfb", &prepack_sfb, "Pre-remap SFB weight scales into CUTLASS layout");
|
||||
}
|
||||
|
||||
@@ -89,13 +89,46 @@ MEGA_MOE_DEBUG = int(os.environ.get("MEGA_MOE_DEBUG", "0"))
|
||||
# Main kernel entry points
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _prepack_weight_sf(weight_sf, N, K, tag):
|
||||
"""Lazily prepack SFB weight scales into CUTLASS layout (once per tag).
|
||||
|
||||
Returns a tensor of shape (E, sfb_size) with SFB already in CUTLASS
|
||||
interleaved layout, skipping the per-call remap+memset+alloc.
|
||||
"""
|
||||
cache_attr = f"_prepacked_{tag}"
|
||||
cached = getattr(_prepack_weight_sf, cache_attr, None)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
from nvfp4_megamoe_kernel.cutlass_nvfp4_gemm.kernel import prepack_sfb
|
||||
|
||||
E = weight_sf.shape[0]
|
||||
# M for layout sizing. Test with different M to confirm SFB is M-independent.
|
||||
# If SFB size changes with M, bucket by M and cache per-bucket.
|
||||
M_for_layout = 128
|
||||
|
||||
packed = []
|
||||
for e in range(E):
|
||||
packed.append(prepack_sfb(weight_sf[e], M_for_layout, N, K))
|
||||
|
||||
packed = torch.stack(packed, dim=0).contiguous()
|
||||
setattr(_prepack_weight_sf, cache_attr, packed)
|
||||
|
||||
if MEGA_MOE_DEBUG:
|
||||
print(f"[PREPACK] {tag}: E={E} N={N} K={K} packed_shape={packed.shape} "
|
||||
f"(was {weight_sf.shape})")
|
||||
|
||||
return packed
|
||||
|
||||
|
||||
def nvfp4_mega_moe_l1(
|
||||
x_fp4, # (num_tokens, K//2) int8 packed E2M1
|
||||
x_sf, # (num_tokens, sf_k_groups) float8_e4m3fn
|
||||
l1_weights, # (E_per_rank, K//2, 2*INTER) int8, column-major for CUTLASS
|
||||
l1_scales, # (E_per_rank, sf_k_groups, 2*INTER) float8_e4m3fn, column-major
|
||||
l1_scales, # (E_per_rank, sf_k_groups, 2*INTER) float8_e4m3fn, column-major — or prepacked
|
||||
topk_ids, # (num_tokens, NUM_TOPK) int32 — local expert IDs
|
||||
alpha=1.0, # fp32 scalar from stage_activation global scale
|
||||
sfb_prepacked=False, # True if l1_scales is prepacked CUTLASS layout
|
||||
):
|
||||
"""L1 GEMM: gate_up_proj — slot-based, no routing weights.
|
||||
|
||||
@@ -110,13 +143,17 @@ def nvfp4_mega_moe_l1(
|
||||
print(f"[nvfp4_moe_l1] tokens={x_fp4.shape[0]} K={K} N={N} native=1")
|
||||
|
||||
x_sf_fp8 = unpack_ue4m3_u32(x_sf) if x_sf.dtype == torch.uint32 else x_sf
|
||||
w_sf_fp8 = unpack_ue4m3_u32(l1_scales) if l1_scales.dtype == torch.uint32 else l1_scales
|
||||
if not sfb_prepacked:
|
||||
w_sf_fp8 = unpack_ue4m3_u32(l1_scales) if l1_scales.dtype == torch.uint32 else l1_scales
|
||||
else:
|
||||
w_sf_fp8 = l1_scales # already prepacked, skip unpack
|
||||
|
||||
slot_out, slot_token = cutlass_grouped_nvfp4_gemm(
|
||||
x_fp4, x_sf_fp8,
|
||||
l1_weights, w_sf_fp8,
|
||||
topk_ids,
|
||||
alpha=alpha,
|
||||
sfb_prepacked=sfb_prepacked,
|
||||
)
|
||||
print(f"[L1-GEMM-OUT] slots={slot_out.shape[0]} N={N} amax={slot_out.abs().max().item():.4e} mean={slot_out.float().mean().item():.4e}")
|
||||
return slot_out, slot_token
|
||||
@@ -126,10 +163,11 @@ def nvfp4_mega_moe_l2(
|
||||
x_fp4, # (num_slots, INTER//2) int8 packed E2M1
|
||||
x_sf, # (num_slots, sf_k_groups) float8_e4m3fn
|
||||
l2_weights, # (E_per_rank, INTER//2, HIDDEN) int8, column-major for CUTLASS
|
||||
l2_scales, # (E_per_rank, sf_k_groups, HIDDEN) float8_e4m3fn, column-major
|
||||
l2_scales, # (E_per_rank, sf_k_groups, HIDDEN) float8_e4m3fn, column-major — or prepacked
|
||||
topk_ids, # (num_tokens, NUM_TOPK) int32 — local expert IDs (for slot mapping)
|
||||
slot_token, # (num_slots,) int64 — token index per slot (from L1)
|
||||
alpha=1.0, # fp32 scalar from stage_activation global scale
|
||||
sfb_prepacked=False, # True if l2_scales is prepacked CUTLASS layout
|
||||
):
|
||||
"""L2 GEMM: down_proj — slot-based, no routing weights.
|
||||
|
||||
@@ -144,7 +182,10 @@ def nvfp4_mega_moe_l2(
|
||||
print(f"[nvfp4_moe_l2] slots={x_fp4.shape[0]} K={K} N={N} native=1")
|
||||
|
||||
x_sf_fp8 = unpack_ue4m3_u32(x_sf) if x_sf.dtype == torch.uint32 else x_sf
|
||||
w_sf_fp8 = unpack_ue4m3_u32(l2_scales) if l2_scales.dtype == torch.uint32 else l2_scales
|
||||
if not sfb_prepacked:
|
||||
w_sf_fp8 = unpack_ue4m3_u32(l2_scales) if l2_scales.dtype == torch.uint32 else l2_scales
|
||||
else:
|
||||
w_sf_fp8 = l2_scales # already prepacked
|
||||
|
||||
# Build local expert IDs per slot (same mapping as L1)
|
||||
num_topk = topk_ids.shape[1]
|
||||
@@ -156,8 +197,9 @@ def nvfp4_mega_moe_l2(
|
||||
slot_out, _ = cutlass_grouped_nvfp4_gemm(
|
||||
x_fp4, x_sf_fp8,
|
||||
l2_weights, w_sf_fp8,
|
||||
slot_expert_ids, # per-slot expert IDs
|
||||
slot_expert_ids,
|
||||
alpha=alpha,
|
||||
sfb_prepacked=sfb_prepacked,
|
||||
)
|
||||
return slot_out # (num_slots, HIDDEN) bfloat16
|
||||
|
||||
@@ -318,11 +360,21 @@ def nvfp4_mega_moe_full(
|
||||
y.zero_()
|
||||
return
|
||||
|
||||
# Step 2: L1 GEMM — slot-based, no routing weights
|
||||
# Prepack SFB weight scales into CUTLASS layout (lazy, once per layer)
|
||||
l1_N = l1_w.shape[2]
|
||||
l1_K = l1_w.shape[1] * 2
|
||||
l1_sf_prepacked = _prepack_weight_sf(l1_sf, l1_N, l1_K, "l1")
|
||||
|
||||
l2_N = l2_w.shape[2]
|
||||
l2_K = l2_w.shape[1] * 2
|
||||
l2_sf_prepacked = _prepack_weight_sf(l2_sf, l2_N, l2_K, "l2")
|
||||
|
||||
# Step 2: L1 GEMM — slot-based, no routing weights, prepacked SFB
|
||||
l1_slots, _ = nvfp4_mega_moe_l1(
|
||||
x_fp4, x_sf, l1_w, l1_sf,
|
||||
x_fp4, x_sf, l1_w, l1_sf_prepacked,
|
||||
topk_ids_local,
|
||||
alpha=l1_global_scale,
|
||||
sfb_prepacked=True,
|
||||
) # (num_slots, 2*INTER) bfloat16
|
||||
|
||||
if MEGA_MOE_DEBUG:
|
||||
@@ -347,11 +399,12 @@ def nvfp4_mega_moe_full(
|
||||
_l2gs = l2_global_scale if isinstance(l2_global_scale, float) else l2_global_scale.item()
|
||||
print(f"[ALPHA L2] alpha={_l2gs:.4e} l1_sf range [{_l1sf_f32.min().item():.4e}, {_l1sf_f32.max().item():.4e}]")
|
||||
|
||||
# Step 5: L2 GEMM — slot-based, no routing weights
|
||||
# Step 5: L2 GEMM — slot-based, no routing weights, prepacked SFB
|
||||
l2_slots = nvfp4_mega_moe_l2(
|
||||
l1_fp4, l1_sf_out, l2_w, l2_sf,
|
||||
l1_fp4, l1_sf_out, l2_w, l2_sf_prepacked,
|
||||
topk_ids_local, slot_token,
|
||||
alpha=l2_global_scale,
|
||||
sfb_prepacked=True,
|
||||
) # (num_slots, HIDDEN) bfloat16
|
||||
|
||||
if MEGA_MOE_DEBUG:
|
||||
|
||||
Reference in New Issue
Block a user