A high-throughput and memory-efficient inference and serving engine for LLMs
[XPU] Enable sequence parallel support for XPU (#38608)
Signed-off-by: chaojun-zhang <chaojun.zhang@intel.com> Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com> Signed-off-by: Chaojun,Zhang <chaojun.zhang@intel.com>
Chaojun Zhang committed
Jun 15, 2026 at 10:26 UTC
2725c84aaed1dd27085655f224b3b3c4ff2e8f1e
5 files changed
+128
-19
tests/compile/conftest.py
+23
@@ -24,6 +24,7 @@ def mock_cuda_platform():
24
def _mock_platform(is_cuda: bool = True, capability: tuple[int, int] | None = None):
25
mock_platform = MagicMock()
26
mock_platform.is_cuda.return_value = is_cuda
27
+ mock_platform.is_xpu.return_value = False
28
device_capability = (
29
DeviceCapability(*capability) if capability is not None else None
30
)
@@ -46,3 +47,25 @@ def mock_cuda_platform():
47
yield mock_platform
48
49
return _mock_platform
50
+
51
+
52
+@pytest.fixture
53
+def mock_xpu_platform():
54
+ """
55
+ Fixture that returns a factory for creating mocked XPU platforms.
56
+
57
+ Usage:
58
+ def test_something(mock_xpu_platform):
59
+ with mock_xpu_platform():
60
+ # test code
61
+ """
62
+
63
+ @contextmanager
64
+ def _mock_platform():
65
+ mock_platform = MagicMock()
66
+ mock_platform.is_cuda.return_value = False
67
+ mock_platform.is_xpu.return_value = True
68
+ with patch("vllm.platforms.current_platform", mock_platform):
69
+ yield mock_platform
70
+
71
+ return _mock_platform
tests/compile/test_sequence_parallelism_threshold.py
+82
@@ -108,3 +108,85 @@ class TestGetSequenceParallelismThreshold:
108
element_size=2,
109
)
110
assert result is not None
111
+
112
+
113
+# XPU-specific constants (must match sequence_parallelism.py values)
114
+_XPU_MIN_HIDDEN_SIZE = 4096
115
+_XPU_MIN_PER_GPU_SIZE_MB = 8.0
116
+
117
+
118
+class TestGetSequenceParallelismThresholdXPU:
119
+ """Tests for get_sequence_parallelism_threshold on XPU platform."""
120
+
121
+ def test_xpu_small_hidden_size_returns_none(self, mock_xpu_platform):
122
+ """XPU with hidden_size below threshold should return None."""
123
+ with mock_xpu_platform():
124
+ result = get_sequence_parallelism_threshold(
125
+ hidden_size=_XPU_MIN_HIDDEN_SIZE - 1,
126
+ tp_size=2,
127
+ element_size=2,
128
+ )
129
+ assert result is None
130
+
131
+ def test_xpu_large_model_returns_threshold(self, mock_xpu_platform):
132
+ """XPU with hidden_size >= threshold should return calculated value."""
133
+ with mock_xpu_platform():
134
+ hidden_size = _XPU_MIN_HIDDEN_SIZE
135
+ tp_size = 2
136
+ element_size = 2
137
+ result = get_sequence_parallelism_threshold(
138
+ hidden_size=hidden_size,
139
+ tp_size=tp_size,
140
+ element_size=element_size,
141
+ )
142
+ # (8 * 2 * 1024 * 1024) // (4096 * 2) = 2048
143
+ MiB = 1024 * 1024
144
+ expected = int(
145
+ (_XPU_MIN_PER_GPU_SIZE_MB * tp_size * MiB) // (hidden_size * element_size)
146
+ )
147
+ assert result == expected
148
+ assert result == 2048
149
+
150
+ @pytest.mark.parametrize(
151
+ "hidden_size,tp_size,element_size,expected",
152
+ [
153
+ # (8 * 1 * 1024 * 1024) // (4096 * 2) = 1024
154
+ (4096, 1, 2, 1024),
155
+ # (8 * 4 * 1024 * 1024) // (4096 * 2) = 4096
156
+ (4096, 4, 2, 4096),
157
+ # (8 * 2 * 1024 * 1024) // (8192 * 2) = 1024
158
+ (8192, 2, 2, 1024),
159
+ # (8 * 2 * 1024 * 1024) // (4096 * 4) = 1024
160
+ (4096, 2, 4, 1024),
161
+ ],
162
+ )
163
+ def test_xpu_threshold_calculation_variations(
164
+ self, mock_xpu_platform, hidden_size, tp_size, element_size, expected
165
+ ):
166
+ """Test XPU threshold calculation with various parameter combinations."""
167
+ with mock_xpu_platform():
168
+ result = get_sequence_parallelism_threshold(
169
+ hidden_size=hidden_size,
170
+ tp_size=tp_size,
171
+ element_size=element_size,
172
+ )
173
+ assert result == expected
174
+
175
+ def test_xpu_hidden_size_boundary(self, mock_xpu_platform):
176
+ """Test behavior at the exact XPU hidden_size boundary."""
177
+ with mock_xpu_platform():
178
+ # Just below threshold
179
+ result = get_sequence_parallelism_threshold(
180
+ hidden_size=_XPU_MIN_HIDDEN_SIZE - 1,
181
+ tp_size=2,
182
+ element_size=2,
183
+ )
184
+ assert result is None
185
+
186
+ # Exactly at threshold
187
+ result = get_sequence_parallelism_threshold(
188
+ hidden_size=_XPU_MIN_HIDDEN_SIZE,
189
+ tp_size=2,
190
+ element_size=2,
191
+ )
192
+ assert result is not None
vllm/compilation/passes/fusion/sequence_parallelism.py
+20
-17
@@ -72,24 +72,27 @@ def get_sequence_parallelism_threshold(
72
"""
73
from vllm.platforms import current_platform
74
75
- if not current_platform.is_cuda():
76
- return None
77
-
78
- capability = current_platform.get_device_capability()
79
- if capability is None:
80
- return None
81
-
82
- # Collapse Blackwell variants (sm100/sm103/...) into one policy bucket.
83
- if current_platform.is_device_capability_family(100):
84
- device_capability = 100
75
+ if current_platform.is_xpu():
76
+ min_hidden_size = 4096
77
+ min_per_gpu_size_mb = 8.0
78
+ elif current_platform.is_cuda():
79
+ capability = current_platform.get_device_capability()
80
+ if capability is None:
81
+ return None
82
+
83
+ # Collapse Blackwell variants (sm100/sm103/...) into one policy bucket.
84
+ if current_platform.is_device_capability_family(100):
85
+ device_capability = 100
86
+ else:
87
+ device_capability = capability.to_int()
88
+
89
+ # Check if device has configured thresholds
90
+ _hidden = SP_MIN_HIDDEN_SIZE.get(device_capability)
91
+ _gpu_mb = SP_MIN_PER_GPU_SIZE_MB.get(device_capability)
92
+ if _hidden is None or _gpu_mb is None:
93
+ return None
94
+ min_hidden_size, min_per_gpu_size_mb = _hidden, _gpu_mb
95
else:
86
- device_capability = capability.to_int()
87
-
88
- # Check if device has configured thresholds
89
- min_hidden_size = SP_MIN_HIDDEN_SIZE.get(device_capability)
90
- min_per_gpu_size_mb = SP_MIN_PER_GPU_SIZE_MB.get(device_capability)
91
-
92
- if min_hidden_size is None or min_per_gpu_size_mb is None:
96
return None
97
98
# Only apply sequence parallelism for models meeting the size threshold
vllm/compilation/passes/pass_manager.py
+3
-1
@@ -29,6 +29,9 @@ if rocm_aiter_ops.is_enabled():
29
RocmAiterTritonAddRMSNormPadFusionPass,
30
)
31
32
+if current_platform.is_cuda_alike() or current_platform.is_xpu():
33
+ from .fusion.sequence_parallelism import SequenceParallelismPass
34
+
35
if current_platform.is_cuda_alike():
36
from .fusion.act_quant_fusion import ActivationQuantFusionPass
37
from .fusion.attn_quant_fusion import AttnQuantFusionPass
@@ -37,7 +40,6 @@ if current_platform.is_cuda_alike():
40
from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass
41
from .fusion.rms_quant_fusion import RMSNormQuantFusionPass
42
from .fusion.rope_kvcache_fusion import RopeKVCacheFusionPass
40
- from .fusion.sequence_parallelism import SequenceParallelismPass
43
from .utility.scatter_split_replace import ScatterSplitReplacementPass
44
from .utility.split_coalescing import SplitCoalescingPass
45
vllm/platforms/xpu.py
-1
@@ -208,7 +208,6 @@ class XPUPlatform(Platform):
208
209
pass_config = compilation_config.pass_config
210
fusion_passes_to_disable = {
211
- "enable_sp": "Sequence parallelism",
211
"fuse_gemm_comms": "Async TP",
212
"fuse_allreduce_rms": "AllReduce + RMSNorm fusion",
213
"fuse_attn_quant": "Attention + quant fusion",