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 8, 16:30 UTC
[Release 2.10] Update to Torch 2.10 - final release (#30525)
Andrey Talman committed
Feb 8, 2026 at 16:51 UTC
f97ca671766c5201404e9fc812e35bf2c4e95a01
17 files changed
+130
-78
.buildkite/image_build/image_build.yaml
+2
-1
@@ -3,6 +3,7 @@ steps:
3
- label: ":docker: Build image"
4
key: image-build
5
depends_on: []
6
+ timeout_in_minutes: 600
7
commands:
8
- if [[ "$BUILDKITE_BRANCH" != "main" ]]; then .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $VLLM_USE_PRECOMPILED $VLLM_MERGE_BASE_COMMIT $IMAGE_TAG; fi
9
- if [[ "$BUILDKITE_BRANCH" == "main" ]]; then .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $VLLM_USE_PRECOMPILED $VLLM_MERGE_BASE_COMMIT $IMAGE_TAG $IMAGE_TAG_LATEST; fi
@@ -41,7 +42,7 @@ steps:
42
limit: 2
43
- exit_status: -10 # Agent was lost
44
limit: 2
44
-
45
+
46
- label: ":docker: Build CPU arm64 image"
47
key: cpu-arm64-image-build
48
depends_on: []
CMakeLists.txt
+5
-5
@@ -56,8 +56,8 @@ endif()
56
# requirements.txt files and should be kept consistent. The ROCm torch
57
# versions are derived from docker/Dockerfile.rocm
58
#
59
-set(TORCH_SUPPORTED_VERSION_CUDA "2.9.1")
60
-set(TORCH_SUPPORTED_VERSION_ROCM "2.9.1")
59
+set(TORCH_SUPPORTED_VERSION_CUDA "2.10.0")
60
+set(TORCH_SUPPORTED_VERSION_ROCM "2.10.0")
61
62
#
63
# Try to find python package with an executable that exactly matches
@@ -433,7 +433,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
433
list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_BF16_KERNEL_SRC})
434
endif()
435
436
- if (MARLIN_SM75_ARCHS)
436
+ if (MARLIN_SM75_ARCHS)
437
file(GLOB MARLIN_TEMPLATE_SM75_KERNEL_SRC "csrc/quantization/marlin/sm75_kernel_*.cu")
438
set_gencode_flags_for_srcs(
439
SRCS "${MARLIN_TEMPLATE_SM75_KERNEL_SRC}"
@@ -445,7 +445,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
445
list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_SM75_KERNEL_SRC})
446
endif()
447
448
- if (MARLIN_FP8_ARCHS)
448
+ if (MARLIN_FP8_ARCHS)
449
file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/quantization/marlin/sm89_kernel_*.cu")
450
set_gencode_flags_for_srcs(
451
SRCS "${MARLIN_TEMPLATE_FP8_KERNEL_SRC}"
@@ -1042,7 +1042,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
1042
list(APPEND VLLM_MOE_EXT_SRC ${MARLIN_MOE_SRC})
1043
endif()
1044
1045
- if (MARLIN_MOE_SM75_ARCHS)
1045
+ if (MARLIN_MOE_SM75_ARCHS)
1046
file(GLOB MARLIN_MOE_SM75_SRC "csrc/moe/marlin_moe_wna16/sm75_kernel_*.cu")
1047
set_gencode_flags_for_srcs(
1048
SRCS "${MARLIN_MOE_SM75_SRC}"
cmake/external_projects/triton_kernels.cmake
+4
-4
@@ -1,9 +1,9 @@
1
# Install OpenAI triton_kernels from https://github.com/triton-lang/triton/tree/main/python/triton_kernels
2
3
-set(DEFAULT_TRITON_KERNELS_TAG "v3.5.0")
3
+set(DEFAULT_TRITON_KERNELS_TAG "v3.6.0")
4
5
# Set TRITON_KERNELS_SRC_DIR for use with local development with vLLM. We expect TRITON_KERNELS_SRC_DIR to
6
-# be directly set to the triton_kernels python directory.
6
+# be directly set to the triton_kernels python directory.
7
if (DEFINED ENV{TRITON_KERNELS_SRC_DIR})
8
message(STATUS "[triton_kernels] Fetch from $ENV{TRITON_KERNELS_SRC_DIR}")
9
FetchContent_Declare(
@@ -24,7 +24,7 @@ else()
24
)
25
endif()
26
27
-# Fetch content
27
+# Fetch content
28
FetchContent_MakeAvailable(triton_kernels)
29
30
if (NOT triton_kernels_SOURCE_DIR)
@@ -47,7 +47,7 @@ install(CODE "file(MAKE_DIRECTORY \"\${CMAKE_INSTALL_PREFIX}/vllm/third_party/tr
47
## Copy .py files to install directory.
48
install(DIRECTORY
49
${TRITON_KERNELS_PYTHON_DIR}
50
- DESTINATION
50
+ DESTINATION
51
vllm/third_party/triton_kernels/
52
COMPONENT triton_kernels
53
FILES_MATCHING PATTERN "*.py")
pyproject.toml
+1
-1
@@ -6,7 +6,7 @@ requires = [
6
"packaging>=24.2",
7
"setuptools>=77.0.3,<81.0.0",
8
"setuptools-scm>=8.0",
9
- "torch == 2.9.1",
9
+ "torch == 2.10.0",
10
"wheel",
11
"jinja2",
12
"grpcio-tools==1.78.0",
requirements/build.txt
+1
-1
@@ -4,7 +4,7 @@ ninja
4
packaging>=24.2
5
setuptools>=77.0.3,<81.0.0
6
setuptools-scm>=8
7
-torch==2.9.1
7
+torch==2.10.0
8
wheel
9
jinja2>=3.1.6
10
regex
requirements/cuda.txt
+3
-3
@@ -5,9 +5,9 @@ numba == 0.61.2 # Required for N-gram speculative decoding
5
6
# Dependencies for NVIDIA GPUs
7
ray[cgraph]>=2.48.0
8
-torch==2.9.1
9
-torchaudio==2.9.1
8
+torch==2.10.0
9
+torchaudio==2.10.0
10
# These must be updated alongside torch
11
-torchvision==0.24.1 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
11
+torchvision==0.25.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
12
# FlashInfer should be updated together with the Dockerfile
13
flashinfer-python==0.6.3
requirements/rocm-build.txt
+5
-6
@@ -1,12 +1,11 @@
1
# Common dependencies
2
-r common.txt
3
4
---extra-index-url https://download.pytorch.org/whl/rocm6.4
5
-torch==2.9.1
6
-torchvision==0.24.1
7
-torchaudio==2.9.1
8
-
9
-triton==3.5.1
4
+--extra-index-url https://download.pytorch.org/whl/test/rocm7.0
5
+torch==2.10.0
6
+torchvision==0.25.0
7
+torchaudio==2.10.0
8
+triton==3.6.0
9
cmake>=3.26.1,<4
10
packaging>=24.2
11
setuptools>=77.0.3,<80.0.0
requirements/test.in
+4
-4
@@ -24,10 +24,10 @@ sentence-transformers>=5.2.0 # required for embedding tests
24
soundfile # required for audio tests
25
jiwer # required for audio tests
26
tblib # for pickling test exceptions
27
-timm==1.0.17 # required for internvl and gemma3n-mm test
28
-torch==2.9.1
29
-torchaudio==2.9.1
30
-torchvision==0.24.1
27
+timm >=1.0.17 # required for internvl and gemma3n-mm test
28
+torch==2.10.0
29
+torchaudio==2.10.0
30
+torchvision==0.25.0
31
transformers_stream_generator # required for qwen-vl test
32
matplotlib # required for qwen-vl test
33
mistral_common[image,audio] >= 1.9.0 # required for voxtral test
requirements/test.txt
+9
-5
@@ -155,6 +155,10 @@ coverage==7.10.6
155
# via pytest-cov
156
cramjam==2.9.0
157
# via fastparquet
158
+cuda-bindings==12.9.4
159
+ # via torch
160
+cuda-pathfinder==1.3.3
161
+ # via cuda-bindings
162
cupy-cuda12x==13.6.0
163
# via ray
164
cycler==0.12.1
@@ -631,7 +635,7 @@ nvidia-nvjitlink-cu12==12.9.86
635
# nvidia-cusolver-cu12
636
# nvidia-cusparse-cu12
637
# torch
634
-nvidia-nvshmem-cu12==3.3.20
638
+nvidia-nvshmem-cu12==3.4.5
639
# via torch
640
nvidia-nvtx-cu12==12.9.79
641
# via torch
@@ -1163,7 +1167,7 @@ tomli==2.2.1
1167
# via schemathesis
1168
tomli-w==1.2.0
1169
# via schemathesis
1166
-torch==2.9.1+cu129
1170
+torch==2.10.0+cu129
1171
# via
1172
# -r requirements/test.in
1173
# accelerate
@@ -1192,7 +1196,7 @@ torch==2.9.1+cu129
1196
# torchvision
1197
# vector-quantize-pytorch
1198
# vocos
1195
-torchaudio==2.9.1+cu129
1199
+torchaudio==2.10.0+cu129
1200
# via
1201
# -r requirements/test.in
1202
# encodec
@@ -1205,7 +1209,7 @@ torchmetrics==1.7.4
1209
# pytorch-lightning
1210
# terratorch
1211
# torchgeo
1208
-torchvision==0.24.1+cu129
1212
+torchvision==0.25.0+cu129
1213
# via
1214
# -r requirements/test.in
1215
# lightly
@@ -1247,7 +1251,7 @@ transformers==4.57.5
1251
# transformers-stream-generator
1252
transformers-stream-generator==0.0.5
1253
# via -r requirements/test.in
1250
-triton==3.5.1
1254
+triton==3.6.0
1255
# via torch
1256
tritonclient==2.64.0
1257
# via -r requirements/test.in
tests/compile/test_aot_compile.py
+10
-30
@@ -90,9 +90,7 @@ def use_vllm_config(vllm_config: VllmConfig):
90
yield
91
92
93
-@pytest.mark.skipif(
94
- not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
95
-)
93
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
94
def test_no_dynamo_cache_entry(monkeypatch: pytest.MonkeyPatch):
95
with monkeypatch.context() as m:
96
vllm_config = make_vllm_config()
@@ -116,9 +114,7 @@ def test_no_dynamo_cache_entry(monkeypatch: pytest.MonkeyPatch):
114
assert torch.allclose(actual, expected)
115
116
119
-@pytest.mark.skipif(
120
- not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
121
-)
117
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
118
def test_force_aot_load(monkeypatch: pytest.MonkeyPatch):
119
with tempfile.TemporaryDirectory() as tmpdirname, monkeypatch.context() as m:
120
args = (torch.randn(10, 10),)
@@ -132,9 +128,7 @@ def test_force_aot_load(monkeypatch: pytest.MonkeyPatch):
128
CompiledMod(vllm_config=vllm_config)(*args)
129
130
135
-@pytest.mark.skipif(
136
- not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
137
-)
131
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
132
def test_save_and_load(monkeypatch: pytest.MonkeyPatch):
133
with monkeypatch.context() as m:
134
args = (torch.randn(10, 10),)
@@ -162,9 +156,7 @@ def test_save_and_load(monkeypatch: pytest.MonkeyPatch):
156
assert torch.allclose(ret, expected)
157
158
165
-@pytest.mark.skipif(
166
- not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
167
-)
159
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
160
def test_cache_load_returns_tuple_consistency(monkeypatch: pytest.MonkeyPatch):
161
"""
162
Test that cache loading correctly handles the returns_tuple logic.
@@ -223,9 +215,7 @@ def test_cache_load_returns_tuple_consistency(monkeypatch: pytest.MonkeyPatch):
215
)
216
217
226
-@pytest.mark.skipif(
227
- not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
228
-)
218
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
219
def test_cache_load_returns_tuple_consistency_tuple_output(
220
monkeypatch: pytest.MonkeyPatch,
221
):
@@ -294,9 +284,7 @@ def test_cache_load_returns_tuple_consistency_tuple_output(
284
)
285
286
297
-@pytest.mark.skipif(
298
- not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
299
-)
287
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
288
def test_shape_env(monkeypatch: pytest.MonkeyPatch):
289
"""
290
Test that the shape environment is correctly serialized and preserved
@@ -333,9 +321,7 @@ def test_shape_env(monkeypatch: pytest.MonkeyPatch):
321
assert guards_string == " - s77 <= 42\n - Eq(Mod(s77, 2), 0)"
322
323
336
-@pytest.mark.skipif(
337
- not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
338
-)
324
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
325
def test_partition_wrapper_applied_on_aot_load(
326
monkeypatch: pytest.MonkeyPatch, vllm_tmp_cache: Path, mocker
327
):
@@ -426,9 +412,7 @@ def test_partition_wrapper_applied_on_aot_load(
412
)
413
414
429
-@pytest.mark.skipif(
430
- not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
431
-)
415
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
416
@create_new_process_for_each_test("spawn")
417
def test_gpt2_cache_hit(monkeypatch: pytest.MonkeyPatch):
418
"""
@@ -492,9 +476,7 @@ def test_gpt2_cache_hit(monkeypatch: pytest.MonkeyPatch):
476
symbolic_shapes_module.make_symbol = original_make_symbol
477
478
495
-@pytest.mark.skipif(
496
- not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
497
-)
479
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
480
class TestStandaloneCompiledArtifacts:
481
def test_init(self):
482
cache = StandaloneCompiledArtifacts()
@@ -668,9 +650,7 @@ class TestStandaloneCompiledArtifacts:
650
assert len(restored_cache.loaded_submodule_store) == 0
651
652
671
-@pytest.mark.skipif(
672
- not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
673
-)
653
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
654
class TestStandaloneCompiledArtifactsIntegration:
655
def test_add_pickle_unpickle(self):
656
cache = StandaloneCompiledArtifacts()
tests/compile/test_dynamic_shapes_compilation.py
+1
-3
@@ -39,9 +39,7 @@ def get_test_models():
39
@pytest.mark.parametrize("use_aot_compile", ["0", "1"])
40
@pytest.mark.parametrize("use_bytecode_hook", [True, False])
41
@pytest.mark.parametrize("evaluate_guards", [False, True])
42
-@pytest.mark.skipif(
43
- not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
44
-)
42
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
43
def test_dynamic_shapes_compilation(
44
monkeypatch,
45
model_name,
tests/kernels/moe/test_shared_fused_moe_routed_transform.py
+5
@@ -14,6 +14,7 @@ import torch.nn as nn
14
from vllm.config import VllmConfig, set_current_vllm_config
15
from vllm.forward_context import set_forward_context
16
from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE
17
+from vllm.utils.torch_utils import is_torch_equal_or_newer
18
19
20
class SimpleLinear(nn.Module):
@@ -60,6 +61,10 @@ def setup_cuda():
61
@pytest.mark.parametrize("num_tokens", [1, 32])
62
@pytest.mark.parametrize("hidden_size,latent_size", [(256, 128), (128, 64)])
63
@pytest.mark.parametrize("dtype", [torch.bfloat16])
64
+@pytest.mark.skipif(
65
+ is_torch_equal_or_newer("2.10.0"),
66
+ reason="Test fails with PyTorch 2.10.0 see: https://github.com/vllm-project/vllm/issues/33995",
67
+)
68
def test_routed_input_transform_inside_vs_outside(
69
num_tokens: int,
70
hidden_size: int,
vllm/compilation/compiler_interface.py
+1
-1
@@ -233,7 +233,7 @@ class InductorStandaloneAdaptor(CompilerInterface):
233
234
from torch._inductor import standalone_compile
235
236
- supports_aot = is_torch_equal_or_newer("2.10.0.dev")
236
+ supports_aot = is_torch_equal_or_newer("2.10.0")
237
238
if not supports_aot and envs.VLLM_USE_MEGA_AOT_ARTIFACT:
239
logger.error(
vllm/compilation/decorators.py
+4
-4
@@ -333,7 +333,7 @@ def _support_torch_compile(
333
) -> None:
334
def mark_dynamic(arg: torch.Tensor, dims: list[int]) -> None:
335
if ds_type == DynamicShapesType.UNBACKED:
336
- if is_torch_equal_or_newer("2.10.0.dev"):
336
+ if is_torch_equal_or_newer("2.10.0"):
337
for dim in dims:
338
torch._dynamo.decorators.mark_unbacked(
339
arg, dim, hint_override=arg.size()[dim]
@@ -373,7 +373,7 @@ def _support_torch_compile(
373
if isinstance(arg, torch.Tensor):
374
# In case dims is specified with negative indexing
375
dims = [arg.ndim + dim if dim < 0 else dim for dim in dims]
376
- if is_torch_equal_or_newer("2.10.0.dev"):
376
+ if is_torch_equal_or_newer("2.10.0"):
377
for dim in dims:
378
torch._dynamo.decorators.mark_unbacked(
379
arg, dim, hint_override=arg.size()[dim]
@@ -525,9 +525,9 @@ def _support_torch_compile(
525
fx_config_patches["backed_size_oblivious"] = True
526
527
# Prepare inductor config patches
528
- # assume_32bit_indexing is only available in torch 2.10.0.dev+
528
+ # assume_32bit_indexing is only available in torch 2.10.0+
529
inductor_config_patches = {}
530
- if is_torch_equal_or_newer("2.10.0.dev"):
530
+ if is_torch_equal_or_newer("2.10.0"):
531
inductor_config_patches["assume_32bit_indexing"] = (
532
self.compilation_config.dynamic_shapes_config.assume_32_bit_indexing
533
)
vllm/envs.py
+1
-1
@@ -271,7 +271,7 @@ def use_aot_compile() -> bool:
271
272
default_value = (
273
"1"
274
- if is_torch_equal_or_newer("2.10.0.dev") and not disable_compile_cache()
274
+ if is_torch_equal_or_newer("2.11.0.dev") and not disable_compile_cache()
275
else "0"
276
)
277
vllm/model_executor/layers/batch_invariant.py
+1
-1
@@ -974,7 +974,7 @@ def enable_batch_invariant_mode():
974
)
975
976
reduced_precision_val = (
977
- (False, False) if is_torch_equal_or_newer("2.10.0.dev") else False
977
+ (False, False) if is_torch_equal_or_newer("2.10.0") else False
978
)
979
torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = (
980
reduced_precision_val
vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py
+73
-8
@@ -27,9 +27,21 @@ logger = init_logger(__name__)
27
if has_triton_kernels():
28
try:
29
import triton_kernels.swiglu
30
- from triton_kernels.matmul_ogs import FnSpecs, FusedActivation, matmul_ogs
31
- from triton_kernels.routing import RoutingData, routing, routing_from_bitmatrix
32
- from triton_kernels.tensor import Bitmatrix
30
+ from triton_kernels.matmul_ogs import (
31
+ FnSpecs,
32
+ FusedActivation,
33
+ GatherIndx,
34
+ RoutingData,
35
+ ScatterIndx,
36
+ matmul_ogs,
37
+ )
38
+ from triton_kernels.tensor import (
39
+ BIT,
40
+ Bitmatrix,
41
+ SparseMatrix,
42
+ make_ragged_tensor_metadata,
43
+ )
44
+ from triton_kernels.topk import topk
45
except (AttributeError, ImportError) as e:
46
logger.error(
47
"Failed to import Triton kernels. Please make sure your triton "
@@ -78,6 +90,58 @@ def pack_bitmatrix(
90
tl.store(bitmatrix_ptrs, y, mask=offsets_m[:, None] < n_rows)
91
92
93
+def legacy_routing_from_bitmatrix(
94
+ bitmatrix: "Bitmatrix",
95
+ expt_scal: torch.Tensor,
96
+ expt_indx: torch.Tensor,
97
+ n_expts_tot: int,
98
+ n_expts_act: int,
99
+) -> tuple["RoutingData", "GatherIndx", "ScatterIndx"]:
100
+ """
101
+ Replacement for the removed triton_kernels.routing.routing_from_bitmatrix.
102
+ Creates routing data from a bitmatrix representation.
103
+ """
104
+ sparse_logits = SparseMatrix(indx=expt_indx, vals=expt_scal, mask=bitmatrix)
105
+ dispatch_indx = sparse_logits.mask_metadata.row_sorted_indx
106
+ combine_indx = sparse_logits.mask_metadata.col_sorted_indx
107
+ ragged_batch_metadata = make_ragged_tensor_metadata(
108
+ sparse_logits.mask_metadata.col_sum,
109
+ dispatch_indx.shape[0],
110
+ )
111
+ gate_scal = sparse_logits.vals.flatten()[combine_indx]
112
+ routing_data = RoutingData(
113
+ gate_scal,
114
+ ragged_batch_metadata.block_sizes,
115
+ n_expts_tot,
116
+ n_expts_act,
117
+ ragged_batch_metadata,
118
+ )
119
+ gather_idx = GatherIndx(combine_indx, dispatch_indx)
120
+ scatter_idx = ScatterIndx(dispatch_indx, combine_indx)
121
+ return routing_data, gather_idx, scatter_idx
122
+
123
+
124
+def legacy_routing(
125
+ logits: torch.Tensor,
126
+ n_expts_act: int,
127
+ sm_first: bool = False,
128
+) -> tuple["RoutingData", "GatherIndx", "ScatterIndx"]:
129
+ """
130
+ Replacement for the removed triton_kernels.routing.routing function.
131
+ Computes routing data from gating logits.
132
+ """
133
+ if sm_first:
134
+ logits = torch.softmax(logits, dim=-1)
135
+ sparse_logits = topk(logits, n_expts_act, apply_softmax=not sm_first)
136
+ return legacy_routing_from_bitmatrix(
137
+ sparse_logits.mask,
138
+ sparse_logits.vals,
139
+ sparse_logits.indx,
140
+ logits.shape[-1],
141
+ n_expts_act,
142
+ )
143
+
144
+
145
def triton_kernel_moe_forward(
146
hidden_states: torch.Tensor,
147
w1, # Tensor or triton_kernels.Tensor
@@ -91,7 +155,7 @@ def triton_kernel_moe_forward(
155
global_num_experts: int = -1,
156
expert_map: torch.Tensor | None = None,
157
) -> torch.Tensor:
94
- routing_data, gather_idx, scatter_idx = routing(
158
+ routing_data, gather_idx, scatter_idx = legacy_routing(
159
gating_output, topk, sm_first=not renormalize
160
)
161
@@ -168,9 +232,10 @@ def triton_kernel_fused_experts(
232
output_tensor = _resize_cache(output_tensor, (batch_dim, M, K))
233
234
act = FusedActivation(
171
- FnSpecs("swiglu", triton_kernels.swiglu.swiglu_fn, ("alpha", "limit")),
235
+ FnSpecs(
236
+ "swiglu", triton_kernels.swiglu.swiglu_fn, ("alpha", "limit"), reduction_n=2
237
+ ),
238
(swiglu_alpha, swiglu_limit),
173
- 2,
239
)
240
gammas = routing_data.gate_scal if routing_data else None
241
@@ -232,12 +297,12 @@ def make_routing_data(
297
bitmatrix_shape = [n_rows, bm_cols * 32]
298
bitmatrix_shape_max = [n_rows, None]
299
bitmatrix = Bitmatrix(
235
- bitmatrix, shape=bitmatrix_shape, shape_max=bitmatrix_shape_max, scratchpad=None
300
+ bitmatrix, dtype=BIT, shape=bitmatrix_shape, shape_max=bitmatrix_shape_max
301
)
302
303
# matmul_ogs expects invalid topk_weights to be -1s
304
topk_weights = torch.where(topk_ids == -1, -1.0, topk_weights)
240
- routing_data, gather_indx, scatter_indx = routing_from_bitmatrix(
305
+ routing_data, gather_indx, scatter_indx = legacy_routing_from_bitmatrix(
306
bitmatrix, topk_weights, topk_ids, num_local_experts, num_topk
307
)
308