A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 16:00 UTC
[NIXL][BUG] Fix Triton heterogeneous TP (#37940)
Signed-off-by: Yifan <yzong@redhat.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>
yzong-rh committed
Apr 1, 2026 at 11:23 UTC
dc0428ebb879af3224e5d9ef5ab93cb04d8d1b27
5 files changed
+62
-22
tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+7
@@ -23,6 +23,10 @@ hybrid_ssm_configs=(
23
# TODO: (NickLucche) Address async scheduling issue with TP>1 separately as this may impact other models.
24
"ENABLE_HMA_FLAG=1 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code,--no-async-scheduling"
25
)
26
+sw_attn_configs=(
27
+ "ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=google/gemma-3-4b-it PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192"
28
+ "ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=google/gemma-3-4b-it PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=1 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192"
29
+)
30
31
# Select config array based on DP_EP env var
32
if [[ -n "${DP_EP:-}" ]]; then
@@ -31,6 +35,9 @@ if [[ -n "${DP_EP:-}" ]]; then
35
elif [[ -n "${HYBRID_SSM:-}" ]]; then
36
configs=("${hybrid_ssm_configs[@]}")
37
echo "HYBRID_SSM is set, using hybrid_ssm_configs."
38
+elif [[ -n "${SW_ATTN:-}" ]]; then
39
+ configs=("${sw_attn_configs[@]}")
40
+ echo "SW_ATTN is set, using sw_attn_configs."
41
else
42
configs=("${tp_configs[@]}")
43
fi
tests/v1/kv_connector/unit/test_nixl_connector.py
+17
-15
@@ -1586,10 +1586,18 @@ def test_register_kv_caches(
1586
expected_base_addrs: list[int]
1587
expected_num_entries: int
1588
kv_caches: dict[str, torch.Tensor]
1589
- assert str(enable_cross_layers).lower() != "true" or (
1590
- (attn_backend not in ("FLASH_ATTN", "FLASHINFER"))
1591
- or connector.prefer_cross_layer_blocks
1589
+ if str(enable_cross_layers).lower() == "true":
1590
+ assert connector.prefer_cross_layer_blocks == (
1591
+ attn_backend in ("FLASH_ATTN", "FLASHINFER", "TRITON_ATTN")
1592
+ )
1593
+ else:
1594
+ assert not connector.prefer_cross_layer_blocks
1595
+
1596
+ test_shape = backend_cls.get_kv_cache_shape(
1597
+ num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
1598
)
1599
+ is_blocks_first = len(test_shape) == 5 and test_shape[0] == 1
1600
+
1601
if connector.prefer_cross_layer_blocks:
1602
with set_current_vllm_config(vllm_config):
1603
_, cross_layers_kv_cache, _ = (
@@ -1619,7 +1627,7 @@ def test_register_kv_caches(
1627
]
1628
expected_num_entries = 1
1629
1622
- expected_blocks_count = 8
1630
+ expected_blocks_count = num_blocks * (2 if is_blocks_first else 1)
1631
1632
kv_caches = {"all-layers": cross_layers_kv_cache}
1633
else:
@@ -1639,12 +1647,6 @@ def test_register_kv_caches(
1647
}
1648
1649
# Store tensor info for validation
1642
-
1643
- test_shape = backend_cls.get_kv_cache_shape(
1644
- num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
1645
- )
1646
- is_blocks_first = len(test_shape) == 5 and test_shape[0] == 1
1647
-
1650
if is_blocks_first:
1651
expected_tensor_size = (
1652
shared_tensor.element_size() * shared_tensor.numel()
@@ -1696,13 +1698,13 @@ def test_register_kv_caches(
1698
1699
if connector.prefer_cross_layer_blocks:
1700
num_blocks = 8
1699
- expected_block_len = expected_tensor_size // num_blocks
1701
else:
1702
num_blocks = kv_cache_config.num_blocks
1702
- if is_blocks_first:
1703
- expected_block_len = expected_tensor_size // num_blocks // 2
1704
- else:
1705
- expected_block_len = expected_tensor_size // num_blocks
1703
+
1704
+ if is_blocks_first:
1705
+ expected_block_len = expected_tensor_size // num_blocks // 2
1706
+ else:
1707
+ expected_block_len = expected_tensor_size // num_blocks
1708
1709
for i, block_entry in enumerate(blocks_data):
1710
block_start_addr, block_len, tp_rank = block_entry
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+16
@@ -330,6 +330,7 @@ class NixlConnector(KVConnectorBase_V1, SupportsHMA):
330
if backend.get_name() not in (
331
"FLASH_ATTN",
332
"FLASHINFER",
333
+ "TRITON_ATTN",
334
):
335
return False
336
@@ -2118,6 +2119,21 @@ class NixlConnectorWorker:
2119
"setting 'enable_permute_local_kv'=True in --kv-transfer-config."
2120
)
2121
2122
+ # Heterogeneous TP requires head-splitting, which only works with
2123
+ # HND layout. MLA and replicated-KV cases don't split on heads.
2124
+ # Mamba doesn't support heterogeneous TP.
2125
+ if (
2126
+ abs(tp_ratio) != 1
2127
+ and not self.use_mla
2128
+ and not self.kv_topo.is_kv_replicated(remote_engine_id)
2129
+ and kv_cache_layout != "HND"
2130
+ and not self.enable_permute_local_kv
2131
+ ):
2132
+ raise RuntimeError(
2133
+ "Heterogeneous TP head-dimension splitting requires contiguous heads. "
2134
+ "Use HND layout on the prefill side."
2135
+ )
2136
+
2137
# Block len can only vary across layers when using MLA.
2138
remote_block_len = nixl_agent_meta.block_lens[0]
2139
if self.use_mla or self.kv_topo.is_kv_replicated(remote_engine_id):
vllm/v1/attention/backends/triton_attn.py
+13
-4
@@ -29,6 +29,7 @@ from vllm.v1.attention.backend import (
29
CommonAttentionMetadata,
30
MultipleOf,
31
)
32
+from vllm.v1.attention.backends.utils import get_kv_cache_layout
33
from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd
34
from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
35
triton_reshape_and_cache_flash,
@@ -309,12 +310,20 @@ class TritonAttentionBackend(AttentionBackend):
310
) -> tuple[int, ...]:
311
# `stride_order` indicates the permutation that gets
312
# us from `get_kv_cache_shape` to the actual memory layout we want.
312
- if include_num_layers_dimension:
313
+ cache_layout = get_kv_cache_layout()
314
+ if cache_layout == "NHD" and include_num_layers_dimension:
315
# (num_blocks, num_layers, 2, block_size, num_kv_heads, head_size)
316
return (1, 0, 2, 3, 4, 5)
315
-
316
- # (num_blocks, 2, block_size, num_kv_heads, head_size)
317
- return (0, 1, 2, 3, 4)
317
+ elif cache_layout == "NHD":
318
+ stride_order = (0, 1, 2, 3, 4)
319
+ elif cache_layout == "HND" and include_num_layers_dimension:
320
+ # (num_blocks, 2, num_kv_heads, num_layers, block_size, head_size)
321
+ return (1, 2, 4, 0, 3, 5)
322
+ elif cache_layout == "HND":
323
+ stride_order = (0, 1, 3, 2, 4)
324
+ else:
325
+ raise ValueError(f"Unknown cache layout: {cache_layout}")
326
+ return stride_order
327
328
@staticmethod
329
def use_cascade_attention(*args, **kwargs) -> bool:
vllm/v1/attention/ops/triton_reshape_and_cache_flash.py
+9
-3
@@ -70,9 +70,15 @@ def reshape_and_cache_kernel_flash(
70
+ (cur_dim % x)
71
)
72
else:
73
- tgt_base = block_idx * block_stride + block_offset * page_stride
74
- tgt_idx_k = tgt_base + tile_pos
75
- tgt_idx_v = tgt_base + tile_pos
73
+ cur_head = tile_pos // head_size
74
+ cur_dim = tile_pos % head_size
75
+ tgt_idx_k = (
76
+ block_idx * block_stride
77
+ + block_offset * page_stride
78
+ + cur_head * head_stride
79
+ + cur_dim
80
+ )
81
+ tgt_idx_v = tgt_idx_k
82
83
# [TILE_SIZE]
84
key_load = tl.load(