Repository navigation
Conversation
Contributor
CI recapDashboard: View test results in Grafana |
Collaborator
|
I'd like to avoid adding new kernels until #47773 is merged. Is that ok or do we need to update the existing kernels for the atlas kernels? Don't hesitate to ping me tho |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
On AMD Strix Halo (Ryzen AI Max,
gfx1151, RDNA 3.5) themamba_ssmandcausal_conv1dpackages have no ROCm wheel, and_KERNEL_MAPPINGonly defines"cuda"and"xpu"entries for the Mamba-2 and causal-conv1d ops, soNemotronHandMamba2models fall back to the slow PyTorcheinsumsegment-sum reference (mamba2_chunk_scan/mamba2_selective_state_update). On NVIDIA GB10 (SM121, aarch64, torch 2.12 + CUDA 13.0),kernels-community/mamba-ssmv4now backfills thetorch212-cu130-aarch64-linuxcell.This PR:
kernels-community/mamba-ssmentries inintegrations/hub_kernels.pyfromversion=3toversion=4so Torch 2.12 aarch64 (SM121/ GH200 / GB200) users resolve the workingv4build (_mamba_ssm_cuda_4a743b8.abi3.so).Device(type="rocm", properties=ROCMProperties(min_capability=115, max_capability=115))inference entries forcausal_conv1d_fn,causal_conv1d_update,mamba_chunk_scan_combined,mamba_split_conv1d_scan_combined, andselective_state_updatepointing at the newly publishedgfx1151ROCm revision ofAtlas-Inference/gdn.Qwen3_5GatedDeltaNetgfx1151pin to the same unified revision and adds a usage note with measured numbers todocs/source/en/model_doc/nemotron_h.md.On
gfx1151, the new Mamba-2 ops run the register-residentSTATE=128HIP kernels from Atlas Inference (mamba2_ssm_scan_regwithTB=16LDS staging for prefill andmamba2_ssm_decode_regfor single-step decode) alongside the existingcausal_conv1d_fwdandcausal_conv1d_updatekernels.What does this PR do?
The gated
rocm(gfx1151) entries for the Mamba-2 andcausal_conv1dops plus themamba-ssmversion=4bump inintegrations/hub_kernels.pyand the docs note indocs/source/en/model_doc/nemotron_h.md.NemotronHMamba2MixerandMamba2Mixerare already decorated with@use_kernelized_func, so no modeling changes are needed.Verified on Strix Halo (
gfx1151,torch.cuda.get_device_capability() == (11, 5)), bf16, torch 2.14.1 + ROCm 7.2, and on DGX Spark GB10 (SM121), torch 2.12.1 + CUDA 13.0,kernels0.17.2, transformersmain:1, 33, 65, 128, 257, 1024) atNVIDIA-Nemotron-3.5-Super-VL-120BMamba-2 dimensions (H=128, P=64, N=128, G=8, D_conv=10240, K=4)]Nemotron-3.5-Super-VL-120BMamba-2 mixer layer ongfx1151, prefill is 3.56x faster at 4,096 tokens (718.9 ms -> 202.1 ms) and 3.07x faster at 8,192 tokens (1,439.6 ms -> 468.4 ms).use_kernels=FalseTTFTuse_kernels=TrueTTFTuse_kernels=FalseDecodeuse_kernels=TrueDecodeFollow-up to #46423, #49127, #49452, and #49453. Just bumping the
trust_remote_code=Truenote untilAtlas-Inferenceis on thekernelstrusted-publisher allowlist!Code Agent Policy
Before submitting
Who can review?
@vasqu (merged #46423, #49127, #49452, #49453), @danieldk (
mamba-ssmv4), @stevhliu (docs)