diff --git a/docs/source/en/model_doc/nemotron_h.md b/docs/source/en/model_doc/nemotron_h.md index 24d95c4a75e9..8b74f3a1a55a 100644 --- a/docs/source/en/model_doc/nemotron_h.md +++ b/docs/source/en/model_doc/nemotron_h.md @@ -67,6 +67,13 @@ print(tokenizer.decode(output[0], skip_special_tokens=True)) ## Notes - Use left padding for batched generation. See [Padding side](../llm_tutorial#padding-side). +- On NVIDIA GB10 (`SM121`) and AMD Strix Halo (`gfx1151`), `mamba-ssm` and `causal-conv1d` have no prebuilt wheels, so `NemotronHMamba2Mixer` falls back to the slow PyTorch `einsum` segment-sum reference. Pass `use_kernels=True` (`pip install -U kernels`) to [`~PreTrainedModel.from_pretrained`] to use the Hub kernels instead (`kernels-community/mamba-ssm` `v4` on CUDA and [`Atlas-Inference/gdn`](https://huggingface.co/kernels/Atlas-Inference/gdn) on `gfx1151` ROCm). Measured on Strix Halo (`gfx1151`, torch 2.14 + ROCm 7.2, `Nemotron-3.5-Super-VL-120B` Mamba-2 mixer dimensions, bf16): + + | Prompt Tokens | `use_kernels=False` TTFT | `use_kernels=True` TTFT | Decode (`False` -> `True`) | + | --- | --- | --- | --- | + | 1,024 | 0.64 s | 0.38 s (1.67x faster) | 27.8 -> 31.5 tok/s | + | 4,096 | 2.71 s | 1.65 s (1.64x faster) | 24.7 -> 29.2 tok/s | + | 16,384 | 18.32 s | 9.38 s (1.95x faster) | 20.9 -> 23.0 tok/s | ## NemotronHConfig diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index 7aba365f3fcd..6c40ceb0c719 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -178,7 +178,7 @@ def _build_kernel_mapping() -> dict: ): LayerRepository( repo_id="Atlas-Inference/gdn", layer_name="Qwen3_5GatedDeltaNet", - revision="13462bf5e936315f341204ae06bf09068392d227", + revision="f9704654e7c18d3149dfa314fed44eaee52f6ca9", # TODO: drop once Atlas-Inference is an allow-listed trusted publisher trust_remote_code=True, ), @@ -188,24 +188,36 @@ def _build_kernel_mapping() -> dict: Mode.TRAINING: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="causal_conv1d_fn", - version=3, + version=4, ), Mode.INFERENCE: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="causal_conv1d_fn", - version=3, + version=4, + ), + }, + Device( + type="rocm", + properties=ROCMProperties(min_capability=115, max_capability=115), + ): { + Mode.INFERENCE: LayerRepository( + repo_id="Atlas-Inference/gdn", + layer_name="causal_conv1d_fn", + revision="f9704654e7c18d3149dfa314fed44eaee52f6ca9", + # TODO: drop once Atlas-Inference is an allow-listed trusted publisher + trust_remote_code=True, ), }, "xpu": { Mode.TRAINING: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="causal_conv1d_fn", - version=3, + version=4, ), Mode.INFERENCE: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="causal_conv1d_fn", - version=3, + version=4, ), }, }, @@ -214,24 +226,36 @@ def _build_kernel_mapping() -> dict: Mode.TRAINING: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="causal_conv1d_update", - version=3, + version=4, ), Mode.INFERENCE: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="causal_conv1d_update", - version=3, + version=4, + ), + }, + Device( + type="rocm", + properties=ROCMProperties(min_capability=115, max_capability=115), + ): { + Mode.INFERENCE: LayerRepository( + repo_id="Atlas-Inference/gdn", + layer_name="causal_conv1d_update", + revision="f9704654e7c18d3149dfa314fed44eaee52f6ca9", + # TODO: drop once Atlas-Inference is an allow-listed trusted publisher + trust_remote_code=True, ), }, "xpu": { Mode.TRAINING: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="causal_conv1d_update", - version=3, + version=4, ), Mode.INFERENCE: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="causal_conv1d_update", - version=3, + version=4, ), }, }, @@ -284,24 +308,36 @@ def _build_kernel_mapping() -> dict: Mode.TRAINING: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="mamba_chunk_scan_combined", - version=3, + version=4, ), Mode.INFERENCE: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="mamba_chunk_scan_combined", - version=3, + version=4, + ), + }, + Device( + type="rocm", + properties=ROCMProperties(min_capability=115, max_capability=115), + ): { + Mode.INFERENCE: LayerRepository( + repo_id="Atlas-Inference/gdn", + layer_name="mamba_chunk_scan_combined", + revision="f9704654e7c18d3149dfa314fed44eaee52f6ca9", + # TODO: drop once Atlas-Inference is an allow-listed trusted publisher + trust_remote_code=True, ), }, "xpu": { Mode.TRAINING: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="mamba_chunk_scan_combined", - version=3, + version=4, ), Mode.INFERENCE: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="mamba_chunk_scan_combined", - version=3, + version=4, ), }, }, @@ -310,24 +346,36 @@ def _build_kernel_mapping() -> dict: Mode.TRAINING: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="mamba_split_conv1d_scan_combined", - version=3, + version=4, ), Mode.INFERENCE: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="mamba_split_conv1d_scan_combined", - version=3, + version=4, + ), + }, + Device( + type="rocm", + properties=ROCMProperties(min_capability=115, max_capability=115), + ): { + Mode.INFERENCE: LayerRepository( + repo_id="Atlas-Inference/gdn", + layer_name="mamba_split_conv1d_scan_combined", + revision="f9704654e7c18d3149dfa314fed44eaee52f6ca9", + # TODO: drop once Atlas-Inference is an allow-listed trusted publisher + trust_remote_code=True, ), }, "xpu": { Mode.TRAINING: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="mamba_split_conv1d_scan_combined", - version=3, + version=4, ), Mode.INFERENCE: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="mamba_split_conv1d_scan_combined", - version=3, + version=4, ), }, }, @@ -336,24 +384,24 @@ def _build_kernel_mapping() -> dict: Mode.TRAINING: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="mamba_inner_fn", - version=3, + version=4, ), Mode.INFERENCE: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="mamba_inner_fn", - version=3, + version=4, ), }, "xpu": { Mode.TRAINING: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="mamba_inner_fn", - version=3, + version=4, ), Mode.INFERENCE: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="mamba_inner_fn", - version=3, + version=4, ), }, }, @@ -362,24 +410,24 @@ def _build_kernel_mapping() -> dict: Mode.TRAINING: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="selective_scan_fn", - version=3, + version=4, ), Mode.INFERENCE: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="selective_scan_fn", - version=3, + version=4, ), }, "xpu": { Mode.TRAINING: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="selective_scan_fn", - version=3, + version=4, ), Mode.INFERENCE: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="selective_scan_fn", - version=3, + version=4, ), }, }, @@ -388,24 +436,36 @@ def _build_kernel_mapping() -> dict: Mode.TRAINING: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="selective_state_update", - version=3, + version=4, ), Mode.INFERENCE: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="selective_state_update", - version=3, + version=4, + ), + }, + Device( + type="rocm", + properties=ROCMProperties(min_capability=115, max_capability=115), + ): { + Mode.INFERENCE: LayerRepository( + repo_id="Atlas-Inference/gdn", + layer_name="selective_state_update", + revision="f9704654e7c18d3149dfa314fed44eaee52f6ca9", + # TODO: drop once Atlas-Inference is an allow-listed trusted publisher + trust_remote_code=True, ), }, "xpu": { Mode.TRAINING: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="selective_state_update", - version=3, + version=4, ), Mode.INFERENCE: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="selective_state_update", - version=3, + version=4, ), }, },