biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Restructure: cutedsl/ -> dsv4/ with proper layering

Browse Source 
- Split bridge.py -> ops/quantize.py, ops/layouts.py, ops/gemm_runner.py
- Renamed classes: CuTeDSLNvfp4Linear -> Nvfp4Linear, etc.
- Moved kernel code to dsv4/kernels/ (gemm, attention, compressor, decode, cuda)
- Moved PyTorch bridges to dsv4/ops/
- Moved nn.Module layers to dsv4layers/
- Moved reference implementations to dsv4/reference/
- Moved vendored CUTLASS code to vendored/
- Archived ~190 debug tests to tests/archive/
- Kept ~15 canonical tests in tests/unit/
- Updated all import paths
- Added stubs for future components (model/, cache/, loader/)
- Updated pyproject.toml: dsv4-inference package name
This commit is contained in:
biondizzle
2026-05-21 17:30:44 +00:00
parent 99e143dd0e
commit 3fb3c925af
274 changed files with 715 additions and 556 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
dsv4/kernels/compressor/__init__.py Normal file
View File

419
dsv4/kernels/compressor/csa_hca.py Normal file
View File

@@ -0,0 +1,419 @@
"""
DeepSeek-V4 CSA/HCA compressor kernels for CUTLASS CuTe DSL / Blackwell.
This is a production-oriented fusion boundary for the PyTorch reference in
`Pasted markdown.md`:
1. Run the projection stage as one or two packed Blackwell GEMMs:
CSA main: H @ [W_a_KV | W_a_Z | W_b_KV | W_b_Z] -> (N, 4*C)
CSA indexer: H @ [W_I_a_KV | W_I_a_Z | W_I_b_KV | W_I_b_Z] -> (N, 4*C_I)
HCA: H @ [W_KV | W_Z] -> (N, 2*C)
Use your tcgen05 / NVFP4 blockscaled GEMM here. Keep the packed projection
outputs in BF16/FP32 as you prefer; the compressor reads them as tensor
elements and accumulates the softmax reduction in FP32.
2. These native CuTe DSL kernels fuse:
bias add + column-wise softmax over token positions + weighted C sum
+ partial RoPE for KV output.
The full projection+compression single-kernel tcgen05 variant is possible, but it
needs your exact NVFP4 weight/scale layout and preferred tile shape. This file is
therefore the safe fusion seam: the expensive D x C math stays in your Blackwell
GEMM, while the small-position softmax/reduction/rope path avoids PyTorch ops and
extra materialization after projection.
Target dimensions used by DeepSeek-V4-Pro reference:
D=7168, C=512, C_I=128, CSA_M=4, HCA_M=128, NOPE=448, ROPE=64.
Assumptions:
* Tensors are contiguous row-major from PyTorch/DLPack.
* Projection buffers are laid out as described above.
* One sequence at a time. State/tail management remains on the caller side.
* For CSA continuation across calls, provide external previous-block B-side
projections for the first committed block. For fresh prefill, set
has_external_prev=False.
NOTE: I cannot compile this in this sandbox because CUTLASS CuTe DSL is not
installed here. It follows the CuTe DSL @jit/@kernel launch style, but you may
need tiny API-name edits if you are pinned to a specific CUTLASS 4.x commit.
"""
from __future__ import annotations
import torch
import cutlass
import cutlass.cute as cute
import cuda.bindings.driver as cuda
# -----------------------------------------------------------------------------
# Small helpers
# -----------------------------------------------------------------------------
LOG2_E = 1.44269504088896340736
def _ceil_div(a: int, b: int) -> int:
return (a + b - 1) // b
@cute.jit
def _expf(x: cutlass.Float32) -> cutlass.Float32:
# CuTe DSL exposes exp2 in cute.math; exp(x) = exp2(x * log2(e)).
return cute.math.exp2(x * cutlass.Float32(LOG2_E))
@cute.jit
def _read_csa_current_packed(
proj: cute.Tensor,
token_idx: cutlass.Int64,
col: cutlass.Int32,
kind: cutlass.Constexpr, # 0 Ca, 1 Za, 2 Cb, 3 Zb
OUT: cutlass.Constexpr,
) -> cutlass.Float32:
return proj[token_idx, kind * OUT + col].to(cutlass.Float32)
@cute.jit
def _read_hca_packed(
proj: cute.Tensor,
token_idx: cutlass.Int64,
col: cutlass.Int32,
kind: cutlass.Constexpr, # 0 C, 1 Z
C: cutlass.Constexpr,
) -> cutlass.Float32:
return proj[token_idx, kind * C + col].to(cutlass.Float32)
# -----------------------------------------------------------------------------
# CSA fused compressor from packed projections
# -----------------------------------------------------------------------------
@cute.jit
def _csa_raw_reduce_one_col(
proj: cute.Tensor, # (N_tokens, 4*OUT): Ca, Za, Cb, Zb
prev_b_proj: cute.Tensor, # (M, 2*OUT): Cb_prev, Zb_prev for first block, may be dummy
B_a: cute.Tensor, # (M, OUT)
B_b: cute.Tensor, # (M, OUT)
block_i: cutlass.Int64,
col: cutlass.Int32,
start_token_in_proj: cutlass.Int64,
has_external_prev: cutlass.Constexpr,
M: cutlass.Constexpr,
OUT: cutlass.Constexpr,
) -> cutlass.Float32:
"""Column-wise CSA softmax+weighted-sum for either main KV or indexer.
This implements exactly:
softmax([Z_a_cur + B_a ; Z_b_prev + B_b], dim=position)
sum S_a*C_a + sum S_b*C_b
with the block-0 no-prev case reducing over M current positions only.
"""
# First pass: max logit.
max_logit = cutlass.Float32(-3.4028234663852886e38)
# Current/current-a side always exists.
for p in cutlass.range_constexpr(M):
tok = start_token_in_proj + block_i * M + p
za = _read_csa_current_packed(proj, tok, col, 1, OUT) + B_a[p, col].to(cutlass.Float32)
max_logit = cute.math.fmax(max_logit, za)
# Previous/b side exists if this is not the first fresh block.
use_prev = (block_i > 0) or has_external_prev
if use_prev:
for p in cutlass.range_constexpr(M):
if block_i > 0:
tok_prev = start_token_in_proj + (block_i - 1) * M + p
zb = _read_csa_current_packed(proj, tok_prev, col, 3, OUT)
else:
# External previous block is packed as [Cb_prev | Zb_prev]
zb = prev_b_proj[p, OUT + col].to(cutlass.Float32)
zb = zb + B_b[p, col].to(cutlass.Float32)
max_logit = cute.math.fmax(max_logit, zb)
# Second pass: exp denominator and weighted value.
denom = cutlass.Float32(0.0)
acc = cutlass.Float32(0.0)
for p in cutlass.range_constexpr(M):
tok = start_token_in_proj + block_i * M + p
za = _read_csa_current_packed(proj, tok, col, 1, OUT) + B_a[p, col].to(cutlass.Float32)
ca = _read_csa_current_packed(proj, tok, col, 0, OUT)
e = _expf(za - max_logit)
denom += e
acc += e * ca
if use_prev:
for p in cutlass.range_constexpr(M):
if block_i > 0:
tok_prev = start_token_in_proj + (block_i - 1) * M + p
cb = _read_csa_current_packed(proj, tok_prev, col, 2, OUT)
zb = _read_csa_current_packed(proj, tok_prev, col, 3, OUT)
else:
cb = prev_b_proj[p, col].to(cutlass.Float32)
zb = prev_b_proj[p, OUT + col].to(cutlass.Float32)
zb = zb + B_b[p, col].to(cutlass.Float32)
e = _expf(zb - max_logit)
denom += e
acc += e * cb
return acc / denom
@cute.kernel
def csa_compress_projected_kernel(
proj_main: cute.Tensor, # (N_tokens, 4*C)
proj_indexer: cute.Tensor, # (N_tokens, 4*C_I)
prev_main_b: cute.Tensor, # (M, 2*C), Cb_prev|Zb_prev; dummy if no ext prev
prev_indexer_b: cute.Tensor, # (M, 2*C_I), Cb_prev|Zb_prev; dummy if no ext prev
B_a: cute.Tensor, # (M, C)
B_b: cute.Tensor, # (M, C)
B_I_a: cute.Tensor, # (M, C_I)
B_I_b: cute.Tensor, # (M, C_I)
cos_sin_cache: cute.Tensor, # (max_pos, ROPE), cos first half, sin second half
kv_out: cute.Tensor, # (n_blocks, C)
indexer_out: cute.Tensor, # (n_blocks, C_I)
n_blocks: cutlass.Int64,
start_token_in_proj: cutlass.Int64,
start_abs_pos: cutlass.Int64,
has_external_prev: cutlass.Constexpr,
M: cutlass.Constexpr,
C: cutlass.Constexpr,
C_I: cutlass.Constexpr,
NOPE: cutlass.Constexpr,
ROPE: cutlass.Constexpr,
COLS_PER_CTA: cutlass.Constexpr,
):
tx, _, _ = cute.arch.thread_idx()
bx, by, bz = cute.arch.block_idx()
block_i = bx.to(cutlass.Int64)
base_col = by * COLS_PER_CTA
col = base_col + tx
# bz == 0: main KV output with RoPE. bz == 1: indexer output, no RoPE.
if bz == 0:
if block_i < n_blocks and col < C:
# RoPE dims need even/odd pair. Let even lane compute/store both.
if col < NOPE:
val = _csa_raw_reduce_one_col(
proj_main, prev_main_b, B_a, B_b,
block_i, col, start_token_in_proj, has_external_prev, M, C,
)
kv_out[block_i, col] = val.to(kv_out.element_type)
else:
rope_col = col - NOPE
if (rope_col % 2) == 0:
# Compute pair and rotate by block end position.
x0 = _csa_raw_reduce_one_col(
proj_main, prev_main_b, B_a, B_b,
block_i, col, start_token_in_proj, has_external_prev, M, C,
)
x1 = _csa_raw_reduce_one_col(
proj_main, prev_main_b, B_a, B_b,
block_i, col + 1, start_token_in_proj, has_external_prev, M, C,
)
block_end_pos = start_abs_pos + block_i * M + (M - 1)
half_idx = rope_col // 2
cosv = cos_sin_cache[block_end_pos, half_idx].to(cutlass.Float32)
sinv = cos_sin_cache[block_end_pos, half_idx + ROPE // 2].to(cutlass.Float32)
kv_out[block_i, col] = (x0 * cosv - x1 * sinv).to(kv_out.element_type)
kv_out[block_i, col + 1] = (x0 * sinv + x1 * cosv).to(kv_out.element_type)
else:
if block_i < n_blocks and col < C_I:
val_i = _csa_raw_reduce_one_col(
proj_indexer, prev_indexer_b, B_I_a, B_I_b,
block_i, col, start_token_in_proj, has_external_prev, M, C_I,
)
indexer_out[block_i, col] = val_i.to(indexer_out.element_type)
@cute.jit
def launch_csa_compress_projected(
proj_main: cute.Tensor,
proj_indexer: cute.Tensor,
prev_main_b: cute.Tensor,
prev_indexer_b: cute.Tensor,
B_a: cute.Tensor,
B_b: cute.Tensor,
B_I_a: cute.Tensor,
B_I_b: cute.Tensor,
cos_sin_cache: cute.Tensor,
kv_out: cute.Tensor,
indexer_out: cute.Tensor,
n_blocks: int,
start_token_in_proj: int,
start_abs_pos: int,
has_external_prev: cutlass.Constexpr,
stream: cuda.CUstream,
M: cutlass.Constexpr = 4,
C: cutlass.Constexpr = 512,
C_I: cutlass.Constexpr = 128,
NOPE: cutlass.Constexpr = 448,
ROPE: cutlass.Constexpr = 64,
COLS_PER_CTA: cutlass.Constexpr = 128,
):
grid_y = _ceil_div(C, COLS_PER_CTA) # enough for main; indexer just masks col < C_I
csa_compress_projected_kernel(
proj_main,
proj_indexer,
prev_main_b,
prev_indexer_b,
B_a,
B_b,
B_I_a,
B_I_b,
cos_sin_cache,
kv_out,
indexer_out,
n_blocks,
start_token_in_proj,
start_abs_pos,
has_external_prev,
M,
C,
C_I,
NOPE,
ROPE,
COLS_PER_CTA,
).launch(
grid=[n_blocks, grid_y, 2],
block=[COLS_PER_CTA, 1, 1],
stream=stream,
)
# -----------------------------------------------------------------------------
# HCA fused compressor from packed projections
# -----------------------------------------------------------------------------
@cute.jit
def _hca_raw_reduce_one_col(
proj: cute.Tensor, # (N_tokens, 2*C): C, Z
B: cute.Tensor, # (M, C)
block_i: cutlass.Int64,
col: cutlass.Int32,
start_token_in_proj: cutlass.Int64,
M: cutlass.Constexpr,
C: cutlass.Constexpr,
) -> cutlass.Float32:
max_logit = cutlass.Float32(-3.4028234663852886e38)
for p in cutlass.range_constexpr(M):
tok = start_token_in_proj + block_i * M + p
z = _read_hca_packed(proj, tok, col, 1, C) + B[p, col].to(cutlass.Float32)
max_logit = cute.math.fmax(max_logit, z)
denom = cutlass.Float32(0.0)
acc = cutlass.Float32(0.0)
for p in cutlass.range_constexpr(M):
tok = start_token_in_proj + block_i * M + p
z = _read_hca_packed(proj, tok, col, 1, C) + B[p, col].to(cutlass.Float32)
c = _read_hca_packed(proj, tok, col, 0, C)
e = _expf(z - max_logit)
denom += e
acc += e * c
return acc / denom
@cute.kernel
def hca_compress_projected_kernel(
proj: cute.Tensor, # (N_tokens, 2*C)
B: cute.Tensor, # (M, C)
cos_sin_cache: cute.Tensor, # (max_pos, ROPE)
kv_out: cute.Tensor, # (n_blocks, C)
n_blocks: cutlass.Int64,
start_token_in_proj: cutlass.Int64,
start_abs_pos: cutlass.Int64,
M: cutlass.Constexpr,
C: cutlass.Constexpr,
NOPE: cutlass.Constexpr,
ROPE: cutlass.Constexpr,
COLS_PER_CTA: cutlass.Constexpr,
):
tx, _, _ = cute.arch.thread_idx()
bx, by, _ = cute.arch.block_idx()
block_i = bx.to(cutlass.Int64)
col = by * COLS_PER_CTA + tx
if block_i < n_blocks and col < C:
if col < NOPE:
val = _hca_raw_reduce_one_col(proj, B, block_i, col, start_token_in_proj, M, C)
kv_out[block_i, col] = val.to(kv_out.element_type)
else:
rope_col = col - NOPE
if (rope_col % 2) == 0:
x0 = _hca_raw_reduce_one_col(proj, B, block_i, col, start_token_in_proj, M, C)
x1 = _hca_raw_reduce_one_col(proj, B, block_i, col + 1, start_token_in_proj, M, C)
block_end_pos = start_abs_pos + block_i * M + (M - 1)
half_idx = rope_col // 2
cosv = cos_sin_cache[block_end_pos, half_idx].to(cutlass.Float32)
sinv = cos_sin_cache[block_end_pos, half_idx + ROPE // 2].to(cutlass.Float32)
kv_out[block_i, col] = (x0 * cosv - x1 * sinv).to(kv_out.element_type)
kv_out[block_i, col + 1] = (x0 * sinv + x1 * cosv).to(kv_out.element_type)
@cute.jit
def launch_hca_compress_projected(
proj: cute.Tensor,
B: cute.Tensor,
cos_sin_cache: cute.Tensor,
kv_out: cute.Tensor,
n_blocks: int,
start_token_in_proj: int,
start_abs_pos: int,
stream: cuda.CUstream,
M: cutlass.Constexpr = 128,
C: cutlass.Constexpr = 512,
NOPE: cutlass.Constexpr = 448,
ROPE: cutlass.Constexpr = 64,
COLS_PER_CTA: cutlass.Constexpr = 128,
):
grid_y = _ceil_div(C, COLS_PER_CTA)
hca_compress_projected_kernel(
proj,
B,
cos_sin_cache,
kv_out,
n_blocks,
start_token_in_proj,
start_abs_pos,
M,
C,
NOPE,
ROPE,
COLS_PER_CTA,
).launch(
grid=[n_blocks, grid_y, 1],
block=[COLS_PER_CTA, 1, 1],
stream=stream,
)
# -----------------------------------------------------------------------------
# PyTorch-side packing helpers
# -----------------------------------------------------------------------------
def pack_csa_main_weights(W_a_KV, W_a_Z, W_b_KV, W_b_Z):
"""Return W packed as [W_a_KV | W_a_Z | W_b_KV | W_b_Z]."""
return torch.cat([W_a_KV, W_a_Z, W_b_KV, W_b_Z], dim=1).contiguous()
def pack_csa_indexer_weights(W_I_a_KV, W_I_a_Z, W_I_b_KV, W_I_b_Z):
"""Return W_I packed as [W_I_a_KV | W_I_a_Z | W_I_b_KV | W_I_b_Z]."""
return torch.cat([W_I_a_KV, W_I_a_Z, W_I_b_KV, W_I_b_Z], dim=1).contiguous()
def pack_hca_weights(W_KV, W_Z):
"""Return W packed as [W_KV | W_Z]."""
return torch.cat([W_KV, W_Z], dim=1).contiguous()
def make_dummy_prev_b(M: int, OUT: int, *, device, dtype):
"""Dummy external previous-block projection for fresh prefill."""
return torch.empty((M, 2 * OUT), device=device, dtype=dtype)