vrr/kvcache-ai-ktransformers
2
0
Fork 0
mirror of https://github.com/kvcache-ai/ktransformers.git synced 2025-09-14 09:09:42 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Mock triton mla due to precision issue

Browse source
This commit is contained in:
Azure 2025-02-16 06:03:12 +00:00
parent c5f036e8a4
commit ff6b265e53
1 changed files with 3 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

4
ktransformers/operators/attention.py
View file

@ -43,11 +43,13 @@ class KDeepseekV2Attention(BaseInjectedModule, DeepseekV2Attention):
orig_module: nn.Module,
device: str = "cuda",
chunck_size: int = 1000,
use_triton: bool = False,
**kwargs):
BaseInjectedModule.__init__(self, key, gguf_loader, config, orig_module, device, **kwargs)
self.orig_module.__init__(orig_module.config,
orig_module.layer_idx)
self.chunck_size = chunck_size # TODO, generate chunck_size automatically.
self.use_triton = use_triton
def get_absorbed(self) -> Tuple[torch.Tensor, torch.Tensor]:
if not (hasattr(self, 'q_absorb') and hasattr(self, 'out_absorb')):
@ -401,7 +403,7 @@ class KDeepseekV2Attention(BaseInjectedModule, DeepseekV2Attention):
cache_position: Optional[torch.LongTensor] = None,
**kwargs,
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
if os.name == 'nt':
if not self.use_triton: # os.name == 'nt'
return self.forward_windows(
hidden_states,
attention_mask,