diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index 8a0676db..989e0997 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -109,7 +109,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): + 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() @@ -152,10 +152,13 @@ class FmhaKernel: # else: sink_bias is already a CuTe tensor (caller must pass via ct.from_dlpack) # Grid: (M_tiles, 1, batch) where M = n_h * T packed into M dimension # For single-head (n_h=1): grid=(1,1,1) — backward compatible - 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).launch(grid=(1,1,self.batch_size),block=[self.threads_per_cta,1,1],stream=stream) + if const_expr(row_sums is None): + row_sums = cute.make_tensor(lse.iterator, cute.make_layout((1,), stride=(0,))) + + 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): + 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: @@ -583,6 +586,8 @@ class FmhaKernel: _ln2 = Float32(0.6931471805599453) # ln(2) lse_val = cute.math.log(row_sum, fastmath=True) + _row_max_safe * _ln2 mLSE[sfw_idx, Int32(0), Int32(0)] = lse_val + # Also output row_sum for external normalization (D5c) + mRowSums[sfw_idx, Int32(0), Int32(0)] = row_sum tmem.relinquish_alloc_permit() tmem.free(tmem_ptr) diff --git a/tests/unit/test_d5c_fused.py b/tests/unit/test_d5c_fused.py index 42ecf58e..3fc549a7 100644 --- a/tests/unit/test_d5c_fused.py +++ b/tests/unit/test_d5c_fused.py @@ -162,6 +162,7 @@ def test_d5c_combined(): # Allocate output c_out = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') lse_out = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') + row_sum_out = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') # Prepare CuTe tensors def to_cute(t): @@ -172,31 +173,32 @@ def test_d5c_combined(): mV = to_cute(v_combined.unsqueeze(-1).contiguous()) mC = to_cute(c_out) mLSE = to_cute(lse_out) + mRowSums = to_cute(row_sum_out) # Compile print('Compiling D5c kernel (combined KV + sink bias)...', flush=True) mSinkBias = to_cute(attn_sink) compiled = cute.compile( kernel, mQ, mK, mV, mC, stream, mLSE, - swa_len=swa_len, sink_bias=mSinkBias, + swa_len=swa_len, sink_bias=mSinkBias, row_sums=mRowSums, ) # Run print('Running D5c kernel...', flush=True) compiled( mQ, mK, mV, mC, stream, mLSE, - swa_len=swa_len, sink_bias=mSinkBias, + swa_len=swa_len, sink_bias=mSinkBias, row_sums=mRowSums, ) torch.cuda.synchronize() # Check results - # Kernel outputs UN-NORMALIZED O (normalize=False). Normalize using per-row LSE. - # O_norm[i] = O_unnorm[i] * exp(-lse[i]) + # Kernel outputs UN-NORMALIZED O (normalize=False). Normalize using per-row row_sum. + # O_norm[i] = O_unnorm[i] / row_sum[i] o_kernel_unnorm = c_out[:, :, 0].float() # (m, hd) - lse_kernel = lse_out[:, 0, 0].float() # (m,) + row_sums = row_sum_out[:, 0, 0].float() # (m,) - # Normalize each row - o_kernel = o_kernel_unnorm * (-lse_kernel.unsqueeze(1)).exp() # (m, hd) + # Normalize each row by its row_sum + o_kernel = o_kernel_unnorm / row_sums.unsqueeze(1).clamp(min=1e-30) cos = torch.nn.functional.cosine_similarity( o_kernel.flatten().unsqueeze(0), @@ -210,7 +212,8 @@ def test_d5c_combined(): if cos < 0.99: print(f' kernel[0,:4]={o_kernel[0,:4].tolist()}') print(f' ref[0,:4]={ref_combined[0,:4].tolist()}') - print(f' LSE range: {lse_kernel.min().item():.4f} to {lse_kernel.max().item():.4f}') + print(f' row_sum range: {row_sums.min().item():.4f} to {row_sums.max().item():.4f}') + print(f' LSE range: {lse_out[:,0,0].min().item():.4f} to {lse_out[:,0,0].max().item():.4f}') def test_d5c_with_causal(): @@ -258,6 +261,7 @@ def test_d5c_with_causal(): c_out = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') lse_out = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') + row_sum_out = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') def to_cute(t): return ct.from_dlpack(t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(t)) @@ -267,22 +271,21 @@ def test_d5c_with_causal(): mV = to_cute(v_combined.unsqueeze(-1).contiguous()) mC = to_cute(c_out) mLSE = to_cute(lse_out) - - print('Compiling D5c kernel (causal + sink bias)...', flush=True) + mRowSums = to_cute(row_sum_out) mSinkBias = to_cute(attn_sink) compiled = cute.compile( kernel, mQ, mK, mV, mC, stream, mLSE, - swa_len=swa_len, sink_bias=mSinkBias, + swa_len=swa_len, sink_bias=mSinkBias, row_sums=mRowSums, ) compiled( mQ, mK, mV, mC, stream, mLSE, - swa_len=swa_len, sink_bias=mSinkBias, + swa_len=swa_len, sink_bias=mSinkBias, row_sums=mRowSums, ) torch.cuda.synchronize() o_kernel_unnorm = c_out[:, :, 0].float() - lse_kernel = lse_out[:, 0, 0].float() - o_kernel = o_kernel_unnorm * (-lse_kernel.unsqueeze(1)).exp() + row_sums = row_sum_out[:, 0, 0].float() + o_kernel = o_kernel_unnorm / row_sums.unsqueeze(1).clamp(min=1e-30) cos = torch.nn.functional.cosine_similarity( o_kernel.flatten().unsqueeze(0),