coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-28 14:57:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

neuron: the grouped wv projection takes the bf16 context through the NKI qdq kernel instead of an in-graph fp8 encode (TTFT 1009 -> 942 ms, step 437 -> 414 ms, tokens identical)

Browse Source 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
whn09 2026-10-28 14:32:47 +08:00
parent deb8c5a958
commit d27e8951db
1 changed files with 17 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
deepinfer/model_executor/models/v4_modules/attn.py
View File

@ -200,6 +200,23 @@ class V4AttentionBase(nn.Module):
def _wv_and_o_proj(self, context: torch.Tensor | tuple[torch.Tensor, torch.Tensor], positions: torch.Tensor) -> torch.Tensor:
if not self.use_fused_norm_rope_attn_rope_cast:
self.rotary_emb(positions, context[..., -self.qk_rope_head_dim :], conjugate=True)
if (
PLATFORM.is_neuron
and not isinstance(context, tuple)
and not self.use_fp8_einsum_output
and self.einsum_out_clamp_value is None
and nki_gemm.has_kernel(self.wv_weight)
):
# The NKI kernel quantizes the bf16 context itself (1x32, as per_token_cast): no fp8 encode in the graph.
ctx = context.view(context.shape[0], self.n_wv_group, self.wv_group_dim)
wv = torch.stack([nki_gemm.linear_qdq(ctx[:, h], self.wv_weight, group=h) for h in range(self.n_wv_group)], dim=1)
out = self.out_proj(wv.flatten(1, 2))
tp_group = getattr(self, "tp_group", None)
if tp_group is not None:
import torch.distributed._functional_collectives as funcol
out = funcol.wait_tensor(funcol.all_reduce(out.float(), "sum", tp_group)).to(out.dtype)
return out
if isinstance(context, tuple):
context_fp8, context_sf = context
gran = 32