Clean up wrapper in flashinfer backend (#2638)

python/sglang/srt/models/llama.py

@@ -516,6 +516,17 @@ class LlamaForCausalLM(nn.Module):
             )
             return None
 
+    def get_embed_and_head(self):
+        return self.model.embed_tokens.weight, self.lm_head.weight
+
+    def set_embed_and_head(self, embed, head):
+        del self.model.embed_tokens.weight
+        del self.lm_head.weight
+        self.model.embed_tokens.weight = embed
+        self.lm_head.weight = head
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+
 
 class Phi3ForCausalLM(LlamaForCausalLM):
     pass