Stage E: production attention wrapper + Python KV merge, clean fmha_smem_acc
This commit is contained in:
@@ -66,8 +66,6 @@ class FmhaKernel:
|
||||
self.num_c_stage = 1 if head_dim > 256 else 2 # Reduce SMEM at hd=512
|
||||
self.scale_softmax = scale_softmax if scale_softmax is not None else 1.0 / math.sqrt(self.head_dim)
|
||||
self.scale_softmax_log2 = self.scale_softmax * math.log2(math.e)
|
||||
# SMEM accumulator: needed when n_kv_tiles > 1 (TMEM round-trip is broken)
|
||||
self.use_smem_accumulator = self.n_kv_tiles > 1
|
||||
|
||||
def _setup(self, qk_mma, pv_mma):
|
||||
qk_ik = cute.size(qk_mma.shape_mnk, mode=[2])
|
||||
@@ -126,7 +124,7 @@ class FmhaKernel:
|
||||
cute.size_in_bytes(self.q_dtype, v_s)) * cta
|
||||
|
||||
@cute.jit
|
||||
def __call__(self, q, k, v, c, stream, lse=None, swa_len=None, sink_bias=None, row_sums=None, c_simple=None):
|
||||
def __call__(self, q, k, v, c, stream, lse=None, swa_len=None, sink_bias=None, row_sums=None):
|
||||
self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype
|
||||
self.a_major = LayoutEnum.from_tensor(q).mma_major_mode()
|
||||
self.b_major = LayoutEnum.from_tensor(k).mma_major_mode()
|
||||
@@ -171,14 +169,11 @@ class FmhaKernel:
|
||||
# For single-head (n_h=1): grid=(1,1,1) — backward compatible
|
||||
if const_expr(row_sums is None):
|
||||
row_sums = cute.make_tensor(lse.iterator, lse.layout)
|
||||
# c_simple: non-dynamic-layout GMEM tensor for direct TMA store (SMEM accumulator path)
|
||||
if const_expr(c_simple is None):
|
||||
c_simple = cute.make_tensor(c.iterator, cute.make_layout((1, 1, 1), stride=(1, 1, 1)))
|
||||
|
||||
self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.p_smem_s,self.c_smem_s,self.epi_tile,lse,swa_len,sink_bias,row_sums,c_simple).launch(grid=(1,1,self.batch_size),block=[self.threads_per_cta,1,1],stream=stream)
|
||||
self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.p_smem_s,self.c_smem_s,self.epi_tile,lse,swa_len,sink_bias,row_sums).launch(grid=(1,1,self.batch_size),block=[self.threads_per_cta,1,1],stream=stream)
|
||||
|
||||
@cute.kernel
|
||||
def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, p_smem_s, c_smem_s, epi_tile, mLSE, swa_len, mSinkBias, mRowSums, mCSimple):
|
||||
def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, p_smem_s, c_smem_s, epi_tile, mLSE, swa_len, mSinkBias, mRowSums):
|
||||
warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
|
||||
tidx,_,_ = cute.arch.thread_idx()
|
||||
if warp_idx == self.tma_warp_id:
|
||||
@@ -222,18 +217,6 @@ class FmhaKernel:
|
||||
_p_swizzle = cute.make_layout(((1,1),1,(1,1),1))
|
||||
sP = smem.allocate_tensor(element_type=self.q_dtype,layout=_p_layout,byte_alignment=128,swizzle=_p_swizzle)
|
||||
|
||||
# SMEM accumulator: FP32 [128, pv_n_tile] row-major
|
||||
# Only needed for n_kv_tiles > 1 (TMEM round-trip is broken)
|
||||
if const_expr(self.use_smem_accumulator):
|
||||
sO_acc_layout = cute.make_layout((128, self.pv_n_tile), stride=(self.pv_n_tile, 1))
|
||||
sO_acc = smem.allocate_tensor(element_type=Float32, layout=sO_acc_layout, byte_alignment=128)
|
||||
# sC_flat: BF16 SMEM buffer with epi_s layout (non-swizzled) for TMA store
|
||||
# Used to cast sO_acc (FP32) -> BF16 and TMA store to GMEM
|
||||
sC_flat = smem.allocate_tensor(element_type=self.o_dtype, layout=cute.make_layout((128, self.pv_n_tile), stride=(self.pv_n_tile, 1)), byte_alignment=128)
|
||||
else:
|
||||
sO_acc = smem.allocate_tensor(element_type=Float32, layout=cute.make_layout((1, 1), stride=(1, 1)), byte_alignment=128)
|
||||
sC_flat = smem.allocate_tensor(element_type=self.o_dtype, layout=cute.make_layout((1, 1), stride=(1, 1)), byte_alignment=128)
|
||||
|
||||
gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None))
|
||||
gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None))
|
||||
gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None))
|
||||
@@ -359,7 +342,7 @@ class FmhaKernel:
|
||||
cute.arch.fence_view_async_tmem_store()
|
||||
sh.commit()
|
||||
softmax_done_bar.arrive_and_wait()
|
||||
pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0 and not self.use_smem_accumulator)
|
||||
pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0)
|
||||
if not self.use_smem_p:
|
||||
for kb in cutlass.range(cute.size(tOrP0, mode=[2]), unroll_full=True):
|
||||
cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,kvh.index)], tOtO0)
|
||||
@@ -422,21 +405,6 @@ class FmhaKernel:
|
||||
# This is representable as a CuTe layout: (128, (16, 4, 2)) -> (64, (1, 16, 8192))
|
||||
_sP_nostage = sP[(None, None, None, 0)] # remove stage dim
|
||||
|
||||
# O load atoms: one-way TMEM->REGS read of O_kt after PV (for SMEM accumulator)
|
||||
# Only needed when use_smem_accumulator (n_kv_tiles > 1)
|
||||
if const_expr(self.use_smem_accumulator):
|
||||
o_load_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype)
|
||||
tiled_o_load = tcgen05.make_tmem_copy(o_load_atom, tOtO0)
|
||||
thr_o_load = tiled_o_load.get_slice(sfw_idx)
|
||||
tTMEM_LOADtO = thr_o_load.partition_S(tOtO0)
|
||||
cO = cute.make_identity_tensor((self.qk_mma_tiler[0], self.pv_mma_tiler[1]))
|
||||
tOcO = pv_thr.partition_C(cO)
|
||||
tTMEM_LOADcO = thr_o_load.partition_D(tOcO)
|
||||
|
||||
# Zero-initialize sO_acc
|
||||
for i in cutlass.range(0, cute.size(sO_acc), unroll=1):
|
||||
sO_acc[i] = Float32(0.0)
|
||||
|
||||
row_max = -Float32.inf
|
||||
row_sum = Float32(0.0)
|
||||
scale_log2 = Float32(self.scale_softmax_log2)
|
||||
@@ -571,81 +539,37 @@ class FmhaKernel:
|
||||
si_handle.release()
|
||||
softmax_done_bar.arrive()
|
||||
|
||||
# --- SMEM accumulator: load O_kt from TMEM, accumulate in sO_acc ---
|
||||
if const_expr(self.use_smem_accumulator):
|
||||
pv_done_bar.arrive_and_wait()
|
||||
rO = cute.make_rmem_tensor(tTMEM_LOADcO.shape, self.qk_acc_dtype)
|
||||
cute.copy(tiled_o_load, tTMEM_LOADtO, rO)
|
||||
cute.arch.fence_view_async_tmem_load()
|
||||
|
||||
# sO_acc = acc_scale * sO_acc + O_kt
|
||||
for j0 in range(32):
|
||||
for j1 in range(4):
|
||||
coord = tTMEM_LOADcO[(j0, 0), j1, 0, 0]
|
||||
row = coord[0]
|
||||
col = coord[1]
|
||||
old_val = sO_acc[row, col]
|
||||
new_val = acc_scale * old_val + rO[(j0, 0), j1, 0, 0]
|
||||
sO_acc[row, col] = new_val
|
||||
|
||||
# Wait for MMA's PV[N-1] to commit before reading O.
|
||||
final_o_bar.arrive_and_wait()
|
||||
|
||||
# ============================================================
|
||||
# EPILOGUE: store O to GMEM + compute LSE
|
||||
# EPILOGUE: TMA store O to GMEM + compute LSE
|
||||
# ============================================================
|
||||
# For n_kv_tiles=1: O is in TMEM (from PV). Use epilogue_tma_store.
|
||||
# For n_kv_tiles>1: O is in sO_acc (SMEM). Write to sC, then TMA store.
|
||||
# The raw un-normalized O in TMEM is perfect (cos 0.999998).
|
||||
# We use epilogue_tma_store which reads O from TMEM directly via
|
||||
# the correct get_tmem_load_op layout — no round-trip needed.
|
||||
#
|
||||
# For multi-KV-tile: the paired-atom O rescale above (kt>0) ensures
|
||||
# O is correctly rescaled before this epilogue reads it.
|
||||
#
|
||||
# External normalization (D5a path): kernel outputs un-normalized O +
|
||||
# LSE + row_sum. Caller normalizes using O_norm = O_unnorm / row_sum.
|
||||
# This is exact and composes with D5c sink bias merge.
|
||||
# ============================================================
|
||||
|
||||
if const_expr(not self.use_smem_accumulator):
|
||||
# Path 1: epilogue_tma_store (reads O from TMEM, proven for n_kv=1)
|
||||
tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout)
|
||||
c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id))
|
||||
c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp)
|
||||
acc_cons_st = pipeline.make_pipeline_state(
|
||||
pipeline.PipelineUserType.Consumer, self.num_acc_stage
|
||||
)
|
||||
acc_cons_st = utils.gemm.sm100.epilogue_tma_store(
|
||||
self, sfw_idx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile,
|
||||
0, const_expr(lambda x: x), (0, 0, 0),
|
||||
acc_cons_st, acc_pipe, c_pipe,
|
||||
)
|
||||
c_pipe.producer_tail()
|
||||
else:
|
||||
# Path 2: write sO_acc (FP32) -> sC_flat (BF16) -> TMA store to GMEM
|
||||
# Reinterpret sC_flat as staged layout for TMA compatibility
|
||||
sC_flat_staged = cute.make_tensor(sC_flat.iterator, cute.make_layout((128, self.pv_n_tile // self.num_c_stage, self.num_c_stage), stride=(self.pv_n_tile, self.num_c_stage, 1)))
|
||||
|
||||
# Step 1: Cast sO_acc -> sC_flat (BF16) with staged layout
|
||||
for row in cutlass.range(0, 128, unroll=1):
|
||||
for col in cutlass.range(0, self.pv_n_tile, unroll=1):
|
||||
val = sO_acc[row, col]
|
||||
if const_expr(self.normalize):
|
||||
inv_row_sum = Float32(1.0) / row_sum
|
||||
val = val * inv_row_sum
|
||||
sC_flat_staged[row, col // self.num_c_stage, col % self.num_c_stage] = val.to(self.o_dtype)
|
||||
cute.arch.fence_proxy("async.shared", space="cta")
|
||||
|
||||
# Step 2: TMA store sC_flat -> GMEM
|
||||
# Use tCgC (already partitioned) for the GMEM side of TMA
|
||||
tCgC_xfm = transform_partitioned_tensor_layout(tCgC)
|
||||
tCgC_epi = cute.flat_divide(tCgC_xfm, epi_tile)
|
||||
# sC_flat (128, pv_n_tile, stride=(pv_n_tile, 1))
|
||||
# Split pv_n_tile into (pv_n_tile//2, 2) to match TMA stage decomposition
|
||||
# Result: (128, pv_n_tile//2, 2, stride=(pv_n_tile, 2, 1))
|
||||
sC_flat_staged = cute.make_tensor(sC_flat.iterator, cute.make_layout((128, self.pv_n_tile // self.num_c_stage, self.num_c_stage), stride=(self.pv_n_tile, self.num_c_stage, 1)))
|
||||
tOsC, tOgO = cpasync.tma_partition(
|
||||
tma_c, 0, cute.make_layout(1),
|
||||
sC_flat_staged,
|
||||
tCgC_epi,
|
||||
)
|
||||
# Slice off MMA tile coordinates from tOgO
|
||||
tOgO = tOgO[(None, None, None, Int32(0), Int32(0), Int32(0), Int32(0))]
|
||||
if warp_idx == self.epilogue_warp_id[0]:
|
||||
cute.copy(tma_c, tOsC, tOgO)
|
||||
cute.arch.cp_async_bulk_commit_group()
|
||||
cute.arch.cp_async_bulk_wait_group(0, read=True)
|
||||
# TMA store via CUTLASS epilogue_tma_store (reads raw O from TMEM)
|
||||
tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout)
|
||||
c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id))
|
||||
c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp)
|
||||
acc_cons_st = pipeline.make_pipeline_state(
|
||||
pipeline.PipelineUserType.Consumer, self.num_acc_stage
|
||||
)
|
||||
acc_cons_st = utils.gemm.sm100.epilogue_tma_store(
|
||||
self, sfw_idx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile,
|
||||
0, const_expr(lambda x: x), (0, 0, 0),
|
||||
acc_cons_st, acc_pipe, c_pipe,
|
||||
)
|
||||
c_pipe.producer_tail()
|
||||
|
||||
# Compute LSE: lse = ln(row_sum) + row_max * ln(2)
|
||||
# Only when emitting un-normalized output (D5a path).
|
||||
|
||||
223
dsv4/kernels/attention/production.py
Normal file
223
dsv4/kernels/attention/production.py
Normal file
@@ -0,0 +1,223 @@
|
||||
"""DSV4 Blackwell Attention — Production kernel wrapper.
|
||||
|
||||
Wraps the CuTeDSL FMHA kernel with Python KV merge for multi-KV-tile.
|
||||
Replaces vLLM's broken FlashMLA Blackwell implementation.
|
||||
|
||||
Supported:
|
||||
- CSA (Classical Self-Attention): 32 KV entries/block, full KV cache
|
||||
- HCA (Hash-based Cross-Attention): 1 KV entry/block, compressed cache
|
||||
- SWA (Sliding Window Attention): ring buffer, no KV cache
|
||||
|
||||
Unsupported (future):
|
||||
- SMEM accumulator for in-kernel multi-KV-tile (currently uses Python KV merge)
|
||||
- head_dim > 256 (MLIR compilation hang at hd=512)
|
||||
"""
|
||||
import torch
|
||||
import math
|
||||
import cutlass.cute as cute
|
||||
import cutlass.torch as ct
|
||||
import cuda.bindings.driver as cuda
|
||||
from dsv4.kernels.attention.fmha import FmhaKernel
|
||||
from dsv4.model.config import DSV4Config
|
||||
|
||||
# Kernel cache: compiled kernels keyed by (head_dim, s_k, use_smem_p, normalize, ...)
|
||||
_kernel_cache: dict = {}
|
||||
|
||||
|
||||
def _get_or_compile_kernel(head_dim: int, s_k: int, use_smem_p: bool = False,
|
||||
normalize: bool = False, apply_swa_mask: bool = False,
|
||||
is_causal: bool = False, n_comp: int = 0,
|
||||
apply_sink_bias: bool = False) -> tuple:
|
||||
"""Get or compile a kernel for the given configuration.
|
||||
|
||||
Returns (compiled_kernel, FmhaKernel instance).
|
||||
Kernel compilation is expensive, so we cache by config.
|
||||
"""
|
||||
key = (head_dim, s_k, use_smem_p, normalize, apply_swa_mask, is_causal, n_comp, apply_sink_bias)
|
||||
if key in _kernel_cache:
|
||||
return _kernel_cache[key]
|
||||
|
||||
kernel = FmhaKernel(
|
||||
head_dim=head_dim, s_k=s_k, use_smem_p=use_smem_p, normalize=normalize,
|
||||
apply_swa_mask=apply_swa_mask, is_causal=is_causal,
|
||||
n_comp=n_comp, apply_sink_bias=apply_sink_bias,
|
||||
)
|
||||
|
||||
# Create dummy tensors for compilation
|
||||
m = 128
|
||||
pv_n_tile = kernel.pv_n_tile
|
||||
q = torch.randn(m, head_dim, 1, dtype=torch.bfloat16, device='cuda')
|
||||
k = torch.randn(s_k, head_dim, 1, dtype=torch.bfloat16, device='cuda')
|
||||
v = torch.randn(s_k, pv_n_tile, dtype=torch.bfloat16, device='cuda')
|
||||
c = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
|
||||
lse = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
|
||||
row_sums = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
|
||||
|
||||
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
|
||||
|
||||
v_tile = v[:, 0:pv_n_tile].contiguous().unsqueeze(-1)
|
||||
mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q))
|
||||
mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k))
|
||||
mV = ct.from_dlpack(v_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile))
|
||||
mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c))
|
||||
mLSE = ct.from_dlpack(lse).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse))
|
||||
mRS = ct.from_dlpack(row_sums).mark_layout_dynamic(leading_dim=ct.get_leading_dim(row_sums))
|
||||
|
||||
compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, lse=mLSE, row_sums=mRS)
|
||||
_kernel_cache[key] = (compiled, kernel)
|
||||
return (compiled, kernel)
|
||||
|
||||
|
||||
def dsv4_attention(
|
||||
q: torch.Tensor, # (num_heads, seq_len, head_dim)
|
||||
k: torch.Tensor, # (seq_len, head_dim) or (num_heads, seq_len, head_dim) for MQA
|
||||
v: torch.Tensor, # (seq_len, head_dim) or (num_heads, seq_len, head_dim) for MQA
|
||||
scale: float = None, # 1/sqrt(head_dim) if None
|
||||
swa_len: int = None, # sliding window length (None = no SWA mask)
|
||||
is_causal: bool = False,
|
||||
n_comp: int = 0, # compressed KV length for D5c sink bias
|
||||
sink_bias: torch.Tensor = None, # per-head logit bias for D5c
|
||||
) -> torch.Tensor:
|
||||
"""Production DSV4 attention using CuTeDSL FMHA kernel.
|
||||
|
||||
Args:
|
||||
q: Query tensor (num_heads, seq_len, head_dim) BF16
|
||||
k: Key tensor (seq_len, head_dim) or (num_heads, seq_len, head_dim) BF16
|
||||
v: Value tensor (seq_len, head_dim) or (num_heads, seq_len, head_dim) BF16
|
||||
scale: Softmax scale (default: 1/sqrt(head_dim))
|
||||
swa_len: Sliding window length (None = no mask)
|
||||
is_causal: Apply causal mask
|
||||
n_comp: Compressed KV length for D5c sink bias merge
|
||||
sink_bias: Per-head FP32 logit bias for sink attention
|
||||
|
||||
Returns:
|
||||
Output tensor (num_heads, seq_len, head_dim) BF16
|
||||
"""
|
||||
n_h, T, hd = q.shape
|
||||
N = k.shape[-2] # KV sequence length
|
||||
scale = scale or (1.0 / math.sqrt(hd))
|
||||
use_smem_p = hd > 64
|
||||
apply_swa_mask = swa_len is not None
|
||||
apply_sink_bias = sink_bias is not None
|
||||
|
||||
# Per-head launch: one kernel call per head
|
||||
output = torch.zeros_like(q)
|
||||
|
||||
for h in range(n_h):
|
||||
q_h = q[h:h+1] # (1, T, hd)
|
||||
k_h = k if k.dim() == 2 else k[h] # shared K (MQA) or per-head
|
||||
if k_h.dim() == 2:
|
||||
k_h = k_h.unsqueeze(0) # (1, N, hd)
|
||||
v_h = v if v.dim() == 2 else v[h]
|
||||
if v_h.dim() == 2:
|
||||
v_h = v_h.unsqueeze(0)
|
||||
|
||||
o_h = _attention_single_head(
|
||||
q_h, k_h, v_h, scale=scale,
|
||||
swa_len=swa_len, is_causal=is_causal,
|
||||
n_comp=n_comp, sink_bias=sink_bias[h] if sink_bias is not None else None,
|
||||
use_smem_p=use_smem_p,
|
||||
)
|
||||
output[h] = o_h[0]
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def _attention_single_head(
|
||||
q: torch.Tensor, # (1, T, hd)
|
||||
k: torch.Tensor, # (1, N, hd)
|
||||
v: torch.Tensor, # (1, N, hd)
|
||||
scale: float,
|
||||
swa_len: int = None,
|
||||
is_causal: bool = False,
|
||||
n_comp: int = 0,
|
||||
sink_bias: torch.Tensor = None,
|
||||
use_smem_p: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""Run FMHA for a single head, with Python KV merge for s_k > 128."""
|
||||
_, T, hd = q.shape
|
||||
N = k.shape[1]
|
||||
apply_swa_mask = swa_len is not None
|
||||
apply_sink_bias = sink_bias is not None
|
||||
|
||||
# Reference output
|
||||
qf = q[0].float() # (T, hd)
|
||||
kf = k[0].float() # (N, hd)
|
||||
vf = v[0].float() # (N, hd)
|
||||
attn_max = (qf @ kf.T * scale).max(dim=-1, keepdim=True)[0]
|
||||
attn_exp = torch.exp(qf @ kf.T * scale - attn_max)
|
||||
attn_sum = attn_exp.sum(dim=-1, keepdim=True)
|
||||
ref_norm = (attn_exp / attn_sum) @ vf
|
||||
|
||||
# Kernel execution with Python KV merge for multi-KV-tile
|
||||
# Each segment is s_k=128 (one KV tile)
|
||||
s_k_per_seg = 128
|
||||
n_segments = (N + s_k_per_seg - 1) // s_k_per_seg
|
||||
|
||||
# Use un-normalized output + LSE for Python KV merge
|
||||
compiled, kernel = _get_or_compile_kernel(
|
||||
head_dim=hd, s_k=s_k_per_seg, use_smem_p=use_smem_p,
|
||||
normalize=False, apply_swa_mask=apply_swa_mask,
|
||||
is_causal=is_causal, n_comp=n_comp if n_comp > 0 else None,
|
||||
apply_sink_bias=apply_sink_bias,
|
||||
)
|
||||
pv_n_tile = kernel.pv_n_tile
|
||||
n_pv_tiles = kernel.n_pv_tiles
|
||||
|
||||
# Accumulate across KV segments using log-sum-exp merge
|
||||
o_accum = torch.zeros(T, hd, dtype=torch.float32, device='cuda')
|
||||
lse_accum = torch.full((T, 1), float('-inf'), dtype=torch.float32, device='cuda')
|
||||
|
||||
for seg in range(n_segments):
|
||||
k_start = seg * s_k_per_seg
|
||||
k_end = min(k_start + s_k_per_seg, N)
|
||||
k_seg = k[:, k_start:k_end]
|
||||
v_seg = v[:, k_start:k_end]
|
||||
|
||||
# If last segment is shorter than s_k_per_seg, pad
|
||||
if k_end - k_start < s_k_per_seg:
|
||||
pad_len = s_k_per_seg - (k_end - k_start)
|
||||
k_seg = torch.cat([k_seg, torch.zeros(1, pad_len, hd, dtype=k.dtype, device='cuda')], dim=1)
|
||||
v_seg = torch.cat([v_seg, torch.zeros(1, pad_len, hd, dtype=v.dtype, device='cuda')], dim=1)
|
||||
|
||||
seg_o = torch.zeros(T, hd, dtype=torch.float32, device='cuda')
|
||||
seg_lse = torch.zeros(T, 1, dtype=torch.float32, device='cuda')
|
||||
|
||||
for nt in range(n_pv_tiles):
|
||||
v_start = nt * pv_n_tile
|
||||
v_end = v_start + pv_n_tile
|
||||
v_tile = v_seg[0, :, v_start:v_end].contiguous()
|
||||
v_kernel = v_tile.unsqueeze(-1)
|
||||
c_tile = torch.zeros(T, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
|
||||
lse_tensor = torch.zeros(T, 1, 1, dtype=torch.float32, device='cuda')
|
||||
row_sums_tensor = torch.zeros(T, 1, 1, dtype=torch.float32, device='cuda')
|
||||
|
||||
# Prepare CuTe tensors
|
||||
q_input = q[0].unsqueeze(-1) # (T, hd, 1)
|
||||
k_input = k_seg[0].unsqueeze(-1) # (s_k_per_seg, hd, 1)
|
||||
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
|
||||
|
||||
mQ = ct.from_dlpack(q_input).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q_input))
|
||||
mK = ct.from_dlpack(k_input).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_input))
|
||||
mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel))
|
||||
mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile))
|
||||
mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor))
|
||||
mRS = ct.from_dlpack(row_sums_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(row_sums_tensor))
|
||||
|
||||
compiled(mQ, mK, mV, mC, stream, lse=mLSE, row_sums=mRS)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
seg_o[:, v_start:v_end] = c_tile[:, :, 0].float()
|
||||
if nt == 0:
|
||||
seg_lse[:, 0] = lse_tensor[:, 0, 0].float()
|
||||
|
||||
# Merge with accumulator using log-sum-exp
|
||||
e_old = torch.exp(lse_accum)
|
||||
e_new = torch.exp(seg_lse)
|
||||
e_sum = e_old + e_new
|
||||
o_accum = (e_old * o_accum + e_new * seg_o) / e_sum
|
||||
lse_accum = torch.log(e_sum)
|
||||
|
||||
# Normalize
|
||||
output = o_accum.to(torch.bfloat16).unsqueeze(0) # (1, T, hd)
|
||||
return output
|
||||
80
tests/unit/test_production.py
Normal file
80
tests/unit/test_production.py
Normal file
@@ -0,0 +1,80 @@
|
||||
"""Test production DSV4 attention wrapper."""
|
||||
import torch
|
||||
import math
|
||||
from dsv4.kernels.attention.production import dsv4_attention
|
||||
|
||||
|
||||
def test_production_basic():
|
||||
"""Test basic single-head attention."""
|
||||
torch.manual_seed(42)
|
||||
hd = 64
|
||||
n_h = 1
|
||||
T = 128
|
||||
N = 128
|
||||
|
||||
q = torch.randn(n_h, T, hd, dtype=torch.bfloat16, device='cuda')
|
||||
k = torch.randn(N, hd, dtype=torch.bfloat16, device='cuda')
|
||||
v = torch.randn(N, hd, dtype=torch.bfloat16, device='cuda')
|
||||
|
||||
# PyTorch reference
|
||||
qf = q[0].float()
|
||||
kf = k.float()
|
||||
vf = v.float()
|
||||
scale = 1.0 / math.sqrt(hd)
|
||||
attn = qf @ kf.T * scale
|
||||
attn_max = attn.max(dim=-1, keepdim=True)[0]
|
||||
attn_exp = torch.exp(attn - attn_max)
|
||||
attn_sum = attn_exp.sum(dim=-1, keepdim=True)
|
||||
ref = ((attn_exp / attn_sum) @ vf).unsqueeze(0)
|
||||
|
||||
out = dsv4_attention(q, k, v)
|
||||
|
||||
cos = torch.nn.functional.cosine_similarity(
|
||||
out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)
|
||||
).item()
|
||||
status = "PASS" if cos >= 0.99 else "FAIL"
|
||||
print(f" hd={hd}, n_h={n_h}, N={N}: cos {cos:.6f} {status}")
|
||||
|
||||
|
||||
def test_production_multi_head():
|
||||
"""Test multi-head attention (per-head launch)."""
|
||||
torch.manual_seed(42)
|
||||
hd = 64
|
||||
n_h = 4
|
||||
T = 128
|
||||
N = 256
|
||||
|
||||
q = torch.randn(n_h, T, hd, dtype=torch.bfloat16, device='cuda')
|
||||
k = torch.randn(n_h, N, hd, dtype=torch.bfloat16, device='cuda')
|
||||
v = torch.randn(n_h, N, hd, dtype=torch.bfloat16, device='cuda')
|
||||
|
||||
# PyTorch reference
|
||||
scale = 1.0 / math.sqrt(hd)
|
||||
ref = torch.zeros_like(q)
|
||||
for h in range(n_h):
|
||||
qf = q[h].float()
|
||||
kf = k[h].float()
|
||||
vf = v[h].float()
|
||||
attn = qf @ kf.T * scale
|
||||
attn_max = attn.max(dim=-1, keepdim=True)[0]
|
||||
attn_exp = torch.exp(attn - attn_max)
|
||||
attn_sum = attn_exp.sum(dim=-1, keepdim=True)
|
||||
ref[h] = ((attn_exp / attn_sum) @ vf).bfloat16()
|
||||
|
||||
out = dsv4_attention(q, k, v)
|
||||
|
||||
cos = torch.nn.functional.cosine_similarity(
|
||||
out.flatten().unsqueeze(0), ref.float().flatten().unsqueeze(0)
|
||||
).item()
|
||||
status = "PASS" if cos >= 0.99 else "FAIL"
|
||||
print(f" hd={hd}, n_h={n_h}, N={N}: cos {cos:.6f} {status}")
|
||||
|
||||
|
||||
def test():
|
||||
print("=== Production DSV4 Attention Wrapper ===\n")
|
||||
test_production_basic()
|
||||
test_production_multi_head()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test()
|
||||
Reference in New Issue
Block a user