[Perf][DSv4.1] Shard the Engram wkv projection across TP ranks (#58678)
Signed-off-by: Shuolei Wang <shuoleiwang123@gmail.com>
Shuolei Wang committed
Sep 26, 2026 at 09:15 UTC
4bb804cc5bf0de81ce8344238b5b4522630874fd
2 files changed
+7
-3
tests/kernels/test_engram.py
+1
-1
@@ -827,7 +827,7 @@ def test_engram_constructor_honors_offload(monkeypatch, backend, cpu_offload):
827
monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_world_size", lambda: 1)
828
monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: 0)
829
monkeypatch.setattr(
830
- engram_ops, "ReplicatedLinear", lambda *a, **k: torch.nn.Identity()
830
+ engram_ops, "ColumnParallelLinear", lambda *a, **k: torch.nn.Identity()
831
)
832
stream = torch.cuda.Stream()
833
kwargs = {"prefetch_stream": stream} if backend == "nvidia" else {}
vllm/models/deepseek_v41/common/engram.py
+6
-2
@@ -47,7 +47,7 @@ from vllm.distributed import (
47
tensor_model_parallel_all_gather,
48
)
49
from vllm.logger import init_logger
50
-from vllm.model_executor.layers.linear import ReplicatedLinear
50
+from vllm.model_executor.layers.linear import ColumnParallelLinear
51
from vllm.model_executor.layers.quantization import QuantizationConfig
52
from vllm.model_executor.utils import set_weight_attrs
53
from vllm.triton_utils import tl, triton
@@ -923,13 +923,17 @@ class Engram(nn.Module):
923
# suffix rule.
924
self.embed_tokens = self._create_embedding(layout, layer_hash_index)
925
n_hash_cols = (layout.max_ngram_size - 1) * layout.n_heads
926
- self.wkv = ReplicatedLinear(
926
+ # Without sequence parallelism every TP rank holds every token, so
927
+ # shard the output columns instead of replicating the projection.
928
+ self.wkv = ColumnParallelLinear(
929
n_hash_cols * layout.head_dim,
930
self.dim * (self.hc_mult + 1),
931
bias=False,
932
+ gather_output=True,
933
quant_config=quant_config,
934
return_bias=False,
935
prefix=f"{prefix}.wkv",
936
+ disable_tp=use_sequence_parallel,
937
)
938
self.q_weight = nn.Parameter(
939
torch.empty(self.hc_mult, self.dim, dtype=torch.bfloat16),