TRF020

MLA models must isolate the KV LoRA expansion (kv_b_proj) in a dedicated method that forward() calls.

Default	Enabled
Scope	All models
Source	mlinter/trf020.py
Show in terminal	mlinter --rule TRF020

What it does

In model directories whose configuration declares kv_lora_rank (Multi-head Latent Attention), checks the attention class owning the KV LoRA expansion projection (kv_b_proj, or any nn.Linear(config.kv_lora_rank, ...)): the expansion must live in a dedicated method (e.g. expand_kv) that forward() calls, not inline. In modular files a method inherited from an imported base counts.

Why is this bad?

External backends (vLLM/SGLang) override the expansion to consume the compressed KV cache directly. Inlined in forward() there is nothing to override, so the backend must materialize the full cache – losing the memory savings MLA exists for.

Example

+    def expand_kv(self, k_nope, k_pe):
+        key_shape = (*k_nope.shape[:-1], -1, self.qk_nope_head_dim + self.v_head_dim)
+        k_nope = self.kv_b_proj(k_nope).view(key_shape).transpose(1, 2)
+        k_nope, value_states = torch.split(k_nope, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+        k_pe = k_pe.expand(*k_nope.shape[:-1], -1)
+        key_states = torch.cat((k_nope, k_pe), dim=-1)
+        return key_states, value_states
+
     def forward(self, hidden_states, ...):
         ...
-        k_nope = self.kv_b_proj(k_pass).view(key_shape).transpose(1, 2)
-        k_nope, value_states = torch.split(k_nope, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
-        k_pe = k_rot.expand(*k_nope.shape[:-1], -1)
-        key_states = torch.cat((k_nope, k_pe), dim=-1)
+        key_states, value_states = self.expand_kv(k_pass, k_rot)

Suppressing this rule

Add a # trf-ignore: TRF020 comment on the flagged line or the line directly above it. See Suppressing rules for whole-file directives and when a suppression is the wrong answer.