mawei/StarPilot
1
0
Fork 0
mirror of https://github.com/firestar5683/StarPilot.git synced 2026-10-27 16:43:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

bump tinygrad (#39085)

Browse Source
This commit is contained in:
Daniel Koepping
2026-10-20 03:55:45 +00:00
committed by firestar5683
parent 926ba1b700
commit 68984ff07b
444 changed files with 15637 additions and 12588 deletions
Download Patch File Download Diff File Expand all files Collapse all files
openpilot/selfdrive/modeld/modeld.py
+2
View File
@@ -231,6 +231,8 @@ class ModelState:
self.input_device.copy_from(self.input_host)
self.input_queues['new_img'] = self.run_warp(**self.warp_inputs)
self.run_model(output_buffers=self.outputs, **self.input_queues)
if self.chestnut and (err := Device[self.outputs['outputs'].device].error_state.host.view(fmt='q')[0]):
raise RuntimeError(f"chestnut failed with {err}")
if after_enqueue is not None:
after_enqueue()
model_output = self.outputs['outputs'].numpy()[0]
tinygrad_repo/.github/actions/setup-tinygrad/action.yml
Vendored
+12 -8
View File
@@ -63,7 +63,10 @@ runs:
echo "OMP_NUM_THREADS=1" >> "$GITHUB_ENV"
# no buffers should be over 300MB in CI
echo "MAX_BUFFER_SIZE=300000000" >> "$GITHUB_ENV"
if [[ "$RUNNER_OS" == "Linux" ]]; then
if [[ "$RUNNER_OS" != "Windows" ]]; then
if [[ "$RUNNER_OS" == "macOS" ]]; then
sudo install -d -o "$USER" /opt/venv /opt/python
fi
echo "VIRTUAL_ENV=/opt/venv/${{ inputs.python-version }}" >> "$GITHUB_ENV"
echo "UV_PYTHON_INSTALL_DIR=/opt/python" >> "$GITHUB_ENV"
else
@@ -78,7 +81,7 @@ runs:
# **** Caching packages ****
- name: Cache Python packages (PR)
if: github.event_name == 'pull_request'
if: github.event_name == 'pull_request' && env.RUNNER_ENVIRONMENT == 'github-hosted'
id: restore-venv-pr
uses: actions/cache/restore@v5
with:
@@ -95,7 +98,7 @@ runs:
# **** Caching downloads ****
- name: Cache downloads (PR)
if: inputs.key != '' && github.event_name == 'pull_request'
if: inputs.key != '' && github.event_name == 'pull_request' && env.RUNNER_ENVIRONMENT == 'github-hosted'
uses: actions/cache/restore@v5
with:
path: ${{ runner.os == 'Linux' && '~/.cache/tinygrad/downloads/' || '~/Library/Caches/tinygrad/downloads/' }}
@@ -115,7 +118,7 @@ runs:
run: |
uv venv --allow-existing --python ${{ inputs.python-version }} "$VIRTUAL_ENV"
DEPS="${{ inputs.deps }}"
uv pip install --python "$VIRTUAL_ENV" -e ".[${DEPS// /,}]" ${{ inputs.pydeps }} --torch-backend cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
uv pip install --python "$VIRTUAL_ENV" -e ".[${DEPS// /,}]" ${{ inputs.pydeps }} --torch-backend cpu
- name: Install dependencies in venv (without extra)
if: inputs.deps == ''
shell: bash
@@ -211,7 +214,7 @@ runs:
echo "installed=$installed" >> "$GITHUB_OUTPUT"
- name: Cache apt (PR)
if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true' || inputs.autogen == 'true') && github.event_name == 'pull_request' && steps.apt-pkgs.outputs.installed == 'false'
if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true' || inputs.autogen == 'true') && github.event_name == 'pull_request' && steps.apt-pkgs.outputs.installed == 'false' && env.RUNNER_ENVIRONMENT == 'github-hosted'
uses: actions/cache/restore@v5
with:
path: /var/cache/apt/archives/
@@ -267,10 +270,11 @@ runs:
if: inputs.cuda == 'true'
shell: bash
run: |
sudo mkdir -p /usr/local/cuda/targets/x86_64-linux
sudo mkdir -p /usr/local/cuda-11.5/targets/x86_64-linux
curl -fL https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvrtc/linux-x86_64/cuda_nvrtc-linux-x86_64-11.5.119-archive.tar.xz \
| sudo tar -xJ -C /usr/local/cuda/targets/x86_64-linux --strip-components=1
echo /usr/local/cuda/targets/x86_64-linux/lib | sudo tee /etc/ld.so.conf.d/cuda-nvrtc.conf
| sudo tar -xJ -C /usr/local/cuda-11.5/targets/x86_64-linux --strip-components=1
sudo ln -sfnT /usr/local/cuda-11.5 /usr/local/cuda
echo /usr/local/cuda-11.5/targets/x86_64-linux/lib | sudo tee /etc/ld.so.conf.d/cuda-11.5.conf
sudo ldconfig
# **** gpuocelot ****
tinygrad_repo/.github/workflows/autogen.yml
Vendored
+1 -1
View File
@@ -1,7 +1,7 @@
name: Autogen
env:
# increment this when downloads substantially change to avoid the internet
CACHE_VERSION: '13'
CACHE_VERSION: '14'
CAPTURE_PROCESS_REPLAY: 1
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
PYTHONPATH: ${{ github.workspace }}
tinygrad_repo/.github/workflows/benchmark.yml
Vendored
+25 -40
View File
@@ -95,7 +95,6 @@ jobs:
shell: bash -e -o pipefail {0}
env:
DEV: ${{ matrix.dev }}
HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
if: github.repository_owner == 'tinygrad'
steps:
- name: Checkout Code
@@ -107,7 +106,7 @@ jobs:
./extra/hcq/hcq_smi.py amd kill_pids --sudoless
- name: Setup (NV)
if: ${{ matrix.dev == 'NV' }}
run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
run: ./extra/hcq/hcq_smi.py nv kill_pids --sudoless
- name: setup staging db
if: github.ref == 'refs/heads/update_benchmark_staging'
run: |
@@ -135,7 +134,6 @@ jobs:
shell: bash -e -o pipefail {0}
env:
DEV: ${{ matrix.dev }}
HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
if: github.repository_owner == 'tinygrad'
steps:
- name: Checkout Code
@@ -147,7 +145,7 @@ jobs:
./extra/hcq/hcq_smi.py amd kill_pids --sudoless
- name: Setup (NV)
if: ${{ matrix.dev == 'NV' }}
run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
run: ./extra/hcq/hcq_smi.py nv kill_pids --sudoless
- name: setup staging db
if: github.ref == 'refs/heads/update_benchmark_staging'
run: |
@@ -183,7 +181,6 @@ jobs:
shell: bash -e -o pipefail {0}
env:
DEV: ${{ matrix.dev }}
HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
if: github.repository_owner == 'tinygrad'
steps:
- name: Checkout Code
@@ -195,7 +192,7 @@ jobs:
./extra/hcq/hcq_smi.py amd kill_pids --sudoless
- name: Setup (NV)
if: ${{ matrix.dev == 'NV' }}
run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
run: ./extra/hcq/hcq_smi.py nv kill_pids --sudoless
- name: Symlink models and datasets
run: |
mkdir -p extra/datasets
@@ -225,7 +222,6 @@ jobs:
shell: bash -e -o pipefail {0}
env:
DEV: ${{ matrix.dev }}
HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
if: github.repository_owner == 'tinygrad'
steps:
- name: Checkout Code
@@ -237,7 +233,7 @@ jobs:
./extra/hcq/hcq_smi.py amd kill_pids --sudoless
- name: Setup (NV)
if: ${{ matrix.dev == 'NV' }}
run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
run: ./extra/hcq/hcq_smi.py nv kill_pids --sudoless
- name: setup staging db
if: github.ref == 'refs/heads/update_benchmark_staging'
run: |
@@ -270,7 +266,6 @@ jobs:
shell: bash -e -o pipefail {0}
env:
DEV: ${{ matrix.dev }}
HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
if: github.repository_owner == 'tinygrad'
steps:
- name: Checkout Code
@@ -282,7 +277,7 @@ jobs:
./extra/hcq/hcq_smi.py amd kill_pids --sudoless
- name: Setup (NV)
if: ${{ matrix.dev == 'NV' }}
run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
run: ./extra/hcq/hcq_smi.py nv kill_pids --sudoless
- name: Symlink models and datasets
run: |
mkdir -p weights
@@ -334,7 +329,7 @@ jobs:
./extra/hcq/hcq_smi.py amd kill_pids --sudoless
- name: Setup (NV)
if: ${{ matrix.dev == 'NV' }}
run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
run: ./extra/hcq/hcq_smi.py nv kill_pids --sudoless
- name: setup staging db
if: github.ref == 'refs/heads/update_benchmark_staging'
run: |
@@ -351,14 +346,14 @@ jobs:
- name: Test tensor cores
run: |
if [[ "${{ matrix.dev }}" == "METAL" ]]; then
python3 test/opt/test_tensor_cores.py
python3 test/runtime/test_tensor_cores.py
DEBUG=2 SHOULD_USE_TC=1 python3 extra/gemm/simple_matmul.py
DEBUG=2 SHOULD_USE_TC=1 HALF=1 python3 extra/gemm/simple_matmul.py
DEBUG=2 SHOULD_USE_TC=1 BFLOAT16=1 python3 extra/gemm/simple_matmul.py
M_START=6 M_STOP=10 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=6 K_STOP=24 K_STEP=1 TC_OPT=2 DEBUG=2 python3 ./extra/gemm/fuzz_matmul.py
elif [[ "${{ matrix.dev }}" == "NV" ]]; then
ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
DEV=NV:PTX ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
ALLOW_TF32=1 python3 test/runtime/test_tensor_cores.py
DEV=NV:PTX ALLOW_TF32=1 python3 test/runtime/test_tensor_cores.py
DEV=CUDA SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
DEV=CUDA SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
DEV=CUDA SHOULD_USE_TC=1 ALLOW_TF32=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
@@ -369,9 +364,9 @@ jobs:
# M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
# DEV=NV:PTX M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
else
python3 test/opt/test_tensor_cores.py
python3 test/runtime/test_tensor_cores.py
# TODO: this is flaky
# DEV=AMD:LLVM python3 test/opt/test_tensor_cores.py
# DEV=AMD:LLVM python3 test/runtime/test_tensor_cores.py
SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
# TODO: AMD compiler bug causes this to fail
@@ -548,9 +543,9 @@ jobs:
- name: Test copy speeds
run: SIZE=64000000 GMMU=0 DEV=USB+AMD taskset -c 4-7 python3 test/external/external_test_usb_asm24.py
driverbenchmarks:
name: PCI Driver Benchmark (DEV=${{ matrix.dev }})
runs-on: [self-hosted, Linux, tinyboxrandom]
drivertests:
name: PCI Driver Tests (DEV=${{ matrix.dev }})
runs-on: ${{ matrix.dev == 'AMD' && '9060xt' || '3080' }}
strategy:
fail-fast: false
matrix:
@@ -565,12 +560,6 @@ jobs:
steps:
- name: Checkout Code
uses: actions/checkout@v6
- name: Setup
run: |
./extra/hcq/hcq_smi.py ${{ matrix.dev }} rmmod --expect
./extra/hcq/hcq_smi.py ${{ matrix.dev }} kill_pids --sudoless
mkdir -p extra/datasets
ln -s /raid/datasets/imagenet extra/datasets/imagenet
- name: setup staging db
if: github.ref == 'refs/heads/update_benchmark_staging'
run: |
@@ -581,24 +570,26 @@ jobs:
- name: Test driver cold start time
run: time DEBUG=3 AM_RESET=1 python3 test/test_tiny.py TestTiny.test_plus
- name: Test driver warm start time
if: ${{ matrix.dev == 'AMD' }}
if: ${{ matrix.dev != 'NV' }}
run: time DEBUG=3 python3 test/test_tiny.py TestTiny.test_plus
- name: Test GPU crash recovery
if: ${{ matrix.dev == 'AMD' }}
if: ${{ matrix.dev != 'NV' }}
run: python3 -m pytest -rA test/external/external_test_gpu_crash.py
- name: Test tensor cores
run: |
if [[ "${{ matrix.dev }}" == "AMD" ]]; then
if [[ "${{ matrix.dev }}" == "NV" ]]; then
ALLOW_TF32=1 python3 test/runtime/test_tensor_cores.py
else
# Fails on 9070
# python3 test/test_linearizer.py test/opt/test_tensor_cores.py
# DEV=AMD:LLVM python3 test/test_linearizer.py test/opt/test_tensor_cores.py
# python3 test/test_linearizer.py test/runtime/test_tensor_cores.py
# DEV=AMD:LLVM python3 test/test_linearizer.py test/runtime/test_tensor_cores.py
# SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
else
ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
fi
- name: Test DISK copy time
run: TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
run: |
truncate -s 4G disk-benchmark
TESTFILE=disk-benchmark python3 test/external/external_benchmark_disk_raw.py
- name: Test CPU copy time
run: |
GRAPH_ONE_KERNEL=1 NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
@@ -607,14 +598,8 @@ jobs:
# - name: HEVC Decode Benchmark
# if: ${{ matrix.dev == 'NV' }}
# run: IGNORE_BEAM_CACHE=1 VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 PYTHONPATH=. python3 extra/hevc/decode.py
- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
if: ${{ matrix.dev == 'NV' }}
run: BENCHMARK_LOG=resnet_10steps MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
- name: Run 10 MLPerf Bert training steps (1 gpu)
# TODO: remove BERT_LAYERS once scheduler is fast
run: BENCHMARK_LOG=bert_10steps CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
- name: Remote
if: ${{ matrix.dev == 'AMD' }}
if: ${{ matrix.dev != 'NV' }}
env:
DEV: PCI+AMD
run: |
tinygrad_repo/.github/workflows/platform.yml
Vendored
+17 -20
View File
@@ -1,7 +1,7 @@
name: Platform Tests
env:
# increment this when downloads substantially change to avoid the internet
CACHE_VERSION: '19'
CACHE_VERSION: '20'
CAPTURE_PROCESS_REPLAY: ${{ github.event_name == 'pull_request' && contains(github.event.pull_request.title, '[pr]') && '1' || '0' }}
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
PYTHONPATH: ${{ github.workspace }}
@@ -23,7 +23,7 @@ jobs:
# ****** OSX Tests ******
unittestmacos:
name: MacOS (unit)
name: MacOS (opt, device)
runs-on: macos-26
timeout-minutes: 20
steps:
@@ -34,18 +34,15 @@ jobs:
with:
key: unittest-macos
deps: testing_unit
- name: Run unit tests
run: DEV=METAL python -m pytest -n=auto test/unit/ --durations=20
- name: Run opt tests
run: DEV=METAL python -m pytest -n=auto test/opt --durations=20
llvm: 'true'
- name: Test tensor core ops (fake)
run: DEV=METAL DEBUG=3 TC=2 python test/backend/test_ops.py TestOps.test_gemm
run: DEV=METAL DEBUG=3 TC=2 python test/runtime/test_ops.py TestOps.test_gemm
- name: Test tensor core ops (real)
run: DEV=METAL DEBUG=3 python test/backend/test_ops.py TestOps.test_big_gemm
run: DEV=METAL DEBUG=3 python test/runtime/test_ops.py TestOps.test_big_gemm
- name: Test Beam Search
run: DEV=METAL IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
- name: Test Device Specific
run: DEV=METAL python3 -m pytest test/device/test_metal.py
- name: Run opt and device tests
run: DEV=METAL python3 -m pytest -n=auto test/runtime/test_kernel_opts.py test/runtime/test_tensor_cores.py test/runtime/test_search.py test/device/metal --durations=20
#- name: Fuzz Test linearizer
# run: DEV=METAL DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
- name: Run process replay tests
@@ -63,6 +60,7 @@ jobs:
with:
key: unittest-macos-mock
deps: testing_unit
llvm: 'true'
amd: 'true'
ocelot: 'true'
- name: Run NULL backend tests
@@ -70,19 +68,17 @@ jobs:
- name: Run pytest (amd)
env:
DEV: MOCKKFD+AMD
HCQ_RUNTIME_DEV: PYTHON
FORWARD_ONLY: 1
run: |
python3 -m pytest -n=auto test/device/test_hcq2.py test/test_tiny.py --durations=20
python3 -m pytest -n=auto test/null/test_hcq2.py test/runtime/test_hcq2.py test/test_tiny.py test/device/amd --durations=20
- name: Run pytest (ptx)
env:
DEV: "MOCK+NV:PTX"
HCQ_RUNTIME_DEV: PYTHON
FORWARD_ONLY: 1
# TODO: failing due to library loading error
CAPTURE_PROCESS_REPLAY: 0
run: |
python3 -m pytest -n=auto test/device/test_hcq2.py test/test_tiny.py \
python3 -m pytest -n=auto test/null/test_hcq2.py test/runtime/test_hcq2.py test/test_tiny.py test/device/nv \
test/testextra/test_hevc.py::TestHevc::test_hevc_decode_compile --durations=20
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -105,12 +101,13 @@ jobs:
with:
key: macos-metal
deps: testing_unit
llvm: 'true'
- name: Check Device.DEFAULT and print some source
run: |
python -c "from tinygrad import Device; assert Device.DEFAULT == 'METAL'"
DEBUG=4 python test/test_tiny.py TestTiny.test_plus
- name: Run backend tests
run: python -m pytest -n=auto test/backend --durations=20 --splits 2 --group ${{ matrix.group }}
- name: Run runtime tests
run: python -m pytest -n=auto test/runtime test/device/metal --durations=20 --splits 2 --group ${{ matrix.group }}
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -143,8 +140,8 @@ jobs:
run: |
python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
DEBUG=4 python test/test_tiny.py TestTiny.test_plus
- name: Run test_tiny
run: python -m pytest -n=auto test/test_tiny.py --durations=20
- name: Run test_tiny and device tests
run: python -m pytest -n=auto test/test_tiny.py test/device/$(python -c 'from tinygrad import Device; print(Device.DEFAULT.lower())') --durations=20
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -180,6 +177,6 @@ jobs:
run: |
python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
DEBUG=4 python test/test_tiny.py TestTiny.test_plus
- name: Run test_tiny
- name: Run test_tiny and device tests
shell: bash
run: python -m pytest -n=auto test/test_tiny.py --durations=20
run: python -m pytest -n=auto test/test_tiny.py test/device/$(python -c 'from tinygrad import Device; print(Device.DEFAULT.lower())') --durations=20
tinygrad_repo/.github/workflows/test.yml
Vendored
+97 -78
View File
@@ -1,11 +1,12 @@
name: Unit Tests
env:
# increment this when downloads substantially change to avoid the internet
CACHE_VERSION: '19'
CACHE_VERSION: '20'
CAPTURE_PROCESS_REPLAY: ${{ github.event_name == 'pull_request' && contains(github.event.pull_request.title, '[pr]') && '1' || '0' }}
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
PYTHONPATH: ${{ github.workspace }}
CHECK_OOB: 1
PYTHONFAULTHANDLER: 1
on:
push:
@@ -21,10 +22,11 @@ concurrency:
jobs:
docs:
name: Docs
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
runs-on: &linux ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'turin' || 'ubuntu-24.04' }}
timeout-minutes: 10
env:
CHECK_OOB: 0
DEV: CPU
steps:
- name: Checkout Code
uses: actions/checkout@v6
@@ -60,7 +62,7 @@ jobs:
torchbackend:
name: Torch Backend Tests
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
runs-on: *linux
timeout-minutes: 15
steps:
- name: Checkout Code
@@ -74,17 +76,17 @@ jobs:
llvm: 'true'
ninja: 'true'
- name: Test ResNet-18
run: DEBUG=2 python3 extra/torch_backend/example.py
run: DEV=CPU DEBUG=2 python3 extra/torch_backend/example.py
- name: Test one op in torch tests
run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
run: DEV=CPU DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
- name: Test Ops with TINY_BACKEND
run: DEV=CPU:LLVM LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/backend/test_ops.py --durations=20
run: DEV=CPU:LLVM LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/runtime/test_ops.py --durations=20
- name: Custom tests
run: DEV=CPU:LLVM GPUS=4 TINY_BACKEND=1 python3 -m pytest -nauto extra/torch_backend/test.py extra/torch_backend/test_inplace.py extra/torch_backend/test_multigpu.py extra/torch_backend/test_kernel_fusion.py --durations=20
torchbackendtrain:
name: Torch Backend Training
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
runs-on: *linux
timeout-minutes: 15
steps:
- name: Checkout Code
@@ -101,7 +103,7 @@ jobs:
bepython:
name: Python Backend
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
runs-on: *linux
timeout-minutes: 15
steps:
- name: Checkout Code
@@ -111,10 +113,10 @@ jobs:
with:
key: be-minimal
deps: testing_unit
- name: Run backend tests
run: SKIP_SLOW_TEST=1 DEV=PYTHON python3 -m pytest -n=auto test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_ops.py test/backend/test_uops.py test/backend/test_symbolic_ops.py test/backend/test_renderer_failures.py::TestRendererFailures --durations=20
- name: Run runtime tests
run: SKIP_SLOW_TEST=1 DEV=PYTHON python3 -m pytest -n=auto test/runtime/test_dtype.py test/runtime/test_dtype_alu.py test/runtime/test_ops.py test/runtime/test_uops.py test/runtime/test_symbolic_ops.py test/runtime/test_renderer_failures.py::TestRendererFailures test/device/python --durations=20
- name: Test IMAGE support
run: IMAGE=1 DEV=PYTHON python3 test/backend/test_ops.py TestOps.test_gemm TestOps.test_simple_conv2d
run: IMAGE=1 DEV=PYTHON python3 test/runtime/test_ops.py TestOps.test_gemm TestOps.test_simple_conv2d
- name: Test emulated tensor cores
env:
DEBUG: 2
@@ -130,16 +132,17 @@ jobs:
'DEV=PYTHON::sm_75 HALF=1' 'DEV=PYTHON::sm_80 HALF=1' 'DEV=PYTHON::sm_80 ALLOW_TF32=1'
- name: Run additional tensor core tests
run: |
DEV=PYTHON::METAL python3 -m pytest -nauto test/opt/test_tensor_cores.py test/null/test_uops_stats.py::TestUOpsStatsMatmulHalf
DEV=PYTHON::gfx1100 python3 -m pytest -nauto test/opt/test_tensor_cores.py test/null/test_uops_stats.py::TestUOpsStatsMatmulHalf
DEV=PYTHON::gfx950 python3 -m pytest -nauto test/opt/test_tensor_cores.py
DEV=PYTHON::gfx1201 python3 -m pytest -nauto test/opt/test_tensor_cores.py
ALLOW_TF32=1 DEV=PYTHON::sm_89 python3 -m pytest -nauto test/opt/test_tensor_cores.py
DEV=PYTHON::METAL python3 -m pytest -nauto test/runtime/test_tensor_cores.py test/null/test_uops_stats.py::TestUOpsStatsMatmulHalf
DEV=PYTHON::gfx1100 python3 -m pytest -nauto test/runtime/test_tensor_cores.py test/null/test_uops_stats.py::TestUOpsStatsMatmulHalf
DEV=PYTHON::gfx942 python3 -m pytest -nauto test/runtime/test_tensor_cores.py
DEV=PYTHON::gfx950 python3 -m pytest -nauto test/runtime/test_tensor_cores.py
DEV=PYTHON::gfx1201 python3 -m pytest -nauto test/runtime/test_tensor_cores.py
ALLOW_TF32=1 DEV=PYTHON::sm_89 python3 -m pytest -nauto test/runtime/test_tensor_cores.py
DEBUG=2 DEV=PYTHON::sm_80 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
linter:
name: Linters
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
runs-on: *linux
timeout-minutes: 10
steps:
@@ -170,7 +173,7 @@ jobs:
nulltest:
name: Null Tests
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
runs-on: *linux
timeout-minutes: 15
steps:
@@ -188,7 +191,7 @@ jobs:
run: SPEC=2 DEV=NULL python -m pytest -n=auto test/null/ --durations=20
- name: Run targeted tests on NULL backend
run: |
DEV=NULL python3 -m unittest test.backend.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
DEV=NULL python3 -m unittest test.runtime.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
DEV=NULL VIZ=1 python3 -m pytest -n=auto test/null/test_viz.py
DEBUG=7 python -m tinygrad.viz.cli --json | jq empty
# TODO: too slow
@@ -203,8 +206,8 @@ jobs:
# run: DEV=NULL python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
unittest:
name: Unit Tests
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
name: Backend Tests (CPU)
runs-on: *linux
timeout-minutes: 15
steps:
@@ -222,10 +225,8 @@ jobs:
run: SKIP=ruff,mypy,tests pre-commit run --all-files
- name: Check Device.DEFAULT
run: python -c "from tinygrad import Device; assert Device.DEFAULT == 'CPU', Device.DEFAULT"
- name: Run unit tests
run: |
DEV=CPU python test/null/test_device.py TestRunAsModule.test_module_runs
DEV=CPU python -m pytest -n=auto test/unit/ --durations=20
- name: Run test as module
run: DEV=CPU python test/null/test_device.py TestRunAsModule.test_module_runs
- name: Run GC tests
run: python test/external/external_uop_gc.py
- name: External Benchmark Schedule
@@ -241,7 +242,7 @@ jobs:
matrix:
group: [1, 2]
name: SPEC=2 (${{ matrix.group }})
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
runs-on: *linux
timeout-minutes: 15
steps:
- name: Checkout Code
@@ -253,12 +254,14 @@ jobs:
deps: testing_unit
llvm: 'true'
- name: Test SPEC=2
run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 test/unit test/backend test/opt --ignore test/backend/test_custom_kernel.py --ignore test/unit/test_hashing.py --splits 2 --group ${{ matrix.group }}
run: DEV=CPU SPEC=2 pytest --maxfail=10 -n auto --durations=30 test/runtime test/device/cpu --splits 2 --group ${{ matrix.group }}
fuzzing:
name: Fuzzing
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
runs-on: *linux
timeout-minutes: 10
env:
DEV: CPU
steps:
- name: Checkout Code
uses: actions/checkout@v6
@@ -273,7 +276,7 @@ jobs:
testopenclimage:
name: CL IMAGE Tests
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
runs-on: *linux
timeout-minutes: 15
steps:
- name: Checkout Code
@@ -286,7 +289,7 @@ jobs:
opencl: 'true'
- name: Test CL IMAGE=1 ops
run: |
DEV=CL IMAGE=1 python -m pytest -n=auto test/backend/test_ops.py --durations=20
DEV=CL IMAGE=1 python -m pytest -n=auto test/runtime/test_ops.py --durations=20
DEV=CL IMAGE=1 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -295,7 +298,7 @@ jobs:
name: openpilot compile tests
env:
LFS_SHA: "51b12d20a25b76fa8a5717c2e5af01f148e469e2"
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
runs-on: *linux
timeout-minutes: 15
steps:
- name: Checkout Code
@@ -330,7 +333,7 @@ jobs:
DEV: CL::IMAGE_PITCH_ALIGNMENT=64
run: |
VIZ=1 IMAGE=1 FLOAT16=1 python examples/openpilot/compile_onnx.py "https://huggingface.co/commaai/openpilot-lfs/resolve/$LFS_SHA/openpilot/selfdrive/modeld/models/driving_supercombo.onnx" /tmp/openpilot_viz.pkl --benchmark-runs 1
ALLOWED_KERNEL_COUNT=162 ALLOWED_READ_IMAGE=1530 ALLOWED_GATED_READ_IMAGE=47 python examples/openpilot/check_image_reads.py
ALLOWED_KERNEL_COUNT=162 ALLOWED_READ_IMAGE=1530 ALLOWED_GATED_READ_IMAGE=44 python examples/openpilot/check_image_reads.py
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -338,7 +341,7 @@ jobs:
testonnxcpu:
name: ONNX (CPU) Tests
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
runs-on: *linux
timeout-minutes: 20
steps:
@@ -351,13 +354,13 @@ jobs:
deps: testing
llvm: 'true'
- name: Test ONNX (CPU)
run: DEV=CPU python -m pytest -n=auto test/external/external_test_onnx_backend.py test/external/external_test_onnx_runner.py test/external/external_test_onnx_ops.py test/backend/test_quantize_onnx.py --durations=20
run: DEV=CPU python -m pytest -n=auto test/external/external_test_onnx_backend.py test/external/external_test_onnx_runner.py test/external/external_test_onnx_ops.py test/device/cpu/test_quantize_onnx.py --durations=20
- name: Run process replay tests
uses: ./.github/actions/process-replay
testoptim:
name: Optimization Tests
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
runs-on: *linux
timeout-minutes: 20
steps:
- name: Checkout Code
@@ -389,10 +392,11 @@ jobs:
testllm:
name: Test LLM
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
runs-on: *linux
timeout-minutes: 15
env:
CHECK_OOB: 0
DEV: CPU
steps:
- name: Checkout Code
uses: actions/checkout@v6
@@ -406,17 +410,17 @@ jobs:
run: |
parallel --link --tagstring '[{1}]' '{2}' \
::: llama 'llama q4' qwen3.5 qwen \
::: $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model llama3.2:1b | tee /dev/stderr | grep -i rooster' \
$'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model llama3.2:1b-q4 | tee /dev/stderr | grep -i rooster' \
$'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model qwen3.5:0.8b | tee /dev/stderr | grep -i rooster' \
$'echo "What\'s a female chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model qwen3:0.6b | tee /dev/stderr | grep -i hen'
::: $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model llama3.2:1b | tee -a /dev/stderr | grep -i rooster' \
$'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model llama3.2:1b-q4 | tee -a /dev/stderr | grep -i rooster' \
$'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model qwen3.5:0.8b | tee -a /dev/stderr | grep -i rooster' \
$'echo "What\'s a female chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model qwen3:0.6b | tee -a /dev/stderr | grep -i hen'
# NOTE: qwen is dumb and only knows about female chickens
# ****** Models Tests ******
testmodels:
name: Models
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
runs-on: *linux
timeout-minutes: 15
steps:
- name: Checkout Code
@@ -436,7 +440,7 @@ jobs:
testdsp:
name: Linux (DSP)
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
runs-on: *linux
timeout-minutes: 15
steps:
- name: Checkout Code
@@ -449,7 +453,7 @@ jobs:
llvm: "true"
qemu: "true"
- name: Run tests
run: MOCKDSP=1 DEV=DSP python -m pytest -n=auto test/test_tiny.py test/backend/test_transcendental.py::TestTranscendentalVectorized test/backend/test_quantize_onnx.py
run: MOCKDSP=1 DEV=DSP python -m pytest -n=auto test/test_tiny.py test/runtime/test_transcendental.py::TestTranscendentalVectorized test/device/dsp
testlinux:
strategy:
@@ -464,7 +468,7 @@ jobs:
- 'WEBGPU'
name: Linux (DEV=${{ matrix.dev }})
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
runs-on: *linux
timeout-minutes: 20
steps:
- name: Checkout Code
@@ -483,18 +487,17 @@ jobs:
run: |
python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
DEBUG=4 python test/test_tiny.py TestTiny.test_plus
- name: Run backend tests
run: python -m pytest -n=auto test/backend --durations=20
- name: Run runtime tests
run: python -m pytest -n=auto test/runtime test/device/$(python -c 'from tinygrad import Device; print(Device.DEFAULT.lower())') --durations=20
- name: Run process replay tests
uses: ./.github/actions/process-replay
testamdasm:
name: AMD ASM IDE
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
runs-on: *linux
timeout-minutes: 20
env:
DEV: MOCKKFD+AMD
HCQ_RUNTIME_DEV: PYTHON
steps:
- name: Checkout Code
uses: actions/checkout@v6
@@ -533,16 +536,14 @@ jobs:
PYTHONPATH="." DEV=MOCKKFD+AMD N=256 python3 extra/gemm/amd_asm_matmul.py
PYTHONPATH="." DEV=MOCKKFD+AMD N=256 python3 extra/gemm/amd_copy_matmul.py
- name: Run LLVM test
run: DEV=MOCKKFD+AMD:LLVM python test/device/test_amd_llvm.py
run: DEV=MOCKKFD+AMD:LLVM python test/device/amd/test_llvm.py
testmockam:
name: Linux (am)
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
runs-on: *linux
timeout-minutes: 15
env:
DEV: MOCKPCI+AMD
HCQ2: 1
HCQ_RUNTIME_DEV: PYTHON
PYTHONPATH: .
steps:
- name: Checkout Code
@@ -554,13 +555,13 @@ jobs:
deps: testing_unit
amd: 'true'
- name: Run tests on MOCKAM
run: python -m pytest test/test_tiny.py test/unit/test_jit.py
run: python -m pytest test/test_tiny.py test/runtime/test_jit.py
- name: Run test_tiny on MOCKUSB
run: HCQ2=0 GMMU=0 DEV=MOCKUSB+AMD python test/test_tiny.py
run: CHECK_OOB=0 GMMU=0 DEV=MOCKUSB+AMD python -m pytest test/test_tiny.py test/runtime/test_hcq2.py
- name: Run test_hcq2 on MOCKPCI
run: python -m pytest test/device/test_hcq2.py
run: python -m pytest test/null/test_hcq2.py test/runtime/test_hcq2.py
- name: Run disk copy tests on MOCKPCI
run: python -m pytest test/unit/test_disk_tensor.py -k test_copy_from_disk
run: python -m pytest test/runtime/test_disk_tensor.py -k test_copy_from_disk
testamd:
strategy:
@@ -570,13 +571,11 @@ jobs:
arch: [gfx1100, gfx1201, gfx950]
name: Linux (${{ matrix.backend }} ${{ matrix.arch }})
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
timeout-minutes: 15
runs-on: *linux
timeout-minutes: 20
env:
DEV: MOCKKFD+AMD:${{ matrix.backend == 'amdllvm' && 'LLVM' || '' }}:${{ matrix.arch }}
SKIP_SLOW_TEST: 1
HCQ2: 1
HCQ_RUNTIME_DEV: PYTHON
PYTHONPATH: .
steps:
- name: Checkout Code
@@ -592,20 +591,36 @@ jobs:
run: |
python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['AMD'], Device.DEFAULT"
DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
- name: Run MXFP4 Llama training on NULL backend
if: ${{ matrix.backend == 'amd' && matrix.arch == 'gfx950' }}
run: PYTHONPATH=. DEV=NULL:HIP:gfx950 MXFP4=1 LLAMA_LAYERS=2 BENCHMARK=3 NULL_ALLOW_COPYOUT=1 NO_HIPCC=1 ROCM_PATH=/opt/rocm JITBEAM=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh
- name: Run pytest (amd)
run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq2.py test/external/external_test_am.py test/backend/test_asm_gemm.py::TestAsmGEMM --durations=20
- name: Run opt tests
run: python -m pytest -n=auto test/opt --durations=20
- name: Run disk copy tests
run: python -m pytest test/unit/test_disk_tensor.py -k test_copy_from_disk
- name: Run TRANSCENDENTAL math
run: TRANSCENDENTAL=2 python -m pytest -n=auto test/backend/test_ops.py::TestOps::test_sin test/backend/test_ops.py::TestOps::test_cos test/backend/test_ops.py::TestOps::test_tan test/backend/test_ops.py::TestOps::test_exp test/backend/test_ops.py::TestOps::test_log --durations=20
run: >-
python -m pytest -n=auto test/runtime/test_ops.py test/runtime/test_dtype.py test/runtime/test_dtype_alu.py
test/runtime/test_kernel_opts.py test/runtime/test_linearizer.py test/runtime/test_randomness.py test/runtime/test_jit.py
test/runtime/test_multitensor.py test/runtime/test_search.py test/runtime/test_tensor_cores.py test/null/test_hcq2.py
test/runtime/test_hcq2.py test/external/external_test_am.py test/device/amd
test/runtime/test_disk_tensor.py::TestDiskTensor::test_copy_from_disk
test/runtime/test_disk_tensor.py::TestDiskTensor::test_copy_from_disk_offset
test/runtime/test_disk_tensor.py::TestDiskTensor::test_copy_from_disk_huge --durations=20
- name: Run process replay tests
uses: ./.github/actions/process-replay
testamdmxfp4:
name: AMD MXFP4 Llama
runs-on: *linux
timeout-minutes: 15
env:
DEV: NULL:HIP:gfx950
steps:
- name: Checkout Code
uses: actions/checkout@v6
- name: Setup Environment
uses: ./.github/actions/setup-tinygrad
with:
key: gfx950-mxfp4
deps: testing_unit
amd: 'true'
- name: Run MXFP4 Llama training on NULL backend
run: MXFP4=1 LLAMA_LAYERS=2 BENCHMARK=3 NULL_ALLOW_COPYOUT=1 NO_HIPCC=1 ROCM_PATH=/opt/rocm JITBEAM=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh
testnvidia:
strategy:
fail-fast: false
@@ -613,7 +628,7 @@ jobs:
backend: [ptx, nv]
name: Linux (${{ matrix.backend }})
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
runs-on: *linux
timeout-minutes: 20
env:
FORWARD_ONLY: 1
@@ -628,16 +643,20 @@ jobs:
cuda: 'true'
ocelot: 'true'
- name: Set env
run: printf "${{ matrix.backend == 'ptx' && 'DEV=MOCK+CUDA:PTX' || matrix.backend == 'nv' && 'DEV=MOCK+NV\nSKIP_SLOW_TEST=1\nHCQ_RUNTIME_DEV=PYTHON' }}" >> $GITHUB_ENV
run: printf "${{ matrix.backend == 'ptx' && 'DEV=MOCK+CUDA:PTX' || matrix.backend == 'nv' && 'DEV=MOCK+NV\nSKIP_SLOW_TEST=1' }}" >> $GITHUB_ENV
- name: Check Device.DEFAULT and print some source
run: |
python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CUDA','NV'], Device.DEFAULT"
DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
- name: Run pytest (cuda)
# skip multitensor because it's slow
run: python -m pytest -n=auto test/backend --ignore test/backend/test_multitensor.py --durations=20
# NV and CUDA share the NVIDIA device tests.
run: python -m pytest -n=auto test/runtime test/device/nv --ignore test/runtime/test_multitensor.py --durations=20
- name: Run NV PTX device tests
if: matrix.backend == 'nv'
run: DEV=MOCK+NV:PTX python -m pytest -n=auto test/device/nv --durations=20
- name: Run TestOps.test_add with PMA
run: VIZ=-1 PMA=1 PMA_BUFFER_SIZE=8 DEBUG=5 python3 test/backend/test_ops.py TestOps.test_add
run: VIZ=-1 PMA=1 PMA_BUFFER_SIZE=8 DEBUG=5 python3 test/runtime/test_ops.py TestOps.test_add
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -652,7 +671,7 @@ jobs:
- 'NULL:QCOMCL:a630'
- 'NULL:NAK:sm_120'
name: Compile-only (DEV=${{ matrix.dev }})
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
runs-on: *linux
timeout-minutes: 15
env:
NULL_ALLOW_COPYOUT: 1
@@ -670,11 +689,11 @@ jobs:
- name: Test IMAGE
shell: bash
if: contains(matrix.dev, 'a630')
run: DEBUG=7 python3 test/backend/test_ops.py TestOps.test_gemm | grep isam
run: DEBUG=7 python3 test/runtime/test_ops.py TestOps.test_gemm | grep isam
- name: Run test_ops
shell: bash
run: |
python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
DEBUG=4 python3 test/runtime/test_ops.py TestOps.test_add
# QCOMCL compiles in qemu, too slow for parallel workers
${{ contains(matrix.dev, 'QCOMCL') && 'PARALLEL=0' || '' }} python -m pytest -n=auto test/backend/test_ops.py --durations=20
${{ contains(matrix.dev, 'QCOMCL') && 'PARALLEL=0' || '' }} python -m pytest -n=auto test/runtime/test_ops.py --durations=20
tinygrad_repo/.pre-commit-config.yaml
+1 -1
View File
@@ -28,7 +28,7 @@ repos:
pass_filenames: false
- id: tests
name: comprehensive test suite
entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/backend/test_ops.py test/backend/test_schedule.py test/backend/test_assign.py test/backend/test_tensor.py test/backend/test_jit.py test/unit/test_schedule_cache.py test/null/test_pattern_matcher.py test/null/test_uop_symbolic.py test/unit/test_helpers.py
entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/runtime/test_ops.py test/runtime/test_schedule.py test/runtime/test_assign.py test/runtime/test_tensor.py test/runtime/test_jit.py test/runtime/test_schedule_cache.py test/null/test_pattern_matcher.py test/null/test_uop_symbolic.py test/runtime/test_helpers.py
language: system
always_run: true
pass_filenames: false
tinygrad_repo/README.md
+1 -1
View File
@@ -194,7 +194,7 @@ For more examples on how to run the full test suite please refer to the [CI work
Some examples of running tests locally:
```sh
python3 -m pip install -e '.[testing]' # install extra deps for testing
python3 test/backend/test_ops.py # just the ops tests
python3 test/runtime/test_ops.py # just the ops tests
python3 -m pytest test/ # whole test suite
```
tinygrad_repo/conftest.py
+7 -1
View File
@@ -8,6 +8,12 @@ def pytest_configure(config):
@pytest.hookimpl(wrapper=True)
def pytest_runtest_call(item):
faulthandler.dump_traceback_later(int(os.getenv("TEST_TIMEOUT", 120)), file=item.config.stash[stderr_fd_key], exit=True)
faulthandler.dump_traceback_later(int(os.getenv("TEST_TIMEOUT", 180)), file=item.config.stash[stderr_fd_key], exit=True)
try: yield
finally: faulthandler.cancel_dump_traceback_later()
@pytest.fixture(autouse=True)
def clear_runtime_cache():
yield
from tinygrad.engine.realize import runtime_cache
runtime_cache.clear()
tinygrad_repo/docs/abstractions4.py
+2 -2
View File
@@ -81,7 +81,7 @@ def example_3_custom_uop(a:Tensor, correct):
lane = UOp.range(LCLS, 1, AxisType.LOCAL)
# accumulate the globals into a per lane accumulator
reduce_loop = UOp.range(buf.shape[1], 2, AxisType.REDUCE)
reduce_loop = UOp.range(buf.shape[1], 2, AxisType.LOOP)
acc = UOp.placeholder((1,), dtypes.float, slot=6, addrspace=AddrSpace.REG)
acc = acc.after(acc.store(0))
acc = acc.after(acc[0].store(acc.after(reduce_loop)[0] + buf[glbl, reduce_loop, lane]).end(reduce_loop))
@@ -91,7 +91,7 @@ def example_3_custom_uop(a:Tensor, correct):
local_accs = local_accs.after(local_accs[lane].store(acc[0]))
# accumulate LOCALs into a single per CU accumulator
late_reduce_loop = UOp.range(LCLS, 3, AxisType.REDUCE)
late_reduce_loop = UOp.range(LCLS, 3, AxisType.LOOP)
acc2 = UOp.placeholder((1,), dtypes.float, slot=7, addrspace=AddrSpace.REG)
acc2 = acc2.after(acc2.store(0))
acc2 = acc2.after(acc2[0].store(acc2.after(late_reduce_loop)[0] + local_accs[late_reduce_loop]).end(late_reduce_loop))[0]
tinygrad_repo/docs/developer/am.md
+2
View File
@@ -16,6 +16,8 @@ Optional requirements:
|----------|------------------|-------------|
| AM_RESET | [1] | Performs a full GPU reset (reloading all firmware and IP blocks) |
| AM_DEBUG | [0-4] | Sets the level of additional debugging information |
| AM_POWER_LIMIT | [watts] | Sets the GPU power limit |
| AM_PERF_LEVEL | [auto, high, low] | Sets the clock policy (default: auto with AM_POWER_LIMIT, high without) |
## AM Driver Details
tinygrad_repo/docs/developer/developer.md
+1 -1
View File
@@ -33,7 +33,7 @@ Then we render the UOps into code with a `Renderer`, then we compile the code to
## Execution
`run_linear` walks the `LINEAR` UOp, dispatching each `CALL` to a runner (kernel, copy, view, encdec, or graph).
`run_linear` walks the `LINEAR` UOp, dispatching each `CALL` to a runner (kernel, hcq batch, copy, encdec, or validate).
## Runtime
tinygrad_repo/docs/env_vars.md
+1 -1
View File
@@ -36,7 +36,7 @@ BEAM | [#] | number of beams in kernel beam search
DEFAULT_FLOAT | [HALF, ...]| specify the default float dtype (FLOAT32, HALF, BFLOAT16, FLOAT64, ...), default to FLOAT32
IMAGE | [1] | enable 2d specific optimizations
FLOAT16 | [1] | use float16 for images instead of float32
JIT | [0-2] | 0=disabled, 1=[jit enabled](quickstart.md#jit) (default), 2=jit enabled, but graphs are disabled
JIT | [0-2] | 0=disabled, 1=[jit enabled](quickstart.md#jit) (default), 2=same as 1 (kept for compatibility)
VIZ | [1] | 0=disabled, 1=[viz enabled](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/viz)
ALLOW_TF32 | [1] | enable TensorFloat-32 tensor cores on Ampere or newer GPUs.
WEBGPU_BACKEND | [WGPUBackendType_Metal, ...] | Force select a backend for WebGPU (Metal, DirectX, OpenGL, Vulkan...)
tinygrad_repo/docs/tensor/ops.md
+1
View File
@@ -52,6 +52,7 @@
::: tinygrad.Tensor.linear
::: tinygrad.Tensor.sequential
::: tinygrad.Tensor.layernorm
::: tinygrad.Tensor.rmsnorm
::: tinygrad.Tensor.batchnorm
::: tinygrad.Tensor.dropout
::: tinygrad.Tensor.one_hot
tinygrad_repo/examples/llm.c/export.py
+1 -1
View File
@@ -99,6 +99,6 @@ if __name__ == "__main__":
#print(srcs[si.ast][1])
main.append("}")
mallocs = [f" {b.dtype.name}* {n} = ({b.dtype.name}*)malloc({b.nbytes});" for n,b in dedup(all_bufs)]
mallocs = [f" {(dt:=Device['CPU'].renderer.render_dtype(b.dtype))}* {n} = ({dt}*)malloc({b.nbytes});" for n,b in dedup(all_bufs)]
with open("out.c", "w") as f: f.write('\n'.join(c_code+premain+mallocs+main))
tinygrad_repo/examples/mlperf/model_train.py
+5 -32
View File
@@ -1282,7 +1282,7 @@ def train_bert():
previous_step = i
def train_llama3():
from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad, FP8_DTYPE, MXFP8, MXFP4
from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad, MXFP4
from examples.llama3 import MODEL_PARAMS
from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
from examples.mlperf.optim import GradAccClipAdamW, clip_grads
@@ -1418,8 +1418,7 @@ def train_llama3():
eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device)
for p in optim.params:
grad_dtype = dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype
p.grad = p.zeros_like(dtype=grad_dtype).contiguous()
p.grad = p.zeros_like().contiguous()
grads = [p.grad for p in optim.params]
scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
@@ -1433,39 +1432,15 @@ def train_llama3():
print(f"loading optim checkpoint from {fn}")
load_state_dict(scheduler, safe_load(fn), realize=False)
fp8_amax = [t for ts in model._fp8_amax.values() for t in ts]
fp8_next_amax = [t for ts in model._fp8_next_amax.values() for t in ts]
fp8_grad_amax = [t for ts in model._fp8_grad_amax.values() for t in ts]
fp8_next_grad_amax = [t for ts in model._fp8_next_grad_amax.values() for t in ts]
fp8_inv_scales = list(model._fp8_inv_scale.values()) + list(model._fp8_next_inv_scale.values())
from tinygrad.nn.state import get_state_dict
model_state = get_state_dict(model)
for wname in model._fp8_inv_scale:
w = model_state[wname]
w._inv_scale = model._fp8_inv_scale[wname]
w._next_inv_scale = model._fp8_next_inv_scale[wname]
if optim.master_params:
idx = next(j for j, p in enumerate(optim.params) if p is w)
master = optim.master_params[idx]
inv = w._inv_scale if w._inv_scale.device == master.device else w._inv_scale.to(master.device)
if MXFP8:
from extra.gemm.cdna_asm_gemm import _mx_block_scale
bs = _mx_block_scale(inv.reshape(-1, inv.shape[-1])).reshape(w.shape)
master.assign((master * bs).contiguous())
else:
master.assign((master * inv.reshape(*inv.shape, *([1]*(w.ndim-inv.ndim)))).contiguous())
# realize everything here
if optim.master_params: Tensor.realize(*optim.master_params)
loss_acc = Tensor.zeros(1, dtype=dtypes.float32, device=device)
Tensor.realize(loss_acc, *optim.params, *fp8_inv_scales, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax)
Tensor.realize(loss_acc, *optim.params)
mxfp4_weights = model.create_mxfp4_weight_cache() if MXFP4 else None
if mxfp4_weights is not None: Tensor.realize(*[x for layers in mxfp4_weights.values() for outputs in layers for x in outputs])
@TinyJit
def minibatch(tokens:Tensor):
model.reset_amax()
if is_dp: tokens = tokens.to(None).shard(device, 0)
if is_mp: tokens = tokens.shard(device)
if not is_sharding: tokens = tokens.to(None)
@@ -1480,7 +1455,7 @@ def train_llama3():
apply_grad(g, new_g.uop)
loss_acc.assign(loss_acc + loss.flatten().float())
return loss_acc.realize(*grads, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax)
return loss_acc.realize(*grads)
@TinyJit
def optim_step():
@@ -1489,14 +1464,12 @@ def train_llama3():
scheduler.step()
for g in grads: g.assign(0)
model.update_amax()
new_mxfp4_w = model.update_mxfp4_weight_cache(mxfp4_weights) if mxfp4_weights is not None else []
lr_cpu = optim.lr.float().to("CPU")
grad_norm_cpu = grad_norm.float().to("CPU")
loss_cpu = loss_acc.to("CPU")
Tensor.realize(lr_cpu, grad_norm_cpu, loss_cpu, loss_acc.assign(0), *grads, *fp8_inv_scales, *fp8_amax, *fp8_grad_amax,
*new_mxfp4_w)
Tensor.realize(lr_cpu, grad_norm_cpu, loss_cpu, loss_acc.assign(0), *grads, *new_mxfp4_w)
return lr_cpu, grad_norm_cpu, loss_cpu
tinygrad_repo/examples/mlperf/models/flat_llama.py
+40 -223
View File
@@ -16,123 +16,39 @@ from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker, ro
from tinygrad.uop.ops import Ops, UOp
from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
from extra.llama_kernels.rmsnorm import rmsnorm
from extra.llama_kernels import FP8_MAX, local_abs_max
ASM_GEMM = getenv("ASM_GEMM", 0)
FUSED_INPUT_QUANTIZE = getenv("FUSED_INPUT_QUANTIZE", 0)
FUSED_ADD_NORM_MUL_QUANTIZE = getenv("FUSED_ADD_NORM_MUL_QUANTIZE", 0)
FUSED_SILU_W13 = getenv("FUSED_SILU_W13", 0)
SPLIT_W13 = getenv("SPLIT_W13", 0)
COLUMNWISE_WEIGHT_SCALE = getenv("COLUMNWISE_WEIGHT_SCALE", 0)
MXFP8 = getenv("MXFP8", 0)
MXFP4 = getenv("MXFP4", 0)
FP8_DTYPE = dtypes.fp8e4m3
FP8_GRAD_DTYPE = dtypes.fp8e5m2
def quantize_fp8(x:Tensor, amax_state:Tensor|None=None):
new_amax = (local_abs_max(x) if isinstance(x.device, tuple) else x.abs().max()).detach().cast(dtypes.float32)
scale = FP8_MAX / ((amax_state if amax_state is not None else new_amax) + 1e-8)
x_scaled = x * scale
x_clamped = x_scaled + (x_scaled.detach().clamp(-FP8_MAX, FP8_MAX) - x_scaled.detach()) # STE
return x_clamped.cast(FP8_DTYPE), scale.float().reciprocal(), new_amax
def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_scale:Tensor|None=None,
x_fp8:Tensor|None=None, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None, x_prequant_mx:tuple|None=None,
next_amax_x:Tensor|None=None, mxfp4_w:tuple[Tensor, Tensor, Tensor, Tensor]|None=None,
def matmul(x:Tensor, w:Tensor, mxfp4:bool=bool(MXFP4), mxfp4_w:tuple[Tensor, Tensor, Tensor, Tensor]|None=None,
x_prequant_mxfp4:tuple[Tensor|None, Tensor|None, Tensor|None, Tensor|None]|None=None) -> tuple[Tensor,...]:
if not fp8:
if ASM_GEMM:
from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T),)
return (x @ w.T,)
if MXFP4:
assert x is not None, "MXFP4 matmul requires an unquantized input"
if mxfp4 or ASM_GEMM:
from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm
if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T, mxfp4=True, mxfp4_w=mxfp4_w, mxfp4_x=x_prequant_mxfp4),)
return (x @ w.T,)
assert w_inv_scale is not None, "fp8 matmul requires w_inv_scale (weights must be stored in fp8 with per-tensor scale)"
if MXFP8:
from extra.gemm.cdna_asm_gemm import asm_gemm, quantize_mxfp8, mx_pack, can_use_asm_gemm, _mx_block_scale
if x_prequant_mx is not None: x_q, x_e8, x_si = x_prequant_mx # fused producer already quantized (2d)
else: x_q, x_e8, x_si = quantize_mxfp8(x.reshape(-1, x.shape[-1]))
l_shape = x.shape[:-1] if x is not None else x_q.shape[:-1]
if can_use_asm_gemm(x_q, w.T):
out = asm_gemm(x_q, w.T, mx=True, mx_scales=(x_si, x_e8, mx_pack(w_inv_scale), w_inv_scale),
mx_w_stored=True).reshape(*l_shape, w.shape[0])
else:
x_phys = (x_q.cast(dtypes.bfloat16) * _mx_block_scale(x_e8)).reshape(*l_shape, x_q.shape[-1])
out = x_phys @ (w.cast(dtypes.bfloat16) * _mx_block_scale(w_inv_scale)).T
return out, x_q
if x_fp8 is None:
if FUSED_INPUT_QUANTIZE:
from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed
x_fp8, _ = quantize_fp8_delayed(x, amax_x, next_amax_x, FP8_DTYPE)
else:
x_fp8, _, new_amax_x = quantize_fp8(x, amax_state=amax_x)
next_amax_x.assign(new_amax_x)
if ASM_GEMM:
from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
if can_use_asm_gemm(x_fp8, w.T):
assert amax_x is not None
if COLUMNWISE_WEIGHT_SCALE:
out = asm_gemm(x_fp8, w.T, x_scale=amax_x, grad_amax_state=grad_amax_state,
next_grad_amax_state=next_grad_amax_state, w_post_scale=w_inv_scale)
else:
out = asm_gemm(x_fp8, w.T, x_scale=amax_x, w_scale=w_inv_scale, grad_amax_state=grad_amax_state,
next_grad_amax_state=next_grad_amax_state)
return out, x_fp8
return (x_fp8.dot(w.T, dtype=dtypes.float) * ((amax_x.float() + 1e-8) / FP8_MAX) * w_inv_scale).cast(dtypes.bfloat16), x_fp8
if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T, mxfp4=mxfp4, mxfp4_w=mxfp4_w, mxfp4_x=x_prequant_mxfp4),)
return (x @ w.T,)
def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor|None,
next_amax_x:Tensor|None, grad_amax_state:Tensor|None, next_grad_amax_state:Tensor|None, mxfp4_w=None):
if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4:
from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_rmsnorm_mul_quantize_fp8
x_fp8, x_normed, rrms = fused_rmsnorm_mul_quantize_fp8(x, norm, amax_x, eps, FP8_DTYPE, next_amax_x)
out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x,
grad_amax_state=grad_amax_state, next_grad_amax_state=next_grad_amax_state)
return out, x_normed, rrms, ret
def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, eps:float, mxfp4_w=None):
x_normed, rrms = rmsnorm(x, eps)
out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale, grad_amax_state=grad_amax_state,
next_grad_amax_state=next_grad_amax_state, next_amax_x=next_amax_x, mxfp4_w=mxfp4_w)
out, *ret = matmul(x_normed * norm, w, mxfp4_w=mxfp4_w)
return out, x_normed, rrms, ret
def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor|None,
next_amax_x:Tensor|None, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None, mxfp4_w=None):
if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4:
from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_add_rmsnorm_mul_quantize_fp8
x_fp8, h, x_normed, rrms = fused_add_rmsnorm_mul_quantize_fp8(x, residual, norm, amax_x, eps, FP8_DTYPE, next_amax_x)
out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x,
grad_amax_state=grad_amax_state, next_grad_amax_state=next_grad_amax_state)
return out, h, x_normed, rrms, ret
def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, eps:float, mxfp4_w=None):
h = x + residual
x_normed, rrms = rmsnorm(h, eps)
out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale, grad_amax_state=grad_amax_state,
next_grad_amax_state=next_grad_amax_state, next_amax_x=next_amax_x, mxfp4_w=mxfp4_w)
out, *ret = matmul(x_normed * norm, w, mxfp4_w=mxfp4_w)
return out, h, x_normed, rrms, ret
def silu_w13_quantize_matmul(x_w13:Tensor, w2:Tensor, s_2:Tensor,
amax_x2:Tensor|None, next_amax_x2:Tensor|None,
grad_amax_xw13:Tensor|None, next_grad_amax_xw13:Tensor|None,
grad_amax_xout:Tensor|None, next_grad_amax_xout:Tensor|None, mxfp4_w=None):
def silu_w13_quantize_matmul(x_w13:Tensor, w2:Tensor, mxfp4_w=None):
if FUSED_SILU_W13 and MXFP4:
from extra.llama_kernels.swiglu import swiglu_mxfp4
x2, x2_mxfp4 = swiglu_mxfp4(x_w13)
out, *ret = matmul(x2, w2, amax_x=amax_x2, w_inv_scale=s_2, grad_amax_state=grad_amax_xout,
next_grad_amax_state=next_grad_amax_xout, next_amax_x=next_amax_x2, mxfp4_w=mxfp4_w,
x_prequant_mxfp4=x2_mxfp4)
return out, ret
if FUSED_SILU_W13 and not MXFP4:
from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13
x2_fp8 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13,
next_grad_amax_state=next_grad_amax_xw13, amax_out=next_amax_x2)
out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, amax_x=amax_x2,
grad_amax_state=grad_amax_xout, next_grad_amax_state=next_grad_amax_xout)
out, *ret = matmul(x2, w2, mxfp4_w=mxfp4_w, x_prequant_mxfp4=x2_mxfp4)
return out, ret
hidden = x_w13.shape[-1] // 2
x_w1, x_w3 = x_w13[..., :hidden], x_w13[..., hidden:]
out, *ret = matmul(x_w1.silu() * x_w3, w2, amax_x=amax_x2, w_inv_scale=s_2, grad_amax_state=grad_amax_xout,
next_grad_amax_state=next_grad_amax_xout, next_amax_x=next_amax_x2, mxfp4_w=mxfp4_w)
out, *ret = matmul(x_w1.silu() * x_w3, w2, mxfp4_w=mxfp4_w)
return out, ret
class FlatTransformer:
@@ -149,16 +65,16 @@ class FlatTransformer:
scaled_std = 0.02 / math.sqrt(2 * n_layers)
# Attention
self.wqkv, s_qkv = self.lin_per_layer(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2)
self.wo, s_o = self.lin_per_layer(self.n_heads * self.head_dim, dim, std=scaled_std)
self.wqkv = self.lin_per_layer(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2)
self.wo = self.lin_per_layer(self.n_heads * self.head_dim, dim, std=scaled_std)
# FeedForward
if SPLIT_W13:
self.w1, s_1 = self.lin_per_layer(dim, hidden_dim)
self.w3, s_3 = self.lin_per_layer(dim, hidden_dim)
self.w1 = self.lin_per_layer(dim, hidden_dim)
self.w3 = self.lin_per_layer(dim, hidden_dim)
else:
self.w13, s_13 = self.lin_per_layer(dim, hidden_dim * 2)
self.w2, s_2 = self.lin_per_layer(hidden_dim, dim, std=scaled_std)
self.w13 = self.lin_per_layer(dim, hidden_dim * 2)
self.w2 = self.lin_per_layer(hidden_dim, dim, std=scaled_std)
self.norm_eps = norm_eps
self.attention_norm = Tensor.ones(n_layers, dim).contiguous()
@@ -171,50 +87,17 @@ class FlatTransformer:
self.output = Tensor.normal(1, vocab_size, dim, mean=0.0, std=0.02, dtype=dtypes.bfloat16)
self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).clone().is_param_(False)
def _amax(): return Tensor.full((), FP8_MAX, dtype=dtypes.float32).contiguous().is_param_(False)
n_amax = 0 if MXFP4 else n_layers
names = ["xqkv", "xo", "x2"]
names += ["x1", "x3"] if SPLIT_W13 else ["x13"]
self._fp8_amax = {name: [_amax() for _ in range(n_amax)] for name in names}
self._fp8_next_amax = {name: [_amax() for _ in range(n_amax)] for name in names}
grad_names = ["xqkv", "xo", "xout"]
grad_names += ["xw1", "xw3"] if SPLIT_W13 else ["xw13"]
self._fp8_grad_amax = {name: [_amax() for _ in range(n_amax)] for name in grad_names}
self._fp8_next_grad_amax = {name: [_amax() for _ in range(n_amax)] for name in grad_names}
w_scales = [("wqkv", s_qkv), ("wo", s_o), ("w2", s_2)]
w_scales += [("w1", s_1), ("w3", s_3)] if SPLIT_W13 else [("w13", s_13)]
self._fp8_inv_scale = {name: (s if MXFP8 else s.float()).contiguous().is_param_(False) for name, s in w_scales}
self._fp8_next_inv_scale = {name: (s if MXFP8 else s.float()).contiguous().is_param_(False) for name, s in w_scales}
def lin_per_layer(self, in_features:int, out_features:int, std:float=0.02, w:Tensor|None=None):
if w is None:
if getenv("ZEROS"): w = Tensor.zeros(self.n_layers, out_features, in_features)
else: w = Tensor.normal(self.n_layers, out_features, in_features, mean=0.0, std=std)
if MXFP8:
from extra.gemm.cdna_asm_gemm import quantize_mxfp8
w_q, w_e8, _ = quantize_mxfp8(w.reshape(self.n_layers * out_features, in_features))
return w_q.reshape(self.n_layers, out_features, in_features), w_e8.reshape(self.n_layers, out_features, in_features // 32)
if MXFP4:
# FP4 is produced dynamically so optimizer updates always start from the current BF16 weight.
return w.cast(dtypes.bfloat16), Tensor.ones(self.n_layers)
amax = (w.abs().max(axis=2) if COLUMNWISE_WEIGHT_SCALE else w.abs().flatten(1).max(1)).detach()
scale = FP8_MAX / (amax + 1e-8)
inv_scale = (amax + 1e-8) / FP8_MAX
scale_b = scale.reshape(self.n_layers, out_features, 1) if COLUMNWISE_WEIGHT_SCALE else scale.reshape(-1, 1, 1)
return (w * scale_b).clamp(-FP8_MAX, FP8_MAX).cast(FP8_DTYPE), inv_scale
return w.cast(dtypes.bfloat16)
def attention(self, x:Tensor, freqs_cis:Tensor, *, attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
amax_xqkv:Tensor|None, amax_xo:Tensor|None, s_qkv:Tensor, s_o:Tensor,
next_amax_xqkv:Tensor|None, next_amax_xo:Tensor|None,
grad_amax_xqkv:Tensor|None, grad_amax_xo:Tensor|None,
next_grad_amax_xqkv:Tensor|None, next_grad_amax_xo:Tensor|None, mxfp4_wqkv=None, mxfp4_wo=None):
def attention(self, x:Tensor, freqs_cis:Tensor, *, attention_norm:Tensor, wqkv:Tensor, wo:Tensor, mxfp4_wqkv=None, mxfp4_wo=None):
bsz, seqlen, _ = x.shape
saves = []
xqkv, x_normed, rrms, s = norm_quantize_matmul(x, attention_norm, wqkv, s_qkv, self.norm_eps,
amax_x=amax_xqkv, grad_amax_state=grad_amax_xqkv,
next_grad_amax_state=next_grad_amax_xqkv, next_amax_x=next_amax_xqkv,
mxfp4_w=mxfp4_wqkv)
xqkv, x_normed, rrms, s = norm_quantize_matmul(x, attention_norm, wqkv, self.norm_eps, mxfp4_w=mxfp4_wqkv)
saves.extend([x_normed, rrms, *s, xqkv])
if getenv("HK_FLASH_ATTENTION"):
from extra.thunder.amd.fa import flash_attention, fused_qkv_rope
@@ -232,8 +115,7 @@ class FlatTransformer:
attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
attn = attn.reshape(bsz, seqlen, -1)
out, *s = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o, grad_amax_state=grad_amax_xo,
next_grad_amax_state=next_grad_amax_xo, next_amax_x=next_amax_xo, mxfp4_w=mxfp4_wo)
out, *s = matmul(attn, wo, mxfp4_w=mxfp4_wo)
saves.extend([*s, out])
return out, saves
@@ -245,41 +127,17 @@ class FlatTransformer:
x_normed, rrms = rmsnorm(h, self.norm_eps)
saves.extend([x_normed, rrms])
inp = x_normed * kwargs["ffn_norm"]
x_w1, *s = matmul(inp, kwargs["w1"], amax_x=kwargs["amax_x1"], w_inv_scale=kwargs["s_1"],
grad_amax_state=kwargs["grad_amax_xw1"], next_grad_amax_state=kwargs["next_grad_amax_xw1"],
next_amax_x=kwargs["next_amax_x1"], mxfp4_w=kwargs.get("mxfp4_w1"))
x_w1, *s = matmul(inp, kwargs["w1"], mxfp4_w=kwargs.get("mxfp4_w1"))
saves.extend([*s, x_w1])
x_w3, *s = matmul(inp, kwargs["w3"], amax_x=kwargs["amax_x3"], w_inv_scale=kwargs["s_3"],
grad_amax_state=kwargs["grad_amax_xw3"], next_grad_amax_state=kwargs["next_grad_amax_xw3"],
next_amax_x=kwargs["next_amax_x3"], mxfp4_w=kwargs.get("mxfp4_w3"))
x_w3, *s = matmul(inp, kwargs["w3"], mxfp4_w=kwargs.get("mxfp4_w3"))
saves.extend([*s, x_w3])
if FUSED_SILU_W13 and MXFP8:
from extra.llama_kernels.fused_silu_mul_quantize_mxfp8 import fused_silu_mul_quantize_mxfp8
aq, ae8, asi = fused_silu_mul_quantize_mxfp8(x_w1.reshape(-1, x_w1.shape[-1]), x_w3.reshape(-1, x_w3.shape[-1]))
out, *s = matmul(None, kwargs["w2"], x_prequant_mx=(aq, ae8, asi), amax_x=kwargs["amax_x2"],
w_inv_scale=kwargs["s_2"], grad_amax_state=kwargs["grad_amax_xout"],
next_grad_amax_state=kwargs["next_grad_amax_xout"], next_amax_x=kwargs["next_amax_x2"])
out = out.reshape(*x_w1.shape[:-1], kwargs["w2"].shape[0])
else:
out, *s = matmul(x_w1.silu() * x_w3, kwargs["w2"], amax_x=kwargs["amax_x2"], w_inv_scale=kwargs["s_2"],
grad_amax_state=kwargs["grad_amax_xout"], next_grad_amax_state=kwargs["next_grad_amax_xout"],
next_amax_x=kwargs["next_amax_x2"], mxfp4_w=kwargs.get("mxfp4_w2"))
out, *s = matmul(x_w1.silu() * x_w3, kwargs["w2"], mxfp4_w=kwargs.get("mxfp4_w2"))
saves.extend([*s, out])
else:
x_w13, h, x_normed, rrms, s = add_norm_quantize_matmul(x, residual, kwargs["ffn_norm"], kwargs["w13"], kwargs["s_13"],
self.norm_eps, amax_x=kwargs["amax_x13"],
next_amax_x=kwargs["next_amax_x13"],
grad_amax_state=kwargs["grad_amax_xw13"],
next_grad_amax_state=kwargs["next_grad_amax_xw13"],
mxfp4_w=kwargs.get("mxfp4_w13"))
x_w13, h, x_normed, rrms, s = add_norm_quantize_matmul(x, residual, kwargs["ffn_norm"], kwargs["w13"],
self.norm_eps, mxfp4_w=kwargs.get("mxfp4_w13"))
saves.extend([x_normed, rrms, *s, x_w13])
out, s = silu_w13_quantize_matmul(x_w13, kwargs["w2"], kwargs["s_2"], amax_x2=kwargs["amax_x2"],
next_amax_x2=kwargs["next_amax_x2"],
grad_amax_xw13=kwargs["grad_amax_xw13"],
next_grad_amax_xw13=kwargs["next_grad_amax_xw13"],
grad_amax_xout=kwargs["grad_amax_xout"],
next_grad_amax_xout=kwargs["next_grad_amax_xout"],
mxfp4_w=kwargs.get("mxfp4_w2"))
out, s = silu_w13_quantize_matmul(x_w13, kwargs["w2"], mxfp4_w=kwargs.get("mxfp4_w2"))
saves.extend([*s, out])
return out, h, saves
@@ -297,40 +155,20 @@ class FlatTransformer:
for v in get_parameters(self): v.shard_(device, axis=None)
else:
# flat per-layer weights: axis 0 is n_layers, so shard axes are +1 vs per-layer Transformer
def _shard_fp8(name:str, axis:int, std:float=0.02):
w = getattr(self, name)
if MXFP8:
from extra.gemm.cdna_asm_gemm import quantize_mxfp8
w_bf16 = Tensor.empty(self.n_layers, w.shape[1], w.shape[2], dtype=dtypes.bfloat16).shard(device, axis=axis).randn_like() * std
w_q, w_e8, _ = quantize_mxfp8(w_bf16)
w.replace(w_q)
self._fp8_inv_scale[name].replace(w_e8.contiguous()).is_param_(False)
self._fp8_next_inv_scale[name].replace(w_e8.contiguous()).is_param_(False)
else:
w.shard_(device, axis=axis)
scale_axis = (1 if axis == 1 else None) if COLUMNWISE_WEIGHT_SCALE else None
self._fp8_inv_scale[name] = self._fp8_inv_scale[name].shard(device, axis=scale_axis).contiguous().is_param_(False)
self._fp8_next_inv_scale[name] = self._fp8_next_inv_scale[name].shard(device, axis=scale_axis).contiguous().is_param_(False)
Tensor.realize(w, self._fp8_inv_scale[name], self._fp8_next_inv_scale[name])
sstd = 0.02 / math.sqrt(2 * self.n_layers)
_shard_fp8("wqkv", 1) # (n_layers, out, dim) shard out
_shard_fp8("wo", 2, sstd) # (n_layers, dim, in) shard in
self.wqkv.shard_(device, axis=1).realize() # (n_layers, out, dim) shard out
self.wo.shard_(device, axis=2).realize() # (n_layers, dim, in) shard in
if SPLIT_W13:
_shard_fp8("w1", 1)
_shard_fp8("w3", 1)
self.w1.shard_(device, axis=1).realize()
self.w3.shard_(device, axis=1).realize()
else:
_shard_fp8("w13", 1) # (n_layers, hidden*2, dim) shard out
_shard_fp8("w2", 2, sstd) # (n_layers, dim, hidden) shard in
self.w13.shard_(device, axis=1).realize() # (n_layers, hidden*2, dim) shard out
self.w2.shard_(device, axis=2).realize() # (n_layers, dim, hidden) shard in
self.attention_norm.shard_(device, axis=None).realize()
self.ffn_norm.shard_(device, axis=None).realize()
self.norm.weight.shard_(device, axis=None).realize()
self.tok_embeddings.weight.shard_(device, axis=0).realize()
self.output.shard_(device, axis=1).realize()
self.freqs_cis.shard_(device, axis=None).realize()
for amax_dict in (self._fp8_amax, self._fp8_next_amax, self._fp8_grad_amax, self._fp8_next_grad_amax):
for name in amax_dict:
for i in range(len(amax_dict[name])):
amax_dict[name][i] = amax_dict[name][i].to(device).contiguous().is_param_(False)
def create_mxfp4_weight_cache(self) -> dict[str, list[tuple[Tensor, Tensor, Tensor, Tensor]]]:
assert MXFP4
@@ -343,42 +181,26 @@ class FlatTransformer:
return [x for name, layers in cache.items() for weight, outputs in zip(getattr(self, name), layers)
for x in quantize_mxfp4(weight, shuffle_row=True, shuffle_col=True, out=outputs)]
def reset_amax(self):
for st in (self._fp8_next_amax, self._fp8_next_grad_amax):
for ts in st.values():
for t in ts: t.assign(0)
def update_amax(self):
for cur, nxt in ((self._fp8_amax, self._fp8_next_amax), (self._fp8_grad_amax, self._fp8_next_grad_amax)):
for name in cur:
for c, n in zip(cur[name], nxt[name]): c.assign(n)
def __call__(self, tokens:Tensor, save:bool=True,
mxfp4_weights:dict[str, list[tuple[Tensor, Tensor, Tensor, Tensor]]]|None=None):
h = self.tok_embeddings(tokens)
freqs_cis = self.freqs_cis.cast(h.dtype)
if not getenv("HK_FLASH_ATTENTION"): freqs_cis = freqs_cis[:, :tokens.shape[1], :, :, :]
a, na, ga, nga, s = self._fp8_amax, self._fp8_next_amax, self._fp8_grad_amax, self._fp8_next_grad_amax, self._fp8_inv_scale
def amax_kwargs(i:int, act_names:tuple[str, ...], grad_names:tuple[str, ...]) -> dict[str, Tensor|None]:
specs = (("amax_", a, act_names), ("next_amax_", na, act_names), ("grad_amax_", ga, grad_names), ("next_grad_amax_", nga, grad_names))
if MXFP4: return dict.fromkeys(f"{prefix}{name}" for prefix, _, names in specs for name in names)
return {f"{prefix}{name}":val[name][i] for prefix, val, names in specs for name in names}
for i in range(self.n_layers):
attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wo=self.wo[i], s_qkv=s["wqkv"][i], s_o=s["wo"][i],
**amax_kwargs(i, ("xqkv", "xo"), ("xqkv", "xo")))
ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], w2=self.w2[i], s_2=s["w2"][i], **amax_kwargs(i, ("x2",), ("xout",)))
attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wo=self.wo[i])
ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], w2=self.w2[i])
if mxfp4_weights is not None:
attn_kwargs.update(mxfp4_wqkv=mxfp4_weights["wqkv"][i], mxfp4_wo=mxfp4_weights["wo"][i])
ffn_kwargs.update(mxfp4_w2=mxfp4_weights["w2"][i])
if SPLIT_W13:
ffn_kwargs.update(w1=self.w1[i], w3=self.w3[i], s_1=s["w1"][i], s_3=s["w3"][i], **amax_kwargs(i, ("x1", "x3"), ("xw1", "xw3")))
ffn_kwargs.update(w1=self.w1[i], w3=self.w3[i])
if mxfp4_weights is not None: ffn_kwargs.update(mxfp4_w1=mxfp4_weights["w1"][i], mxfp4_w3=mxfp4_weights["w3"][i])
else:
ffn_kwargs.update(w13=self.w13[i], s_13=s["w13"][i], **amax_kwargs(i, ("x13",), ("xw13",)))
ffn_kwargs.update(w13=self.w13[i])
if mxfp4_weights is not None: ffn_kwargs.update(mxfp4_w13=mxfp4_weights["w13"][i])
h, *_ = self.run_layer(h, freqs_cis, attn_kwargs, ffn_kwargs, save=save)
logits = matmul(self.norm(h), self.output[0], fp8=False)[0]
logits = matmul(self.norm(h), self.output[0], mxfp4=False)[0]
return logits
def _get_pads(uop:UOp) -> list[UOp]:
@@ -437,11 +259,7 @@ if __name__ == "__main__":
if is_mp: vocab_mask.shard_(device, axis=2).realize()
# preallocate all the grad buffers and zero them out
grad_dtype = lambda x: dtypes.bfloat16 if x.dtype in dtypes.fp8s else x.dtype
grads = {x:x.zeros_like(dtype=grad_dtype(x)).contiguous() for x in state.values() if x.is_param}
fp8_amax = [t for ts in model._fp8_amax.values() for t in ts]
fp8_grad_amax = [t for ts in model._fp8_grad_amax.values() for t in ts]
grads = {x:x.zeros_like().contiguous() for x in state.values() if x.is_param}
# print model size
sz = 0
@@ -459,13 +277,12 @@ if __name__ == "__main__":
@TinyJit
def fwd_bwd(tokens:Tensor):
with Timing("python forward: "):
model.reset_amax()
logits = model(tokens[:, :-1], save=llama_size=="8B")
loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
with Timing("python backward: "):
for t,g in zip(grads, loss.gradient(*grads)):
apply_grad(grads[t], g.uop)
with Timing("run fwd_bwd: "): loss.realize(*grads.values(), *fp8_amax, *fp8_grad_amax)
with Timing("run fwd_bwd: "): loss.realize(*grads.values())
@TinyJit
def optim_step():
tinygrad_repo/examples/mlperf/models/gpt_oss.py
+64 -21
View File
@@ -14,7 +14,7 @@ from extra.models.llama import apply_rotary_emb
from extra.llama_kernels.rmsnorm import rmsnorm
from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8, asm_gemm, can_use_asm_gemm, mx_pack
from extra.gemm.moe_gemm import grouped_mx_gemm
from extra.gemm.moe_routing import route, dispatch, combine, router_mfma, Routing, BLOCK_ROW
from extra.gemm.moe_routing import route, dispatch, dispatch_fp8, combine, router_mfma, Routing, BLOCK_ROW
from extra.gptoss_kernels.embedding import GPTOSSEmbedding
FP8_DTYPE = dtypes.fp8e4m3
@@ -61,7 +61,7 @@ def dequant_weight(w_q:Tensor, w_scale:Tensor) -> Tensor:
fxn = _dequant_fwd_fxn(w_q.as_param(0).uop, w_scale.as_param(1).uop, w_q.device)
return Tensor(fxn.uop.call_with_output(w_q.uop, w_scale.uop, grad_fxn=_dequant_bwd))
def matmul_mx(x:Tensor|tuple[Tensor, Tensor], w_q:Tensor, w_scale:Tensor) -> Tensor:
def matmul_mx(x:Tensor|tuple[Tensor, Tensor], w_q:Tensor, w_scale:Tensor, *, padded_output:bool=False) -> Tensor:
if isinstance(x, tuple):
assert ASM_GEMM, "pre-quantized MXFP8 input requires ASM_GEMM"
from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm, mx_pack
@@ -79,6 +79,7 @@ def matmul_mx(x:Tensor|tuple[Tensor, Tensor], w_q:Tensor, w_scale:Tensor) -> Ten
ws = ws.pad(((0, npad), (0, 0)), value=127).cast(dtypes.uint8)
assert can_use_asm_gemm(x_q, wq.T)
out = asm_gemm(x_q, wq.T, mx=True, mx_scales=(mx_pack(x_e8), x_e8, mx_pack(ws), ws), mx_w_stored=True)
if padded_output: return out
return (out[:, :N] if npad else out).reshape(*l_shape, N).cast(dtypes.bfloat16)
l_shape = x.shape[:-1]
if ASM_GEMM:
@@ -100,7 +101,9 @@ def matmul_mx(x:Tensor|tuple[Tensor, Tensor], w_q:Tensor, w_scale:Tensor) -> Ten
else: x_q, x_e8, x_si = quantize_mxfp8(x2)
if x_si is not None and can_use_asm_gemm(x_q, wq.T):
out = asm_gemm(x_q, wq.T, mx=True, mx_scales=(x_si, x_e8, mx_pack(ws), ws), mx_w_stored=True)
if padded_output: return out
return (out[:, :N] if npad else out).reshape(*l_shape, N).cast(dtypes.bfloat16)
assert not padded_output, "padded output requires an ASM GEMM"
x_phys = quant_dequant_mx(x.reshape(-1, x.shape[-1])).reshape(*l_shape, x.shape[-1])
w_phys = dequant_weight(w_q, w_scale)
return (x_phys @ w_phys.T).cast(dtypes.bfloat16)
@@ -223,7 +226,7 @@ class GPTOSS:
x_q, x_e8, rrms = rmsnorm_mul_quantize_mxfp8(x, attention_norm, self.norm_eps)
qkv = matmul_mx((x_q, x_e8), wqkv, wqkv_scale) + wqkv_bias
norm_saves = [x_q, x_e8, rrms]
if getenv("FUSED_RMSNORM_MUL", 0):
elif getenv("FUSED_RMSNORM_MUL", 0):
from extra.gptoss_kernels.rmsnorm import rmsnorm_mul
x_normed, rrms = rmsnorm_mul(x, attention_norm, self.norm_eps)
qkv = matmul_mx(x_normed, wqkv, wqkv_scale) + wqkv_bias
@@ -233,11 +236,15 @@ class GPTOSS:
qkv = matmul_mx(x_normed * attention_norm, wqkv, wqkv_scale) + wqkv_bias
norm_saves = [x_normed, rrms]
qkv = qkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
xq = qkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
xk, xv = qkv[:, :, :, self.n_rep], qkv[:, :, :, self.n_rep + 1]
xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16) # (B,N,H,D)/(B,N,KV,D)
if getenv("FUSED_QKV_ROPE", 0):
from extra.gptoss_kernels.qkv_rope import fused_qkv_rope
xq, xk, xv = fused_qkv_rope(qkv, freqs_cis)
else:
qkv = qkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
xq = qkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
xk, xv = qkv[:, :, :, self.n_rep], qkv[:, :, :, self.n_rep + 1]
xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16) # (B,N,H,D)/(B,N,KV,D)
fa_saves = []
if getenv("HK_FLASH_ATTENTION"):
@@ -257,8 +264,13 @@ class GPTOSS:
w = (e / (e.sum(-1, keepdim=True) + (sink - m).exp())).cast(dtypes.bfloat16)
attn = (w @ xvm).permute(0, 3, 1, 2, 4).reshape(bsz, seqlen, self.n_heads * self.head_dim)
out = matmul_mx(attn, wo, wo_scale) + wo_bias
return out, [attn] + norm_saves + fa_saves
proj = matmul_mx(attn, wo, wo_scale, padded_output=bool(getenv("GPTOSS_RESIDUAL_HIP", 0)))
if getenv("GPTOSS_RESIDUAL_HIP", 0):
from extra.gptoss_kernels.residual import wo_bias_add
out = wo_bias_add(proj[:, :self.dim].reshape(x.shape), wo_bias)
else:
out = proj + wo_bias
return out, [attn] + norm_saves + fa_saves, proj
def feed_forward(self, x:Tensor, *, ffn_norm:Tensor, gate:Tensor, gate_bias:Tensor,
w_gate_up:Tensor, w_gate_up_scale:Tensor, w_gate_up_bias:Tensor,
@@ -271,21 +283,47 @@ class GPTOSS:
x_normed, rrms = rmsnorm(x, self.norm_eps)
inp = x_normed * ffn_norm
logits = router_mfma(inp, gate, gate_bias) if getenv("ROUTER_MFMA", 0) else inp.float() @ gate.float().T + gate_bias.float()
dim, inter = self.dim, self.intermediate_size
if getenv("GROUPED_MOE", 0):
bsz, seqlen = x.shape[:2]
inp, logits = inp.reshape(-1, dim), logits.reshape(-1, self.n_experts)
r = route(logits, self.experts_per_tok, self.n_experts)
xg = dispatch(_pad_cols(inp.cast(dtypes.bfloat16)), r)
quantized = None
if getenv("FUSED_ROUTER_DGRAD", 0): assert getenv("FUSED_ROUTER_TOPK", 0) and getenv("FP8_DISPATCH", 0)
if getenv("FUSED_ROUTER_TOPK", 0):
from extra.gptoss_kernels.router_topk import fused_router, fused_router_quantize
from extra.gemm.moe_routing import route_topk
if getenv("FUSED_ROUTER_DGRAD", 0):
weights, topi, q, scales = fused_router_quantize(inp, gate, gate_bias)
quantized = (q, scales)
else:
weights, topi = fused_router(inp, gate, gate_bias)
r = route_topk(weights, topi, self.n_experts)
else:
logits = router_mfma(inp, gate, gate_bias) if getenv("ROUTER_MFMA", 0) else inp.float() @ gate.float().T + gate_bias.float()
r = route(logits.reshape(-1, self.n_experts), self.experts_per_tok, self.n_experts)
inp = inp.reshape(-1, dim)
xg = dispatch_fp8(quantized, r) if quantized is not None else \
(dispatch_fp8 if getenv("FP8_DISPATCH", 0) else dispatch)(_pad_cols(inp.cast(dtypes.bfloat16)), r)
h = grouped_mx_gemm(xg, (w_gate_up, w_gate_up_scale), r.off)[:, :2*inter] + _moe_bias_tile(w_gate_up_bias, r).cast(dtypes.bfloat16)
y = swiglu(h, self.swiglu_limit)
z = grouped_mx_gemm(_pad_cols(y.cast(dtypes.bfloat16)), (w_down, w_down_scale), r.off)[:, :dim] \
+ _moe_bias_tile(w_down_bias, r).cast(dtypes.bfloat16)
out = combine(z, r, inp.shape[0], self.experts_per_tok).reshape(bsz, seqlen, dim)
return out, [x_normed, rrms, xg, h, y, z, r.weights, r.dest_row, r.off]
if getenv("FUSED_SWIGLU", 0):
from extra.gptoss_kernels.swiglu import fused_swiglu_quantize
y = fused_swiglu_quantize(h)
else:
y = swiglu(h, self.swiglu_limit)
if getenv("FUSED_COMBINE", 0):
from extra.gptoss_kernels.combine import down_combine
yq, ye8 = y if isinstance(y, tuple) else quantize_mxfp8(_pad_cols(y.cast(dtypes.bfloat16)))[:2]
out, z = down_combine(yq, ye8, w_down, w_down_scale, r.off, w_down_bias, r.dest_row, r.weights)
out = out.reshape(bsz, seqlen, dim)
else:
z = grouped_mx_gemm(y if isinstance(y, tuple) else _pad_cols(y.cast(dtypes.bfloat16)), (w_down, w_down_scale), r.off)[:, :dim] \
+ _moe_bias_tile(w_down_bias, r).cast(dtypes.bfloat16)
out = combine(z, r, inp.shape[0], self.experts_per_tok).reshape(bsz, seqlen, dim)
return out, [x_normed, rrms, *(xg if isinstance(xg, tuple) else (xg,)), h, *(y if isinstance(y, tuple) else (y,)), z,
r.weights, r.topi, r.dest_row, r.off,
*([quantized[1]] if quantized is not None else [])]
else:
logits = router_mfma(inp, gate, gate_bias) if getenv("ROUTER_MFMA", 0) else inp.float() @ gate.float().T + gate_bias.float()
thresh = logits.topk(self.experts_per_tok)[0][..., -1:]
weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1)
@@ -301,10 +339,15 @@ class GPTOSS:
@function(precompile=True, precompile_backward=True)
def run_layer(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, sliding:bool, attn_kwargs:dict, ffn_kwargs:dict, save:bool=True):
attn, attn_saves = self.attention(x, freqs_cis, mask, sliding, **attn_kwargs)
attn, attn_saves, proj = self.attention(x, freqs_cis, mask, sliding, **attn_kwargs)
h = x + attn
ffn, ffn_saves = self.feed_forward(h, **ffn_kwargs)
h = h + ffn
if save: ffn_saves.append(h)
if getenv("GPTOSS_RESIDUAL_HIP", 0):
from extra.gptoss_kernels.residual import residual_join
h = residual_join(h, x, proj, attn_kwargs["wo_bias"], ffn)
else:
h = h + ffn
if save: return (h, *attn_saves, *ffn_saves)
return (h,)
tinygrad_repo/examples/mlperf/models/test_flat_llama.py
+1 -23
View File
@@ -2,7 +2,7 @@ import os
os.environ["WQKV"] = "1"
import unittest
import numpy as np
from tinygrad import Tensor, nn, dtypes
from tinygrad import Tensor, nn
from tinygrad.device import Device
from examples.mlperf.models.llama import Transformer
from examples.mlperf.models.flat_llama import FlatTransformer
@@ -111,27 +111,5 @@ class TestFlatLlama(unittest.TestCase):
self.assertEqual(ref_logits.shape, flat_logits.shape)
np.testing.assert_allclose(flat_logits, ref_logits, atol=1e-4, rtol=1e-4)
@unittest.skipUnless(dtypes.fp8e4m3 in Device[Device.DEFAULT].renderer.supported_dtypes(), "fp8 not supported on this device")
def test_forward_fp8(self):
import examples.mlperf.models.flat_llama as flat_llama_mod
old_fp8 = flat_llama_mod.FP8
try:
flat_llama_mod.FP8 = 1
Tensor.manual_seed(42)
params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
ref = Transformer(**params)
flat = FlatTransformer(**params)
copy_weights(flat, ref)
Tensor.realize(*nn.state.get_state_dict(flat).values())
tokens = Tensor([[1, 50, 100, 999, 2]])
ref_logits = ref(tokens).numpy()
flat_logits = flat(tokens).numpy()
self.assertEqual(ref_logits.shape, flat_logits.shape)
# FP8 has lower precision, allow larger tolerance
np.testing.assert_allclose(flat_logits, ref_logits, atol=1.0, rtol=0.1)
finally:
flat_llama_mod.FP8 = old_fp8
if __name__ == "__main__":
unittest.main()
tinygrad_repo/examples/mlperf/optim.py
+14 -37
View File
@@ -7,8 +7,6 @@ from tinygrad.uop.ops import UOp, Ops, AxisType
STOCHASTIC_ROUND = getenv("STOCHASTIC_ROUND", 0)
MASTER_WEIGHTS = getenv("MASTER_WEIGHTS", 0)
ZERO_OPTIM = getenv("ZERO_OPTIM", 0)
FP8_AMAX_MARGIN = getenv("FP8_AMAX_MARGIN", 1.1)
IMMEDIATE_SCALE = getenv("IMMEDIATE_SCALE", 0)
MXFP8 = getenv("MXFP8", 0)
PRESTORE_WT = getenv("PRESTORE_WT", 0)
@@ -65,8 +63,7 @@ class GradAccClipAdamW(Optimizer):
updates, extra = self._step([], grads)
for i, tt in enumerate(self.params): tt.assign(self._apply_update(tt, updates[i], self.master_params[i] if self.master_params else None))
fp8_inv_scales = [tt._inv_scale for tt in self.params if hasattr(tt, '_inv_scale')]
fp8_next_inv_scales = [tt._next_inv_scale for tt in self.params if hasattr(tt, '_next_inv_scale')]
return extra + self.params + self.buffers + (self.master_params or []) + fp8_inv_scales + fp8_next_inv_scales
return extra + self.params + self.buffers + (self.master_params or []) + fp8_inv_scales
def fstep(self, grads:list[Tensor], grad_norm:Tensor|None=None):
Tensor.realize(*([grad_norm] if grad_norm is not None else []), *self.fschedule_step(grads))
@@ -102,39 +99,19 @@ class GradAccClipAdamW(Optimizer):
if STOCHASTIC_ROUND and t.dtype == dtypes.bfloat16:
out = stochastic_round_bf16(new_w)
return out.shard_like(t) if offloaded else out
if t.dtype in dtypes.fp8s:
if MXFP8:
from extra.gemm.cdna_asm_gemm import quantize_mxfp8
w_q, w_e8, _ = quantize_mxfp8(new_w.reshape(-1, new_w.shape[-1]))
if self.zero: w_q, w_e8 = self._zero_gather(w_q), self._zero_gather(w_e8)
new_e8 = w_e8.reshape(t._inv_scale.shape)
t._inv_scale.assign(new_e8.shard_like(t._inv_scale) if offloaded else new_e8)
ret = w_q.reshape(t.shape)
if PRESTORE_WT and hasattr(t, '_wT_q'):
from extra.gemm.cdna_asm_gemm import _mx_block_scale_3d
w_phys = ret.cast(dtypes.bfloat16) * _mx_block_scale_3d(new_e8).cast(dtypes.bfloat16)
wT_q, wT_e8, _ = quantize_mxfp8(w_phys.transpose(1, 2))
t._wT_q.assign(wT_q.shard_like(t._wT_q) if offloaded else wT_q)
t._wT_e8.assign(wT_e8.shard_like(t._wT_e8) if offloaded else wT_e8)
return ret.shard_like(t) if offloaded else ret
from examples.mlperf.models.flat_llama import FP8_MAX
if IMMEDIATE_SCALE:
amax_axis = tuple(range(t._inv_scale.ndim, new_w.ndim))
new_inv = ((new_w.float().abs().max(axis=amax_axis).detach() + 1e-8) / FP8_MAX).cast(t._inv_scale.dtype)
t._inv_scale.assign(new_inv.shard_like(t._inv_scale) if offloaded else new_inv)
scale = new_inv.reciprocal().reshape(*new_inv.shape, *([1]*(new_w.ndim-new_inv.ndim)))
ret = (new_w * scale).clamp(-FP8_MAX, FP8_MAX).cast(t.dtype)
return ret.shard_like(t) if offloaded else ret
# delayed scaling: reuse previous step's inv_scale
t._inv_scale.assign(t._next_inv_scale)
inv_scale = t._inv_scale.to(new_w.device) if offloaded else t._inv_scale
scale = inv_scale.reciprocal().reshape(*inv_scale.shape, *([1]*(new_w.ndim-inv_scale.ndim)))
scaled = (new_w * scale).clamp(-FP8_MAX, FP8_MAX)
ret = scaled.cast(t.dtype)
# update inv_scale for next step from quantized result
new_amax = (ret.float().abs().max(axis=tuple(range(inv_scale.ndim, ret.ndim))) * inv_scale * FP8_AMAX_MARGIN).detach()
new_inv = ((new_amax + 1e-8) / FP8_MAX).cast(t._inv_scale.dtype)
t._next_inv_scale.assign(new_inv.shard_like(t._next_inv_scale) if offloaded else new_inv)
if MXFP8 and t.dtype in dtypes.fp8s:
from extra.gemm.cdna_asm_gemm import quantize_mxfp8
w_q, w_e8, _ = quantize_mxfp8(new_w.reshape(-1, new_w.shape[-1]))
if self.zero: w_q, w_e8 = self._zero_gather(w_q), self._zero_gather(w_e8)
new_e8 = w_e8.reshape(t._inv_scale.shape)
t._inv_scale.assign(new_e8.shard_like(t._inv_scale) if offloaded else new_e8)
ret = w_q.reshape(t.shape)
if PRESTORE_WT and hasattr(t, '_wT_q'):
from extra.gemm.cdna_asm_gemm import _mx_block_scale_3d
w_phys = ret.cast(dtypes.bfloat16) * _mx_block_scale_3d(new_e8).cast(dtypes.bfloat16)
wT_q, wT_e8, _ = quantize_mxfp8(w_phys.transpose(1, 2))
t._wT_q.assign(wT_q.shard_like(t._wT_q) if offloaded else wT_q)
t._wT_e8.assign(wT_e8.shard_like(t._wT_e8) if offloaded else wT_e8)
return ret.shard_like(t) if offloaded else ret
out = new_w.cast(t.dtype)
return out.shard_like(t) if offloaded else out
tinygrad_repo/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_405b/implementations/tinybox_8xMI350X/dev_beam.sh
-4
View File
@@ -16,12 +16,8 @@ export USE_ATOMICS=${USE_ATOMICS:-1}
export ASM_GEMM=${ASM_GEMM:-1}
export WQKV=${WQKV:-1}
export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
export FP8=${FP8:-1}
export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
export FAST_CE=${FAST_CE:-0}
export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
export SPLIT_W13=${SPLIT_W13:-1}
export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
tinygrad_repo/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_405b/implementations/tinybox_8xMI350X/dev_run.sh
-4
View File
@@ -16,12 +16,8 @@ export USE_ATOMICS=${USE_ATOMICS:-1}
export ASM_GEMM=${ASM_GEMM:-1}
export WQKV=${WQKV:-1}
export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
export FP8=${FP8:-1}
export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
export FAST_CE=${FAST_CE:-0}
export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
export SPLIT_W13=${SPLIT_W13:-1}
export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
tinygrad_repo/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh
-3
View File
@@ -19,9 +19,6 @@ export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
export MXFP4=${MXFP4:-1}
export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
export FAST_CE=${FAST_CE:-1}
export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1}
export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
export SPLIT_W13=${SPLIT_W13:-0}
export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
tinygrad_repo/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam_mp.sh
-4
View File
@@ -16,12 +16,8 @@ export USE_ATOMICS=${USE_ATOMICS:-1}
export ASM_GEMM=${ASM_GEMM:-1}
export WQKV=${WQKV:-1}
export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
export FP8=${FP8:-1}
export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
export FAST_CE=${FAST_CE:-0}
export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
export SPLIT_W13=${SPLIT_W13:-1}
export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
tinygrad_repo/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh
-3
View File
@@ -19,9 +19,6 @@ export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
export MXFP4=${MXFP4:-1}
export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
export FAST_CE=${FAST_CE:-1}
export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1}
export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
export SPLIT_W13=${SPLIT_W13:-0}
export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
tinygrad_repo/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run_mp.sh
-4
View File
@@ -16,12 +16,8 @@ export USE_ATOMICS=${USE_ATOMICS:-1}
export ASM_GEMM=${ASM_GEMM:-1}
export WQKV=${WQKV:-1}
export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
export FP8=${FP8:-1}
export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
export FAST_CE=${FAST_CE:-0}
export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
export SPLIT_W13=${SPLIT_W13:-1}
export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
tinygrad_repo/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/run_and_time.sh
-3
View File
@@ -20,9 +20,6 @@ export MASTER_WEIGHTS=1
export MXFP4=1
export ALLREDUCE_CAST=1
export FAST_CE=1
export FUSED_INPUT_QUANTIZE=1
export FUSED_GRAD_QUANTIZE=1
export FUSED_ADD_NORM_MUL_QUANTIZE=1
export FUSED_SILU_W13=1
export SPLIT_W13=0
tinygrad_repo/examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/llama31_8b/implementations/tinybox_2x8xMI350X/dev_beam.sh
Executable
+5
View File
@@ -0,0 +1,5 @@
#!/usr/bin/env bash
export NODES=${NODES:-"tinyamd3 tinyamd4"}
export REMOTE=${REMOTE:-"local,$(echo $NODES | tr ' ' '\n' | grep -vx "$(hostname)" | head -1):6667"}
export DEV=PCI+AMD REMOTE_TIMEOUT=600 ALLREDUCE_NODE_NDEVS=8 DP=16 BS=32 EVAL_BS=16
exec bash "$(dirname "$0")/../tinybox_8xMI350X/dev_beam.sh"
tinygrad_repo/examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/llama31_8b/implementations/tinybox_2x8xMI350X/dev_run.sh
Executable
+5
View File
@@ -0,0 +1,5 @@
#!/usr/bin/env bash
export NODES=${NODES:-"tinyamd3 tinyamd4"}
export REMOTE=${REMOTE:-"local,$(echo $NODES | tr ' ' '\n' | grep -vx "$(hostname)" | head -1):6667"}
export DEV=PCI+AMD REMOTE_TIMEOUT=600 ALLREDUCE_NODE_NDEVS=8 DP=16 BS=32 EVAL_BS=16
exec bash "$(dirname "$0")/../tinybox_8xMI350X/dev_run.sh"
tinygrad_repo/examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/llama31_8b/implementations/tinybox_2x8xMI350X/setup.sh
Executable
+3
View File
@@ -0,0 +1,3 @@
#!/usr/bin/env bash
for p in $(lspci -D -d ::0604 | cut -d" " -f1); do sudo -n setpci -s $p ECAP_ACS+0x6.w=0000:000c 2>/dev/null || true; done
tinygrad_repo/examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh
-4
View File
@@ -16,12 +16,8 @@ export USE_ATOMICS=${USE_ATOMICS:-1}
export ASM_GEMM=${ASM_GEMM:-1}
export WQKV=${WQKV:-1}
export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
export FP8=${FP8:-1}
export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
export FAST_CE=${FAST_CE:-1}
export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1}
export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
export SPLIT_W13=${SPLIT_W13:-0}
export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
tinygrad_repo/examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam_mp.sh
-4
View File
@@ -16,12 +16,8 @@ export USE_ATOMICS=${USE_ATOMICS:-1}
export ASM_GEMM=${ASM_GEMM:-1}
export WQKV=${WQKV:-1}
export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
export FP8=${FP8:-1}
export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
export FAST_CE=${FAST_CE:-0}
export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
export SPLIT_W13=${SPLIT_W13:-1}
export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
tinygrad_repo/examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh
-4
View File
@@ -16,12 +16,8 @@ export USE_ATOMICS=${USE_ATOMICS:-1}
export ASM_GEMM=${ASM_GEMM:-1}
export WQKV=${WQKV:-1}
export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
export FP8=${FP8:-1}
export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
export FAST_CE=${FAST_CE:-1}
export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1}
export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
export SPLIT_W13=${SPLIT_W13:-0}
export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
tinygrad_repo/examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run_mp.sh
-4
View File
@@ -16,12 +16,8 @@ export USE_ATOMICS=${USE_ATOMICS:-1}
export ASM_GEMM=${ASM_GEMM:-1}
export WQKV=${WQKV:-1}
export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
export FP8=${FP8:-1}
export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
export FAST_CE=${FAST_CE:-0}
export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
export SPLIT_W13=${SPLIT_W13:-1}
export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
tinygrad_repo/examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/run_and_time.sh
-4
View File
@@ -17,12 +17,8 @@ export USE_ATOMICS=1
export ASM_GEMM=1
export WQKV=1
export MASTER_WEIGHTS=1
export FP8=1
export ALLREDUCE_CAST=1
export FAST_CE=1
export FUSED_INPUT_QUANTIZE=1
export FUSED_GRAD_QUANTIZE=1
export FUSED_ADD_NORM_MUL_QUANTIZE=1
export FUSED_SILU_W13=1
export SPLIT_W13=0
tinygrad_repo/examples/openpilot/check_image_reads.py
+4 -4
View File
@@ -1,16 +1,16 @@
import pickle, sys
from tinygrad.helpers import getenv, is_image_shape, temp
from tinygrad.uop.ops import Ops
from tinygrad.viz.serve import VizData, _reconstruct
from tinygrad.viz.serve import VizData, _reconstruct, canonicalize_ast
with open(sys.argv[1] if len(sys.argv) > 1 else temp("rewrites.pkl", append_user=True), "rb") as f: data = VizData(pickle.load(f))
steps = [s for key, ctx in zip(data.trace.keys, data.trace.rewrites) if key.display_name.startswith("JIT ")
for s in ctx if s.name == "View captured linear"]
assert len(steps) == 1, f"expected one JIT compilation, found {len(steps)}"
asts = [call.without_after.src[0] for call in _reconstruct(data, steps[0].sink).src]
by_ast = {key.keys[1]: _reconstruct(data, s.sink) for key, ctx in zip(data.trace.keys, data.trace.rewrites)
for s in ctx if s.name == "View Program" and key.keys[1] in asts}
programs = [by_ast[ast] if ast.op is Ops.SINK else ast for ast in asts if ast.op in {Ops.SINK, Ops.PROGRAM}]
by_ast = {canonicalize_ast(_reconstruct(data, b.sink)): _reconstruct(data, p.sink) for ctx in data.trace.rewrites
for b in ctx if b.name == "View Base AST" for p in ctx if p.name == "View Program"}
programs = [by_ast[canonicalize_ast(ast)] if ast.op is Ops.SINK else ast for ast in asts if ast.op in {Ops.SINK, Ops.PROGRAM}]
read_image, gated_read_image = 0, 0
for prg in programs:
uops = prg.src[1].src
tinygrad_repo/examples/openpilot/helpers.py
+12 -22
View File
@@ -2,7 +2,7 @@
import io, pathlib, pickle, shutil, struct, time, tempfile
import numpy as np
from typing import Callable
from tinygrad import Tensor, Device, Context, dtypes
from tinygrad import Tensor, Device, Context
from tinygrad.device import Buffer
from tinygrad.uop.ops import PatternMatcher, UPat, Ops, graph_rewrite
from tinygrad.nn.state import get_parameters
@@ -22,7 +22,7 @@ def dump_pickle(obj, path):
if not isinstance(b, Buffer) or b.device != Device.DEFAULT: return None
offset += buffers.write(bytes((-offset) % 256))
offset = (off:=offset) + buffers.write(b.as_memoryview())
return b.size, b.dtype, off
return b.nbytes, off
p = pickle.Pickler(opcodes)
p.persistent_id = persistent_id
@@ -36,33 +36,23 @@ def dump_pickle(obj, path):
def load_pickle(path, *, out_of_band=False):
with open(path, "rb") as f:
if not out_of_band: return pickle.load(f)
opcodes = f.read(struct.unpack('<q', f.read(8))[0])
if f.tell() == pathlib.Path(path).stat().st_size:
return pickle.loads(opcodes)
buffers = Tensor(pathlib.Path(path))[f.tell():].uop.buffer.ensure_allocated()
opcodes, buffers = f.read(struct.unpack('<q', f.read(8))[0]), Tensor(pathlib.Path(path))[f.tell():].to(Device.DEFAULT).realize().uop.buffer
# FIXME: we load in chunks here because hcq_submit for one large copy is very slow to compile
arena, CHUNK_SIZE = Buffer(Device.DEFAULT, buffers.nbytes, dtypes.uchar, preallocate=True), 32 << 20
for off in range(0, buffers.nbytes, CHUNK_SIZE):
size = min(CHUNK_SIZE, buffers.nbytes-off)
arena.view(size, dtypes.uchar, off).ensure_allocated().copy_from(buffers.view(size, dtypes.uchar, off).ensure_allocated())
def persistent_load(pid): return arena.view(*pid)
def persistent_load(pid): return buffers.view(*pid)
u = pickle.Unpickler(io.BytesIO(opcodes))
u.persistent_load = persistent_load
return u.load()
@Context(OPENPILOT_HACKS=1, **{'AMD': {'TC_OPT': 2, 'TC_MIN_GLOBALS': 32}}.get(Device.DEFAULT, {}))
def benchmark(fxn:Callable, cb=None, **kwargs):
# Select benchmark options when called, without probing devices on import.
with Context(OPENPILOT_HACKS=1, **{'AMD': {'TC_OPT': 2, 'TC_MIN_GLOBALS': 32}}.get(Device.DEFAULT, {})):
Device.default.synchronize()
start = time.perf_counter()
if (output := fxn(**kwargs)) is not None: output.realize()
Device.default.synchronize()
end = time.perf_counter()
if cb: cb(end-start)
return [t.numpy().copy() for t in get_parameters(kwargs.get('output_buffers', output))]
Device.default.synchronize()
start = time.perf_counter()
if (output := fxn(**kwargs)) is not None: output.realize()
Device.default.synchronize()
end = time.perf_counter()
if cb: cb(end-start)
return [t.numpy().copy() for t in get_parameters(kwargs.get('output_buffers', output))]
pm_retargetable = PatternMatcher([
(UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(), UPat()), name="p"), lambda p: p.replace(src=p.src[:-1]) if p.arg.target.device == "CPU" else None)
tinygrad_repo/examples/webgpu/stable_diffusion/compile.py
+1 -1
View File
@@ -114,7 +114,7 @@ if __name__ == "__main__":
linear, output_bufs = jit_model(step, *step.input)
functions, statements, bufs, _ = compile_net(linear, output_bufs)
state = get_state_dict(model)
weights = {(id(b), b.offset, b.size, b.dtype): name for name, x in state.items() if (b:=x.uop.base.realized) is not None}
weights = {(id(b), b.offset, b.nbytes, x.dtype): name for name, x in state.items() if (b:=x.uop.base.realized) is not None}
kernel_code = '\n\n'.join([f"const {key} = `{fixup_code(code, key)}`;" for key, code in functions.items()])
kernel_names = ', '.join([name for (name, _, _, _) in statements])
input_names = [f"input{i}" for i in range(len(step.input))]
tinygrad_repo/extra/amdpci/am_smi.py
+5 -2
View File
@@ -4,7 +4,7 @@ import time, mmap, sys, shutil, os, glob, subprocess, argparse, collections
from tinygrad.helpers import DEBUG, NO_COLOR, colored, ansilen
from tinygrad.runtime.autogen import libc
from tinygrad.runtime.autogen.am import am
from tinygrad.runtime.support.hcq import MMIOInterface
from tinygrad.runtime.support.memory import MMIOInterface
from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager, AMPageTableEntry
from tinygrad.runtime.support.am.ip import AM_SOC, AM_GMC, AM_IH, AM_PSP, AM_SMU, AM_GFX, AM_SDMA
@@ -118,6 +118,9 @@ class SMICtx:
def _open_am_device(self, pcibus):
if pcibus not in self.opened_pci_resources:
if not os.path.exists(f"/sys/bus/pci/devices/{pcibus}"):
if DEBUG >= 2: print(f"/sys/bus/pci/devices/{pcibus} not found")
return
bar_fds = {bar: os.open(f"/sys/bus/pci/devices/{pcibus}/resource{bar}", os.O_RDWR | os.O_SYNC) for bar in [0, 2, 5]}
bar_size = {0: get_bar0_size(pcibus), 2: os.fstat(bar_fds[2]).st_size, 5: os.fstat(bar_fds[5]).st_size}
@@ -161,7 +164,7 @@ class SMICtx:
case (13,0,6): table_t = dev.smu.smu_mod.MetricsTableV0_t
case (13,0,12): table_t = dev.smu.smu_mod.MetricsTable_t
case _: table_t = dev.smu.smu_mod.SmuMetricsExternal_t
tables[dev] = dev.smu.read_table(table_t, dev.smu.smu_mod.SMU_TABLE_SMU_METRICS) if dev.pci_state == "D0" else None
tables[dev] = dev.smu.read_table(table_t, getattr(dev.smu.smu_mod, "TABLE_SMU_METRICS", dev.smu.smu_mod.SMU_TABLE_SMU_METRICS)) if dev.pci_state == "D0" else None
return tables
def _pick_nonzero_avg(self, vals) -> int:
tinygrad_repo/extra/amdpci/hive_reset.py
+1 -1
View File
@@ -3,7 +3,7 @@
import os
from tinygrad.helpers import Context
from tinygrad.runtime.support.system import System, PCIDevice
from tinygrad.runtime.support.hcq import FileIOInterface
from tinygrad.runtime.support.system import FileIOInterface
from tinygrad.runtime.support.am.amdev import AMDev
if __name__ == "__main__":
tinygrad_repo/extra/benchmark_llm.py
+2 -1
View File
@@ -8,10 +8,11 @@ if __name__ == "__main__":
parser.add_argument("--prompt-tokens", type=int, default=1024, help="number of prompt tokens (default: %(default)s)")
parser.add_argument("--decode-tokens", type=int, default=16, help="number of tokens to decode (default: %(default)s)")
parser.add_argument("--chunk-size", type=int, default=32, help="chunk size for prefill (default: %(default)s)")
parser.add_argument("--shard", type=int, default=1, help="tensor parallel device count (default: %(default)s)")
args = parser.parse_args()
st = time.perf_counter()
model, _ = Transformer.from_gguf(args.model, args.max_context)
model, _ = Transformer.from_gguf(args.model, args.max_context, shard=args.shard)
print(f"load {time.perf_counter()-st:.3f}s", flush=True)
st = time.perf_counter()
tinygrad_repo/extra/export_model.py
+2 -2
View File
@@ -26,7 +26,7 @@ def compile_net(linear:UOp, output_bufs:List[Buffer]) -> Tuple[Dict[str,str], Li
if bu.op is Ops.PARAM: key, name, size = ("in", bu.arg.slot), f"input{bu.arg.slot}", prod(bu.shape)*bu.dtype.itemsize
else:
b = bu.buffer
key, size = (id(b.base), b.offset, b.size, b.dtype), b.size*b.dtype.itemsize
key, size = (id(b.base), b.offset, b.nbytes, bu.dtype), b.nbytes
if key in bufs: return bufs[key][0]
if (name:=output_name.get(id(b))) is None:
name, n = f"buf_{n}", n+1
@@ -244,7 +244,7 @@ def export_model(model, target:str, *inputs, model_name: Optional[str] = "model"
with Context(JIT=2): linear, output_bufs = jit_model(model, *inputs)
functions, statements, bufs, bufs_to_save = compile_net(linear, output_bufs)
state = get_state_dict(model)
weight_names = {(id(b), b.offset, b.size, b.dtype): name for name, x in state.items() if (b:=x.uop.base.realized) is not None}
weight_names = {(id(b), b.offset, b.nbytes, x.dtype): name for name, x in state.items() if (b:=x.uop.base.realized) is not None}
input_names = [f"input{i}" for i in range(len(inputs))]
output_names = [f"output{i}" for i in range(len(output_bufs))]
tinygrad_repo/extra/fp8/fp8_linear.py
+2 -2
View File
@@ -1,6 +1,6 @@
from typing import Callable, Any
from tinygrad import Tensor, dtypes, nn, UOp
from tinygrad.uop.ops import KernelInfo, AxisType, Ops
from tinygrad.uop.ops import KernelInfo, Ops
def quantize_to_fp8(x: Tensor, dtype=dtypes.fp8e4m3):
fp8_min = -448.0 if dtype == dtypes.fp8e4m3 else -57344.0
@@ -21,7 +21,7 @@ def custom_matmul(output: UOp, inp: UOp, weight: UOp) -> UOp:
seq_idx = UOp.range(SEQ, 2)
out_idx = UOp.range(OUT, 3)
batch_idx = UOp.range(output.size//SEQ//OUT, 1)
reduce_idx = UOp.range(IN, 0, AxisType.REDUCE)
reduce_idx = UOp.range(IN, 0)
product = (inp.index((seq_idx*IN+reduce_idx+batch_idx*IN*SEQ)) * weight.index((out_idx*IN+reduce_idx))).cast(dtypes.float)
reduced = product.reduce(reduce_idx, arg=Ops.ADD)
store_op = output.index((seq_idx*OUT+out_idx+batch_idx*OUT*SEQ)).store(reduced).end(batch_idx, seq_idx, out_idx)
tinygrad_repo/extra/gemm/amd_call_matmul.py
+3 -3
View File
@@ -9,7 +9,7 @@ def call(*ranges):
def wrapper(*args):
ret = fxn(*ranges, *args).end(*ranges)
# closing the GLOBAL ranges finishes the kernel
if any(r.arg[-1] is AxisType.GLOBAL for r in ranges): ret = ret.sink(arg=KernelInfo(opts_to_apply=()))
if any(r.axis_type is AxisType.GLOBAL for r in ranges): ret = ret.sink(arg=KernelInfo(opts_to_apply=()))
return ret
return wrapper
return decorator
@@ -48,7 +48,7 @@ def block_128x128_gemm(lane:UOp, wave_m:UOp, wave_n:UOp, c:UOp, a:UOp, b:UOp) ->
acc = UOp.placeholder((TM, TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
acc = acc.after(acc.store(acc.zeros_like(buffer=False)))
@call(UOp.range(K // BLOCK_K, 100, AxisType.REDUCE))
@call(UOp.range(K // BLOCK_K, 100, AxisType.LOOP))
def tile_reduce(k_tile:UOp):
A_local = UOp.placeholder((BLOCK_K, BLOCK_M), a.dtype, slot=0, addrspace=AddrSpace.LOCAL)
B_local = UOp.placeholder((BLOCK_K, BLOCK_N), b.dtype, slot=1, addrspace=AddrSpace.LOCAL)
@@ -60,7 +60,7 @@ def block_128x128_gemm(lane:UOp, wave_m:UOp, wave_n:UOp, c:UOp, a:UOp, b:UOp) ->
# NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
A_local, B_local = A_local.after(A_store, B_store), B_local.after(A_store, B_store)
@call(UOp.range(BLOCK_K, 101, AxisType.REDUCE))
@call(UOp.range(BLOCK_K, 101, AxisType.LOOP))
def inner_reduce(k:UOp):
# registers for LOCAL -> REG
a_frag = UOp.placeholder((TM//UNROLL_M, UNROLL_M), dtypes.float, slot=0, addrspace=AddrSpace.REG)
tinygrad_repo/extra/gemm/amd_copy_matmul.py
+3 -3
View File
@@ -51,7 +51,7 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
a = a.reshape(K // BLOCK_K, BLOCK_K, BLOCK_M)
b = b.reshape(K // BLOCK_K, BLOCK_K, BLOCK_N)
k_tile = UOp.range(K // BLOCK_K, 100, AxisType.REDUCE)
k_tile = UOp.range(K // BLOCK_K, 100, AxisType.LOOP)
# copy with transpose for wmma (input is k×spatial, LDS is spatial×k)
A_copy = A_local.permute((1,0)) if use_wmma else A_local
@@ -69,7 +69,7 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
acc = acc.after(acc.store(acc.zeros_like(buffer=False)))
if use_wmma:
k = UOp.range(BLOCK_K // WMMA_K, 101, AxisType.REDUCE)
k = UOp.range(BLOCK_K // WMMA_K, 101, AxisType.LOOP)
tile_m = UOp.range(TM // WMMA_ACC, 200)
tile_n = UOp.range(TN, 201)
@@ -87,7 +87,7 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
a_frag = UOp.placeholder((TM//UNROLL_M, UNROLL_M), dtypes.float, slot=0, addrspace=AddrSpace.REG)
b_frag = UOp.placeholder((TN//UNROLL_N, UNROLL_N), dtypes.float, slot=1, addrspace=AddrSpace.REG)
k = UOp.range(BLOCK_K, 101, AxisType.REDUCE)
k = UOp.range(BLOCK_K, 101, AxisType.LOOP)
a_frag = a_frag.after(a_frag.store(A_local[k].reshape(WAVES_M, TM//UNROLL_M, LANES_PER_WAVE_M, UNROLL_M)[wave_m, :, lane_m, :]))
b_frag = b_frag.after(b_frag.store(B_local[k].reshape(WAVES_N, TN//UNROLL_N, LANES_PER_WAVE_N, UNROLL_N)[wave_n, :, lane_n, :]))
tinygrad_repo/extra/gemm/amd_uop_matmul.py
+2 -2
View File
@@ -45,7 +45,7 @@ def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[block_id_m, :, block_id_n, :]
# open the main reduction range
k_tile_range = UOp.range(K // BLOCK_K, 0, AxisType.REDUCE)
k_tile_range = UOp.range(K // BLOCK_K, 0, AxisType.LOOP)
a = a.reshape(M // BLOCK_M, BLOCK_M, K // BLOCK_K, BLOCK_K)[block_id_m, :, k_tile_range, :]
b = b.reshape(K // BLOCK_K, BLOCK_K, N // BLOCK_N, BLOCK_N)[k_tile_range, :, block_id_n, :]
@@ -70,7 +70,7 @@ def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
A_local, B_local = A_local.after(A_local_store, B_local_store), B_local.after(A_local_store, B_local_store)
# open inner k range
k = UOp.range(BLOCK_K, 3, AxisType.REDUCE)
k = UOp.range(BLOCK_K, 3, AxisType.LOOP)
# ---------------------------
# LOCAL -> REG (per-wave tiles)
tinygrad_repo/extra/gemm/cdna_asm_gemm.py
+26 -161
View File
@@ -5,87 +5,12 @@ from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
from tinygrad.renderer import Estimates
from tinygrad.helpers import getenv, all_same, DEBUG, ceildiv
from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8
from extra.llama_kernels.quantize_mxfp4 import quantize_mxfp4
FP8_DTYPE = dtypes.fp8e4m3
TILE_M, TILE_N, TILE_K = 256, 256, 64
# ** FP8 GEMM custom kernel
@functools.cache
def custom_hk_fp8_gemm(C:UOp, A:UOp, B:UOp, *args:UOp, dname:str, scale_mode:int=3) -> UOp:
# scale_mode: 0=no scale, 1=x only, 2=w only, 3=both
n_scales = (1 if scale_mode & 1 else 0) + (1 if scale_mode & 2 else 0) + (1 if scale_mode & 4 else 0)
scales, extra = args[:n_scales], args[n_scales:]
M, K = A.shape[0]*A.shape[1], A.shape[2]
N, K2 = B.shape[(1 if B.ndim == 3 else 0):]
assert K == K2, f"{A.shape} {B.shape}"
block_size = 256
threads = UOp.special(64 * 8, "lidx0")
workgroups = UOp.special((M // block_size) * (N // block_size), "gidx0")
sink_inputs = (C.base, A.base, B.base) + tuple(s.base for s in scales) + (threads, workgroups)
sink = UOp.sink(*sink_inputs,
arg=KernelInfo(f"hk_fp8_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K, mem=(M*K+N*K)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
src = (kittens_path/"gemm_fp8.cpp").read_text()
lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
"-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}",
f"-DSCALE_MODE={scale_mode}"]).compile_cached(src)
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
UOp(Ops.BINARY, arg=lib)))
# ** FP8 AtB GEMM custom kernel
@functools.cache
def custom_hk_fp8_atb_gemm(C:UOp, A:UOp, B:UOp, *args:UOp, dname:str, scale_mode:int=5) -> UOp:
# C = A.T @ B, A and B are physically [K, M] and [K, N].
n_scales = (1 if scale_mode & 1 else 0) + (1 if scale_mode & 2 else 0) + (1 if scale_mode & 4 else 0)
scales = args[:n_scales]
K, M = A.shape[0]*A.shape[1], A.shape[2]
K2, N = B.shape[0]*B.shape[1], B.shape[2]
assert K == K2, f"{A.shape} {B.shape}"
block_m, block_n, block_k, num_warps = 256, 256, 128, 8
assert M % block_m == 0 and N % block_n == 0 and K % block_k == 0, f"invalid fp8 atb tile {(block_m, block_n, block_k)} for {(M, N, K)}"
threads = UOp.special(64 * num_warps, "lidx0")
workgroups = UOp.special((M // block_m) * (N // block_n), "gidx0")
sink_inputs = (C.base, A.base, B.base) + tuple(s.base for s in scales) + (threads, workgroups)
sink = UOp.sink(*sink_inputs,
arg=KernelInfo(f"hk_fp8_atb_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K, mem=(M*K+N*K)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
src = (kittens_path/"gemm_fp8_atb.cpp").read_text()
lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
"-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}",
f"-DSCALE_MODE={scale_mode}"]).compile_cached(src)
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
UOp(Ops.BINARY, arg=lib)))
def hk_fp8_atb_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, g_amax:Tensor|None=None) -> Tensor:
assert a.dtype == b.dtype == FP8_DTYPE, f"expected fp8, got {a.dtype} {b.dtype}"
assert a.ndim == b.ndim == 3 and a.shape[:2] == b.shape[:2], f"{a.shape} {b.shape}"
batch, rows, M = a.shape
N = b.shape[2]
assert M % TILE_M == 0 and N % TILE_N == 0 and (batch * rows) % 128 == 0, \
f"fp8 atb shape {a.shape} {b.shape} must produce (M,N,K) multiples of ({TILE_M},{TILE_N},128)"
is_multi = isinstance(a.device, tuple)
reduce_out = False
if is_multi:
ndev = len(a.device)
if a.uop.axis in (0, 1) or b.uop.axis in (0, 1): inv, out_axis, reduce_out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a.device), 0, True
elif b.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M, N // ndev, dtype=dtypes.bfloat16, device=a.device), 2
elif a.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M // ndev, N, dtype=dtypes.bfloat16, device=a.device), 1
else: inv, out_axis, reduce_out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a.device), 0, True
out = Tensor(inv.uop.unshard(out_axis), device=a.device)
dname = a.device[0]
else:
out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a.device)
dname = a.device
dname = dname.split(":")[0]
scales = tuple(s for s in (x_scale, g_amax) if s is not None)
scale_mode = (1 if x_scale is not None else 0) | (4 if g_amax is not None else 0)
out = Tensor.custom_kernel(out, a, b, *scales, fxn=functools.partial(custom_hk_fp8_atb_gemm, dname=dname, scale_mode=scale_mode))[0]
if reduce_out: out = out.sum(0)
return out.squeeze(0) if out.ndim == 3 else out
# ** MXFP8 GEMM custom kernel
@functools.cache
@@ -221,7 +146,7 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
assert K == K2
m = UOp.range(M, 1)
n = UOp.range(N, 2)
k = UOp.range(K, 0, AxisType.REDUCE)
k = UOp.range(K, 0, AxisType.LOOP)
mul = (A.flatten().index((m*UOp.const(K)+k))*
B.flatten().index((k*UOp.const(N)+n))).cast(dtypes.float32)
red = mul.reduce(k, arg=Ops.ADD).cast(C.dtype)
@@ -294,79 +219,29 @@ def hk_bf16_atb_gemm(a:Tensor, b:Tensor) -> Tensor:
# ** backward gemm, might use the asm gemm
def custom_gemm_bw(gradient:UOp, kernel:UOp, n_scales:int=2, has_grad_amax:bool=False, has_w_post:bool=False):
def custom_gemm_bw(gradient:UOp, kernel:UOp):
inputs = kernel.src[1:]
if inputs[1].dtype == FP8_DTYPE:
out, a, b = inputs[:3]
i = 3
s_x = inputs[i]; i += 1
has_w = n_scales >= 2
s_w = inputs[i] if has_w else None; i += has_w
s_g_amax = inputs[i] if n_scales == 3 else None; i += (n_scales == 3)
grad_amax_state = inputs[i] if has_grad_amax else None; i += has_grad_amax
next_grad_amax_state = inputs[i] if has_grad_amax else None; i += has_grad_amax
w_post = inputs[i] if has_w_post else None
a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
s_x_t = Tensor(s_x, device=a.device)
s_w_t = Tensor(s_w, device=a.device) if has_w else None
s_g_amax_t = Tensor(s_g_amax, device=a.device) if s_g_amax is not None else None
w_post_t = Tensor(w_post, device=a.device) if has_w_post else None
g_t = g_t[:a.shape[0]]
from extra.llama_kernels.cast_amax import _grad_fp8_mailbox
from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed
gbase = gradient.base if hasattr(gradient, "base") else gradient
mailbox_entry = _grad_fp8_mailbox.pop(gbase, None) or _grad_fp8_mailbox.pop(gradient, None)
if mailbox_entry is not None:
g_fp8_u, grad_amax_u = mailbox_entry
g_fp8 = Tensor(g_fp8_u, device=a.device)[:a.shape[0]]
g_amax = Tensor(grad_amax_u, device=a.device)
else:
assert grad_amax_state is not None, "fp8 matmul bwd needs either a mailbox entry or a grad_amax_state"
if getenv("CURRENT_GRAD_SCALE", 0):
g_fp8, _, g_amax = quantize_fp8(g_t, amax_state=None)
elif getenv("FUSED_GRAD_QUANTIZE", 0):
grad_amax_t = Tensor(grad_amax_state, device=a.device)
g_amax = grad_amax_t
g_fp8, _ = quantize_fp8_delayed(g_t, g_amax, Tensor(next_grad_amax_state, device=a.device))
else:
grad_amax_t = Tensor(grad_amax_state, device=a.device)
g_amax = grad_amax_t
g_fp8, _, new_grad_amax = quantize_fp8(g_t, amax_state=g_amax)
store_effect = next_grad_amax_state.store(new_grad_amax.uop)
g_fp8 = Tensor(g_fp8.contiguous().uop.after(store_effect), device=a.device)
# dgrad: applies grad/activation amax scales in the GEMM epilogue; w_scale is already inverse.
assert s_g_amax_t is None, "fp8 GEMM bwd through g_amax scaling is unsupported"
grad_a = asm_gemm(g_fp8, b_t, x_scale=s_x_t, w_scale=s_w_t, g_amax=g_amax) if has_w else asm_gemm(g_fp8, b_t, x_scale=s_x_t, g_amax=g_amax)
# wgrad: no w_scale
grad_b = hk_fp8_atb_gemm(g_fp8, a_t, x_scale=s_x_t, g_amax=g_amax)
# wgrad: rescale if not scalar
if w_post_t is not None:
grad_b = grad_b / w_post_t.reshape(*w_post_t.shape, *([1]*(grad_b.ndim - w_post_t.ndim)))
# one None per input: (out, a, b, x_scale[, w_scale][, grad_amax][, w_post_scale])
ret = (None, grad_a.uop, grad_b.uop) + tuple(None for _ in inputs[3:])
return ret
hk_bf16 = len(inputs) == 4 and inputs[1].dtype == dtypes.bfloat16
if hk_bf16:
out, a, b_t, b = inputs
assert all_same([gradient.device, a.device, b_t.device, b.device, out.device])
else:
hk_bf16 = len(inputs) == 4 and inputs[1].dtype == dtypes.bfloat16
if hk_bf16:
out, a, b_t, b = inputs
assert all_same([gradient.device, a.device, b_t.device, b.device, out.device])
else:
assert len(inputs) == 3, f"regular gemm must have exactly 3 sources, got: {len(inputs)}"
out, a, b = inputs
assert all_same([gradient.device, a.device, b.device, out.device])
a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
g_t = g_t[:a.shape[0]]
if hk_bf16 and g_t.dtype != b_t.dtype: g_t = g_t.cast(b_t.dtype)
if can_use_asm_gemm(g_t, b_t.T): grad_a = asm_gemm(g_t, b_t.T).uop
else: grad_a = (g_t @ b_t.T).uop
if hk_bf16:
grad_b = hk_bf16_atb_gemm(a_t, g_t).uop
else:
a_t_flat, g_t_flat = a_t.permute(2, 0, 1).reshape(a_t.shape[2], -1), g_t.reshape(-1, g_t.shape[-1])
if can_use_asm_gemm(a_t_flat, g_t_flat): grad_b = asm_gemm(a_t_flat, g_t_flat).uop
else: grad_b = (a_t_flat @ g_t_flat).uop
# hk_bf16 uses b.T, writes gradients only for a and b
return (None, grad_a, None, grad_b) if hk_bf16 else (None, grad_a, grad_b)
assert len(inputs) == 3, f"regular gemm must have exactly 3 sources, got: {len(inputs)}"
out, a, b = inputs
assert all_same([gradient.device, a.device, b.device, out.device])
a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
g_t = g_t[:a.shape[0]]
if hk_bf16 and g_t.dtype != b_t.dtype: g_t = g_t.cast(b_t.dtype)
if can_use_asm_gemm(g_t, b_t.T): grad_a = asm_gemm(g_t, b_t.T).uop
else: grad_a = (g_t @ b_t.T).uop
if hk_bf16:
grad_b = hk_bf16_atb_gemm(a_t, g_t).uop
else:
a_t_flat, g_t_flat = a_t.permute(2, 0, 1).reshape(a_t.shape[2], -1), g_t.reshape(-1, g_t.shape[-1])
if can_use_asm_gemm(a_t_flat, g_t_flat): grad_b = asm_gemm(a_t_flat, g_t_flat).uop
else: grad_b = (a_t_flat @ g_t_flat).uop
# hk_bf16 uses b.T, writes gradients only for a and b
return (None, grad_a, None, grad_b) if hk_bf16 else (None, grad_a, grad_b)
# ** mxfp8 gemm backward
@@ -418,12 +293,11 @@ def custom_mxfp4_gemm_bw(gradient:UOp, kernel:UOp):
# ** main gemm function
def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=None, grad_amax_state:Tensor|None=None,
next_grad_amax_state:Tensor|None=None,
w_post_scale:Tensor|None=None, mx:bool=False, mx_scales:tuple|None=None, mx_w_stored:bool=False, g_amax:Tensor|None=None,
def asm_gemm(a:Tensor, b:Tensor, w_post_scale:Tensor|None=None, mx:bool=False, mx_scales:tuple|None=None, mx_w_stored:bool=False,
a_pretranspose:Tensor|None=None, mxfp4:bool=False, mxfp4_w:tuple[Tensor, Tensor, Tensor, Tensor]|None=None,
mxfp4_x:tuple[Tensor|None, Tensor|None, Tensor|None, Tensor|None]|None=None) -> Tensor:
assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
assert mx or a.dtype != FP8_DTYPE, "FP8 GEMM requires MXFP8 block scaling"
if mxfp4:
assert not mx and mx_scales is None, "mxfp4 owns quantization; mx/mx_scales are for mxfp8"
assert a.dtype == dtypes.bfloat16, f"cannot quantize {a.dtype} to mxfp4"
@@ -490,15 +364,6 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N
grad_fxn = functools.partial(custom_mx_gemm_bw, has_w_post=has_w_post, w_stored=mx_w_stored)
extra = [w_post_scale] if w_post_scale is not None else []
out = Tensor.custom_kernel(out, a_q.reshape(a.shape), b_q, a_si, b_si, a_e8, b_e8, *extra, fxn=fxn, grad_fxn=grad_fxn)[0]
# fp8 gemm computes a@b.T, kernel multiplies output by x_scale * w_scale before bf16 store
elif a.dtype == FP8_DTYPE:
scales = tuple(s for s in (x_scale, w_scale, g_amax) if s is not None)
scale_mode = (1 if x_scale is not None else 0) | (2 if w_scale is not None else 0) | (4 if g_amax is not None else 0)
assert (grad_amax_state is None) == (next_grad_amax_state is None)
extra = ([grad_amax_state, next_grad_amax_state] if grad_amax_state is not None else []) + ([w_post_scale] if w_post_scale is not None else [])
fxn = functools.partial(custom_hk_fp8_gemm, dname=dname, scale_mode=scale_mode)
bw = functools.partial(custom_gemm_bw, n_scales=len(scales), has_grad_amax=grad_amax_state is not None, has_w_post=w_post_scale is not None)
out = Tensor.custom_kernel(out, a, b.T, *scales, *extra, fxn=fxn, grad_fxn=bw)[0]
elif a.dtype == dtypes.bfloat16:
out = Tensor.custom_kernel(out, a, b.T, b, fxn=functools.partial(custom_hk_bf16_gemm, dname=dname), grad_fxn=custom_gemm_bw)[0]
else:
tinygrad_repo/extra/gemm/gemm_mxfp4.py
+119 -116
View File
@@ -8,6 +8,10 @@ class Kernel:
self.instructions.append(inst)
inst._target, inst._pos = target, self.pos
self.pos += inst.size()
def waitcnt(self, lgkm=None, vm=None, exp=None):
vmcnt, lgkmcnt, expcnt = vm if vm is not None else 63, lgkm if lgkm is not None else 15, exp if exp is not None else 7
waitcnt = (vmcnt & 0xf) | ((expcnt & 0x7) << 4) | ((lgkmcnt & 0xf) << 8) | ((vmcnt & 0x30) << 10)
self.emit(s_waitcnt(simm16=waitcnt))
def finalize(self):
for inst in self.instructions:
if inst._target is not None: inst.simm16 = (self.labels[inst._target] - inst._pos - inst.size()) // 4
@@ -50,7 +54,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(s_mov_b32(s[49], s[2]))
k.emit(s_mov_b32(s[47], s[3]))
k.emit(v_readfirstlane_b32_e32(v[46], v[3]))
k.emit(s_waitcnt(49279))
k.waitcnt(lgkm=0)
if (tile_m, tile_n) == (128, 512):
for i in range(2):
@@ -271,7 +275,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(buffer_load_dwordx4(v[0:3], v[214 + i * 1], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1))
for i in range(24):
k.emit(v_accvgpr_write(v[232 + i * 1], 0))
k.emit(s_waitcnt(20347))
k.waitcnt(vm=27)
k.emit(s_barrier())
k.emit(ds_read_b128(v[8:11], v[216]))
k.emit(ds_read_b128(v[24:27], v[216], v[0], v[0], 0, 64))
@@ -316,11 +320,10 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(s_mul_i32(s[62], s[36], 64))
k.emit(v_add_u32_e32(v[242], s[62], v[240]))
k.emit(v_add_u32_e32(v[243], s[62], v[241]))
k.label('L0_0FAC')
k.emit(s_cmp_lt_i32(s[46], 2))
k.emit(s_cbranch_scc0(1283), target='L0_23C0')
k.label('L0_0FB4')
k.emit(s_waitcnt(127))
k.emit(s_cbranch_scc0(), target='T128_LOOP_LDS1')
k.label('T128_LOOP_LDS0')
k.waitcnt(lgkm=0, vm=15)
k.emit(v_mfma_fp4(v[0:3], v[72:75], v[8:11], 0, 0, v[204], v[200]))
k.emit(s_barrier())
for i in range(2):
@@ -351,7 +354,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(buffer_load_dwordx4(v[164:167], v[227], s[16:19], 0, 0, 1))
k.emit(v_mfma_fp4(v[56:59], v[100:103], v[32:35], 1, 3, v[205], v[201]))
k.emit(v_mfma_fp4(v[60:63], v[100:103], v[36:39], 3, 3, v[205], v[201]))
k.emit(s_waitcnt(3965))
k.waitcnt(vm=13)
k.emit(v_mfma_fp4(v[64:67], v[104:107], v[8:11], 0, 0, v[206], v[200]))
k.emit(s_add_u32(s[62], LIT, s[60], 512))
k.emit(v_mfma_fp4(v[68:71], v[104:107], v[12:15], 2, 0, v[206], v[200]))
@@ -398,7 +401,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(buffer_load_dwordx4(v[188:191], v[233], s[16:19], 0, 0, 1))
k.emit(v_mfma_fp4(v[120:123], v[132:135], v[32:35], 1, 3, v[207], v[201]))
k.emit(v_mfma_fp4(v[124:127], v[132:135], v[36:39], 3, 3, v[207], v[201]))
k.emit(s_waitcnt(16498))
k.waitcnt(lgkm=0, vm=18)
k.emit(v_mfma_fp4(v[128:131], v[72:75], v[40:43], 0, 0, v[204], v[202]))
k.emit(s_barrier())
k.emit(v_mfma_fp4(v[132:135], v[72:75], v[44:47], 2, 0, v[204], v[202]))
@@ -481,8 +484,8 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_mfma_fp4(v[248:251], v[132:135], v[64:67], 1, 3, v[207], v[203]))
k.emit(s_cmp_lt_i32(s[60], s[61]))
k.emit(v_mfma_fp4(v[252:255], v[132:135], v[68:71], 3, 3, v[207], v[203]))
k.emit(s_cbranch_scc0(1926), target='L0_37CC')
k.emit(s_waitcnt(127))
k.emit(s_cbranch_scc0(), target='T128_EPILOGUE')
k.waitcnt(lgkm=0, vm=15)
k.emit(v_mfma_fp4(v[0:3], v[136:139], v[8:11], 0, 0, v[208], v[200]))
k.emit(s_barrier())
for i in range(2):
@@ -513,7 +516,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(buffer_load_dwordx4(v[100:103], v[227], s[16:19], 0, 0, 1))
k.emit(v_mfma_fp4(v[56:59], v[164:167], v[32:35], 1, 3, v[209], v[201]))
k.emit(v_mfma_fp4(v[60:63], v[164:167], v[36:39], 3, 3, v[209], v[201]))
k.emit(s_waitcnt(3965))
k.waitcnt(vm=13)
k.emit(v_mfma_fp4(v[64:67], v[168:171], v[8:11], 0, 0, v[210], v[200]))
k.emit(s_add_u32(s[62], LIT, s[60], 512))
k.emit(v_mfma_fp4(v[68:71], v[168:171], v[12:15], 2, 0, v[210], v[200]))
@@ -560,7 +563,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(buffer_load_dwordx4(v[124:127], v[233], s[16:19], 0, 0, 1))
k.emit(v_mfma_fp4(v[120:123], v[196:199], v[32:35], 1, 3, v[211], v[201]))
k.emit(v_mfma_fp4(v[124:127], v[196:199], v[36:39], 3, 3, v[211], v[201]))
k.emit(s_waitcnt(16498))
k.waitcnt(lgkm=0, vm=18)
k.emit(v_mfma_fp4(v[128:131], v[136:139], v[40:43], 0, 0, v[208], v[202]))
k.emit(s_barrier())
k.emit(v_mfma_fp4(v[132:135], v[136:139], v[44:47], 2, 0, v[208], v[202]))
@@ -643,10 +646,10 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_mfma_fp4(v[248:251], v[196:199], v[64:67], 1, 3, v[211], v[203]))
k.emit(s_cmp_lt_i32(s[60], s[61]))
k.emit(v_mfma_fp4(v[252:255], v[196:199], v[68:71], 3, 3, v[211], v[203]))
k.emit(s_cbranch_scc0(1284), target='L0_37CC')
k.emit(s_branch(64253), target='L0_0FB4')
k.label('L0_23C0')
k.emit(s_waitcnt(127))
k.emit(s_cbranch_scc0(), target='T128_EPILOGUE')
k.emit(s_branch(), target='T128_LOOP_LDS0')
k.label('T128_LOOP_LDS1')
k.waitcnt(lgkm=0, vm=15)
k.emit(v_mfma_fp4(v[0:3], v[72:75], v[8:11], 0, 0, v[204], v[200]))
k.emit(s_barrier())
for i in range(2):
@@ -679,7 +682,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_mfma_fp4(v[56:59], v[100:103], v[32:35], 1, 3, v[205], v[201]))
k.emit(buffer_load_dwordx4(v[164:167], v[227], s[16:19], 0, 0, 1))
k.emit(v_mfma_fp4(v[60:63], v[100:103], v[36:39], 3, 3, v[205], v[201]))
k.emit(s_waitcnt(3965))
k.waitcnt(vm=13)
k.emit(v_mfma_fp4(v[64:67], v[104:107], v[8:11], 0, 0, v[206], v[200]))
k.emit(s_add_u32(s[62], LIT, s[60], 512))
k.emit(v_mfma_fp4(v[68:71], v[104:107], v[12:15], 2, 0, v[206], v[200]))
@@ -724,7 +727,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_mfma_fp4(v[120:123], v[132:135], v[32:35], 1, 3, v[207], v[201]))
k.emit(v_mfma_fp4(v[124:127], v[132:135], v[36:39], 3, 3, v[207], v[201]))
k.emit(buffer_load_dwordx4(v[188:191], v[233], s[16:19], 0, 0, 1))
k.emit(s_waitcnt(16498))
k.waitcnt(lgkm=0, vm=18)
k.emit(v_mfma_fp4(v[128:131], v[72:75], v[40:43], 0, 0, v[204], v[202]))
k.emit(s_barrier())
k.emit(v_mfma_fp4(v[132:135], v[72:75], v[44:47], 2, 0, v[204], v[202]))
@@ -803,8 +806,8 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_mfma_fp4(v[248:251], v[132:135], v[64:67], 1, 3, v[207], v[203]))
k.emit(s_cmp_lt_i32(s[60], s[61]))
k.emit(v_mfma_fp4(v[252:255], v[132:135], v[68:71], 3, 3, v[207], v[203]))
k.emit(s_cbranch_scc0(643), target='L0_37CC')
k.emit(s_waitcnt(127))
k.emit(s_cbranch_scc0(), target='T128_EPILOGUE')
k.waitcnt(lgkm=0, vm=15)
k.emit(v_mfma_fp4(v[0:3], v[136:139], v[8:11], 0, 0, v[208], v[200]))
k.emit(s_barrier())
for i in range(2):
@@ -837,7 +840,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_mfma_fp4(v[56:59], v[164:167], v[32:35], 1, 3, v[209], v[201]))
k.emit(buffer_load_dwordx4(v[100:103], v[227], s[16:19], 0, 0, 1))
k.emit(v_mfma_fp4(v[60:63], v[164:167], v[36:39], 3, 3, v[209], v[201]))
k.emit(s_waitcnt(3965))
k.waitcnt(vm=13)
k.emit(v_mfma_fp4(v[64:67], v[168:171], v[8:11], 0, 0, v[210], v[200]))
k.emit(s_add_u32(s[62], LIT, s[60], 512))
k.emit(v_mfma_fp4(v[68:71], v[168:171], v[12:15], 2, 0, v[210], v[200]))
@@ -882,7 +885,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_mfma_fp4(v[120:123], v[196:199], v[32:35], 1, 3, v[211], v[201]))
k.emit(v_mfma_fp4(v[124:127], v[196:199], v[36:39], 3, 3, v[211], v[201]))
k.emit(buffer_load_dwordx4(v[124:127], v[233], s[16:19], 0, 0, 1))
k.emit(s_waitcnt(16498))
k.waitcnt(lgkm=0, vm=18)
k.emit(v_mfma_fp4(v[128:131], v[136:139], v[40:43], 0, 0, v[208], v[202]))
k.emit(s_barrier())
k.emit(v_mfma_fp4(v[132:135], v[136:139], v[44:47], 2, 0, v[208], v[202]))
@@ -961,10 +964,10 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_mfma_fp4(v[248:251], v[196:199], v[64:67], 1, 3, v[211], v[203]))
k.emit(s_cmp_lt_i32(s[60], s[61]))
k.emit(v_mfma_fp4(v[252:255], v[196:199], v[68:71], 3, 3, v[211], v[203]))
k.emit(s_cbranch_scc0(1), target='L0_37CC')
k.emit(s_branch(64253), target='L0_23C0')
k.label('L0_37CC')
k.emit(s_waitcnt())
k.emit(s_cbranch_scc0(), target='T128_EPILOGUE')
k.emit(s_branch(), target='T128_LOOP_LDS1')
k.label('T128_EPILOGUE')
k.waitcnt(lgkm=0, vm=0, exp=0)
k.emit(s_barrier())
k.emit(s_mul_i32(s[62], s[36], 0))
k.emit(v_add_u32_e32(v[244], s[62], v[240]))
@@ -1215,7 +1218,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(s_nop(1))
k.emit(buffer_store_dwordx4(v[16:19], v[244], s[4:7], 0, 0, 1))
k.emit(v_add_i32(v[244], v[244], 64))
k.emit(s_waitcnt())
k.waitcnt(lgkm=0, vm=0, exp=0)
k.emit(s_endpgm())
elif (tile_m, tile_n) == (192, 256):
k.emit(s_mul_i32(s[63], LIT, 8, 192))
@@ -1246,7 +1249,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(s_nop(3))
k.emit(s_lshl_b32(s[62], s[62], 3))
k.emit(s_cmp_lt_i32(s[3], s[62]))
k.emit(s_cbranch_scc0(33), target='L1_01C0')
k.emit(s_cbranch_scc0(), target='T192_REMAP_IDENTITY')
k.emit(s_add_u32(s[49], s[44], LIT, 255))
k.emit(s_lshr_b32(s[48], s[49], 8))
k.emit(s_mul_i32(s[49], s[48], s[3]))
@@ -1255,38 +1258,38 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(s_lshl_b32(s[47], s[63], 5))
k.emit(s_mul_i32(s[62], s[62], s[47]))
k.emit(s_cmp_lt_i32(s[49], s[62]))
k.emit(s_cbranch_scc0(13), target='L1_0198')
k.emit(s_cbranch_scc0(), target='T192_REMAP_NARROW')
k.emit(s_and_b32(s[62], s[49], LIT, 255))
k.emit(s_and_b32(s[47], s[62], 31))
k.emit(s_lshr_b32(s[48], s[62], 5))
k.emit(s_lshr_b32(s[49], s[49], 8))
k.label('L1_0178')
k.label('T192_REMAP_WIDE_LOOP')
k.emit(s_cmp_lt_i32(s[49], s[63]))
k.emit(s_cbranch_scc1(3), target='L1_018C')
k.emit(s_cbranch_scc1(), target='T192_REMAP_WIDE_DONE')
k.emit(s_sub_i32(s[49], s[49], s[63]))
k.emit(s_add_i32(s[48], s[48], 8))
k.emit(s_branch(65531), target='L1_0178')
k.label('L1_018C')
k.emit(s_branch(), target='T192_REMAP_WIDE_LOOP')
k.label('T192_REMAP_WIDE_DONE')
k.emit(s_mul_i32(s[49], s[49], 32))
k.emit(s_add_i32(s[47], s[47], s[49]))
k.emit(s_branch(12), target='L1_01C8')
k.label('L1_0198')
k.emit(s_branch(), target='T192_REMAP_DONE')
k.label('T192_REMAP_NARROW')
k.emit(s_sub_i32(s[49], s[49], s[62]))
k.emit(s_sub_i32(s[63], s[48], s[47]))
k.emit(s_mov_b32(s[48], 0))
k.label('L1_01A4')
k.label('T192_REMAP_NARROW_LOOP')
k.emit(s_cmp_lt_i32(s[49], s[63]))
k.emit(s_cbranch_scc1(3), target='L1_01B8')
k.emit(s_cbranch_scc1(), target='T192_REMAP_NARROW_DONE')
k.emit(s_sub_i32(s[49], s[49], s[63]))
k.emit(s_add_i32(s[48], s[48], 1))
k.emit(s_branch(65531), target='L1_01A4')
k.label('L1_01B8')
k.emit(s_branch(), target='T192_REMAP_NARROW_LOOP')
k.label('T192_REMAP_NARROW_DONE')
k.emit(s_add_i32(s[47], s[47], s[49]))
k.emit(s_branch(2), target='L1_01C8')
k.label('L1_01C0')
k.emit(s_branch(), target='T192_REMAP_DONE')
k.label('T192_REMAP_IDENTITY')
k.emit(s_mov_b32(s[47], s[2]))
k.emit(s_mov_b32(s[48], s[3]))
k.label('L1_01C8')
k.label('T192_REMAP_DONE')
k.emit(s_lshr_b32(s[37], s[37], 1))
k.emit(s_mul_i32(s[62], s[48], LIT, 192))
k.emit(s_mul_hi_u32(s[63], s[37], s[62]))
@@ -1493,7 +1496,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(buffer_load_dwordx4(v[144:147], v[189], s[16:19], 0, 1024, 1))
k.emit(buffer_load_dwordx4(v[148:151], v[190], s[16:19], 0, 1024, 1))
k.emit(buffer_load_dword(v[176], v[193], s[24:27], 0, 0, 1))
k.emit(s_waitcnt(20347))
k.waitcnt(vm=27)
k.emit(s_barrier())
k.emit(ds_read_b128(v[8:11], v[184]))
k.emit(ds_read_b128(v[16:19], v[184], v[0], v[0], 0, 64))
@@ -1539,12 +1542,12 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
for i in range(23):
k.emit(v_add_u32_e32(v[196 + i * 1], s[62], v[195 + i * 1]))
k.emit(s_cmp_lt_i32(s[46], 2))
k.emit(s_cbranch_scc0(1063), target='L1_1E2C')
k.label('L1_0D90')
k.emit(s_waitcnt(20338))
k.emit(s_cbranch_scc0(), target='T192_LOOP_LDS1')
k.label('T192_LOOP_LDS0')
k.waitcnt(vm=18)
k.emit(s_barrier())
for i in range(2):
k.emit(s_waitcnt(50559))
k.waitcnt(lgkm=5)
for j54 in range(2):
k.emit(v_mfma_fp4(v[0 + i * 32:3 + i * 32], v[104 + j54 * 8:107 + j54 * 8], v[8 + j54 * 8 + i * 16:11 + j54 * 8 + i * 16], 0, 0 + j54 * 3, v[174], v[168 + i * 1]))
k.emit(ds_read_b128(v[40 + j54 * 4 + i * 16:43 + j54 * 4 + i * 16], v[184], v[0], v[0], 0, 0 + i * 128, 33 + j54 * 2 + i * 16))
@@ -1554,7 +1557,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(ds_read_b128(v[48 + j54 * 4 + i * 16:51 + j54 * 4 + i * 16], v[184], v[0], v[0], 0, 64 + i * 128, 33 + j54 * 2 + i * 16))
k.emit(v_mfma_fp4(v[12 + i * 32:15 + i * 32], v[108 + j54 * 8:111 + j54 * 8], v[12 + j54 * 8 + i * 16:15 + j54 * 8 + i * 16], 3, 0 + j54 * 3, v[174], v[168 + i * 1]))
k.emit(ds_read_b32(v[170 + i * 1], v[188], v[0], v[0], 0, 0, 2 + i * 1))
k.emit(s_waitcnt(50559))
k.waitcnt(lgkm=5)
k.emit(v_mfma_fp4(v[64:67], v[104:107], v[40:43], 0, 0, v[174], v[170]))
k.emit(ds_read_b128(v[72:75], v[184], v[0], v[0], 0, 0, 66))
k.emit(v_mfma_fp4(v[68:71], v[104:107], v[44:47], 2, 0, v[174], v[170]))
@@ -1576,7 +1579,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(s_addc_u32(s[17], 0, s[17]))
k.emit(s_sub_u32(s[18], s[18], s[67]))
k.emit(ds_read_b32(v[172], v[188], v[0], v[0], 0, 0, 4))
k.emit(s_waitcnt(50559))
k.waitcnt(lgkm=5)
k.emit(v_mfma_fp4(v[96:99], v[104:107], v[56:59], 0, 0, v[174], v[171]))
k.emit(s_add_u32(s[24], s[24], s[69]))
k.emit(ds_read_b128(v[88:91], v[184], v[0], v[0], 0, 128, 82))
@@ -1592,7 +1595,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_mfma_fp4(v[108:111], v[116:119], v[68:71], 3, 3, v[174], v[171]))
k.emit(ds_read_b32(v[173], v[188], v[0], v[0], 0, 0, 5))
k.emit(s_barrier())
k.emit(s_waitcnt(50559))
k.waitcnt(lgkm=5)
k.emit(v_mfma_fp4(v[128:131], v[104:107], v[72:75], 0, 0, v[174], v[172]))
k.emit(v_mfma_fp4(v[132:135], v[104:107], v[76:79], 2, 0, v[174], v[172]))
k.emit(s_add_u32(NULL, 0, s[65]))
@@ -1605,7 +1608,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(buffer_load_dword(v[0], v[187], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1))
k.emit(v_mfma_fp4(v[136:139], v[116:119], v[80:83], 1, 3, v[174], v[172]))
k.emit(v_mfma_fp4(v[140:143], v[116:119], v[84:87], 3, 3, v[174], v[172]))
k.emit(s_waitcnt(49279))
k.waitcnt(lgkm=0)
k.emit(v_mfma_fp4(v[160:163], v[104:107], v[88:91], 0, 0, v[174], v[173]))
k.emit(v_mfma_fp4(v[164:167], v[104:107], v[92:95], 2, 0, v[174], v[173]))
k.emit(s_add_u32(NULL, 0, s[64]))
@@ -1618,7 +1621,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(buffer_load_dwordx4(v[0:3], v[179], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1))
k.emit(v_mfma_fp4(v[168:171], v[116:119], v[96:99], 1, 3, v[174], v[173]))
k.emit(v_mfma_fp4(v[172:175], v[116:119], v[100:103], 3, 3, v[174], v[173]))
k.emit(s_waitcnt(20338))
k.waitcnt(vm=18)
k.emit(s_barrier())
k.emit(v_mfma_fp4(v[16:19], v[120:123], v[8:11], 0, 0, v[175], v[168]))
k.emit(v_mfma_fp4(v[20:23], v[120:123], v[12:15], 2, 0, v[175], v[168]))
@@ -1689,11 +1692,11 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(ds_read_b128(v[24 + j55 * 8 + i * 4:27 + j55 * 8 + i * 4], v[185], v[0], v[0], 0, 128 + j55 * 64, 16 + i * 2))
k.emit(v_mfma_fp4(v[180 + j55 * 8:183 + j55 * 8], v[120 + j55 * 4 + i * 8:123 + j55 * 4 + i * 8], v[92 + i * 8:95 + i * 8], 2 + j55 * 1, 0 + i * 3, v[175], v[173]))
k.emit(ds_read_b32(v[169], v[188], v[0], v[0], 0, 0, 9))
k.emit(s_cbranch_scc0(1596), target='L1_2EC8')
k.emit(s_waitcnt(20338))
k.emit(s_cbranch_scc0(), target='T192_EPILOGUE')
k.waitcnt(vm=18)
k.emit(s_barrier())
for i in range(2):
k.emit(s_waitcnt(50559))
k.waitcnt(lgkm=5)
for j56 in range(2):
k.emit(v_mfma_fp4(v[0 + i * 32:3 + i * 32], v[136 + j56 * 8:139 + j56 * 8], v[8 + j56 * 8 + i * 16:11 + j56 * 8 + i * 16], 0, 0 + j56 * 3, v[176], v[168 + i * 1]))
k.emit(ds_read_b128(v[40 + j56 * 4 + i * 16:43 + j56 * 4 + i * 16], v[185], v[0], v[0], 0, 0 + i * 128, 33 + j56 * 2 + i * 16))
@@ -1703,7 +1706,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(ds_read_b128(v[48 + j56 * 4 + i * 16:51 + j56 * 4 + i * 16], v[185], v[0], v[0], 0, 64 + i * 128, 33 + j56 * 2 + i * 16))
k.emit(v_mfma_fp4(v[12 + i * 32:15 + i * 32], v[140 + j56 * 8:143 + j56 * 8], v[12 + j56 * 8 + i * 16:15 + j56 * 8 + i * 16], 3, 0 + j56 * 3, v[176], v[168 + i * 1]))
k.emit(ds_read_b32(v[170 + i * 1], v[188], v[0], v[0], 0, 0, 10 + i * 1))
k.emit(s_waitcnt(50559))
k.waitcnt(lgkm=5)
k.emit(v_mfma_fp4(v[64:67], v[136:139], v[40:43], 0, 0, v[176], v[170]))
k.emit(ds_read_b128(v[72:75], v[185], v[0], v[0], 0, 0, 66))
k.emit(v_mfma_fp4(v[68:71], v[136:139], v[44:47], 2, 0, v[176], v[170]))
@@ -1725,7 +1728,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(s_addc_u32(s[17], 0, s[17]))
k.emit(s_sub_u32(s[18], s[18], s[67]))
k.emit(ds_read_b32(v[172], v[188], v[0], v[0], 0, 0, 12))
k.emit(s_waitcnt(50559))
k.waitcnt(lgkm=5)
k.emit(v_mfma_fp4(v[96:99], v[136:139], v[56:59], 0, 0, v[176], v[171]))
k.emit(s_add_u32(s[24], s[24], s[69]))
k.emit(ds_read_b128(v[88:91], v[185], v[0], v[0], 0, 128, 82))
@@ -1741,7 +1744,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_mfma_fp4(v[108:111], v[148:151], v[68:71], 3, 3, v[176], v[171]))
k.emit(ds_read_b32(v[173], v[188], v[0], v[0], 0, 0, 13))
k.emit(s_barrier())
k.emit(s_waitcnt(50559))
k.waitcnt(lgkm=5)
for i in range(2):
k.emit(v_mfma_fp4(v[128:131], v[136 + i * 8:139 + i * 8], v[72 + i * 8:75 + i * 8], 0, 0 + i * 3, v[176], v[172]))
k.emit(v_mfma_fp4(v[132:135], v[136 + i * 8:139 + i * 8], v[76 + i * 8:79 + i * 8], 2, 0 + i * 3, v[176], v[172]))
@@ -1749,7 +1752,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(buffer_load_dword(v[0], v[186 + i * 1], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1))
k.emit(v_mfma_fp4(v[136:139], v[140 + i * 8:143 + i * 8], v[72 + i * 8:75 + i * 8], 1, 0 + i * 3, v[176], v[172]))
k.emit(v_mfma_fp4(v[140:143], v[140 + i * 8:143 + i * 8], v[76 + i * 8:79 + i * 8], 3, 0 + i * 3, v[176], v[172]))
k.emit(s_waitcnt(49279))
k.waitcnt(lgkm=0)
for i in range(2):
k.emit(v_mfma_fp4(v[160:163], v[136 + i * 8:139 + i * 8], v[88 + i * 8:91 + i * 8], 0, 0 + i * 3, v[176], v[173]))
k.emit(v_mfma_fp4(v[164:167], v[136 + i * 8:139 + i * 8], v[92 + i * 8:95 + i * 8], 2, 0 + i * 3, v[176], v[173]))
@@ -1757,7 +1760,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(buffer_load_dwordx4(v[0:3], v[178 + i * 1], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1))
k.emit(v_mfma_fp4(v[168:171], v[140 + i * 8:143 + i * 8], v[88 + i * 8:91 + i * 8], 1, 0 + i * 3, v[176], v[173]))
k.emit(v_mfma_fp4(v[172:175], v[140 + i * 8:143 + i * 8], v[92 + i * 8:95 + i * 8], 3, 0 + i * 3, v[176], v[173]))
k.emit(s_waitcnt(20338))
k.waitcnt(vm=18)
k.emit(s_barrier())
k.emit(v_mfma_fp4(v[16:19], v[152:155], v[8:11], 0, 0, v[177], v[168]))
k.emit(v_mfma_fp4(v[20:23], v[152:155], v[12:15], 2, 0, v[177], v[168]))
@@ -1828,13 +1831,13 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(ds_read_b128(v[24 + j57 * 8 + i * 4:27 + j57 * 8 + i * 4], v[184], v[0], v[0], 0, 128 + j57 * 64, 16 + i * 2))
k.emit(v_mfma_fp4(v[180 + j57 * 8:183 + j57 * 8], v[152 + j57 * 4 + i * 8:155 + j57 * 4 + i * 8], v[92 + i * 8:95 + i * 8], 2 + j57 * 1, 0 + i * 3, v[177], v[173]))
k.emit(ds_read_b32(v[169], v[188], v[0], v[0], 0, 0, 1))
k.emit(s_cbranch_scc0(1064), target='L1_2EC8')
k.emit(s_branch(64473), target='L1_0D90')
k.label('L1_1E2C')
k.emit(s_waitcnt(20338))
k.emit(s_cbranch_scc0(), target='T192_EPILOGUE')
k.emit(s_branch(), target='T192_LOOP_LDS0')
k.label('T192_LOOP_LDS1')
k.waitcnt(vm=18)
k.emit(s_barrier())
for i in range(2):
k.emit(s_waitcnt(50559))
k.waitcnt(lgkm=5)
for j58 in range(2):
k.emit(v_mfma_fp4(v[0 + i * 32:3 + i * 32], v[104 + j58 * 8:107 + j58 * 8], v[8 + j58 * 8 + i * 16:11 + j58 * 8 + i * 16], 0, 0 + j58 * 3, v[174], v[168 + i * 1]))
k.emit(buffer_load_dwordx4(v[152 + j58 * 4 + i * 8:155 + j58 * 4 + i * 8], v[191 + j58 * 1], s[16:19], 0, 0 + i * 1024, 1))
@@ -1844,7 +1847,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_mfma_fp4(v[12 + i * 32:15 + i * 32], v[108 + j58 * 8:111 + j58 * 8], v[12 + j58 * 8 + i * 16:15 + j58 * 8 + i * 16], 3, 0 + j58 * 3, v[174], v[168 + i * 1]))
k.emit(ds_read_b128(v[48 + j58 * 4 + i * 16:51 + j58 * 4 + i * 16], v[184], v[0], v[0], 0, 64 + i * 128, 33 + j58 * 2 + i * 16))
k.emit(ds_read_b32(v[170 + i * 1], v[188], v[0], v[0], 0, 0, 2 + i * 1))
k.emit(s_waitcnt(50559))
k.waitcnt(lgkm=5)
k.emit(v_mfma_fp4(v[64:67], v[104:107], v[40:43], 0, 0, v[174], v[170]))
k.emit(buffer_load_dword(v[177], v[194], s[24:27], 0, 0, 1))
k.emit(v_mfma_fp4(v[68:71], v[104:107], v[44:47], 2, 0, v[174], v[170]))
@@ -1866,7 +1869,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(s_sub_u32(s[18], s[18], s[67]))
k.emit(ds_read_b128(v[84:87], v[184], v[0], v[0], 0, 64, 68))
k.emit(ds_read_b32(v[172], v[188], v[0], v[0], 0, 0, 4))
k.emit(s_waitcnt(50559))
k.waitcnt(lgkm=5)
k.emit(v_mfma_fp4(v[96:99], v[104:107], v[56:59], 0, 0, v[174], v[171]))
k.emit(s_add_u32(s[24], s[24], s[69]))
k.emit(v_mfma_fp4(v[100:103], v[104:107], v[60:63], 2, 0, v[174], v[171]))
@@ -1882,7 +1885,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(ds_read_b128(v[100:103], v[184], v[0], v[0], 0, 192, 84))
k.emit(ds_read_b32(v[173], v[188], v[0], v[0], 0, 0, 5))
k.emit(s_barrier())
k.emit(s_waitcnt(50559))
k.waitcnt(lgkm=5)
k.emit(v_mfma_fp4(v[128:131], v[104:107], v[72:75], 0, 0, v[174], v[172]))
k.emit(s_add_u32(NULL, 0, s[65]))
k.emit(buffer_load_dword(v[0], v[186], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1))
@@ -1895,7 +1898,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_mfma_fp4(v[132:135], v[112:115], v[84:87], 2, 3, v[174], v[172]))
k.emit(v_mfma_fp4(v[136:139], v[116:119], v[80:83], 1, 3, v[174], v[172]))
k.emit(v_mfma_fp4(v[140:143], v[116:119], v[84:87], 3, 3, v[174], v[172]))
k.emit(s_waitcnt(49279))
k.waitcnt(lgkm=0)
k.emit(v_mfma_fp4(v[160:163], v[104:107], v[88:91], 0, 0, v[174], v[173]))
k.emit(s_add_u32(NULL, 0, s[64]))
k.emit(buffer_load_dwordx4(v[0:3], v[178], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1))
@@ -1908,7 +1911,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_mfma_fp4(v[164:167], v[112:115], v[100:103], 2, 3, v[174], v[173]))
k.emit(v_mfma_fp4(v[168:171], v[116:119], v[96:99], 1, 3, v[174], v[173]))
k.emit(v_mfma_fp4(v[172:175], v[116:119], v[100:103], 3, 3, v[174], v[173]))
k.emit(s_waitcnt(20338))
k.waitcnt(vm=18)
k.emit(s_barrier())
k.emit(v_mfma_fp4(v[16:19], v[120:123], v[8:11], 0, 0, v[175], v[168]))
k.emit(s_add_u32(NULL, LIT, s[64], 8448))
@@ -1975,11 +1978,11 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_mfma_fp4(v[180 + j59 * 8:183 + j59 * 8], v[120 + j59 * 4 + i * 8:123 + j59 * 4 + i * 8], v[92 + i * 8:95 + i * 8], 2 + j59 * 1, 0 + i * 3, v[175], v[173]))
k.emit(ds_read_b128(v[24 + j59 * 8 + i * 4:27 + j59 * 8 + i * 4], v[185], v[0], v[0], 0, 128 + j59 * 64, 16 + i * 2))
k.emit(ds_read_b32(v[169], v[188], v[0], v[0], 0, 0, 9))
k.emit(s_cbranch_scc0(533), target='L1_2EC8')
k.emit(s_waitcnt(20338))
k.emit(s_cbranch_scc0(), target='T192_EPILOGUE')
k.waitcnt(vm=18)
k.emit(s_barrier())
for i in range(2):
k.emit(s_waitcnt(50559))
k.waitcnt(lgkm=5)
for j60 in range(2):
k.emit(v_mfma_fp4(v[0 + i * 32:3 + i * 32], v[136 + j60 * 8:139 + j60 * 8], v[8 + j60 * 8 + i * 16:11 + j60 * 8 + i * 16], 0, 0 + j60 * 3, v[176], v[168 + i * 1]))
k.emit(buffer_load_dwordx4(v[120 + j60 * 4 + i * 8:123 + j60 * 4 + i * 8], v[191 + j60 * 1], s[16:19], 0, 0 + i * 1024, 1))
@@ -1989,7 +1992,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_mfma_fp4(v[12 + i * 32:15 + i * 32], v[140 + j60 * 8:143 + j60 * 8], v[12 + j60 * 8 + i * 16:15 + j60 * 8 + i * 16], 3, 0 + j60 * 3, v[176], v[168 + i * 1]))
k.emit(ds_read_b128(v[48 + j60 * 4 + i * 16:51 + j60 * 4 + i * 16], v[185], v[0], v[0], 0, 64 + i * 128, 33 + j60 * 2 + i * 16))
k.emit(ds_read_b32(v[170 + i * 1], v[188], v[0], v[0], 0, 0, 10 + i * 1))
k.emit(s_waitcnt(50559))
k.waitcnt(lgkm=5)
k.emit(v_mfma_fp4(v[64:67], v[136:139], v[40:43], 0, 0, v[176], v[170]))
k.emit(buffer_load_dword(v[175], v[194], s[24:27], 0, 0, 1))
k.emit(v_mfma_fp4(v[68:71], v[136:139], v[44:47], 2, 0, v[176], v[170]))
@@ -2011,7 +2014,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(s_sub_u32(s[18], s[18], s[67]))
k.emit(ds_read_b128(v[84:87], v[185], v[0], v[0], 0, 64, 68))
k.emit(ds_read_b32(v[172], v[188], v[0], v[0], 0, 0, 12))
k.emit(s_waitcnt(50559))
k.waitcnt(lgkm=5)
k.emit(v_mfma_fp4(v[96:99], v[136:139], v[56:59], 0, 0, v[176], v[171]))
k.emit(s_add_u32(s[24], s[24], s[69]))
k.emit(v_mfma_fp4(v[100:103], v[136:139], v[60:63], 2, 0, v[176], v[171]))
@@ -2027,7 +2030,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(ds_read_b128(v[100:103], v[185], v[0], v[0], 0, 192, 84))
k.emit(ds_read_b32(v[173], v[188], v[0], v[0], 0, 0, 13))
k.emit(s_barrier())
k.emit(s_waitcnt(50559))
k.waitcnt(lgkm=5)
for i in range(2):
k.emit(v_mfma_fp4(v[128:131], v[136 + i * 8:139 + i * 8], v[72 + i * 8:75 + i * 8], 0, 0 + i * 3, v[176], v[172]))
k.emit(s_add_u32(NULL, LIT, s[65], 2048 + i * 1024))
@@ -2035,7 +2038,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_mfma_fp4(v[132:135], v[136 + i * 8:139 + i * 8], v[76 + i * 8:79 + i * 8], 2, 0 + i * 3, v[176], v[172]))
k.emit(v_mfma_fp4(v[136:139], v[140 + i * 8:143 + i * 8], v[72 + i * 8:75 + i * 8], 1, 0 + i * 3, v[176], v[172]))
k.emit(v_mfma_fp4(v[140:143], v[140 + i * 8:143 + i * 8], v[76 + i * 8:79 + i * 8], 3, 0 + i * 3, v[176], v[172]))
k.emit(s_waitcnt(49279))
k.waitcnt(lgkm=0)
for i in range(2):
k.emit(v_mfma_fp4(v[160:163], v[136 + i * 8:139 + i * 8], v[88 + i * 8:91 + i * 8], 0, 0 + i * 3, v[176], v[173]))
k.emit(s_add_u32(NULL, LIT, s[64], 25344 + i * 4224))
@@ -2043,7 +2046,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_mfma_fp4(v[164:167], v[136 + i * 8:139 + i * 8], v[92 + i * 8:95 + i * 8], 2, 0 + i * 3, v[176], v[173]))
k.emit(v_mfma_fp4(v[168:171], v[140 + i * 8:143 + i * 8], v[88 + i * 8:91 + i * 8], 1, 0 + i * 3, v[176], v[173]))
k.emit(v_mfma_fp4(v[172:175], v[140 + i * 8:143 + i * 8], v[92 + i * 8:95 + i * 8], 3, 0 + i * 3, v[176], v[173]))
k.emit(s_waitcnt(20338))
k.waitcnt(vm=18)
k.emit(s_barrier())
k.emit(v_mfma_fp4(v[16:19], v[152:155], v[8:11], 0, 0, v[177], v[168]))
k.emit(s_add_u32(NULL, LIT, s[64], 33792))
@@ -2110,10 +2113,10 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_mfma_fp4(v[180 + j61 * 8:183 + j61 * 8], v[152 + j61 * 4 + i * 8:155 + j61 * 4 + i * 8], v[92 + i * 8:95 + i * 8], 2 + j61 * 1, 0 + i * 3, v[177], v[173]))
k.emit(ds_read_b128(v[24 + j61 * 8 + i * 4:27 + j61 * 8 + i * 4], v[184], v[0], v[0], 0, 128 + j61 * 64, 16 + i * 2))
k.emit(ds_read_b32(v[169], v[188], v[0], v[0], 0, 0, 1))
k.emit(s_cbranch_scc0(1), target='L1_2EC8')
k.emit(s_branch(64473), target='L1_1E2C')
k.label('L1_2EC8')
k.emit(s_waitcnt())
k.emit(s_cbranch_scc0(), target='T192_EPILOGUE')
k.emit(s_branch(), target='T192_LOOP_LDS1')
k.label('T192_EPILOGUE')
k.waitcnt(lgkm=0, vm=0, exp=0)
k.emit(s_barrier())
k.emit(v_lshrrev_b32_e32(v[4], 5))
k.emit(v_mul_i32_i24_e32(v[4], 16, v[4]))
@@ -2201,17 +2204,17 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_permlane16_swap_b32_e32(v[17], v[19]))
k.emit(s_nop(1))
k.emit(ds_write_b128(v[0], v[4], v[16:19], v[0], 0, 0 + i * 64, 88))
k.emit(s_waitcnt(49279))
k.waitcnt(lgkm=0)
k.emit(v_mul_i32_i24_e64(v[4], v[0], 16))
k.emit(v_add_i32(v[4], v[4], s[62]))
k.emit(ds_read_b128(v[16:19], v[4]))
for i in range(23):
k.emit(s_waitcnt(49279))
k.waitcnt(lgkm=0)
k.emit(buffer_store_dwordx4(v[16:19], v[195 + i * 1], s[4:7], 0, 0, 1))
k.emit(ds_read_b128(v[16:19], v[4], v[0], v[0], 0, 0, 4 + i * 4))
k.emit(s_waitcnt(49279))
k.waitcnt(lgkm=0)
k.emit(buffer_store_dwordx4(v[16:19], v[218], s[4:7], 0, 0, 1))
k.emit(s_waitcnt())
k.waitcnt(lgkm=0, vm=0, exp=0)
k.emit(s_endpgm())
elif (tile_m, tile_n) == (256, 256):
k.emit(s_add_u32(s[55], s[44], LIT, 255))
@@ -2222,20 +2225,20 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(s_lshr_b32(s[52], s[55], 8))
k.emit(s_lshl_b32(s[52], s[52], 5))
k.emit(s_mov_b32(s[49], 0))
k.label('L2_00E8')
k.label('T256_REMAP_STRIP_LOOP')
k.emit(s_cmp_lt_i32(s[48], s[52]))
k.emit(s_cbranch_scc1(3), target='L2_00FC')
k.emit(s_cbranch_scc1(), target='T256_REMAP_STRIP_DONE')
k.emit(s_sub_i32(s[48], s[48], s[52]))
k.emit(s_add_i32(s[49], s[49], 32))
k.emit(s_branch(65531), target='L2_00E8')
k.label('L2_00FC')
k.emit(s_branch(), target='T256_REMAP_STRIP_LOOP')
k.label('T256_REMAP_STRIP_DONE')
k.emit(s_sub_i32(s[54], s[54], s[49]))
k.emit(s_cmp_lt_i32(s[54], 32))
k.emit(s_cbranch_scc1(3), target='L2_0114')
k.emit(s_cbranch_scc1(), target='T256_REMAP_SMALL_STRIP')
k.emit(s_lshr_b32(s[47], s[48], 5))
k.emit(s_and_b32(s[52], s[48], 31))
k.emit(s_branch(32), target='L2_0194')
k.label('L2_0114')
k.emit(s_branch(), target='T256_REMAP_DONE')
k.label('T256_REMAP_SMALL_STRIP')
k.emit(v_cvt_f32_u32_e32(v[4], s[54]))
k.emit(s_sub_i32(s[47], 0, s[54]))
k.emit(v_rcp_iflag_f32_e32(v[4], v[4]))
@@ -2263,7 +2266,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(s_nop(3))
k.emit(s_mul_i32(s[52], s[54], s[47]))
k.emit(s_sub_i32(s[52], s[48], s[52]))
k.label('L2_0194')
k.label('T256_REMAP_DONE')
k.emit(s_add_i32(s[49], s[52], s[49]))
k.emit(s_mov_b32(s[6], -16))
k.emit(s_mov_b32(s[10], -16))
@@ -2447,7 +2450,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(s_sub_u32(s[14 + i * 8], s[14 + i * 8], s[61 + i * 2]))
for i in range(8):
k.emit(v_accvgpr_write(v[248 + i * 1], 0))
k.emit(s_waitcnt(20345))
k.waitcnt(vm=25)
k.emit(s_barrier())
k.emit(ds_read_b128(v[8:11], v[220]))
k.emit(ds_read_b128(v[40:43], v[220], v[0], v[0], 0, 64))
@@ -2497,9 +2500,9 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(s_cselect_b32(s[61 + i * 1], s[61 + i * 1], 0))
k.emit(s_cselect_b32(s[63 + i * 1], s[63 + i * 1], 0))
k.emit(s_cmp_lt_i32(s[46], 2))
k.emit(s_cbranch_scc0(1367), target='L2_25B8')
k.label('L2_105C')
k.emit(s_waitcnt(122))
k.emit(s_cbranch_scc0(), target='T256_LOOP_LDS1_INIT')
k.label('T256_LOOP_LDS0')
k.waitcnt(lgkm=0, vm=10)
k.emit(v_mfma_fp4(v[0:3], v[136:139], v[8:11], 0, 0, v[208], v[200]))
k.emit(s_barrier())
k.emit(s_nop())
@@ -2571,7 +2574,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_mfma_fp4(v[92:95], v[160:163], v[68:71], 2, 3, v[209], v[203]))
k.emit(v_mfma_fp4(v[120:123], v[164:167], v[64:67], 1, 3, v[209], v[203]))
k.emit(v_mfma_fp4(v[124:127], v[164:167], v[68:71], 3, 3, v[209], v[203]))
k.emit(s_waitcnt(127))
k.waitcnt(lgkm=0, vm=15)
k.emit(v_mfma_fp4(v[128:131], v[136:139], v[72:75], 0, 0, v[208], v[204]))
k.emit(s_barrier())
k.emit(s_nop())
@@ -2666,8 +2669,8 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(s_cmp_lt_i32(s[50], s[51]))
k.emit(v_mfma_fp4(v[248:251], v[164:167], v[128:131], 1, 3, v[209], v[207]))
k.emit(v_mfma_fp4(v[252:255], v[164:167], v[132:135], 3, 3, v[209], v[207]))
k.emit(s_cbranch_scc0(2051), target='L2_3B10')
k.emit(s_waitcnt(122))
k.emit(s_cbranch_scc0(), target='T256_EPILOGUE')
k.waitcnt(lgkm=0, vm=10)
k.emit(v_mfma_fp4(v[0:3], v[168:171], v[8:11], 0, 0, v[210], v[200]))
k.emit(s_barrier())
k.emit(s_nop())
@@ -2739,7 +2742,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_mfma_fp4(v[92:95], v[192:195], v[68:71], 2, 3, v[211], v[203]))
k.emit(v_mfma_fp4(v[120:123], v[196:199], v[64:67], 1, 3, v[211], v[203]))
k.emit(v_mfma_fp4(v[124:127], v[196:199], v[68:71], 3, 3, v[211], v[203]))
k.emit(s_waitcnt(127))
k.waitcnt(lgkm=0, vm=15)
k.emit(v_mfma_fp4(v[128:131], v[168:171], v[72:75], 0, 0, v[210], v[204]))
k.emit(s_barrier())
k.emit(s_nop())
@@ -2834,12 +2837,12 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(s_cmp_lt_i32(s[50], s[51]))
k.emit(v_mfma_fp4(v[248:251], v[196:199], v[128:131], 1, 3, v[211], v[207]))
k.emit(v_mfma_fp4(v[252:255], v[196:199], v[132:135], 3, 3, v[211], v[207]))
k.emit(s_cbranch_scc0(1367), target='L2_3B10')
k.emit(s_branch(64169), target='L2_105C')
k.label('L2_25B8')
k.emit(s_cbranch_scc0(), target='T256_EPILOGUE')
k.emit(s_branch(), target='T256_LOOP_LDS0')
k.label('T256_LOOP_LDS1_INIT')
k.emit(s_nop())
k.label('L2_25BC')
k.emit(s_waitcnt(122))
k.label('T256_LOOP_LDS1')
k.waitcnt(lgkm=0, vm=10)
k.emit(v_mfma_fp4(v[0:3], v[136:139], v[8:11], 0, 0, v[208], v[200]))
k.emit(s_barrier())
k.emit(s_nop())
@@ -2911,7 +2914,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_mfma_fp4(v[92:95], v[160:163], v[68:71], 2, 3, v[209], v[203]))
k.emit(v_mfma_fp4(v[120:123], v[164:167], v[64:67], 1, 3, v[209], v[203]))
k.emit(v_mfma_fp4(v[124:127], v[164:167], v[68:71], 3, 3, v[209], v[203]))
k.emit(s_waitcnt(127))
k.waitcnt(lgkm=0, vm=15)
k.emit(v_mfma_fp4(v[128:131], v[136:139], v[72:75], 0, 0, v[208], v[204]))
k.emit(s_barrier())
k.emit(s_nop())
@@ -3006,8 +3009,8 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(s_cmp_lt_i32(s[50], s[51]))
k.emit(v_mfma_fp4(v[248:251], v[164:167], v[128:131], 1, 3, v[209], v[207]))
k.emit(v_mfma_fp4(v[252:255], v[164:167], v[132:135], 3, 3, v[209], v[207]))
k.emit(s_cbranch_scc0(684), target='L2_3B10')
k.emit(s_waitcnt(122))
k.emit(s_cbranch_scc0(), target='T256_EPILOGUE')
k.waitcnt(lgkm=0, vm=10)
k.emit(v_mfma_fp4(v[0:3], v[168:171], v[8:11], 0, 0, v[210], v[200]))
k.emit(s_barrier())
k.emit(s_nop())
@@ -3079,7 +3082,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_mfma_fp4(v[92:95], v[192:195], v[68:71], 2, 3, v[211], v[203]))
k.emit(v_mfma_fp4(v[120:123], v[196:199], v[64:67], 1, 3, v[211], v[203]))
k.emit(v_mfma_fp4(v[124:127], v[196:199], v[68:71], 3, 3, v[211], v[203]))
k.emit(s_waitcnt(127))
k.waitcnt(lgkm=0, vm=15)
k.emit(v_mfma_fp4(v[128:131], v[168:171], v[72:75], 0, 0, v[210], v[204]))
k.emit(s_barrier())
k.emit(s_nop())
@@ -3173,10 +3176,10 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(s_cmp_lt_i32(s[50], s[51]))
k.emit(v_mfma_fp4(v[248:251], v[196:199], v[128:131], 1, 3, v[211], v[207]))
k.emit(v_mfma_fp4(v[252:255], v[196:199], v[132:135], 3, 3, v[211], v[207]))
k.emit(s_cbranch_scc0(1), target='L2_3B10')
k.emit(s_branch(64171), target='L2_25BC')
k.label('L2_3B10')
k.emit(s_waitcnt(112))
k.emit(s_cbranch_scc0(), target='T256_EPILOGUE')
k.emit(s_branch(), target='T256_LOOP_LDS1')
k.label('T256_EPILOGUE')
k.waitcnt(lgkm=0, vm=0)
k.emit(s_barrier())
for i in range(2):
for j94 in range(4):
@@ -3391,7 +3394,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(s_nop(1))
k.emit(buffer_store_dwordx4(v[16:19], v[250], s[4:7], 0, 0, 1))
k.emit(v_add_i32(v[250], v[250], 64))
k.emit(s_waitcnt())
k.waitcnt(lgkm=0, vm=0, exp=0)
k.emit(s_endpgm())
else:
raise AssertionError(f'unsupported tile {(tile_m, tile_n)}')
tinygrad_repo/extra/gemm/metal_uop_matmul.py
+1 -1
View File
@@ -18,7 +18,7 @@ def hand_spec_tc_cores():
a = UOp.placeholder((N, N), dtypes.float, slot=1).reshape((N//8, 8, N//8, 8))
b = UOp.placeholder((N, N), dtypes.float, slot=2).reshape((N//8, 8, N//8, 8))
gk = UOp.range(N // 8, 0, AxisType.REDUCE)
gk = UOp.range(N // 8, 0, AxisType.LOOP)
a_tc = UOp.stack(*[mat_idx(a, gx, gk, warp, i) for i in range(2)])
b_tc = UOp.stack(*[mat_idx(b, gk, gy, warp, i) for i in range(2)])
tinygrad_repo/extra/gemm/mi350x_uop_matmul.py
+3 -3
View File
@@ -111,7 +111,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
return Asl.after(barrier), Bsl.after(barrier)
def compute_on_locals(acc:UOp, Asl:UOp, Bsl:UOp, rng:int, afters:tuple[UOp, ...]=()) -> UOp:
K_inner_loop = UOp.range(BLOCK_K//TC_K, rng, AxisType.REDUCE)
K_inner_loop = UOp.range(BLOCK_K//TC_K, rng, AxisType.LOOP)
# load from locals into registers
Ar = UOp.placeholder((BLOCK_M//TC_M//WARPGROUP_SIZE,), dtypes.half, slot=1, addrspace=AddrSpace.REG)
@@ -150,7 +150,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
if not getenv("PIPELINE"):
As, Bs = make_locals(slot=0)
K_outer_loop = UOp.range(K//BLOCK_K, 0, AxisType.REDUCE)
K_outer_loop = UOp.range(K//BLOCK_K, 0, AxisType.LOOP)
As, Bs = load_to_locals(K_outer_loop, As, Bs, 1000, barrier=True)
acc_store = compute_on_locals(acc, As, Bs, 1500, afters=(K_outer_loop,))
acc = acc.after(acc_store.barrier().end(K_outer_loop))
@@ -160,7 +160,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
As1, Bs1 = make_locals(slot=2)
As0, Bs0 = load_to_locals(0, As0, Bs0, 1000)
K_outer_loop = UOp.range((K//BLOCK_K-2)//2, 0, AxisType.REDUCE)
K_outer_loop = UOp.range((K//BLOCK_K-2)//2, 0, AxisType.LOOP)
As1, Bs1 = load_to_locals(K_outer_loop+1, As1, Bs1, 2000, barrier=False)
acc_store = compute_on_locals(acc, As0, Bs0, 1500, afters=(K_outer_loop,))
As0, Bs0 = load_to_locals(K_outer_loop+2, As0, Bs0, 3000, barrier=False)
tinygrad_repo/extra/gemm/mi350x_uop_matmul_2.py
+3 -3
View File
@@ -34,7 +34,7 @@ def copy(dest:UOp, src:UOp, rng:int, set=False, upcast=()):
return dest.after(copy) if set else copy
def compute_on_locals(acc:UOp, Asl:UOp, Bsl:UOp, rng:int, afters:tuple[UOp, ...], warpgroup, warp) -> UOp:
K_inner_loop = UOp.range(BLOCK_K//TC_K, rng, AxisType.REDUCE)
K_inner_loop = UOp.range(BLOCK_K//TC_K, rng, AxisType.LOOP)
# load from locals into registers
Ar = UOp.placeholder((BLOCK_M//TC_M//WARPGROUP_SIZE,), dtypes.half, slot=1, addrspace=AddrSpace.REG)
@@ -68,7 +68,7 @@ def compute_on_locals(acc:UOp, Asl:UOp, Bsl:UOp, rng:int, afters:tuple[UOp, ...]
def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
gx, gy = UOp.special(M//BLOCK_M, "gidx0"), UOp.special(N//BLOCK_N, "gidx1")
K_outer_loop = UOp.range(K//BLOCK_K, 0, AxisType.REDUCE)
K_outer_loop = UOp.range(K//BLOCK_K, 0, AxisType.LOOP)
# split out the globals into blocks
C = C.src[0].cast(dtypes.float).reshape((M//BLOCK_M, BLOCK_M, N//BLOCK_N, BLOCK_N))
@@ -114,7 +114,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
C_view = C[gx, :, gy, :].reshape(BLOCK_M//TC_M//WARPGROUP_SIZE, WARPGROUP_SIZE, TC_M, BLOCK_N//TC_N, TC_N)[:, warpgroup, warp%16, :, (warp//16)*4]
sink = copy(C_view, acc.after(sink), rng=300)
else:
sink = C.after(barrier.end(K_outer_loop))[0,0,0,0].store(As[0,0]+Bs[0,0])
sink = C.after(barrier.end(K_outer_loop))[0,0,0,0].store((As[0,0]+Bs[0,0]).float())
return sink.sink(arg=KernelInfo(name="custom_gemm", opts_to_apply=())).simplify()
tinygrad_repo/extra/gemm/moe_gemm.py
+20 -11
View File
@@ -4,7 +4,7 @@ from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
from tinygrad.helpers import getenv
from tinygrad.renderer import Estimates
from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
from extra.gemm.cdna_asm_gemm import quantize_mxfp8, _mx_block_scale, _mx_block_scale_3d
from extra.gemm.cdna_asm_gemm import quantize_mxfp8, mx_pack, _mx_block_scale, _mx_block_scale_3d
ZERO_OPTIM = getenv("ZERO_OPTIM", 0)
@@ -95,31 +95,40 @@ def custom_grouped_mx_gemm_bw(gradient:UOp, kernel:UOp, w_stored:bool=False) ->
bq = Tensor(inputs[2], device=inputs[2].device)
ae8 = Tensor(inputs[5], device=inputs[5].device)
be8 = Tensor(inputs[6], device=inputs[6].device)
E, N = bq.shape[0], bq.shape[1]
M, K = aq.shape
g = Tensor(gradient, device=aq.device).reshape(M, N).cast(dtypes.bfloat16)
expert_off = Tensor(inputs[7], device=inputs[7].device)
grad_xq, grad_wq = grouped_mx_gemm_backward(Tensor(gradient), aq, ae8, bq, be8, expert_off, w_stored=w_stored)
return (None, grad_xq.uop, grad_wq.uop) + tuple(None for _ in inputs[3:])
def grouped_mx_gemm_backward(g:Tensor, aq:Tensor, ae8:Tensor, bq:Tensor, be8:Tensor, expert_off:Tensor, *,
w_stored:bool=False, quantized_gradient:tuple[Tensor, Tensor]|None=None) -> tuple[Tensor, Tensor]:
E, N = bq.shape[:2]
g = g.reshape(aq.shape[0], N).cast(dtypes.bfloat16)
x_phys = (aq.cast(dtypes.bfloat16) * _mx_block_scale(ae8).cast(dtypes.bfloat16))
w_phys = (bq.cast(dtypes.bfloat16) * _mx_block_scale_3d(be8).cast(dtypes.bfloat16))
expert_off = Tensor(inputs[7], device=inputs[7].device)
grad_x = grouped_mx_gemm(g, w_phys.transpose(1, 2), expert_off)
grad_x = grouped_mx_gemm(g if quantized_gradient is None else quantized_gradient, w_phys.transpose(1, 2), expert_off)
grad_w = grouped_mx_wgrad(g, x_phys, expert_off, E)
grad_xq = grad_x * _mx_block_scale(ae8).cast(dtypes.bfloat16)
grad_wq = grad_w.contiguous() if w_stored else (grad_w * _mx_block_scale_3d(be8).cast(dtypes.bfloat16)).contiguous()
return (None, grad_xq.uop, grad_wq.uop) + tuple(None for _ in inputs[3:])
return grad_xq, grad_wq
_grouped_bw_stored = functools.partial(custom_grouped_mx_gemm_bw, w_stored=True)
def grouped_mx_gemm(x:Tensor, w:Tensor|tuple[Tensor, Tensor], expert_off:Tensor) -> Tensor:
def grouped_mx_gemm(x:Tensor|tuple[Tensor, Tensor], w:Tensor|tuple[Tensor, Tensor], expert_off:Tensor) -> Tensor:
if isinstance(x, tuple):
x_q, x_e8 = x
x_si, x = mx_pack(x_e8), x_q
else:
x_q, x_e8, x_si = quantize_mxfp8(x)
if (pre_quantized := isinstance(w, tuple)):
w_q, w_e8 = w
E, N, K2 = w_q.shape
else:
E, N, K2 = w.shape
M, K = x.shape
assert K == K2, f"shape mismatch {x.shape} {w.shape}"
assert M % 256 == 0 and N % 256 == 0 and K % 128 == 0, f"grouped mxfp8 needs M%256,N%256,K%128, got {x.shape} {w.shape}"
assert K == K2, f"shape mismatch K {K} != {K2}"
assert x_e8.shape == (M, K // 32)
assert M % 256 == 0 and N % 256 == 0 and K % 128 == 0, f"grouped mxfp8 needs M%256,N%256,K%128, got {M,K,N}"
dname = (x.device[0] if isinstance(x.device, tuple) else x.device).split(":")[0]
x_q, x_e8, x_si = quantize_mxfp8(x)
if not pre_quantized: w_q, w_e8, _ = quantize_mxfp8(w)
w_si = mx_pack_3d(w_e8)
xe_in, out_shape = x_e8.reshape(M, K // 32), (M, N)
tinygrad_repo/extra/gemm/moe_routing.py
+40 -7
View File
@@ -123,8 +123,14 @@ def _gscatter_bwd(gradient:UOp, kernel:UOp) -> tuple:
dev = src_u.device
G, T_l, D = src_u.shape
k = idx_u.shape[1] // T_l
if getenv("GGATHER_SUM_HIP", 0):
from extra.gptoss_kernels.gather_sum import gather_sum
assert k == 4
return None, gather_sum(Tensor(gradient), Tensor(idx_u)).uop, None
sel = grouped_gather_rows(Tensor(gradient, device=dev), Tensor(idx_u, device=dev), G)
return (None, sel.reshape(G, T_l, k, D).sum(2).cast(src_u.dtype).uop, None)
# FP8 dispatch carries BF16 gradients back to the quantizer, not FP8-rounded gradients.
grad_dtype = dtypes.bfloat16 if src_u.dtype == dtypes.fp8e4m3 else src_u.dtype
return (None, sel.reshape(G, T_l, k, D).sum(2).cast(grad_dtype).uop, None)
def grouped_scatter_rows(src:Tensor, idx:Tensor, m_l:int) -> Tensor:
G, T_l, D = src.shape
@@ -135,9 +141,12 @@ def m_max_for(t_local:int, experts_per_tok:int, n_experts:int) -> int:
return (-(-t_local * experts_per_tok // BLOCK_ROW) + n_experts) * BLOCK_ROW
class Routing:
def __init__(self, weights:Tensor, dest_row:Tensor, off:Tensor, m_l:int, n_groups:int, t_local:int):
def __init__(self, weights:Tensor, dest_row:Tensor, off:Tensor, m_l:int, n_groups:int, t_local:int,
topi:Tensor|None=None, counts:Tensor|None=None):
self.weights, self.dest_row = weights, dest_row
self.off = off
self.topi = topi
self.counts = counts
self.m_l, self.n_groups, self.t_local = m_l, n_groups, t_local
@property
@@ -158,22 +167,46 @@ def route(logits:Tensor, experts_per_tok:int, n_experts:int) -> Routing:
T, E = logits.shape
k, G = experts_per_tok, n_groups_of(logits)
assert T % G == 0, f"tokens {T} must split across {G} devices"
T_l, m_l = T // G, m_max_for(T // G, k, n_experts)
T_l = T // G
topv, topi = logits.reshape(G, T_l, E).topk(k)
weights = topv.softmax(-1)
if getenv("FUSED_ROUTER_TOPK", 0):
from extra.gptoss_kernels.router_topk import fused_router_topk
weights, topi = fused_router_topk(logits.reshape(G, T_l, E))
else:
topv, topi = logits.reshape(G, T_l, E).topk(k)
weights = topv.softmax(-1)
return route_topk(weights, topi, n_experts)
def route_topk(weights:Tensor, topi:Tensor, n_experts:int) -> Routing:
G, T_l, k = weights.shape
E, m_l = n_experts, m_max_for(T_l, k, n_experts)
m = topi.reshape(G, T_l * k).cast(dtypes.int32).one_hot(E).cast(dtypes.int32)
pad = ((m.sum(1) + (BLOCK_ROW - 1)) // BLOCK_ROW) * BLOCK_ROW
counts = m.sum(1)
pad = ((counts + (BLOCK_ROW - 1)) // BLOCK_ROW) * BLOCK_ROW
off = pad.cumsum(1).pad(((0, 0), (1, 0)))
dest_row = ((m.cumsum(1) + off[:, :E].reshape(G, 1, E)) * m).sum(-1).sub(1).cast(dtypes.int32)
return Routing(weights, dest_row, off, m_l, G, T_l)
return Routing(weights, dest_row, off, m_l, G, T_l, topi=topi, counts=counts)
def dispatch(x:Tensor, r:Routing) -> Tensor:
G, D = r.n_groups, x.shape[-1]
return grouped_scatter_rows(x.reshape(G, r.t_local, D), r.dest_row, r.m_l).reshape(G * r.m_l, D)
def dispatch_fp8(x:Tensor|tuple[Tensor, Tensor], r:Routing) -> tuple[Tensor, Tensor]:
from extra.gptoss_kernels.quantize_mxfp8 import quantize_mxfp8_fused_qe8
q, e8 = x if isinstance(x, tuple) else quantize_mxfp8_fused_qe8(x)
if getenv("DISPATCH_GATHER", 0):
from extra.gptoss_kernels.dispatch import inverse_rows, dispatch_gather
assert r.counts is not None
src_row = inverse_rows(r.dest_row, r.counts, r.off, r.m_l)
return tuple(dispatch_gather(t.reshape(r.n_groups, r.t_local, -1), r.dest_row, src_row) for t in (q, e8))
return dispatch(q, r), dispatch(e8, r)
def combine(y:Tensor, r:Routing, n_tokens:int, experts_per_tok:int) -> Tensor:
G, D, k = r.n_groups, y.shape[-1], experts_per_tok
if getenv("FUSED_COMBINE", 0):
from extra.gptoss_kernels.combine import fused_combine
assert k == 4
return fused_combine(y.reshape(G, r.m_l, D), r.dest_row, r.weights.reshape(G, r.t_local, k)).reshape(n_tokens, D)
sel = grouped_gather_rows(y.reshape(G, r.m_l, D), r.dest_row, G).reshape(G, r.t_local, k, D)
return (sel * r.weights.reshape(G, r.t_local, k, 1).cast(sel.dtype)).sum(2).reshape(n_tokens, D).cast(y.dtype)
tinygrad_repo/extra/gemm/tinygrad_nv_matmul.py
+2 -2
View File
@@ -14,8 +14,8 @@ if __name__ == "__main__":
C = A.matmul(B)
if getenv("GEMV"):
opts = [
Opt(op=OptOps.SPLIT, axis=1, arg=(8, AxisType.UNROLL)),
Opt(op=OptOps.SPLIT, axis=1, arg=(32, AxisType.GROUP_REDUCE)),
Opt(op=OptOps.SPLIT, axis=1, arg=(8, AxisType.UPCAST)),
Opt(op=OptOps.SPLIT, axis=1, arg=(32, AxisType.LOCAL)),
]
else:
opts = [
tinygrad_repo/extra/gemm_fragment.py
+1 -1
View File
@@ -38,7 +38,7 @@ def alloc_shared(shape:tuple[int, ...], dtype:DType, slot:int) -> UOp:
def alloc_fragment(shape:tuple[int, ...], dtype:DType, slot:int, axes:tuple[int, ...], rngs:tuple[UOp, ...]) -> UOp:
"""T.alloc_fragment: per-thread REG fragment + UNSHARD over the LOCAL thread grid."""
assert len(axes) == len(rngs)
assert all(tnum.op is Ops.RANGE and tnum.arg[-1] is AxisType.LOCAL for tnum in rngs), "fragments shard over LOCAL ranges"
assert all(tnum.op is Ops.RANGE and tnum.axis_type is AxisType.LOCAL for tnum in rngs), "fragments shard over LOCAL ranges"
by_axis = dict(zip(axes, rngs))
shard_shape = tuple(s // (int(by_axis[i].vmax)+1) if i in by_axis else s for i, s in enumerate(shape))
fragment = UOp.placeholder(shard_shape, dtype, slot, AddrSpace.REG)
tinygrad_repo/extra/gptoss_kernels/combine/__init__.py
+116
View File
@@ -0,0 +1,116 @@
import functools, pathlib
from tinygrad import Tensor, dtypes, function
from tinygrad.uop.ops import UOp, Ops, KernelInfo
from extra.llama_kernels import alloc_like, compile_hip
@functools.cache
def _combine_forward(out:UOp, z:UOp, dest_row:UOp, weights:UOp) -> UOp:
groups, tokens, hidden = out.shape
rows = z.shape[1]
assert hidden == 2880 and z.shape == (groups, rows, hidden)
assert dest_row.shape == (groups, tokens*4) and weights.shape == (groups, tokens, 4)
assert out.dtype == z.dtype == dtypes.bfloat16 and dest_row.dtype == dtypes.int32 and weights.dtype == dtypes.float32
sink = UOp.sink(out.base, z.base, dest_row.base, weights.base, UOp.special(256, "lidx0"),
UOp.special(groups*tokens, "gidx0"), arg=KernelInfo("gptoss_combine_forward"))
src = (pathlib.Path(__file__).parent/"forward.cpp").read_text()
lib = compile_hip(src, [f"-DTOKENS={tokens}", f"-DROWS={rows}", "-fno-fast-math"])
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
@functools.cache
def _combine_weights_backward(out:UOp, z:UOp, gradient:UOp, dest_row:UOp) -> UOp:
groups, tokens, k = out.shape
rows = z.shape[1]
assert k == 4 and z.shape == (groups, rows, 2880) and gradient.shape == (groups, tokens, 2880)
assert dest_row.shape == (groups, tokens*4) and dest_row.dtype == dtypes.int32
assert out.dtype == dtypes.float32 and z.dtype == gradient.dtype == dtypes.bfloat16
sink = UOp.sink(out.base, z.base, gradient.base, dest_row.base, UOp.special(256, "lidx0"),
UOp.special(groups*tokens, "gidx0"), arg=KernelInfo("gptoss_combine_weights_backward"))
src = (pathlib.Path(__file__).parent/"weights_backward.cpp").read_text()
lib = compile_hip(src, [f"-DTOKENS={tokens}", f"-DROWS={rows}", "-fno-fast-math"])
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
@functools.cache
def _inverse_rows(out:UOp, dest_row:UOp) -> UOp:
groups, entries = dest_row.shape
g, i = UOp.range(groups, 0), UOp.range(entries, 1)
row = dest_row.index(g, i).load().cast(dtypes.weakint)
return out.index(g, row).store(i.cast(dtypes.int32)).end(g, i).sink(arg=KernelInfo("gptoss_combine_inverse", opts_to_apply=()))
@functools.cache
def _combine_input_backward(out:UOp, gradient:UOp, src_row:UOp, weights:UOp) -> UOp:
groups, rows, hidden = out.shape
tokens = gradient.shape[1]
assert hidden == 2880 and gradient.shape == (groups, tokens, hidden)
assert src_row.shape == (groups, rows) and weights.shape == (groups, tokens, 4)
assert out.dtype == gradient.dtype == dtypes.bfloat16 and src_row.dtype == dtypes.int32 and weights.dtype == dtypes.float32
sink = UOp.sink(out.base, gradient.base, src_row.base, weights.base, UOp.special(256, "lidx0"),
UOp.special(groups*rows, "gidx0"), arg=KernelInfo("gptoss_combine_input_backward"))
src = (pathlib.Path(__file__).parent/"input_backward.cpp").read_text()
lib = compile_hip(src, [f"-DTOKENS={tokens}", f"-DROWS={rows}", "-fno-fast-math"])
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
def _combine_backward(gradient:UOp, kernel:UOp) -> tuple:
z, dest_row, weights = (Tensor(u) for u in kernel.src[2:5])
src_row = z[:, :, 0].full_like(-1, dtype=dtypes.int32)
src_row = Tensor.custom_kernel(src_row, dest_row, fxn=_inverse_rows)[0]
dz = alloc_like(z.shape, z.dtype, z.device, z.uop.axis)
dz = Tensor.custom_kernel(dz, Tensor(gradient), src_row, weights, fxn=_combine_input_backward)[0]
dw = alloc_like(weights.shape, weights.dtype, weights.device, weights.uop.axis)
dw = Tensor.custom_kernel(dw, z, Tensor(gradient), dest_row, fxn=_combine_weights_backward)[0]
return None, dz.uop, None, dw.uop
def fused_combine(z:Tensor, dest_row:Tensor, weights:Tensor) -> Tensor:
out = alloc_like((*weights.shape[:2], z.shape[-1]), z.dtype, z.device, z.uop.axis)
return Tensor.custom_kernel(out, z, dest_row, weights, fxn=_combine_forward, grad_fxn=_combine_backward)[0]
@functools.cache
def _combine_input_quantize_backward(out:UOp, q:UOp, scales:UOp, gradient:UOp, src_row:UOp, weights:UOp) -> UOp:
groups, rows = src_row.shape
tokens = gradient.shape[1]
assert out.shape == q.shape == (groups*rows, 3072) and scales.shape == (groups*rows, 96)
assert gradient.shape == (groups, tokens, 2880) and weights.shape == (groups, tokens, 4)
assert out.dtype == gradient.dtype == dtypes.bfloat16 and q.dtype == dtypes.fp8e4m3 and scales.dtype == dtypes.uint8
assert src_row.dtype == dtypes.int32 and weights.dtype == dtypes.float32
sink = UOp.sink(out.base, q.base, scales.base, gradient.base, src_row.base, weights.base,
UOp.special(256, "lidx0"), UOp.special(groups*rows, "gidx0"), arg=KernelInfo("gptoss_combine_input_quantize_backward"))
src = (pathlib.Path(__file__).parent/"input_quantize_backward.cpp").read_text()
lib = compile_hip(src, [f"-DTOKENS={tokens}", f"-DROWS={rows}", "-fno-fast-math"])
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
def combine_input_quantize_backward(gradient:Tensor, dest_row:Tensor, weights:Tensor, rows:int) -> tuple[Tensor, Tensor, Tensor]:
groups = weights.shape[0]
src_row = alloc_like((groups, rows), dtypes.int32, gradient.device, gradient.uop.axis).full_like(-1)
src_row = Tensor.custom_kernel(src_row, dest_row, fxn=_inverse_rows)[0]
out = alloc_like((groups*rows, 3072), dtypes.bfloat16, gradient.device, gradient.uop.axis)
q = alloc_like(out.shape, dtypes.fp8e4m3, gradient.device, gradient.uop.axis)
scales = alloc_like((groups*rows, 96), dtypes.uint8, gradient.device, gradient.uop.axis)
return tuple(Tensor.custom_kernel(out, q, scales, gradient, src_row, weights, fxn=_combine_input_quantize_backward)[:3])
def _down_bias(bias:Tensor, off:Tensor, rows:int) -> Tensor:
from extra.gemm.moe_routing import Routing, BLOCK_ROW
r = Routing(None, None, off, rows, off.shape[0], 0)
tile_bias = r.tile_e.one_hot(bias.shape[0]).float() @ bias.float()
return tile_bias.reshape(-1, 1, 2880).expand(-1, BLOCK_ROW, -1).reshape(-1, 2880).cast(dtypes.bfloat16)
def _down_combine_backward(gradient:UOp, call:UOp) -> tuple:
from extra.gemm.moe_gemm import grouped_mx_gemm_backward
aq, ae8, bq, be8, off, bias, dest_row, weights = (Tensor(u) for u in call.src[1:9])
_, z = (Tensor(u) for u in call.unbound_outputs)
groups, tokens, _ = weights.shape
rows = aq.shape[0] // groups
g = Tensor(gradient).reshape(groups, tokens, 2880)
dz, q, scales = combine_input_quantize_backward(g, dest_row, weights, rows)
dx, db = grouped_mx_gemm_backward(dz, aq, ae8, bq, be8, off, w_stored=True, quantized_gradient=(q, scales))
dbias, = _down_bias(bias, off, rows).gradient(bias, gradient=dz[:, :2880])
dw = alloc_like(weights.shape, weights.dtype, weights.device, weights.uop.axis)
dw = Tensor.custom_kernel(dw, z.reshape(groups, rows, 2880), g, dest_row, fxn=_combine_weights_backward)[0]
return dx.uop, None, db.uop, None, None, dbias.uop, None, dw.uop
@function(grad_fxn=_down_combine_backward)
def down_combine(aq:Tensor, ae8:Tensor, bq:Tensor, be8:Tensor, off:Tensor, bias:Tensor,
dest_row:Tensor, weights:Tensor) -> tuple[Tensor, Tensor]:
from extra.gemm.moe_gemm import grouped_mx_gemm
groups = weights.shape[0]
rows = aq.shape[0] // groups
z = grouped_mx_gemm((aq, ae8), (bq, be8), off)[:, :2880] + _down_bias(bias, off, rows)
return fused_combine(z.reshape(groups, rows, 2880), dest_row, weights), z
tinygrad_repo/extra/gptoss_kernels/combine/forward.cpp
+25
View File
@@ -0,0 +1,25 @@
#include <hip/hip_runtime.h>
#include <hip/hip_bf16.h>
// One workgroup gathers and combines the four expert rows for one token.
extern "C" __global__ __launch_bounds__(256) void gptoss_combine_forward(
__hip_bfloat16 *__restrict__ out, const __hip_bfloat16 *__restrict__ z,
const int *__restrict__ dest_row, const float *__restrict__ weights) {
const int token = blockIdx.x, group = token / TOKENS, tid = threadIdx.x;
__shared__ int rows[4];
__shared__ float ws[4];
if (tid < 4) {
rows[tid] = group * ROWS + dest_row[token * 4 + tid];
ws[tid] = weights[token * 4 + tid];
}
__syncthreads();
for (int d = tid; d < 2880; d += 256) {
float acc = 0.0f;
#pragma unroll
for (int j = 0; j < 4; j++) {
acc = fmaf(ws[j], (float)z[(long long)rows[j] * 2880 + d], acc);
}
out[(long long)token * 2880 + d] = (__hip_bfloat16)acc;
}
}
tinygrad_repo/extra/gptoss_kernels/combine/input_backward.cpp
+20
View File
@@ -0,0 +1,20 @@
#include <hip/hip_runtime.h>
#include <hip/hip_bf16.h>
extern "C" __global__ __launch_bounds__(256) void gptoss_combine_input_backward(
__hip_bfloat16 *__restrict__ out, const __hip_bfloat16 *__restrict__ dout,
const int *__restrict__ src_row, const float *__restrict__ weights) {
const int row = blockIdx.x, group = row / ROWS, tid = threadIdx.x;
const int src = src_row[row];
// Return before forming input addresses for unused grouped rows.
if (src < 0) {
for (int d = tid; d < 2880; d += 256) out[(long long)row * 2880 + d] = (__hip_bfloat16)0.0f;
return;
}
const int token = group * TOKENS + src / 4;
const float weight = weights[token * 4 + src % 4];
for (int d = tid; d < 2880; d += 256) {
out[(long long)row * 2880 + d] = (__hip_bfloat16)((float)dout[(long long)token * 2880 + d] * weight);
}
}
tinygrad_repo/extra/gptoss_kernels/combine/input_quantize_backward.cpp
+56
View File
@@ -0,0 +1,56 @@
#include <hip/hip_runtime.h>
#include <hip/hip_bf16.h>
#include <hip/hip_fp8.h>
typedef short short2v __attribute__((ext_vector_type(2)));
typedef unsigned uint4v __attribute__((ext_vector_type(4)));
typedef unsigned short ushort4v __attribute__((ext_vector_type(4)));
extern "C" __global__ __launch_bounds__(256) void gptoss_combine_input_quantize_backward(
__hip_bfloat16 *__restrict__ out, unsigned char *__restrict__ q, unsigned char *__restrict__ scales,
const __hip_bfloat16 *__restrict__ dout, const int *__restrict__ src_row, const float *__restrict__ weights) {
const int row = blockIdx.x, group = row / ROWS, tid = threadIdx.x;
// Four lanes own one MX block, with aligned eight-value loads and stores per lane.
const int src = src_row[row], lane = tid % 4;
for (int block = tid / 4; block < 96; block += 64) {
const int d = block * 32 + lane * 8;
const long long index = (long long)row * 3072 + d;
// Neither padding nor unused rows may read dout or weights.
if (src < 0 || d >= 2880) {
*reinterpret_cast<uint4v*>(&out[index]) = uint4v{};
*reinterpret_cast<ushort4v*>(&q[index]) = ushort4v{};
if (lane == 0) scales[(long long)row * 96 + block] = 0;
continue;
}
const int token = group * TOKENS + src / 4;
const float weight = weights[token * 4 + src % 4];
const uint4v input = *reinterpret_cast<const uint4v*>(&dout[(long long)token * 2880 + d]);
uint4v output;
float2 values[4];
float amax = 0.0f;
#pragma unroll
for (int i = 0; i < 4; i++) {
const unsigned bits = input[i];
const float2 dy = make_float2(__builtin_bit_cast(float, bits << 16), __builtin_bit_cast(float, bits & 0xffff0000u));
const float2 value = __bfloat1622float2(__float22bfloat162_rn(make_float2(dy.x * weight, dy.y * weight)));
values[i] = value;
const __hip_bfloat162 rounded_pair = __float22bfloat162_rn(value);
output[i] = (unsigned)__bfloat16_as_ushort(rounded_pair.x) | ((unsigned)__bfloat16_as_ushort(rounded_pair.y) << 16);
amax = fmaxf(amax, fmaxf(fabsf(value.x), fabsf(value.y)));
}
*reinterpret_cast<uint4v*>(&out[index]) = output;
#pragma unroll
for (int offset = 2; offset > 0; offset >>= 1) amax = fmaxf(amax, __shfl_down(amax, offset, 4));
amax = __shfl(amax, 0, 4);
const unsigned e8 = min((__builtin_bit_cast(unsigned, amax) >> 23) & 255u, 254u);
if (lane == 0) scales[(long long)row * 96 + block] = e8;
ushort4v quantized;
#pragma unroll
for (int i = 0; i < 4; i++) {
short2v packed = {0, 0};
packed = __builtin_amdgcn_cvt_scalef32_pk_fp8_f32(packed, values[i].x, values[i].y, __builtin_bit_cast(float, e8 << 23), false);
quantized[i] = __builtin_bit_cast(unsigned short, packed[0]);
}
*reinterpret_cast<ushort4v*>(&q[index]) = quantized;
}
}
tinygrad_repo/extra/gptoss_kernels/combine/weights_backward.cpp
+44
View File
@@ -0,0 +1,44 @@
#include <hip/hip_runtime.h>
#include <hip/hip_bf16.h>
// One workgroup computes all four routing-weight gradients, reusing dout across expert slots.
extern "C" __global__ __launch_bounds__(256) void gptoss_combine_weights_backward(
float *__restrict__ dw, const __hip_bfloat16 *__restrict__ z,
const __hip_bfloat16 *__restrict__ dout, const int *__restrict__ dest_row) {
const int token = blockIdx.x, group = token / TOKENS, tid = threadIdx.x;
float sums[4] = {};
int rows[4];
#pragma unroll
for (int j = 0; j < 4; j++) rows[j] = group * ROWS + dest_row[token * 4 + j];
for (int d = tid; d < 2880; d += 256) {
const float dy = (float)dout[(long long)token * 2880 + d];
#pragma unroll
for (int j = 0; j < 4; j++)
sums[j] = fmaf(dy, (float)z[(long long)rows[j] * 2880 + d], sums[j]);
}
// Reduce each slot across the four wavefronts.
#pragma unroll
for (int offset = 32; offset > 0; offset >>= 1)
#pragma unroll
for (int j = 0; j < 4; j++) sums[j] += __shfl_down(sums[j], offset, 64);
__shared__ float warp_sums[4][4];
const int lane = tid % 64, warp = tid / 64;
#pragma unroll
for (int j = 0; j < 4; j++) if (lane == 0) warp_sums[warp][j] = sums[j];
__syncthreads();
if (warp == 0) {
#pragma unroll
for (int j = 0; j < 4; j++) sums[j] = lane < 4 ? warp_sums[lane][j] : 0.0f;
#pragma unroll
for (int offset = 32; offset > 0; offset >>= 1)
#pragma unroll
for (int j = 0; j < 4; j++) sums[j] += __shfl_down(sums[j], offset, 64);
if (lane == 0)
#pragma unroll
for (int j = 0; j < 4; j++) dw[token * 4 + j] = sums[j];
}
}
tinygrad_repo/extra/gptoss_kernels/dispatch/__init__.py
+48
View File
@@ -0,0 +1,48 @@
import functools, pathlib
from tinygrad import Tensor, dtypes
from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
from extra.llama_kernels import alloc_like, compile_hip
@functools.cache
def _inverse_rows(out:UOp, dest_row:UOp, counts:UOp, off:UOp) -> UOp:
groups, rows = out.shape
entries, experts = dest_row.shape[1], counts.shape[1]
assert rows >= entries and counts.shape == (groups, experts) and off.shape == (groups, experts+1)
sink = UOp.sink(out.base, dest_row.base, counts.base, off.base, UOp.special(256, "lidx0"),
UOp.special(groups*((rows+255)//256), "gidx0"), arg=KernelInfo(f"dispatch_inverse_{groups}_{rows}_{entries}_{experts}"))
src = (pathlib.Path(__file__).parent/"inverse.cpp").read_text()
lib = compile_hip(src, [f"-DROWS={rows}", f"-DENTRIES={entries}", f"-DEXPERTS={experts}"])
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
def inverse_rows(dest_row:Tensor, counts:Tensor, off:Tensor, rows:int) -> Tensor:
out = alloc_like((dest_row.shape[0], rows), dtypes.int32, dest_row.device, dest_row.uop.axis)
return Tensor.custom_kernel(out, dest_row, counts, off, fxn=_inverse_rows)[0]
@functools.cache
def _dispatch_gather(out:UOp, x:UOp, dest_row:UOp, src_row:UOp) -> UOp:
from extra.gemm.moe_routing import _blk_for
groups, tokens, hidden = x.shape
rows = src_row.shape[1]
assert out.shape == (groups*rows, hidden) and dest_row.shape == (groups, tokens*4)
assert x.dtype == out.dtype and x.dtype.itemsize == 1
block = _blk_for(hidden)
g, row = UOp.range(groups, 0, AxisType.GLOBAL), UOp.range(rows, 1, AxisType.GLOBAL)
outer, inner = UOp.range(hidden//block, 2), UOp.range(block, 3, AxisType.LOCAL)
col = outer*block + inner
source = src_row.index(g, row).load().cast(dtypes.weakint)
anchor = dest_row.index(g, 0).load() >= 0
valid = (source >= 0) & anchor
token = valid.where(source, source.const_like(0)) // 4
value = valid.where(x.index(g, token, col).load(), UOp.const(0, x.dtype))
return out.reshape(groups, rows, hidden).index(g, row, col).store(value).end(g, row, outer, inner).sink(
arg=KernelInfo(f"dispatch_gather_{groups}_{rows}_{tokens}_{hidden}", opts_to_apply=()))
def _dispatch_gather_bwd(gradient:UOp, kernel:UOp) -> tuple:
from extra.gemm.moe_routing import _gscatter_bwd
# The first three arguments match scatter's backward; src_row is forward-only.
x = kernel.src[2]
return (*_gscatter_bwd(gradient.reshape(x.shape[0], -1, x.shape[-1]), kernel), None)
def dispatch_gather(x:Tensor, dest_row:Tensor, src_row:Tensor) -> Tensor:
out = alloc_like((x.shape[0]*src_row.shape[1], x.shape[-1]), x.dtype, x.device, x.uop.axis)
return Tensor.custom_kernel(out, x, dest_row, src_row, fxn=_dispatch_gather, grad_fxn=_dispatch_gather_bwd)[0]
tinygrad_repo/extra/gptoss_kernels/dispatch/gather.cpp
+13
View File
@@ -0,0 +1,13 @@
#include <hip/hip_runtime.h>
// Copy FP8/scale bytes unchanged; every output is written, including padded rows.
// dest is retained as a saved backward input, but forward uses only src_row.
extern "C" __global__ __launch_bounds__(64) void dispatch_gather(
unsigned char* __restrict__ out, const unsigned char* __restrict__ x,
const int* __restrict__ dest, const int* __restrict__ src_row) {
const int row = blockIdx.x;
const int source = src_row[row];
const long long token = (long long)(row / ROWS) * TOKENS + (source >= 0 ? source / 4 : 0);
for (int d = threadIdx.x; d < HIDDEN; d += 64)
out[(long long)row * HIDDEN + d] = source >= 0 ? x[token * HIDDEN + d] : 0;
}
tinygrad_repo/extra/gptoss_kernels/dispatch/inverse.cpp
+16
View File
@@ -0,0 +1,16 @@
#include <hip/hip_runtime.h>
// Valid inverse writes and padding writes are disjoint, so this needs no initialization pass or grid barrier.
extern "C" __global__ __launch_bounds__(256) void dispatch_inverse(
int* __restrict__ out, const int* __restrict__ dest, const int* __restrict__ counts, const int* __restrict__ off) {
const int group = blockIdx.x / ((ROWS + 255) / 256);
const int i = (blockIdx.x % ((ROWS + 255) / 256)) * 256 + threadIdx.x;
if (i < ENTRIES) out[(long long)group * ROWS + dest[(long long)group * ENTRIES + i]] = i;
if (i < ROWS) {
int expert = 0;
#pragma unroll
for (int e = 1; e < EXPERTS; e++) expert += i >= off[group * (EXPERTS + 1) + e];
if (i >= off[group * (EXPERTS + 1) + expert] + counts[group * EXPERTS + expert])
out[(long long)group * ROWS + i] = -1;
}
}
tinygrad_repo/extra/gptoss_kernels/gather_sum/__init__.py
+24
View File
@@ -0,0 +1,24 @@
import functools, pathlib
from tinygrad import Tensor, dtypes
from tinygrad.uop.ops import UOp, Ops, KernelInfo
from extra.llama_kernels import alloc_like, compile_hip
@functools.cache
def _gather_sum(out:UOp, table:UOp, indices:UOp) -> UOp:
groups, tokens, hidden = out.shape
rows = table.shape[1]
assert table.shape == (groups, rows, hidden) and indices.shape == (groups, tokens*4)
assert out.dtype == table.dtype == dtypes.bfloat16 and indices.dtype == dtypes.int32
assert hidden % 64 == 0
sink = UOp.sink(out.base, table.base, indices.base, UOp.special(64, "lidx0"), UOp.special(groups*tokens, "gidx0"),
arg=KernelInfo(f"gather_sum_{groups}_{tokens}_{rows}_{hidden}"))
src = (pathlib.Path(__file__).parent/"gather_sum.cpp").read_text()
lib = compile_hip(src, [f"-DTOKENS={tokens}", f"-DROWS={rows}", f"-DHIDDEN={hidden}", "-fno-fast-math"])
return UOp(Ops.PROGRAM,
src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
def gather_sum(table:Tensor, indices:Tensor) -> Tensor:
groups, _, hidden = table.shape
assert indices.shape[0] == groups and indices.shape[1] % 4 == 0
out = alloc_like((groups, indices.shape[1]//4, hidden), table.dtype, table.device, table.uop.axis)
return Tensor.custom_kernel(out, table, indices, fxn=_gather_sum)[0]
tinygrad_repo/extra/gptoss_kernels/gather_sum/gather_sum.cpp
+23
View File
@@ -0,0 +1,23 @@
#include <hip/hip_runtime.h>
#include <hip/hip_bf16.h>
// One wave per token: gather its four expert gradients and write only their sum.
extern "C" __global__ __launch_bounds__(64) void gather_sum(
__hip_bfloat16* __restrict__ out, const __hip_bfloat16* __restrict__ table, const int* __restrict__ indices) {
#pragma clang fp reassociate(off)
const int token = blockIdx.x;
const long long group_base = (long long)(token / TOKENS) * ROWS * HIDDEN;
const int4 rows = *reinterpret_cast<const int4*>(&indices[(long long)token * 4]);
const long long row0 = group_base + (long long)rows.x * HIDDEN;
const long long row1 = group_base + (long long)rows.y * HIDDEN;
const long long row2 = group_base + (long long)rows.z * HIDDEN;
const long long row3 = group_base + (long long)rows.w * HIDDEN;
#pragma unroll 2
for (int d = threadIdx.x; d < HIDDEN; d += 64) {
float sum = (float)table[row0 + d];
sum = sum + (float)table[row1 + d];
sum = sum + (float)table[row2 + d];
sum = sum + (float)table[row3 + d];
out[(long long)token * HIDDEN + d] = (__hip_bfloat16)sum;
}
}
tinygrad_repo/extra/gptoss_kernels/qkv_rope/__init__.py
+42
View File
@@ -0,0 +1,42 @@
import functools, pathlib
from tinygrad import Tensor, dtypes
from tinygrad.uop.ops import UOp, Ops, KernelInfo
from extra.llama_kernels import alloc_like, compile_hip
@functools.cache
def _forward(q:UOp, k:UOp, v:UOp, xqkv:UOp, freqs:UOp) -> UOp:
assert xqkv.shape == (2, 8192, 5120) and q.shape == (2, 8192, 64, 64) and k.shape == v.shape == (2, 8192, 8, 64)
sink = UOp.sink(q.base, k.base, v.base, xqkv.base, freqs.base,
UOp.special(256, "lidx0"), UOp.special(2, "gidx0"), UOp.special(8192, "gidx1"),
arg=KernelInfo("fused_qkv_rope_forward_gptoss_packed"))
src = (pathlib.Path(__file__).parent/"forward.cpp").read_text()
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
UOp(Ops.BINARY, arg=compile_hip(src, []))))
@functools.cache
def _backward_kernel(out:UOp, dq:UOp, dk:UOp, dv:UOp, freqs:UOp) -> UOp:
assert out.shape == (2, 8192, 5120) and dq.shape == (2, 64, 8192, 64) and dk.shape == dv.shape == (16, 8192, 8, 64)
sink = UOp.sink(out.base, dq.base, dk.base, dv.base, freqs.base,
UOp.special(256, "lidx0"), UOp.special(2, "gidx0"), UOp.special(128, "gidx1"), UOp.special(80, "gidx2"),
arg=KernelInfo("fused_qkv_rope_backward_gptoss"))
src = (pathlib.Path(__file__).parent/"backward.cpp").read_text()
include = pathlib.Path(__file__).parents[2]/"thunder"/"amd"/"include"
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
UOp(Ops.BINARY, arg=compile_hip(src, [f"-I{include}", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS"]))))
def _backward(dq:UOp, dk:UOp, dv:UOp, call:UOp) -> tuple:
from extra.thunder.amd.fa import _fa_native_grads
# Consume FA's shuffled dQ and unreduced dK/dV directly, without materializing their logical views.
native = _fa_native_grads(dq, dk, dv)
assert native is not None, "GPT-OSS QKV/RoPE backward requires native Flash Attention gradients"
xqkv, freqs = (Tensor(u) for u in call.src[4:6])
out = alloc_like(xqkv.shape, xqkv.dtype, xqkv.device, xqkv.uop.axis)
out, *_ = Tensor.custom_kernel(out, *(Tensor(u) for u in native), freqs, fxn=_backward_kernel)
return None, None, None, out.uop, None
def fused_qkv_rope(xqkv:Tensor, freqs:Tensor) -> tuple[Tensor, Tensor, Tensor]:
batch, seq, hidden = xqkv.shape
assert hidden == 5120 and freqs.shape == (1, seq, 1, 32, 2) and xqkv.dtype == freqs.dtype == dtypes.bfloat16
q, k, v = (alloc_like((batch, seq, heads, 64), dtypes.bfloat16, xqkv.device, xqkv.uop.axis) for heads in (64, 8, 8))
q, k, v, *_ = Tensor.custom_kernel(q, k, v, xqkv, freqs, fxn=_forward, grad_fxn=_backward)
return q, k, v
tinygrad_repo/extra/gptoss_kernels/qkv_rope/backward.cpp
+168
View File
@@ -0,0 +1,168 @@
#include "kittens.cuh"
using namespace kittens;
#ifndef ATTN_B
#define ATTN_B 2
#endif
#ifndef ATTN_N
#define ATTN_N 8192
#endif
#ifndef ATTN_H
#define ATTN_H 64
#endif
#ifndef ATTN_H_KV
#define ATTN_H_KV 8
#endif
#ifndef ATTN_D
#define ATTN_D 64
#endif
#ifndef THREADS_PER_BLOCK
#define THREADS_PER_BLOCK 256
#endif
constexpr int GROUP_SIZE = ATTN_H / ATTN_H_KV;
constexpr int HALF_D = ATTN_D / 2;
constexpr int PACKED_H = ATTN_H_KV * (GROUP_SIZE + 2);
// The current D=64 FA emits one dK/dV partial per query head in each GQA group.
constexpr int KV_PARTIALS = GROUP_SIZE;
constexpr int Q_HEADS_PER_WG = 1;
constexpr int Q_FIELDS = ATTN_H / Q_HEADS_PER_WG;
constexpr int NUM_WARPS = 4;
constexpr int TILE_N = 16;
template<typename T> using grad_tile = rt<T, TILE_N, ATTN_D, row_l, rt_16x32_s>;
template<int axis, ducks::rt::row_layout RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<RT>>
__device__ __forceinline__ void load_fa_shuffled(RT &dst, const GL &src, const COORD &idx) {
using U = typename GL::dtype;
using U2 = base_types::packing<U>::packed_type;
U *src_ptr = (U*)&src[(idx.template unit_coord<axis, 3>())];
const int row_stride = src.template stride<axis>();
const int lane = kittens::laneid();
const int tile_row_stride = row_stride * dst.base_tile_rows;
const int tile_stride = dst.base_tile_rows * dst.base_tile_cols;
const uint32_t buffer_size = src.batch() * src.depth() * src.rows() * src.cols() * sizeof(U);
const buffer_resource br = make_buffer_resource(reinterpret_cast<uintptr_t>(src_ptr), buffer_size, 0x00020000);
#pragma unroll
for (int i = 0; i < dst.height; i++) {
#pragma unroll
for (int j = 0; j < dst.width; j++) {
const float4 loaded = std::bit_cast<float4>(llvm_amdgcn_raw_buffer_load_b128(
std::bit_cast<i32x4>(br), (i * tile_row_stride + j * tile_stride + lane * 8) * sizeof(U), 0, 0));
const U2 *packed = reinterpret_cast<const U2*>(&loaded);
#pragma unroll
for (int k = 0; k < dst.packed_per_base_tile; k++) dst.tiles[i][j].data[k] = packed[k];
}
}
}
template<int axis, ducks::rt::row_layout RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<RT>>
__device__ __forceinline__ void store_fa_shuffled(const GL &dst, const RT &src, const COORD &idx) {
using U = typename GL::dtype;
U *dst_ptr = (U*)&dst[(idx.template unit_coord<axis, 3>())];
const int row_stride = dst.template stride<axis>();
const int lane = kittens::laneid();
const int row_offset = (lane % 4) * 4;
const int col_offset = ((lane / 32) * 16) + (((lane % 32) / 16) * 2) + (((lane % 16) / 4) * 4);
const uint32_t buffer_size = dst.batch() * dst.depth() * dst.rows() * dst.cols() * sizeof(U);
const buffer_resource br = make_buffer_resource(reinterpret_cast<uintptr_t>(dst_ptr), buffer_size, 0x00020000);
#pragma unroll
for (int i = 0; i < src.height; i++) {
const int row = src.base_tile_rows * i + row_offset;
#pragma unroll
for (int j = 0; j < src.width; j++) {
const int col = src.base_tile_cols * j + col_offset;
#pragma unroll
for (int k = 0; k < src.packed_per_base_tile; k++) llvm_amdgcn_raw_buffer_store_b32(
*reinterpret_cast<const uint32_t*>(&src.tiles[i][j].data[k]), std::bit_cast<i32x4>(br),
((row + k) * row_stride + col) * sizeof(U), 0, 0);
}
}
}
template<ducks::rt::row_layout RT>
__device__ __forceinline__ void inverse_rope_fa(RT &tile, const bf16_2 *freqs, const int n_base) {
const int lane = kittens::laneid();
const int row_offset = (lane % 4) * 4;
const int col_offset = ((lane / 32) * 16) + (((lane % 32) / 16) * 2) + (((lane % 16) / 4) * 4);
#pragma unroll
for (int i = 0; i < tile.height; i++) {
#pragma unroll
for (int j = 0; j < tile.width; j++) {
const int col = tile.base_tile_cols * j + col_offset;
#pragma unroll
for (int k = 0; k < tile.packed_per_base_tile; k++) {
const int row = tile.base_tile_rows * i + row_offset + k;
const float2 cs = __bfloat1622float2(freqs[(n_base + row) * HALF_D + col / 2]);
const float2 g = __bfloat1622float2(tile.tiles[i][j].data[k]);
tile.tiles[i][j].data[k] = __float22bfloat162_rn(make_float2(g.x * cs.x + g.y * cs.y, -g.x * cs.y + g.y * cs.x));
}
}
}
}
template<ducks::rt::row_layout RT>
__device__ __forceinline__ void inverse_rope(RT &tile, const bf16_2 *freqs, const int n_base) {
const int lane = kittens::laneid();
#pragma unroll
for (int i = 0; i < tile.height; i++) {
const int row = tile.base_tile_rows * i + lane % tile.base_tile_rows;
#pragma unroll
for (int j = 0; j < tile.width; j++) {
#pragma unroll
for (int k = 0; k < tile.packed_per_base_tile; k++) {
const int col = tile.base_tile_cols * j + tile.base_tile_stride * (lane / tile.base_tile_rows) + 2 * k;
const float2 cs = __bfloat1622float2(freqs[(n_base + row) * HALF_D + col / 2]);
const float2 g = __bfloat1622float2(tile.tiles[i][j].data[k]);
tile.tiles[i][j].data[k] = __float22bfloat162_rn(make_float2(g.x * cs.x + g.y * cs.y, -g.x * cs.y + g.y * cs.x));
}
}
}
}
extern "C" __global__ __launch_bounds__(THREADS_PER_BLOCK) void
fused_qkv_rope_backward(
bf16* __restrict__ dxqkv,
const bf16* __restrict__ dq,
const bf16* __restrict__ dk,
const bf16* __restrict__ dv,
const bf16* __restrict__ freqs_cis) {
gl<bf16, -1, -1, -1, -1> out{dxqkv, ATTN_B, ATTN_N, PACKED_H, ATTN_D};
gl<bf16, -1, -1, -1, -1> dqg{const_cast<bf16*>(dq), ATTN_B, ATTN_H, ATTN_N, ATTN_D};
gl<bf16, -1, -1, -1, -1> dkg{const_cast<bf16*>(dk), ATTN_B * KV_PARTIALS, ATTN_N, ATTN_H_KV, ATTN_D};
gl<bf16, -1, -1, -1, -1> dvg{const_cast<bf16*>(dv), ATTN_B * KV_PARTIALS, ATTN_N, ATTN_H_KV, ATTN_D};
const int b = blockIdx.x, n_tile = blockIdx.y * NUM_WARPS + kittens::warpid(), n_base = n_tile * TILE_N;
const int field = blockIdx.z;
if (field < Q_FIELDS) {
#pragma unroll
for (int qj = 0; qj < Q_HEADS_PER_WG; qj++) {
const int qh = field * Q_HEADS_PER_WG + qj;
grad_tile<bf16> tile;
load_fa_shuffled<2>(tile, dqg, {b, qh, n_tile, 0});
inverse_rope_fa(tile, reinterpret_cast<const bf16_2*>(freqs_cis), n_base);
const int out_head = (qh / GROUP_SIZE) * (GROUP_SIZE + 2) + qh % GROUP_SIZE;
store_fa_shuffled<1>(out, tile, {b, n_tile, out_head, 0});
}
} else {
const int kv_field = field - Q_FIELDS;
const bool is_k = kv_field < ATTN_H_KV;
const int kvh = kv_field - (is_k ? 0 : ATTN_H_KV);
const auto &src = is_k ? dkg : dvg;
grad_tile<bf16> partial, tile;
grad_tile<float> partial_f, sum;
zero(sum);
#pragma unroll
for (int p = 0; p < KV_PARTIALS; p++) {
load<1>(partial, src, {b * KV_PARTIALS + p, n_tile, kvh, 0});
copy(partial_f, partial);
add(sum, sum, partial_f);
}
copy(tile, sum);
if (is_k) inverse_rope(tile, reinterpret_cast<const bf16_2*>(freqs_cis), n_base);
const int out_head = kvh * (GROUP_SIZE + 2) + GROUP_SIZE + !is_k;
store<1>(out, tile, {b, n_tile, out_head, 0});
}
}
tinygrad_repo/extra/gptoss_kernels/qkv_rope/forward.cpp
+89
View File
@@ -0,0 +1,89 @@
#include <hip/hip_runtime.h>
#include <hip/hip_bf16.h>
#ifndef ATTN_B
#define ATTN_B 2
#endif
#ifndef ATTN_N
#define ATTN_N 8192
#endif
#ifndef ATTN_H
#define ATTN_H 64
#endif
#ifndef ATTN_H_KV
#define ATTN_H_KV 8
#endif
#ifndef ATTN_D
#define ATTN_D 64
#endif
#ifndef THREADS_PER_BLOCK
#define THREADS_PER_BLOCK 256
#endif
constexpr int GROUP_SIZE = ATTN_H / ATTN_H_KV;
constexpr int HALF_D = ATTN_D / 2;
constexpr int PACKED_D = (GROUP_SIZE + 2) * ATTN_D;
using float2v = float __attribute__((ext_vector_type(2)));
static_assert(ATTN_B == 2 && ATTN_N == 8192 && ATTN_H == 64 && ATTN_H_KV == 8 && ATTN_D == 64,
"this kernel is the exact per-device GPTOSS QKV/RoPE shape");
static_assert(THREADS_PER_BLOCK == ATTN_H_KV * HALF_D, "one thread must own each KV-head/RoPE pair");
extern "C" __global__ __launch_bounds__(THREADS_PER_BLOCK) void
fused_qkv_rope_forward(
__hip_bfloat16* __restrict__ q,
__hip_bfloat16* __restrict__ k,
__hip_bfloat16* __restrict__ v,
const __hip_bfloat16* __restrict__ xqkv,
const __hip_bfloat16* __restrict__ freqs_cis) {
const int b = blockIdx.x;
const int n = blockIdx.y;
const int bn = b * ATTN_N + n;
const int kvh = threadIdx.x / HALF_D;
const int pair = threadIdx.x % HALF_D;
const int even = pair * 2;
const int packed_head = bn * ATTN_H_KV * PACKED_D + kvh * PACKED_D;
const int q_head = bn * ATTN_H * ATTN_D + kvh * GROUP_SIZE * ATTN_D;
const int kv_head = (bn * ATTN_H_KV + kvh) * ATTN_D;
// Cos/sin is shared by every Q/K slot owned by this (KV head, pair) thread.
const __hip_bfloat16 *freq_pair = freqs_cis + (n * HALF_D + pair) * 2;
const uint32_t cs_raw = *reinterpret_cast<const uint32_t*>(freq_pair);
const float2 cs = __bfloat1622float2(*reinterpret_cast<const __hip_bfloat162*>(freq_pair));
uint32_t ab_raw[GROUP_SIZE + 1];
bool exceptional = (cs_raw & 0x7f80u) == 0x7f80u || (cs_raw & 0x7f800000u) == 0x7f800000u;
#pragma unroll
for (int slot = 0; slot < GROUP_SIZE + 1; slot++) {
ab_raw[slot] = *reinterpret_cast<const uint32_t*>(xqkv + packed_head + slot * ATTN_D + even);
exceptional |= (ab_raw[slot] & 0x7f80u) == 0x7f80u || (ab_raw[slot] & 0x7f800000u) == 0x7f800000u;
}
if (__builtin_expect(exceptional, 0)) {
#pragma unroll
for (int slot = 0; slot < GROUP_SIZE + 1; slot++) {
const float2 ab = __bfloat1622float2(*reinterpret_cast<const __hip_bfloat162*>(&ab_raw[slot]));
__hip_bfloat16 *dst = slot < GROUP_SIZE ? q : k;
const int out = slot < GROUP_SIZE ? q_head + slot * ATTN_D + even : kv_head + even;
// Match the original kernel's packed instruction association exactly for NaN/Inf values. The low-lane
// instruction modifier, rather than pre-negating an operand, preserves its NaN sign/payload behavior.
const float2v base = float2v{ab.x, ab.y} * float2v{cs.x, cs.x};
const float2v cross = {ab.y, ab.x}, sine = {cs.y, cs.y};
float2v rotated;
asm volatile("v_pk_fma_f32 %0, %1, %2, %3 neg_lo:[1,0,0]"
: "=v"(rotated) : "v"(sine), "v"(cross), "v"(base));
*reinterpret_cast<__hip_bfloat162*>(dst + out) =
__float22bfloat162_rn(make_float2(rotated[0], rotated[1]));
}
} else {
#pragma unroll
for (int slot = 0; slot < GROUP_SIZE + 1; slot++) {
const float2 ab = __bfloat1622float2(*reinterpret_cast<const __hip_bfloat162*>(&ab_raw[slot]));
__hip_bfloat16 *dst = slot < GROUP_SIZE ? q : k;
const int out = slot < GROUP_SIZE ? q_head + slot * ATTN_D + even : kv_head + even;
*reinterpret_cast<__hip_bfloat162*>(dst + out) = __float22bfloat162_rn(
make_float2(ab.x * cs.x - ab.y * cs.y, ab.x * cs.y + ab.y * cs.x));
}
}
// V is already in output order; copy the aligned BF16 pair without conversion.
*reinterpret_cast<__hip_bfloat162*>(v + kv_head + even) =
*reinterpret_cast<const __hip_bfloat162*>(xqkv + packed_head + (GROUP_SIZE + 1) * ATTN_D + even);
}
tinygrad_repo/extra/gptoss_kernels/quantize_mxfp8/__init__.py
+3 -3
View File
@@ -15,14 +15,14 @@ def _custom_quantize_mxfp8_qe8(fp8_out:UOp, e8_out:UOp, x:UOp) -> UOp:
x, fp8_out, e8_out = x.reshape(n_elems), fp8_out.reshape(n_elems), e8_out.reshape(n_blocks)
wg = UOp.range(n_blocks // THREADS_PER_WG, 0, AxisType.GLOBAL)
tid = UOp.range(THREADS_PER_WG, 1, AxisType.LOCAL)
lane = UOp.range(BLK, 3, AxisType.UNROLL)
lane = UOp.range(BLK, 3, AxisType.UPCAST)
block = wg * THREADS_PER_WG + tid
idx = block * BLK + lane
x_f = x[idx].cast(dtypes.float)
abs_x = (x_f < 0.0).where(-x_f, x_f)
blk_max = abs_x.reduce(lane, arg=Ops.MAX)
e8f = (blk_max.maximum(1e-38).log2().floor() + 127.0).maximum(0.0).minimum(254.0)
qscale = (127.0 - e8f).exp2()
e8f = ((blk_max.bitcast(dtypes.uint32) >> 23) & 0xFF).minimum(254)
qscale = ((e8f.const_like(254) - e8f) << 23).bitcast(dtypes.float32)
scaled = (x_f * qscale).maximum(-FP8_MAX).minimum(FP8_MAX)
fp8_store = fp8_out[idx].store(scaled.cast(fp8_out.dtype)).end(lane)
e8_store = e8_out.after(fp8_store)[block].store(e8f.cast(dtypes.uint8))
tinygrad_repo/extra/gptoss_kernels/residual/__init__.py
+54
View File
@@ -0,0 +1,54 @@
import functools, math, pathlib
from tinygrad import Tensor, dtypes, function
from tinygrad.helpers import ALLREDUCE_CAST
from tinygrad.runtime.support.compiler_amd import HIPCompiler
from tinygrad.uop.ops import UOp, Ops, KernelInfo
from extra.llama_kernels import alloc_like, compile_hip
@functools.cache
def _residual_kernel(out:UOp, x:UOp, proj:UOp, bias:UOp, moe:UOp) -> UOp:
rows = math.prod(x.shape[:-1])
assert out.shape == x.shape == moe.shape and x.shape[-1] == 2880
assert proj.shape == (rows, 3072) and bias.shape == (2880,)
assert all(t.dtype == dtypes.bfloat16 for t in (out, x, proj, bias, moe))
sink = UOp.sink(out.base, x.base, proj.base, bias.base, moe.base,
UOp.special(256, "lidx0"), UOp.special(rows, "gidx0"), arg=KernelInfo("gptoss_residual_join_vec"))
src = (pathlib.Path(__file__).parent/"residual.cpp").read_text()
lib = compile_hip(src, [f"-DROWS={rows}", "-DREAL_D=2880", "-DPAD_D=3072", "-fno-fast-math"])
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
def _residual_backward(gradient:UOp, call:UOp) -> tuple:
return gradient, None, None, None, gradient
@function(grad_fxn=_residual_backward)
def residual_join(h:Tensor, x:Tensor, proj:Tensor, bias:Tensor, moe:Tensor) -> Tensor:
out = alloc_like(x.shape, x.dtype, x.device, x.uop.axis)
return Tensor.custom_kernel(out, x, proj, bias, moe, fxn=_residual_kernel)[0]
@functools.cache
def _wo_bias_kernel(out:UOp, inp:UOp, *, stage:str) -> UOp:
assert stage in ("partial", "final")
partial = stage == "partial"
assert out.shape == ((16, 2880) if partial else (2880,)) and inp.shape == ((16384, 2880) if partial else (16, 2880))
assert out.dtype == (dtypes.float32 if partial else dtypes.bfloat16) and inp.dtype == (dtypes.bfloat16 if partial else dtypes.float32)
sink = UOp.sink(out.base, inp.base, UOp.special(64, "lidx0"), UOp.special(45, "gidx0"),
*((UOp.special(16, "gidx1"),) if partial else ()), arg=KernelInfo(f"gptoss_wo_bias_{stage}"))
src = (pathlib.Path(__file__).parent/f"wo_bias_{stage}.cpp").read_text()
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
UOp(Ops.BINARY, arg=HIPCompiler("gfx950").compile_cached(src))))
def wo_bias_gradient(g:Tensor) -> Tensor:
groups = len(g.device) if isinstance(g.device, tuple) else 1
assert groups == 1 or (g.uop.axis == 0 and ALLREDUCE_CAST), "WO bias sum requires the baseline BF16 aggregation"
partial = alloc_like((groups*16, 2880), dtypes.float32, g.device, g.uop.axis)
local = alloc_like((2880,), dtypes.bfloat16, g.device)
partial = Tensor.custom_kernel(partial, g.reshape(-1, 2880), fxn=functools.partial(_wo_bias_kernel, stage="partial"))[0]
local = Tensor.custom_kernel(local, partial, fxn=functools.partial(_wo_bias_kernel, stage="final"))[0]
return Tensor(local.uop.allreduce(Ops.ADD, g.device)) if groups > 1 else local
def _wo_bias_backward(gradient:UOp, call:UOp) -> tuple:
return gradient, wo_bias_gradient(Tensor(gradient)).uop
@function(grad_fxn=_wo_bias_backward)
def wo_bias_add(x:Tensor, bias:Tensor) -> Tensor:
return x + bias
tinygrad_repo/extra/gptoss_kernels/residual/residual.cpp
+46
View File
@@ -0,0 +1,46 @@
#include <hip/hip_runtime.h>
#include <stdint.h>
using bf16 = __bf16;
using uint4v = uint32_t __attribute__((ext_vector_type(4)));
#if !defined(ROWS) || !defined(REAL_D) || !defined(PAD_D)
#error "ROWS, REAL_D and PAD_D are required"
#endif
// Exact GPT-OSS final residual: bf16(bf16(x + bf16(proj + bias)) + moe).
// proj has the physical GEMM pitch PAD_D while every other matrix is packed REAL_D.
extern "C" __global__ __launch_bounds__(256) void gptoss_residual_join_vec(
bf16 *__restrict__ out, const bf16 *__restrict__ x, const bf16 *__restrict__ proj,
const bf16 *__restrict__ bias, const bf16 *__restrict__ moe) {
const int row = blockIdx.x;
const int tid = threadIdx.x;
#pragma unroll
for (int pass = 0; pass < 2; pass++) {
const int col = (tid + pass * 256) * 8;
if (col >= REAL_D) continue;
const long long live_off = (long long)row * REAL_D + col;
const long long pad_off = (long long)row * PAD_D + col;
const uint4v xv = *reinterpret_cast<const uint4v *>(x + live_off);
const uint4v pv = *reinterpret_cast<const uint4v *>(proj + pad_off);
const uint4v bv = *reinterpret_cast<const uint4v *>(bias + col);
const uint4v mv = *reinterpret_cast<const uint4v *>(moe + live_off);
uint4v ov;
const uint16_t *xh = reinterpret_cast<const uint16_t *>(&xv);
const uint16_t *ph = reinterpret_cast<const uint16_t *>(&pv);
const uint16_t *bh = reinterpret_cast<const uint16_t *>(&bv);
const uint16_t *mh = reinterpret_cast<const uint16_t *>(&mv);
uint16_t *oh = reinterpret_cast<uint16_t *>(&ov);
#pragma unroll
for (int i = 0; i < 8; i++) {
const bf16 xb = __builtin_bit_cast(bf16, xh[i]);
const bf16 pb = __builtin_bit_cast(bf16, ph[i]);
const bf16 bb = __builtin_bit_cast(bf16, bh[i]);
const bf16 mb = __builtin_bit_cast(bf16, mh[i]);
const bf16 attn = (bf16)((float)pb + (float)bb);
const bf16 h = (bf16)((float)xb + (float)attn);
oh[i] = __builtin_bit_cast(uint16_t, (bf16)((float)h + (float)mb));
}
*reinterpret_cast<uint4v *>(out + live_off) = ov;
}
}
tinygrad_repo/extra/gptoss_kernels/residual/wo_bias_final.cpp
+14
View File
@@ -0,0 +1,14 @@
typedef unsigned long size_t;
typedef __bf16 bf16;
extern "C" __attribute__((device, const)) size_t __ockl_get_local_id(unsigned int);
extern "C" __attribute__((device, const)) size_t __ockl_get_group_id(unsigned int);
// The local BF16 boundary must precede the existing data-parallel aggregation.
extern "C" __attribute__((global, amdgpu_flat_work_group_size(64,64)))
void gptoss_wo_bias_final(bf16* out, const float* in) {
int d = __ockl_get_group_id(0)*64 + __ockl_get_local_id(0);
float acc=0;
#pragma unroll
for (int p=0;p<16;p++) acc=acc+in[p*2880+d];
out[d]=(bf16)acc;
}
tinygrad_repo/extra/gptoss_kernels/residual/wo_bias_partial.cpp
+21
View File
@@ -0,0 +1,21 @@
typedef unsigned long size_t;
typedef __bf16 bf16;
extern "C" __attribute__((device, const)) size_t __ockl_get_local_id(unsigned int);
extern "C" __attribute__((device, const)) size_t __ockl_get_group_id(unsigned int);
// Keep the generated GPTOSS reduction's 16 ordered 1024-token partials.
// Prefetching changes only load scheduling; all FP32 additions stay in order.
extern "C" __attribute__((global, amdgpu_flat_work_group_size(64,64)))
void gptoss_wo_bias_partial(float* out, const bf16* in) {
int d = __ockl_get_group_id(0)*64 + __ockl_get_local_id(0);
int p = __ockl_get_group_id(1);
float acc = 0;
for (int r=0; r<1024; r+=64) {
float x[64];
#pragma unroll
for (int j=0; j<64; j++) x[j] = (float)in[(p*1024+r+j)*2880+d];
#pragma unroll
for (int j=0; j<64; j++) acc = acc+x[j];
}
out[p*2880+d] = acc;
}
tinygrad_repo/extra/gptoss_kernels/rmsnorm/__init__.py
+1 -2
View File
@@ -170,8 +170,7 @@ def fast_final_rmsnorm(x:Tensor, weight:Tensor, eps:float) -> Tensor:
axis = x.uop.axis if isinstance(x.device, tuple) else None
local_shape = x.uop.shard_shape if axis is not None else x.shape
if axis not in (None, 0) or math.prod(local_shape[:-1]) != 16384:
xf = x.float()
return (xf * (xf.square().mean(-1, keepdim=True) + eps).rsqrt()).cast(x.dtype) * weight
return x.float().rmsnorm(eps=eps).cast(x.dtype) * weight
fxn = _fast_final_rmsnorm_fwd_fxn(x.as_param(0).uop, weight.as_param(1).uop, eps, x.device)
outputs = UOp.call_with_outputs((fxn[0].uop, fxn[1].uop, fxn[2].uop),
x.uop, weight.uop, grad_fxn=functools.partial(_fast_final_rmsnorm_bwd, eps=eps))
tinygrad_repo/extra/gptoss_kernels/router_topk/__init__.py
+112
View File
@@ -0,0 +1,112 @@
import functools, math, pathlib
from tinygrad import Tensor, dtypes, function
from tinygrad.helpers import getenv
from tinygrad.runtime.support.compiler_amd import HIPCompiler, HIPCCCompiler
from tinygrad.uop.ops import UOp, Ops, KernelInfo
from extra.llama_kernels import alloc_like, compile_hip
@functools.cache
def _router_topk_fwd(weights:UOp, indices:UOp, logits:UOp) -> UOp:
tokens = math.prod(logits.shape[:-1])
sink = UOp.sink(weights.base, indices.base, logits.base,
UOp.special(256, "lidx0"), UOp.special((tokens+255)//256, "gidx0"),
arg=KernelInfo(f"moe_router_topk_{tokens}_32_4"))
src = (pathlib.Path(__file__).parent/"forward.cpp").read_text()
return UOp(Ops.PROGRAM,
src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
UOp(Ops.BINARY, arg=compile_hip(src, [f"-DTOKENS={tokens}"]))))
@functools.cache
def _router_topk_bwd_kernel(grad_logits:UOp, bias_partials:UOp, grad_weights:UOp, weights:UOp, indices:UOp) -> UOp:
tokens = math.prod(grad_logits.shape[:-1])
sink = UOp.sink(grad_logits.base, bias_partials.base, grad_weights.base, weights.base, indices.base,
UOp.special(256, "lidx0"), UOp.special((tokens+255)//256, "gidx0"),
arg=KernelInfo(f"moe_router_topk_bwd_{tokens}_32_4"))
src = (pathlib.Path(__file__).parent/"backward.cpp").read_text()
return UOp(Ops.PROGRAM,
src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
UOp(Ops.BINARY, arg=compile_hip(src, [f"-DTOKENS={tokens}"]))))
def router_topk_backward(gradient:Tensor, weights:Tensor, indices:Tensor) -> tuple[Tensor, Tensor]:
groups, tokens, _ = weights.shape
grad_logits = alloc_like((groups, tokens, 32), dtypes.float32, weights.device, weights.uop.axis)
bias_partials = alloc_like((groups, (tokens+255)//256, 32), dtypes.float32, weights.device, weights.uop.axis)
grad_logits, bias_partials, *_ = Tensor.custom_kernel(grad_logits, bias_partials, gradient.contiguous(), weights, indices,
fxn=_router_topk_bwd_kernel)
return grad_logits, bias_partials
def _router_topk_bwd(gradient:UOp, kernel:UOp) -> tuple:
weights_u, indices_u = kernel.src[1:3]
grad_logits, _ = router_topk_backward(Tensor(gradient), Tensor(weights_u.after(kernel)), Tensor(indices_u.after(kernel)))
return None, None, grad_logits.uop
@functools.cache
def _router_weight_bwd_kernel(out:UOp, x:UOp, gradient:UOp) -> UOp:
rows, hidden = math.prod(x.shape[:-1]), x.shape[-1]
assert rows % 512 == 0 and hidden % 16 == 0
sink = UOp.sink(out.base, x.base, gradient.base, UOp.special(1024, "lidx0"), UOp.special(hidden//16, "gidx0"),
arg=KernelInfo(f"moe_router_fp32_wgrad_{rows}_{hidden}_32"))
src = (pathlib.Path(__file__).parent/"weight_backward.cpp").read_text()
include = pathlib.Path(__file__).parents[2]/"thunder"/"amd"/"include"
lib = HIPCCCompiler("gfx950", [f"-I{include}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS",
f"-DROUTER_M={rows}", f"-DROUTER_K={hidden}"]).compile_cached(src)
return UOp(Ops.PROGRAM,
src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
def router_weight_gradient(x:Tensor, gradient:Tensor) -> Tensor:
assert x.dtype == dtypes.bfloat16 and gradient.dtype == dtypes.float32
groups = len(x.device) if isinstance(x.device, tuple) and x.uop.axis is not None else 1
partial = alloc_like((groups, 32, x.shape[-1]), dtypes.bfloat16, x.device, x.uop.axis)
partial, *_ = Tensor.custom_kernel(partial, x, gradient, fxn=_router_weight_bwd_kernel)
return partial.sum(0, dtype=x.dtype)
def _router_bwd(gradient:UOp, call:UOp) -> tuple:
x, weight, bias = (Tensor(u) for u in call.src[1:4])
weights, indices = (Tensor(u) for u in call.unbound_outputs)
grad_logits, bias_partials = router_topk_backward(Tensor(gradient), weights, indices)
grad_x, = (x.float() @ weight.float().T).gradient(x, gradient=grad_logits.reshape(*x.shape[:-1], 32))
grad_weight = router_weight_gradient(x, grad_logits)
return grad_x.uop, grad_weight.uop, bias_partials.sum(1).cast(bias.dtype).sum(0, dtype=bias.dtype).uop
@function(grad_fxn=_router_bwd)
def fused_router(x:Tensor, weight:Tensor, bias:Tensor) -> tuple[Tensor, Tensor]:
from extra.gemm.moe_routing import router_mfma
logits = router_mfma(x, weight, bias) if getenv("ROUTER_MFMA", 0) else x.float() @ weight.float().T + bias.float()
groups = len(x.device) if isinstance(x.device, tuple) else 1
return fused_router_topk(logits.reshape(groups, -1, 32))
@functools.cache
def _router_input_bwd_kernel(out:UOp, weight:UOp, gradient:UOp, residual:UOp, scales:UOp) -> UOp:
assert math.prod(out.shape[:-1]) == 16384 and out.shape[-1] == 2880 and weight.shape == (32, 2880)
assert math.prod(gradient.shape[:-1]) == 16384 and gradient.shape[-1] == 32
assert residual.shape == (16384, 3072) and scales.shape == (16384, 96)
sink = UOp.sink(out.base, weight.base, gradient.base, residual.base, scales.base,
UOp.special(16, "lidx0"), UOp.special(4, "lidx1"), UOp.special(9, "gidx0"), UOp.special(2048, "gidx1"),
arg=KernelInfo("moe_router_dgrad_fused"))
src = (pathlib.Path(__file__).parent/"input_backward.cpp").read_text()
return UOp(Ops.PROGRAM,
src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
UOp(Ops.BINARY, arg=HIPCompiler("gfx950").compile_cached(src))))
def _router_quantize_bwd(gradient:UOp, quantized_gradient:UOp, *, call:UOp) -> tuple:
x, weight, bias = (Tensor(u) for u in call.src[1:4])
weights, indices, _, scales = (Tensor(u) for u in call.unbound_outputs)
grad_logits, bias_partials = router_topk_backward(Tensor(gradient), weights, indices)
grad_x = alloc_like(x.shape, x.dtype, x.device, x.uop.axis)
grad_x, *_ = Tensor.custom_kernel(grad_x, weight, grad_logits, Tensor(quantized_gradient), scales, fxn=_router_input_bwd_kernel)
return grad_x.uop, router_weight_gradient(x, grad_logits).uop, bias_partials.sum(1).cast(bias.dtype).sum(0, dtype=bias.dtype).uop
@function(grad_fxn=_router_quantize_bwd)
def fused_router_quantize(x:Tensor, weight:Tensor, bias:Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]:
from extra.gptoss_kernels.quantize_mxfp8 import quantize_mxfp8_fused_qe8
weights, indices = fused_router(x, weight, bias)
quantized, scales = quantize_mxfp8_fused_qe8(x.reshape(-1, 2880).pad(((0, 0), (0, 192))))
return weights, indices, quantized, scales
def fused_router_topk(logits:Tensor) -> tuple[Tensor, Tensor]:
assert logits.ndim == 3 and logits.shape[-1] == 32 and logits.dtype == dtypes.float32
axis = logits.uop.axis
weights = alloc_like((*logits.shape[:-1], 4), dtypes.float32, logits.device, axis)
indices = alloc_like(weights.shape, dtypes.int32, logits.device, axis)
weights, indices, *_ = Tensor.custom_kernel(weights, indices, logits, fxn=_router_topk_fwd, grad_fxn=_router_topk_bwd)
return weights, indices
tinygrad_repo/extra/gptoss_kernels/router_topk/backward.cpp
+33
View File
@@ -0,0 +1,33 @@
#include <hip/hip_runtime.h>
extern "C" __global__ __launch_bounds__(256) void moe_router_topk_bwd(
float *__restrict__ grad_logits, float *__restrict__ bias_partials, const float *__restrict__ grad_weights,
const float *__restrict__ weights, const int *__restrict__ indices) {
__shared__ float partial[4][32];
if (threadIdx.x < 128) partial[threadIdx.x / 32][threadIdx.x % 32] = 0.0f;
__syncthreads();
const int token = blockIdx.x * 256 + threadIdx.x;
if (token < TOKENS) {
const long long in = (long long)token * 4;
float dot = 0.0f;
#pragma unroll
for (int i = 0; i < 4; i++) dot += grad_weights[in + i] * weights[in + i];
const long long out = (long long)token * 32;
#pragma unroll
for (int e = 0; e < 32; e++) grad_logits[out + e] = 0.0f;
#pragma unroll
for (int i = 0; i < 4; i++) {
const int e = indices[in + i];
const float dg = weights[in + i] * (grad_weights[in + i] - dot);
grad_logits[out + e] = dg;
atomicAdd(&partial[threadIdx.x / 64][e], dg);
}
}
__syncthreads();
if (threadIdx.x < 32) {
float sum = 0.0f;
#pragma unroll
for (int w = 0; w < 4; w++) sum += partial[w][threadIdx.x];
bias_partials[(long long)blockIdx.x * 32 + threadIdx.x] = sum;
}
}
tinygrad_repo/extra/gptoss_kernels/router_topk/forward.cpp
+37
View File
@@ -0,0 +1,37 @@
#include <hip/hip_runtime.h>
extern "C" __global__ __launch_bounds__(256) void moe_router_topk(
float *__restrict__ weights, int *__restrict__ indices, const float *__restrict__ logits) {
const int token = blockIdx.x * 256 + threadIdx.x;
if (token >= TOKENS) return;
float topv[4];
int topi[4];
#pragma unroll
for (int i = 0; i < 4; i++) { topv[i] = -3.402823466e+38f; topi[i] = -1; }
const long long base = (long long)token * 32;
#pragma unroll
for (int e = 0; e < 32; e++) {
float v = logits[base + e];
int pos = 4;
#pragma unroll
for (int i = 0; i < 4; i++) if (pos == 4 && v > topv[i]) pos = i;
if (pos < 4) {
#pragma unroll
for (int i = 3; i > 0; i--) if (i > pos) { topv[i] = topv[i-1]; topi[i] = topi[i-1]; }
topv[pos] = v;
topi[pos] = e;
}
}
float denom = 0.0f;
float ex[4];
#pragma unroll
for (int i = 0; i < 4; i++) { ex[i] = expf(topv[i] - topv[0]); denom += ex[i]; }
const long long out = (long long)token * 4;
#pragma unroll
for (int i = 0; i < 4; i++) {
weights[out + i] = ex[i] / denom;
indices[out + i] = topi[i];
}
}
tinygrad_repo/extra/gptoss_kernels/router_topk/input_backward.cpp
+68
View File
@@ -0,0 +1,68 @@
typedef long unsigned int size_t;
typedef __bf16 bf16;
typedef float float4_t __attribute__((ext_vector_type(4)));
extern "C" __attribute__((device, const)) size_t __ockl_get_local_id(unsigned int);
extern "C" __attribute__((device, const)) size_t __ockl_get_group_id(unsigned int);
// GPT-OSS router dgrad join, specialized to the production local shard:
// bf16(bf16(sum_e bf16(weight[e,d]) * fp32(grad_logits[t,e]))
// + bf16(bf16(residual[t,d]) * 2**(127-e8[t,d/32])))
extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(64, 64)))
moe_router_dgrad_fused(bf16 *__restrict__ out, const bf16 *__restrict__ weight,
const float *__restrict__ grad_logits, const bf16 *__restrict__ residual,
const unsigned char *__restrict__ e8) {
constexpr int K = 2880, KP = 3072, E = 32;
constexpr int DIM_TILE = 320, DIMS_PER_THREAD = 5;
const int lane_x = __ockl_get_local_id(0);
const int lane_y = __ockl_get_local_id(1);
const int dim_tile = __ockl_get_group_id(0);
const int token8 = __ockl_get_group_id(1);
const int dim0 = dim_tile * DIM_TILE + lane_y * 16 + lane_x;
const int token0 = token8 * 8;
float acc[8][DIMS_PER_THREAD] = {};
#pragma unroll 1
for (int ec = 0; ec < E / 4; ec++) {
const int expert = ec * 4;
const float4_t g0 = *reinterpret_cast<const float4_t *>(grad_logits + (long long)(token0 + 0) * E + expert);
const float4_t g1 = *reinterpret_cast<const float4_t *>(grad_logits + (long long)(token0 + 1) * E + expert);
const float4_t g2 = *reinterpret_cast<const float4_t *>(grad_logits + (long long)(token0 + 2) * E + expert);
const float4_t g3 = *reinterpret_cast<const float4_t *>(grad_logits + (long long)(token0 + 3) * E + expert);
const float4_t g4 = *reinterpret_cast<const float4_t *>(grad_logits + (long long)(token0 + 4) * E + expert);
const float4_t g5 = *reinterpret_cast<const float4_t *>(grad_logits + (long long)(token0 + 5) * E + expert);
const float4_t g6 = *reinterpret_cast<const float4_t *>(grad_logits + (long long)(token0 + 6) * E + expert);
const float4_t g7 = *reinterpret_cast<const float4_t *>(grad_logits + (long long)(token0 + 7) * E + expert);
#pragma unroll
for (int j = 0; j < DIMS_PER_THREAD; j++) {
const int d = dim0 + j * 64;
const float w0 = (float)weight[(long long)(expert + 0) * K + d];
const float w1 = (float)weight[(long long)(expert + 1) * K + d];
const float w2 = (float)weight[(long long)(expert + 2) * K + d];
const float w3 = (float)weight[(long long)(expert + 3) * K + d];
acc[0][j] = acc[0][j] + w0*g0.x + w1*g0.y + w2*g0.z + w3*g0.w;
acc[1][j] = acc[1][j] + w0*g1.x + w1*g1.y + w2*g1.z + w3*g1.w;
acc[2][j] = acc[2][j] + w0*g2.x + w1*g2.y + w2*g2.z + w3*g2.w;
acc[3][j] = acc[3][j] + w0*g3.x + w1*g3.y + w2*g3.z + w3*g3.w;
acc[4][j] = acc[4][j] + w0*g4.x + w1*g4.y + w2*g4.z + w3*g4.w;
acc[5][j] = acc[5][j] + w0*g5.x + w1*g5.y + w2*g5.z + w3*g5.w;
acc[6][j] = acc[6][j] + w0*g6.x + w1*g6.y + w2*g6.z + w3*g6.w;
acc[7][j] = acc[7][j] + w0*g7.x + w1*g7.y + w2*g7.z + w3*g7.w;
}
}
#pragma unroll
for (int t = 0; t < 8; t++) {
#pragma unroll
for (int j = 0; j < DIMS_PER_THREAD; j++) {
const int token = token0 + t, d = dim0 + j * 64;
const unsigned char scale_e = e8[(long long)token * (KP / 32) + d / 32];
// Integer E8M0 exponents produce an exact power of two. e8=254 is 2^-127, a float subnormal.
const unsigned scale_bits = scale_e >= 254 ? (0x00400000u >> ((unsigned)scale_e - 254u)) : (254u - (unsigned)scale_e) << 23;
const float qscale = __builtin_bit_cast(float, scale_bits);
const bf16 router = (bf16)acc[t][j];
const bf16 dequant = (bf16)((float)residual[(long long)token * KP + d] * qscale);
out[(long long)token * K + d] = (bf16)((float)router + (float)dequant);
}
}
}
tinygrad_repo/extra/gptoss_kernels/router_topk/weight_backward.cpp
+69
View File
@@ -0,0 +1,69 @@
#include "kittens.cuh"
using namespace kittens;
constexpr int NUM_WARPS = 16;
constexpr int THREADS = NUM_WARPS * 64;
constexpr int TILE = 16;
constexpr int PREFETCH = 8;
using float4_t = float __attribute__((ext_vector_type(4)));
static_assert(ROUTER_K % TILE == 0);
static_assert(ROUTER_M % (4 * NUM_WARPS * PREFETCH) == 0);
// Each block computes 16 features for all 32 experts, retaining FP32 logit gradients and accumulation.
extern "C" __global__ __launch_bounds__(THREADS, 2) void moe_router_fp32_wgrad(
bf16 *__restrict__ out, const bf16 *__restrict__ x, const float *__restrict__ gradient) {
const int lane = laneid();
const int warp = threadIdx.x / 64;
const int feature_base = blockIdx.x * TILE;
const int k_lane = lane / TILE;
const int mn_lane = lane % TILE;
float4_t accum[2] = {};
#pragma unroll 1
for (int mt = warp; mt < ROUTER_M / 4; mt += NUM_WARPS * PREFETCH) {
float av[PREFETCH], bv0[PREFETCH], bv1[PREFETCH];
#pragma unroll
for (int p = 0; p < PREFETCH; p++) {
const int token = (mt + p * NUM_WARPS) * 4 + k_lane;
av[p] = (float)x[(long long)token * ROUTER_K + feature_base + mn_lane];
bv0[p] = gradient[(long long)token * 32 + mn_lane];
bv1[p] = gradient[(long long)token * 32 + TILE + mn_lane];
}
// Keep the prefetched operands live together instead of serializing load/MFMA pairs.
asm volatile("" :: "v"(av[0]), "v"(bv0[0]), "v"(bv1[0]), "v"(av[1]), "v"(bv0[1]), "v"(bv1[1]),
"v"(av[2]), "v"(bv0[2]), "v"(bv1[2]), "v"(av[3]), "v"(bv0[3]), "v"(bv1[3]),
"v"(av[4]), "v"(bv0[4]), "v"(bv1[4]), "v"(av[5]), "v"(bv0[5]), "v"(bv1[5]),
"v"(av[6]), "v"(bv0[6]), "v"(bv1[6]), "v"(av[7]), "v"(bv0[7]), "v"(bv1[7]) : "memory");
#pragma unroll
for (int p = 0; p < PREFETCH; p++) {
accum[0] = __builtin_amdgcn_mfma_f32_16x16x4f32(av[p], bv0[p], accum[0], 0, 0, 0);
accum[1] = __builtin_amdgcn_mfma_f32_16x16x4f32(av[p], bv1[p], accum[1], 0, 0, 0);
}
}
__shared__ float partial[2][NUM_WARPS][64][4];
#pragma unroll
for (int et = 0; et < 2; et++) {
#pragma unroll
for (int i = 0; i < 4; i++) partial[et][warp][lane][i] = accum[et][i];
}
__syncthreads();
if (warp == 0) {
#pragma unroll
for (int et = 0; et < 2; et++) {
float4_t total = {};
#pragma unroll
for (int w = 0; w < NUM_WARPS; w++) {
#pragma unroll
for (int i = 0; i < 4; i++) total[i] += partial[et][w][lane][i];
}
const int expert = et * TILE + lane % 16;
const int feature0 = feature_base + 4 * (lane / 16);
#pragma unroll
for (int i = 0; i < 4; i++) out[(long long)expert * ROUTER_K + feature0 + i] = (bf16)total[i];
}
}
}
tinygrad_repo/extra/gptoss_kernels/swiglu/__init__.py
+68
View File
@@ -0,0 +1,68 @@
import functools
from tinygrad import Tensor, dtypes
from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
from extra.llama_kernels import alloc_like
ALPHA, LIMIT, LOG2E = 1.702, 7.0, 1.4426950408889634
@functools.cache
def _swiglu_quantize(q:UOp, e8:UOp, h:UOp) -> UOp:
rows = h.shape[0]
wg = UOp.range((rows*24 + 255)//256, 0, AxisType.GLOBAL)
tid = UOp.range(256, 1, AxisType.LOCAL)
sb = UOp.range(4, 2, AxisType.UPCAST)
lane = UOp.range(32, 3, AxisType.UPCAST)
block = (wg*256 + tid)*4 + sb
row, col = block//96, (block%96)*32 + lane
valid = (row < rows) & (col < 2880)
h = h.reshape(rows*5760)
gate = h.index((row*5760 + col*2).valid(valid)).load().cast(dtypes.float32)
up = h.index((row*5760 + col*2 + 1).valid(valid)).load().cast(dtypes.float32)
glu, linear = gate.minimum(LIMIT), up.maximum(-LIMIT).minimum(LIMIT)
sig = (1.0 + (glu * (ALPHA * -LOG2E)).exp2()).reciprocal()
act = glu * sig * (linear + 1.0)
amax = (act < 0.0).where(-act, act).reduce(lane, arg=Ops.MAX)
exponent = (amax.maximum(1e-38).log2().floor() + 127.0).maximum(0.0).minimum(254.0)
scaled = (act * (127.0 - exponent).exp2()).maximum(-448.0).minimum(448.0)
store = q.reshape(rows*3072).index((block*32 + lane).valid(row < rows)).store(scaled.cast(q.dtype)).end(lane)
return e8.reshape(rows*96).after(store).index(block.valid(row < rows)).store(exponent.cast(dtypes.uint8)).end(sb, tid, wg).sink(
arg=KernelInfo("gptoss_swiglu_quantize", opts_to_apply=()))
@functools.cache
def _swiglu_backward(out:UOp, h:UOp, gradient:UOp, e8:UOp) -> UOp:
rows = h.shape[0]
wg = UOp.range((rows*2880 + 256*8 - 1)//(256*8), 0, AxisType.GLOBAL)
tid = UOp.range(256, 1, AxisType.LOCAL)
lane = UOp.range(8, 2, AxisType.UPCAST)
idx = (wg*256 + tid)*8 + lane
valid = idx < rows*2880
row, col = idx//2880, idx%2880
h, out = h.reshape(rows*5760), out.reshape(rows*5760)
gate = h.index((idx*2).valid(valid)).load().cast(dtypes.float32)
up = h.index((idx*2 + 1).valid(valid)).load().cast(dtypes.float32)
scale = (127.0 - e8.reshape(rows*96).index((row*96 + col//32).valid(valid)).load().cast(dtypes.float32)).exp2()
dy = gradient.reshape(rows*3072).index((row*3072 + col).valid(valid)).load().cast(dtypes.float32) * scale
glu, linear = gate.minimum(LIMIT), up.maximum(-LIMIT).minimum(LIMIT)
sig = (1.0 + (glu * (ALPHA * -LOG2E)).exp2()).reciprocal()
sprime = sig * (1.0 + ALPHA * glu * (1.0 - sig))
dgate, dup = dy * sprime * (linear + 1.0), dy * (glu * sig)
gate_ok = (gate < LIMIT).where(1.0, 0.0)
up_ok = (up > -LIMIT).where(1.0, 0.0) * (up < LIMIT).where(1.0, 0.0)
store = out.index((idx*2).valid(valid)).store((dgate * gate_ok).cast(out.dtype))
return out.after(store).index((idx*2 + 1).valid(valid)).store((dup * up_ok).cast(out.dtype)).end(lane, tid, wg).sink(
arg=KernelInfo("gptoss_swiglu_backward", opts_to_apply=()))
def _swiglu_quantize_backward(gradient:UOp, kernel:UOp) -> tuple:
_, e8, h = kernel.src[1:4]
out = alloc_like(h.shape, h.dtype, h.device, h.axis)
out = Tensor.custom_kernel(out, Tensor(h), Tensor(gradient).cast(dtypes.bfloat16), Tensor(e8.after(kernel)), fxn=_swiglu_backward)[0]
return None, None, out.uop
def fused_swiglu_quantize(h:Tensor) -> tuple[Tensor, Tensor]:
assert h.ndim == 2 and h.shape[1] == 5760 and h.dtype == dtypes.bfloat16
assert h.uop.shard_shape[0] % 2 == 0
rows = h.shape[0]
q = alloc_like((rows, 3072), dtypes.fp8e4m3, h.device, h.uop.axis)
e8 = alloc_like((rows, 96), dtypes.uint8, h.device, h.uop.axis)
q, e8, _ = Tensor.custom_kernel(q, e8, h, fxn=_swiglu_quantize, grad_fxn=_swiglu_quantize_backward)
return q, e8
tinygrad_repo/extra/hcq/hcq_smi.py
+7
View File
@@ -65,6 +65,13 @@ def cmd_kill_pids(args):
devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
use_sudo = not getattr(args, "sudoless", False)
# with the kernel driver there are no locks and /proc/<pid>/fd is root-only, so lsof misses them. ask the driver.
for _ in range(128 if args.backend == "nv" and _is_module_loaded("nvidia") else 0):
if not (pids := set(subprocess.check_output(['nvidia-smi', '--query-compute-apps=pid', '--format=csv,noheader']).decode().split())): break
print(f"Killing processes {pids} (which use nvidia)")
subprocess.run((['sudo'] if use_sudo else []) + ['kill', '-9', *pids])
time.sleep(0.2)
for dev in devs:
for i in range(128):
if i > 0: time.sleep(0.2)
tinygrad_repo/extra/hcq1/graph.py
+70 -6
View File
@@ -2,13 +2,77 @@ import collections, time
from typing import Any, cast
from tinygrad.helpers import round_up, PROFILE, ALL2ALL, merge_dicts, getenv, suppress_finalizing, TracingKey, unwrap
from extra.hcq1.hcq import HCQBuffer, HCQCompiled, HCQAllocator, HCQSignal, HWQueue, HCQArgsState
from tinygrad.runtime.support.hcq import BumpAllocator, MMIOInterface
from tinygrad.runtime.support.hcq2 import DepsTracker
from tinygrad.runtime.support.memory import BumpAllocator, MMIOInterface
from tinygrad.device import BufferStorage, Buffer, BufferSpec, Compiled, Device, MultiBuffer, ProfileGraphEntry, ProfileGraphEvent
from tinygrad.dtype import dtypes
from tinygrad.uop.ops import UOp, Ops, Variable
from tinygrad.engine.jit import GraphRunner, MultiGraphRunner
from tinygrad.engine.realize import estimate_uop, get_runtime, unwrap_multi, resolve_params, get_call_arg_uops
from tinygrad.renderer import Estimates
class GraphRunner:
def __init__(self, linear:UOp, input_uops:tuple[UOp, ...]=()):
self.linear = linear.src[0]
self.calls: list[tuple[int, UOp, list[Buffer], dict[str, int]]] = []
self.runtimes: list[Any|None] = []
self.uop_replace: list[list[tuple[int, int]]] = []
for call in self.linear.src:
replace = [(p, b.arg.slot) for p, b in enumerate(get_call_arg_uops(call)) if b.op is Ops.PARAM]
for dev_idx, (bufs, device_vars) in enumerate(unwrap_multi(call, resolve_params(call, input_uops))):
self.calls.append((dev_idx, call.body, [b.ensure_allocated() for b in bufs], device_vars))
self.runtimes.append(get_runtime(bufs[0].device, call.body) if call.body.op is Ops.PROGRAM else None)
self.uop_replace.append(replace)
self.var_vals_replace:dict[int, list[tuple[int, int]]] = {}
self.launch_dims_replace:dict[int, tuple[int|None, int|None]] = {}
self.launch_dims_base:dict[int, tuple[tuple[int|float, ...], tuple[int, ...]]] = {}
def is_sym_dim(dim) -> bool: return not all(isinstance(d, (int, float)) for d in dim)
crs = [(j, self.calls[j][1].arg, self.calls[j][3]) for j in range(len(self.calls)) if self.calls[j][1].op is Ops.PROGRAM]
self.vars = sorted({v.expr for _,p,dv in crs for v in p.vars if v.expr not in dv})
self.symbolic_dims = dedup(tuple(d) for _,p,_ in crs for d in (p.local_size, p.global_size) if is_sym_dim(d))
def find_symbolic_dim(dim:tuple[int,int,int]): return self.symbolic_dims.index(tuple(dim)) if tuple(dim) in self.symbolic_dims else None
for j,p,dv in crs:
if (replace:=[(i, self.vars.index(v.expr)) for i, v in enumerate(p.vars) if v.expr not in dv]):
self.var_vals_replace[j] = replace
global_dim_idx, local_dim_idx = find_symbolic_dim(p.global_size), find_symbolic_dim(p.local_size)
if global_dim_idx is not None or local_dim_idx is not None:
self.launch_dims_replace[j] = (global_dim_idx, local_dim_idx)
self.launch_dims_base[j] = (tuple(p.global_size), tuple(p.local_size))
estimates = sum((estimate_uop(call) for call in self.linear.src), Estimates())
self.device, self.estimates = self.calls[0][2][0].device.split(":")[0], estimates.simplify()
def __call__(self, input_uops:tuple[UOp, ...], var_vals:dict[str, int], wait=False) -> float|None: raise NotImplementedError("override this")
def updated_vars(self, var_vals: dict[str, int]):
vals = [var_vals[v] for v in self.vars]
for j, vidxs in self.var_vals_replace.items():
for i, v in vidxs: yield j, i, vals[v]
def updated_launch_dims(self, var_vals: dict[str, int]):
dims = [tuple(sym_infer(s, var_vals) for s in dim) for dim in self.symbolic_dims]
for j, (gl, lc) in self.launch_dims_replace.items():
yield j, (dims[gl] if gl is not None else self.launch_dims_base[j][0]), (dims[lc] if lc is not None else self.launch_dims_base[j][1])
@staticmethod
def _all_devs(batch_devs:list[Compiled], new_call:UOp) -> list[Compiled]:
return dedup(batch_devs + [Device[x] for b in get_call_arg_uops(new_call)
for x in (b.device if isinstance(b.device, tuple) else (b.device,))])
@staticmethod
def supports_uop(batch_devs:list[Compiled], new_call:UOp) -> bool:
return new_call.op is Ops.CALL and new_call.body.op is Ops.PROGRAM and len(GraphRunner._all_devs(batch_devs, new_call)) == 1
class HCQGraph(GraphRunner):
def _access_resources(self, bufs:list[Buffer], write:list[int], new_dependency:Any):
if not hasattr(self, "deps"): self.deps = DepsTracker()
return self.deps.access_resources([UOp.from_buffer(b.base, dtypes.uint8)[b.offset:b.offset + b.nbytes] for b in bufs], write, new_dependency)
class HCQGraph(MultiGraphRunner):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.devices = list({cast(HCQCompiled, Device[b.device]) for (_,_,bufs,_) in self.calls for b in bufs})
@@ -84,7 +148,7 @@ class HCQGraph(MultiGraphRunner):
self.device_vars: dict[HCQCompiled, dict[str, int]] = {}
for j, ((_, ast, bufs, device_vars), runtime) in enumerate(zip(self.calls, self.runtimes)):
is_xfer = ast.op is Ops.COPY and hasattr(alc:=Device[bufs[0].device].allocator, '_transfer') and alc.supports_transfer \
is_xfer = ast.op is Ops.STORE and hasattr(alc:=Device[bufs[0].device].allocator, '_transfer') and alc.supports_transfer \
and bufs[0].device.split(":")[0] == bufs[1].device.split(":")[0]
ji_devs = [cast(HCQCompiled, Device[b.device]) for b in bufs] if is_xfer else []
is_rdma = len(ji_devs) > 0 and not any(d._is_cpu() for d in ji_devs) and len(set(d.peer_group for d in ji_devs)) > 1
@@ -196,7 +260,7 @@ class HCQGraph(MultiGraphRunner):
dest_queue.signal(dest_out_signal, dest_out_val)
self.num_rdma_ops[(dest_rdma, src_rdma)] += 1
elif ast.op is Ops.COPY:
elif ast.op is Ops.STORE:
dest, src = bufs[0], bufs[1]
uop_replace_j = dict(self.uop_replace[j])
for bufid in range(len(bufs)):
@@ -330,7 +394,7 @@ class HCQGraph(MultiGraphRunner):
try: [d.rdma_dev() for d in all_devs if not d._is_cpu()]
except RuntimeError: return False
if new_call.src[0].op is Ops.COPY:
if new_call.src[0].op is Ops.STORE:
# MOCKGPU is not supported, since it can't execute commands in parallel
is_xfer = len(set(type(d) for d in all_devs)) == 1 and hasattr(alc:=all_devs[0].allocator, '_transfer') and alc.supports_transfer
return is_xfer or (all_devs[0].hw_copy_queue_t is not None and not getattr(all_devs[0], 'iface', None).__class__.__name__.startswith("MOCK"))
tinygrad_repo/extra/hcq1/hcq.py
+5 -21
View File
@@ -361,7 +361,8 @@ class HCQCompiled(Compiled, Generic[SignalType]):
signal_t:Type[SignalType]|None=None, comp_queue_t:Callable[..., HWQueue]|None=None, copy_queue_t:Callable[..., HWQueue]|None=None,
kernargs_size=(16 << 20), sigalloc_size=0x1000, can_recover:bool=False, arch=None):
from extra.hcq1.graph import HCQGraph
super().__init__(device, allocator, compilers, runtime, HCQGraph, arch=arch)
super().__init__(device, allocator, compilers, runtime, arch=arch)
self.graph = HCQGraph
self.peer_group = getattr(getattr(self, 'iface', None), 'peer_group', device.split(":")[0])
HCQCompiled.peer_groups[self.peer_group].append(self)
@@ -475,8 +476,9 @@ class HCQAllocatorBase(Allocator[HCQDeviceType], Generic[HCQDeviceType]):
This class implements basic copy operations following the HCQ API, utilizing both types of `HWQueue`.
"""
def __init__(self, dev:HCQDeviceType, batch_size:int=(2 << 20), batch_cnt:int=32, copy_bufs=None, **kwargs):
super().__init__(dev, **kwargs)
def __init__(self, dev:HCQDeviceType, batch_size:int=(2 << 20), batch_cnt:int=32, copy_bufs=None, supports_transfer:bool=True):
super().__init__(dev)
self.supports_transfer = supports_transfer
self.b = copy_bufs or [self._alloc(batch_size, BufferSpec(host=True)).buf for _ in range(batch_cnt)]
self.b_timeline, self.b_next = [0] * len(self.b), 0
@@ -521,24 +523,6 @@ class HCQAllocator(HCQAllocatorBase, Generic[HCQDeviceType]):
.signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
self.b_timeline[self.b_next] = self.dev.timeline_value - 1
def copy_from_disk(self, dest:HCQBuffer, src, size):
def _get_temp_buf():
# Check if the next buffer is safe to be used (its signal has passed) and reserve it.
if self.b_timeline[(self.b_next + 1) % len(self.b)] <= self.dev.timeline_signal.value:
self.b_timeline[(self.b_next + 1) % len(self.b)], self.b_next = (1 << 64), (self.b_next + 1) % len(self.b)
return (self.b[self.b_next].cpu_view(), self.b_next)
return None
assert self.dev.hw_copy_queue_t is not None
with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=TracingKey(f"DISK -> {self.dev.device}", ret=size), enabled=PROFILE,
dev_suff="SDMA:0"):
for (batch_info, dst_off, src_off, copy_size) in src.device.allocator._copyout_sharded(src, size, _get_temp_buf, seg_len=self.b[0].size,
use_ioring=type(self.b[0].cpu_view()) is MMIOInterface):
self.dev.hw_copy_queue_t().wait(self.dev.timeline_signal, self.dev.timeline_value - 1) \
.copy(dest.offset(dst_off), self.b[batch_info[1]].offset(src_off), copy_size) \
.signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
self.b_timeline[batch_info[1]] = self.dev.timeline_value - 1
def _copyout(self, dest:memoryview, src:HCQBuffer):
self.dev.synchronize()
if self.dev.hw_copy_queue_t is None:
tinygrad_repo/extra/hcq1/ops_amd_old.py
+6 -5
View File
@@ -4,8 +4,9 @@ import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, co
assert sys.platform != 'win32'
from dataclasses import dataclass
from extra.hcq1.hcq import HCQBuffer, HCQCompiled, HCQAllocator, HWQueue, CLikeArgsState, HCQSignal, HCQProgram, hcq_profile
from tinygrad.runtime.support.hcq import FileIOInterface
from tinygrad.runtime.support.hcq import MMIOInterface, BumpAllocator, hcq_filter_visible_devices
from tinygrad.runtime.support.system import FileIOInterface
from tinygrad.runtime.support.memory import MMIOInterface, BumpAllocator
from tinygrad.runtime.support.system import filter_visible_devices
from tinygrad.uop.ops import sint
from tinygrad.device import BufferStorage, Compiled, BufferSpec, TinyELF
from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, lo32, hi32, colored, prod, TracingKey
@@ -641,7 +642,7 @@ class AMDProgram(HCQProgram['AMDDevice']):
class AMDAllocator(HCQAllocator['AMDDevice']):
def __init__(self, dev:AMDDevice):
super().__init__(dev, copy_bufs=getattr(dev.iface, 'copy_bufs', None),
supports_copy_from_disk=dev.has_sdma_queue, supports_transfer=dev.has_sdma_queue and not dev.is_usb())
supports_transfer=dev.has_sdma_queue and not dev.is_usb())
def _alloc(self, size:int, options:BufferSpec) -> BufferStorage:
opaque = self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_sdma_queue)
@@ -780,7 +781,7 @@ class KFDIface:
if KFDIface.kfd is None:
KFDIface.kfd = FileIOInterface("/dev/kfd", os.O_RDWR)
gpus = [g for g in FileIOInterface(kfd_topo_path).listdir() if self._is_usable_gpu(FileIOInterface(f"{kfd_topo_path}/{g}/gpu_id"))]
KFDIface.gpus = hcq_filter_visible_devices(sorted(gpus, key=lambda x: int(x.split('/')[-1])), "AMD")
KFDIface.gpus = filter_visible_devices(sorted(gpus, key=lambda x: int(x.split('/')[-1])), "AMD")
KFDIface.count = len(KFDIface.gpus)
if device_id >= len(KFDIface.gpus): raise RuntimeError(f"No device found for {device_id}. Requesting more devices than the system has?")
@@ -996,7 +997,7 @@ class PCIIface(PCIIfaceBase):
class USBIface(PCIIface):
def __init__(self, dev, dev_id): # pylint: disable=super-init-not-called
if dev_id >= len(visible:=hcq_filter_visible_devices(USB3.list_devices(0xADD1, 0x0001) + USB3.list_devices(0x3801, 0x0001), "AMD")):
if dev_id >= len(visible:=filter_visible_devices(USB3.list_devices(0xADD1, 0x0001) + USB3.list_devices(0x3801, 0x0001), "AMD")):
raise RuntimeError(f"AMD:{dev_id} does not exist ({pluralize('device', len(visible))} available)")
self.dev, self.pci_dev, self.vram_bar, self.count = dev, USBPCIDevice("AM", *visible[dev_id]), 0, len(visible)
self.dev_impl = AMDev(self.pci_dev)
tinygrad_repo/extra/hcq1/ops_rdma.py
+1 -1
View File
@@ -3,7 +3,7 @@ import mmap, struct, functools, atexit
from typing import cast
from tinygrad.uop.ops import sint
from extra.hcq1.hcq import HCQCompiled, HCQAllocatorBase, HCQAllocator, HWQueue, HCQBuffer
from tinygrad.runtime.support.hcq import FileIOInterface
from tinygrad.runtime.support.system import FileIOInterface
from tinygrad.runtime.support.system import System, PCIIfaceBase, PCIAllocationMeta
from tinygrad.runtime.support.memory import VirtMapping, AddrSpace
from extra.mlx_driver.mlxdev import MLXDev, MLXQP
tinygrad_repo/extra/hcq1/remote.py
+2 -1
View File
@@ -1,7 +1,8 @@
from __future__ import annotations
import os, mmap, array, functools, contextlib, itertools, struct, socket, subprocess, time, enum, atexit
from tinygrad.helpers import getenv, temp, ceildiv, unwrap, fetch, system, _ensure_downloads_dir, DEBUG, flatten
from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface
from tinygrad.runtime.support.memory import MMIOInterface
from tinygrad.runtime.support.system import FileIOInterface
from tinygrad.runtime.support.system import PCIDevice, System
class RemoteCmd(enum.IntEnum):
tinygrad_repo/extra/hcq1/test_hcq.py
+36 -36
View File
@@ -32,8 +32,8 @@ class TestHCQ(unittest.TestCase):
def setUp(self):
TestHCQ.d0.synchronize()
TestHCQ.a.uop.buffer.copy_from(Buffer("PYTHON", 2, dtypes.float, opaque=memoryview(bytearray(struct.pack("ff", 0, 1)))))
TestHCQ.b.uop.buffer.copy_from(Buffer("PYTHON", 2, dtypes.float, opaque=memoryview(bytearray(struct.pack("ff", 0, 0)))))
TestHCQ.a.uop.buffer.copy_from(Buffer("PYTHON", 8, initial_value=bytes(struct.pack("ff", 0, 1))))
TestHCQ.b.uop.buffer.copy_from(Buffer("PYTHON", 8, initial_value=bytes(struct.pack("ff", 0, 0))))
TestHCQ.d0.synchronize() # wait for copyins to complete
# Test signals
@@ -171,8 +171,8 @@ class TestHCQ(unittest.TestCase):
prg = to_program(replace_opts(si.src[0], [Opt(op=OptOps.SPLIT, axis=0, arg=(3, AxisType.LOCAL)) for _ in range(3)]), TestHCQ.d0.renderer)
runtime = get_runtime(Device.DEFAULT, prg)
zb = Buffer(Device.DEFAULT, 3 * 3 * 3, dtypes.int, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
zt = Buffer(Device.DEFAULT, 3 * 3 * 3, dtypes.int, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
zb = Buffer(Device.DEFAULT, 3 * 3 * 3 * 4, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
zt = Buffer(Device.DEFAULT, 3 * 3 * 3 * 4, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
ctypes.memset(zb._buf.va_addr, 0, zb.nbytes)
kernargs = runtime.fill_kernargs([zt._buf, zb._buf])
@@ -211,8 +211,8 @@ class TestHCQ(unittest.TestCase):
if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
sz = 64 << 20
buf1 = Buffer(Device.DEFAULT, sz, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(Device.DEFAULT, sz, dtypes.int8, options=BufferSpec(host=True, nolru=True)).ensure_allocated()
buf1 = Buffer(Device.DEFAULT, sz, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(Device.DEFAULT, sz, options=BufferSpec(host=True, nolru=True)).ensure_allocated()
ctypes.memset(buf2._buf.va_addr, 1, sz)
TestHCQ.d0.hw_copy_queue_t().wait(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1) \
@@ -231,8 +231,8 @@ class TestHCQ(unittest.TestCase):
# NOTE: these must be a multiple of 8 for .view(fmt='Q') to work
for sz in [(1 << 32) - 8, (1 << 32), (1 << 32) + 8, (5 << 30), (6 << 30) - 0x4642ee0]:
buf1 = Buffer(Device.DEFAULT, sz, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(Device.DEFAULT, sz, dtypes.int8, options=BufferSpec(host=True, nolru=True)).ensure_allocated()
buf1 = Buffer(Device.DEFAULT, sz, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(Device.DEFAULT, sz, options=BufferSpec(host=True, nolru=True)).ensure_allocated()
ctypes.memset(buf2._buf.va_addr, 0x3e, sz)
buf2_q_view = buf2.host.view(fmt='Q')
@@ -274,8 +274,8 @@ class TestHCQ(unittest.TestCase):
virt_dest_addr = Variable("virt_dest_addr", 0, 0xffffffffffffffff, dtypes.uint64)
sz = 64 << 20
buf1 = Buffer(Device.DEFAULT, sz, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(Device.DEFAULT, sz, dtypes.int8, options=BufferSpec(host=True, nolru=True)).ensure_allocated()
buf1 = Buffer(Device.DEFAULT, sz, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(Device.DEFAULT, sz, options=BufferSpec(host=True, nolru=True)).ensure_allocated()
ctypes.memset(buf2._buf.va_addr, 1, sz)
q = TestHCQ.d0.hw_copy_queue_t().wait(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1) \
@@ -352,8 +352,8 @@ class TestHCQ(unittest.TestCase):
# THEORY: the bandwidth is low here because it's only using one SDMA queue. I suspect it's more stable like this at least.
SZ = 200_000_000
a = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
b = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
a = Buffer(Device.DEFAULT, SZ, options=BufferSpec(nolru=True)).allocate()
b = Buffer(Device.DEFAULT, SZ, options=BufferSpec(nolru=True)).allocate()
sig_st, sig_en = TestHCQ.d0.new_signal(), TestHCQ.d0.new_signal()
st = time.perf_counter()
@@ -380,8 +380,8 @@ class TestHCQ(unittest.TestCase):
except Exception: self.skipTest("no multidevice, test skipped")
SZ = 200_000_000
b = Buffer(f"{Device.DEFAULT}:1", SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
a = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
b = Buffer(f"{Device.DEFAULT}:1", SZ, options=BufferSpec(nolru=True)).allocate()
a = Buffer(Device.DEFAULT, SZ, options=BufferSpec(nolru=True)).allocate()
TestHCQ.d0.allocator._map(b._buf)
sig_st, sig_en = TestHCQ.d0.new_signal(), TestHCQ.d0.new_signal()
@@ -420,8 +420,8 @@ class TestHCQ(unittest.TestCase):
def test_small_copies_from_host_buf(self):
if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
buf1 = Buffer(Device.DEFAULT, 1, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(Device.DEFAULT, 1, dtypes.int8, options=BufferSpec(host=True, nolru=True)).ensure_allocated()
buf1 = Buffer(Device.DEFAULT, 1, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(Device.DEFAULT, 1, options=BufferSpec(host=True, nolru=True)).ensure_allocated()
for i in range(256):
ctypes.memset(buf2._buf.va_addr, i, 1)
@@ -437,9 +437,9 @@ class TestHCQ(unittest.TestCase):
def test_small_copies_from_host_buf_intercopy(self):
if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
buf1 = Buffer(Device.DEFAULT, 1, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(Device.DEFAULT, 1, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
buf3 = Buffer(Device.DEFAULT, 1, dtypes.int8, options=BufferSpec(host=True, nolru=True)).ensure_allocated()
buf1 = Buffer(Device.DEFAULT, 1, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(Device.DEFAULT, 1, options=BufferSpec(nolru=True)).ensure_allocated()
buf3 = Buffer(Device.DEFAULT, 1, options=BufferSpec(host=True, nolru=True)).ensure_allocated()
for i in range(256):
ctypes.memset(buf3._buf.va_addr, i, 1)
@@ -459,9 +459,9 @@ class TestHCQ(unittest.TestCase):
try: _ = Device[f"{Device.DEFAULT}:1"]
except Exception: self.skipTest("no multidevice, test skipped")
buf1 = Buffer(Device.DEFAULT, 1, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(f"{Device.DEFAULT}:1", 1, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
buf3 = Buffer(Device.DEFAULT, 1, dtypes.int8, options=BufferSpec(host=True, nolru=True)).ensure_allocated()
buf1 = Buffer(Device.DEFAULT, 1, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(f"{Device.DEFAULT}:1", 1, options=BufferSpec(nolru=True)).ensure_allocated()
buf3 = Buffer(Device.DEFAULT, 1, options=BufferSpec(host=True, nolru=True)).ensure_allocated()
TestHCQ.d0.allocator._map(buf2._buf)
for i in range(256):
@@ -482,8 +482,8 @@ class TestHCQ(unittest.TestCase):
prg = to_program(b.schedule_linear().src[-1].src[0], TestHCQ.d0.renderer)
runtime = get_runtime(TestHCQ.d0.device, prg)
buf1 = Buffer(Device.DEFAULT, 2, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(Device.DEFAULT, 2, dtypes.int8, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
buf1 = Buffer(Device.DEFAULT, 2, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(Device.DEFAULT, 2, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
kernargs_ptr = runtime.fill_kernargs([buf1._buf, buf2._buf])
@@ -503,9 +503,9 @@ class TestHCQ(unittest.TestCase):
def test_memory_barrier_before_copy(self):
if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
buf1 = Buffer(Device.DEFAULT, 1, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(Device.DEFAULT, 1, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
buf3 = Buffer(Device.DEFAULT, 1, dtypes.int8, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
buf1 = Buffer(Device.DEFAULT, 1, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(Device.DEFAULT, 1, options=BufferSpec(nolru=True)).ensure_allocated()
buf3 = Buffer(Device.DEFAULT, 1, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
for i in range(256):
ctypes.memset(buf3._buf.va_addr, i, 1)
@@ -526,7 +526,7 @@ class TestHCQ(unittest.TestCase):
assert buf2.as_memoryview()[0] == i
def test_write(self):
buf = Buffer(Device.DEFAULT, 4, dtypes.uint32, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
buf = Buffer(Device.DEFAULT, 16, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
try:
TestHCQ.d0.hw_compute_queue_t().write(buf._buf, 0x42) \
@@ -539,7 +539,7 @@ class TestHCQ(unittest.TestCase):
assert buf.as_memoryview().cast("I")[0] == 0x42
def test_poll_bit_set(self):
buf = Buffer(Device.DEFAULT, 4, dtypes.uint32, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
buf = Buffer(Device.DEFAULT, 16, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
try:
TestHCQ.d0.hw_compute_queue_t().write(buf._buf, 0x01000000, b64=False) \
@@ -551,7 +551,7 @@ class TestHCQ(unittest.TestCase):
TestHCQ.d0.timeline_value += 1
def test_poll_bit_clear(self):
buf = Buffer(Device.DEFAULT, 4, dtypes.uint32, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
buf = Buffer(Device.DEFAULT, 16, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
try:
TestHCQ.d0.hw_compute_queue_t().write(buf._buf, 0xFE000000, b64=False) \
@@ -566,7 +566,7 @@ class TestHCQ(unittest.TestCase):
if Device[Device.DEFAULT].hw_copy_queue_t is None: self.skipTest("skip device without copy queue")
sz = 0x2000
cpu_buffer = Buffer("CPU", sz, dtypes.uint8, options=BufferSpec(cpu_access=True)).ensure_allocated()
cpu_buffer = Buffer("CPU", sz, options=BufferSpec(cpu_access=True)).ensure_allocated()
cpu_buffer.host.view(fmt='B')[:] = bytes([x & 0xff for x in range(sz)])
for devid in range(6):
@@ -575,7 +575,7 @@ class TestHCQ(unittest.TestCase):
try: d = Device[f"{Device.DEFAULT}:{devid}"]
except Exception: break
local_buf = Buffer(f"{Device.DEFAULT}:{devid}", sz, dtypes.uint8, options=BufferSpec(cpu_access=True)).ensure_allocated()
local_buf = Buffer(f"{Device.DEFAULT}:{devid}", sz, options=BufferSpec(cpu_access=True)).ensure_allocated()
d.allocator._map(cpu_buffer._buf)
@@ -585,7 +585,7 @@ class TestHCQ(unittest.TestCase):
d.timeline_signal.wait(d.timeline_value)
d.timeline_value += 1
np.testing.assert_equal(cpu_buffer.numpy(), local_buf.numpy(), "failed")
np.testing.assert_equal(np.frombuffer(cpu_buffer.as_memoryview(), np.uint8), np.frombuffer(local_buf.as_memoryview(), np.uint8), "failed")
@unittest.skipUnless(MOCKGPU and not (DEV.device == "AMD" and DEV.interface == "MOCKPCI"), "Emulate this on MOCKGPU to check the path in CI")
def test_on_device_hang(self):
@@ -637,8 +637,8 @@ class TestHCQ(unittest.TestCase):
if TestHCQ.d0.peer_group == d1.peer_group: self.skipTest("devices in same peer group, no RDMA path")
SZ = 200_000_000
a = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
b = Buffer(f"{Device.DEFAULT}:7", SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
a = Buffer(Device.DEFAULT, SZ, options=BufferSpec(nolru=True)).allocate()
b = Buffer(f"{Device.DEFAULT}:7", SZ, options=BufferSpec(nolru=True)).allocate()
# warmup
TestHCQ.d0.allocator._transfer(a._buf, b._buf, SZ, src_dev=d1, dest_dev=TestHCQ.d0)
@@ -655,7 +655,7 @@ class TestHCQ(unittest.TestCase):
print(f"cross device rdma copy: {et_ms:.2f} ms, {gb_s:.2f} GB/s")
assert 1 <= gb_s <= 100
np.testing.assert_equal(a.numpy(), b.numpy(), "failed")
np.testing.assert_equal(np.frombuffer(a.as_memoryview(), np.uint8), np.frombuffer(b.as_memoryview(), np.uint8), "failed")
if __name__ == "__main__":
unittest.main()
tinygrad_repo/extra/hip_large_kernel.py
+3 -3
View File
@@ -12,14 +12,14 @@ extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_si
}
""")
dev.compiler.disassemble(mbin)
buf0 = Buffer(Device.DEFAULT, 1*65537, dtypes.uint8).ensure_allocated()
buf0 = Buffer(Device.DEFAULT, 1*65537).ensure_allocated()
prg = dev.runtime("write_ones", mbin)
prg(buf0._buf, global_size=(1,65537,1), local_size=(1,1,1), wait=True)
import numpy as np
def to_np(buf): return np.frombuffer(buf.as_memoryview().cast(buf.dtype.fmt), dtype=_to_np_dtype(buf.dtype))
def to_np(buf, dtype): return np.frombuffer(buf.as_memoryview(), dtype=_to_np_dtype(dtype))
big = to_np(buf0)
big = to_np(buf0, dtypes.uint8)
print(big)
print((big-1).nonzero())
tinygrad_repo/extra/llama_kernels/cast_amax/__init__.py
-73
View File
@@ -1,73 +0,0 @@
from __future__ import annotations
import functools, pathlib
from tinygrad import Tensor, dtypes
from tinygrad.uop.ops import UOp, Ops, KernelInfo
from tinygrad.renderer import Estimates
from extra.llama_kernels import NUM_WG, THREADS_PER_WG, compile_cpp, alloc_like, dname_of
# module-level mailbox: grad_xw13 UOp -> (grad_xw13_fp8 UOp, delayed amax UOp)
# lets cdna_asm_gemm's bwd reuse the fp8 companion produced by the fused silu_mul bwd kernel
# instead of doing a redundant bf16 -> fp8 quantize.
_grad_fp8_mailbox:dict[UOp, tuple[UOp, UOp]] = {}
@functools.cache
def _custom_fused_bwd_w13(grad_xw13_fp8:UOp, grad_amax_next:UOp, grad_amax:UOp,
xw13:UOp, grad_x2:UOp, amax_state:UOp, grad_amax_state:UOp, dname:str) -> UOp:
hidden = xw13.shape[2] // 2
n_elems = xw13.shape[0] * xw13.shape[1] * hidden
threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
mem = n_elems * 2 * 3 + n_elems * 2 + 4 + 4
sink = UOp.sink(grad_xw13_fp8.base, grad_amax_next.base, grad_amax.base,
xw13.base, grad_x2.base, amax_state.base, grad_amax_state.base, threads, workgroups,
arg=KernelInfo(f"fused_silu_mul_bwd_w13_{n_elems}", estimates=Estimates(ops=10*n_elems, mem=mem)))
src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_bwd_w13.cpp", n_elems, hidden)
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)),
UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
@functools.cache
def _custom_fused_cast_amax_w13(fp8_out:UOp, amax_out:UOp, xw13:UOp, amax_state:UOp, grad_amax_state:UOp,
next_grad_amax_state:UOp, dname:str) -> UOp:
# NOTE: grad_amax_state is plumbed through as an unused fwd input so the bwd kernel can read it via kernel.src
hidden = xw13.shape[2] // 2
n_elems = xw13.shape[0] * xw13.shape[1] * hidden
threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
mem = n_elems * 2 * 2 + n_elems + 4
sink = UOp.sink(fp8_out.base, amax_out.base, xw13.base, amax_state.base, threads, workgroups,
arg=KernelInfo(f"fused_silu_mul_cast_amax_w13_{n_elems}", estimates=Estimates(ops=5*n_elems, mem=mem)))
src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_fwd_w13.cpp", n_elems, hidden)
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)),
UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp):
_, _, xw13, amax_state, grad_amax_state, next_grad_amax_state = kernel.src[1:]
device = xw13.device
axis = xw13.axis if isinstance(device, tuple) else None
grad_xw13_fp8 = alloc_like(xw13.shape, dtypes.fp8e4m3, device, axis)
grad_amax_next = Tensor(next_grad_amax_state, device=device)
grad_amax_state_t = Tensor(grad_amax_state, device=device)
fxn = functools.partial(_custom_fused_bwd_w13, dname=dname_of(device))
grad_amax = grad_amax_state_t.empty_like()
grad_xw13_fp8, grad_amax_next, grad_amax, *_ = Tensor.custom_kernel(
grad_xw13_fp8, grad_amax_next, grad_amax,
Tensor(xw13, device=device), Tensor(gradient, device=device).cast(dtypes.bfloat16),
Tensor(amax_state, device=device), grad_amax_state_t, fxn=fxn)
grad_xw13_uop = grad_xw13_fp8.uop.cast(dtypes.bfloat16)
assert grad_xw13_fp8.uop.op is Ops.AFTER, f"expected AFTER, got {grad_xw13_fp8.uop.op}"
# Stash fp8 companion for cdna_asm_gemm's bwd to attach to grad_a.
_grad_fp8_mailbox[grad_xw13_uop] = (grad_xw13_fp8.uop, grad_amax_state_t.uop)
return (None, None, grad_xw13_uop, None, None, None)
def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype, grad_amax_state:Tensor,
next_grad_amax_state:Tensor, amax_out:Tensor) -> Tensor:
# NOTE: silu(xw1)*xw3 -> fp8 + amax over fused xw13 layout. Returns fp8.
# grad_amax_state: delayed amax for grad_xw13 fp8 quantization in the backward.
assert xw13.dtype == dtypes.bfloat16, f"expected bf16, got {xw13.dtype}"
MBS, SEQ, H2 = xw13.shape
assert H2 % 2 == 0, f"w13 last-axis must be even, got {H2}"
HIDDEN = H2 // 2
axis = xw13.uop.axis if isinstance(xw13.device, tuple) else None
fp8_out = alloc_like((MBS, SEQ, HIDDEN), fp8_dtype, xw13.device, axis)
fxn = functools.partial(_custom_fused_cast_amax_w13, dname=dname_of(xw13.device))
fp8_out, amax_out, *_ = Tensor.custom_kernel(fp8_out, amax_out, xw13, amax_state, grad_amax_state, next_grad_amax_state,
fxn=fxn, grad_fxn=_fused_quantize_bwd_w13)
return fp8_out
tinygrad_repo/extra/llama_kernels/cast_amax/cast_amax_bwd_w13.cpp
-97
View File
@@ -1,97 +0,0 @@
#include <hip/hip_runtime.h>
#include <hip/hip_bf16.h>
#include <hip/hip_fp8.h>
#ifndef N_ELEMS
#define N_ELEMS 234881024
#endif
#ifndef HIDDEN
#define HIDDEN 14336
#endif
#ifndef NUM_WG
#define NUM_WG 1024
#endif
#ifndef THREADS_PER_WG
#define THREADS_PER_WG 256
#endif
constexpr int VEC = 8;
constexpr float FP8_MAX = 448.0f;
static_assert(N_ELEMS % VEC == 0, "N_ELEMS must be divisible by VEC");
static_assert(HIDDEN % VEC == 0, "HIDDEN must be divisible by VEC");
// fused silu*mul backward, three outputs in a single HBM pass:
// 1) fp8 grad_xw13_fp8 — delayed-scale quantize using grad_amax_state (mailbox to matmul bwd)
// 2) fp32 grad_amax_next — scalar |grad_xw13| via global atomic max
// 3) fp32 grad_amax_out — delayed grad amax used for quantize/GEMM epilogue scale
// grad_amax_state is read for the fp8 scale. The store of new_grad_amax into grad_amax_state's
// buffer is built in Python as a separate effect and threaded into grad_a via .after(store).
extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
fused_silu_mul_bwd_w13(
__hip_fp8_storage_t* __restrict__ grad_xw13_fp8_out, // fp8, 2*N_ELEMS
float* __restrict__ grad_amax_next, // fp32 scalar, initialized to 0 before launch
float* __restrict__ grad_amax_out, // fp32 scalar delayed grad amax
const __hip_bfloat16* __restrict__ xw13, // bf16, 2*N_ELEMS
const __hip_bfloat16* __restrict__ grad_x2, // bf16, N_ELEMS
const float* __restrict__ amax_state, // fp32 scalar (fwd x2 amax)
const float* __restrict__ grad_amax_state) // fp32 scalar (delayed grad amax)
{
__shared__ float sdata[THREADS_PER_WG];
const int tid = threadIdx.x;
const int wg = blockIdx.x;
const int gid = wg * THREADS_PER_WG + tid;
const int stride_elems = NUM_WG * THREADS_PER_WG * VEC;
const float scale = FP8_MAX / (static_cast<float>(*amax_state) + 1e-8f);
const float grad_amax = static_cast<float>(*grad_amax_state);
const float g_scale = FP8_MAX / (grad_amax + 1e-8f);
float local_max = 0.0f;
if (wg == 0 && tid == 0) *grad_amax_out = grad_amax;
for (int base = gid * VEC; base < N_ELEMS; base += stride_elems) {
const int outer = base / HIDDEN;
const int inner = base % HIDDEN;
const int xw1_off = outer * 2 * HIDDEN + inner;
const int xw3_off = xw1_off + HIDDEN;
float4 x1_raw = *reinterpret_cast<const float4*>(&xw13[xw1_off]);
float4 x3_raw = *reinterpret_cast<const float4*>(&xw13[xw3_off]);
float4 g_raw = *reinterpret_cast<const float4*>(&grad_x2[base]);
const __hip_bfloat16 *x1 = reinterpret_cast<const __hip_bfloat16*>(&x1_raw);
const __hip_bfloat16 *x3 = reinterpret_cast<const __hip_bfloat16*>(&x3_raw);
const __hip_bfloat16 *gv = reinterpret_cast<const __hip_bfloat16*>(&g_raw);
__hip_fp8_storage_t fp8_1[VEC], fp8_3[VEC];
#pragma unroll
for (int i = 0; i < VEC; i++) {
const float f1 = static_cast<float>(x1[i]);
const float f3 = static_cast<float>(x3[i]);
const float fg = static_cast<float>(gv[i]);
const float sig = 1.0f / (1.0f + __expf(-f1));
const float silu = f1 * sig;
const float silu_prime = sig + silu * (1.0f - sig);
const float gs = fg * scale;
const float g1 = gs * silu_prime * f3;
const float g3 = gs * silu;
local_max = fmaxf(local_max, fmaxf(fabsf(g1), fabsf(g3)));
fp8_1[i] = __hip_cvt_float_to_fp8(fmaxf(-FP8_MAX, fminf(FP8_MAX, g1 * g_scale)), __HIP_SATFINITE, __HIP_E4M3);
fp8_3[i] = __hip_cvt_float_to_fp8(fmaxf(-FP8_MAX, fminf(FP8_MAX, g3 * g_scale)), __HIP_SATFINITE, __HIP_E4M3);
}
*reinterpret_cast<uint64_t*>(&grad_xw13_fp8_out[xw1_off]) = *reinterpret_cast<uint64_t*>(fp8_1);
*reinterpret_cast<uint64_t*>(&grad_xw13_fp8_out[xw3_off]) = *reinterpret_cast<uint64_t*>(fp8_3);
}
sdata[tid] = local_max;
__syncthreads();
for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
if (tid < s) sdata[tid] = fmaxf(sdata[tid], sdata[tid + s]);
__syncthreads();
}
if (tid == 0 && sdata[0] > *grad_amax_next)
atomicMax(reinterpret_cast<int32_t*>(grad_amax_next), __float_as_int(sdata[0]));
}
tinygrad_repo/extra/llama_kernels/cast_amax/cast_amax_fwd_w13.cpp
-79
View File
@@ -1,79 +0,0 @@
#include <hip/hip_runtime.h>
#include <hip/hip_bf16.h>
#include <hip/hip_fp8.h>
#ifndef N_ELEMS
#define N_ELEMS 234881024
#endif
#ifndef HIDDEN
#define HIDDEN 14336
#endif
#ifndef NUM_WG
#define NUM_WG 1024
#endif
#ifndef THREADS_PER_WG
#define THREADS_PER_WG 256
#endif
constexpr int VEC = 8;
constexpr float FP8_MAX = 448.0f;
static_assert(N_ELEMS % VEC == 0, "N_ELEMS must be divisible by VEC");
static_assert(HIDDEN % VEC == 0, "HIDDEN must be divisible by VEC (so VEC loads don't straddle block boundary)");
extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
fused_silu_mul_cast_amax_w13(
__hip_fp8_storage_t* __restrict__ fp8_out, // fp8, N_ELEMS
float* __restrict__ amax_out, // fp32 scalar, initialized to 0 before launch
const __hip_bfloat16* __restrict__ xw13, // bf16, 2*N_ELEMS
const float* __restrict__ amax_state) // fp32 scalar
{
__shared__ float sdata[THREADS_PER_WG];
const int tid = threadIdx.x;
const int wg = blockIdx.x;
const int gid = wg * THREADS_PER_WG + tid;
const int stride_elems = NUM_WG * THREADS_PER_WG * VEC;
const float scale = FP8_MAX / (static_cast<float>(*amax_state) + 1e-8f);
float local_max = 0.0f;
// grid-stride over 8-element groups
for (int base = gid * VEC; base < N_ELEMS; base += stride_elems) {
// interleaved xw13 layout: xw1 and xw3 are not contiguous halves
const int outer = base / HIDDEN;
const int inner = base % HIDDEN;
const int xw1_off = outer * 2 * HIDDEN + inner;
const int xw3_off = xw1_off + HIDDEN;
float4 x1_raw = *reinterpret_cast<const float4*>(&xw13[xw1_off]);
float4 x3_raw = *reinterpret_cast<const float4*>(&xw13[xw3_off]);
const __hip_bfloat16 *x1 = reinterpret_cast<const __hip_bfloat16*>(&x1_raw);
const __hip_bfloat16 *x3 = reinterpret_cast<const __hip_bfloat16*>(&x3_raw);
__hip_fp8_storage_t out[VEC];
#pragma unroll
for (int i = 0; i < VEC; i++) {
const float f1 = static_cast<float>(x1[i]);
const float f3 = static_cast<float>(x3[i]);
const float silu = f1 / (1.0f + __expf(-f1));
const float x2 = silu * f3;
local_max = fmaxf(local_max, fabsf(x2));
const float x_scaled = fmaxf(-FP8_MAX, fminf(FP8_MAX, x2 * scale));
out[i] = __hip_cvt_float_to_fp8(x_scaled, __HIP_SATFINITE, __HIP_E4M3);
}
*reinterpret_cast<uint64_t*>(&fp8_out[base]) = *reinterpret_cast<uint64_t*>(out);
}
// LDS tree reduction: per-workgroup amax, then global atomic into the scalar.
sdata[tid] = local_max;
__syncthreads();
for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
if (tid < s) sdata[tid] = fmaxf(sdata[tid], sdata[tid + s]);
__syncthreads();
}
if (tid == 0 && sdata[0] > *amax_out) atomicMax(reinterpret_cast<int32_t*>(amax_out), __float_as_int(sdata[0]));
}
tinygrad_repo/extra/llama_kernels/fused_ce/__init__.py
+4 -4
View File
@@ -1,6 +1,6 @@
import functools
from tinygrad import Tensor, dtypes
from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
from tinygrad.uop.ops import UOp, Ops, KernelInfo
@functools.cache
def _custom_fused_ce_loss_fwd(loss_out:UOp, max_out:UOp, lse_out:UOp, logits:UOp, targets:UOp,
@@ -9,13 +9,13 @@ def _custom_fused_ce_loss_fwd(loss_out:UOp, max_out:UOp, lse_out:UOp, logits:UOp
b = row // seq
s = row % seq
v_max = UOp.range(vocab, 1, axis_type=AxisType.REDUCE)
v_max = UOp.range(vocab, 1)
row_max = logits[b, s, v_max].cast(dtypes.float).reduce(v_max, arg=Ops.MAX)
v_lse = UOp.range(vocab, 2, axis_type=AxisType.REDUCE)
v_lse = UOp.range(vocab, 2)
row_lse = (logits[b, s, v_lse].cast(dtypes.float) - row_max).exp().reduce(v_lse, arg=Ops.ADD).log() + row_max
v_smooth = UOp.range(vocab, 3, axis_type=AxisType.REDUCE)
v_smooth = UOp.range(vocab, 3)
target = logits[b, s, targets[row].cast(dtypes.weakint)].cast(dtypes.float)
mean_logits = logits[b, s, v_smooth].cast(dtypes.float).reduce(v_smooth, arg=Ops.ADD) / vocab
loss = row_lse - (1.0 - label_smoothing) * target - label_smoothing * mean_logits
tinygrad_repo/extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/__init__.py
-150
View File
@@ -1,150 +0,0 @@
from __future__ import annotations
import functools, pathlib
from tinygrad import Tensor, dtypes
from tinygrad.uop.ops import UOp, Ops, KernelInfo
from tinygrad.renderer import Estimates
from extra.llama_kernels import NUM_WG, THREADS_PER_WG, alloc_like, alloc_local, dname_of, compile_hip
def _src() -> str: return (pathlib.Path(__file__).parent/"fused_rmsnorm_mul_quantize_fp8.cpp").read_text()
def _src_bwd() -> str: return (pathlib.Path(__file__).parent/"fused_rmsnorm_mul_quantize_fp8_bwd.cpp").read_text()
@functools.cache
def _custom_fwd(fp8_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_out:UOp,
x:UOp, weight:UOp, amax_state:UOp, dname:str, eps_val:float) -> UOp:
MBS, SEQ, HIDDEN = x.shape
n_elems = MBS * SEQ * HIDDEN
threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
mem = n_elems * 2 + n_elems + MBS * SEQ * 4 + n_elems + HIDDEN * 2 + 4 + 4
sink = UOp.sink(fp8_out.base, x_normed_out.base, rrms_out.base, amax_out.base,
x.base, weight.base, amax_state.base, threads, workgroups,
arg=KernelInfo(f"fused_rmsnorm_mul_quantize_fp8_{n_elems}_h{HIDDEN}_eps{eps_val:.0e}",
estimates=Estimates(ops=6*n_elems, mem=mem)))
defines = [f"-DN_ELEMS={n_elems}", f"-DHIDDEN={HIDDEN}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}",
f"-DEPS_LITERAL={eps_val}f"]
src = _src()
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)),
UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
@functools.cache
def _custom_fwd_add(fp8_out:UOp, h_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_out:UOp,
x:UOp, residual:UOp, weight:UOp, amax_state:UOp, dname:str, eps_val:float) -> UOp:
MBS, SEQ, HIDDEN = x.shape
n_elems = MBS * SEQ * HIDDEN
threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
mem = n_elems * 2 * 4 + MBS * SEQ * 4 + HIDDEN * 2 + 4 + 4
sink = UOp.sink(fp8_out.base, h_out.base, x_normed_out.base, rrms_out.base, amax_out.base,
x.base, residual.base, weight.base, amax_state.base, threads, workgroups,
arg=KernelInfo(f"fused_add_rmsnorm_mul_quantize_fp8_{n_elems}_h{HIDDEN}_eps{eps_val:.0e}",
estimates=Estimates(ops=7*n_elems, mem=mem)))
defines = [f"-DN_ELEMS={n_elems}", f"-DHIDDEN={HIDDEN}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}",
f"-DEPS_LITERAL={eps_val}f", f"-DHAS_RESIDUAL=1"]
src = _src()
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)),
UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
@functools.cache
def _custom_bwd(grad_x:UOp, grad_weight_partial:UOp,
grad_fp8:UOp, x_normed:UOp, rrms:UOp, weight:UOp, amax_state:UOp, dname:str) -> UOp:
MBS, SEQ, HIDDEN = x_normed.shape
n_elems = MBS * SEQ * HIDDEN
threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
mem = n_elems * 2 * 3 + NUM_WG * HIDDEN * 4 + MBS * SEQ * 4 + HIDDEN * 2 + 4
sink = UOp.sink(grad_x.base, grad_weight_partial.base,
grad_fp8.base, x_normed.base, rrms.base, weight.base, amax_state.base, threads, workgroups,
arg=KernelInfo(f"fused_rmsnorm_mul_quantize_fp8_bwd_{n_elems}_h{HIDDEN}",
estimates=Estimates(ops=8*n_elems, mem=mem)))
defines = [f"-DN_ELEMS={n_elems}", f"-DHIDDEN={HIDDEN}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"]
src = _src_bwd()
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)),
UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
def _bwd_common(fp8_grad_u, h_grad_u, x_u, x_normed_u, rrms_u, weight_u, amax_state_u, kernel:UOp):
device = x_u.device
MBS, SEQ, HIDDEN = x_normed_u.shape
axis = x_normed_u.axis if isinstance(device, tuple) else None
grad_x = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, device, axis)
grad_weight_partial = alloc_local((NUM_WG, HIDDEN), dtypes.float32, device, axis)
grad_h_from_fp8 = None
grad_weight_uop = None
if fp8_grad_u is not None:
fxn = functools.partial(_custom_bwd, dname=dname_of(device))
grad_x_t, grad_weight_partial_t, *_ = Tensor.custom_kernel(
grad_x, grad_weight_partial,
Tensor(fp8_grad_u, device=device).cast(dtypes.bfloat16),
Tensor(x_normed_u.after(kernel), device=device),
Tensor(rrms_u.after(kernel), device=device),
Tensor(weight_u, device=device),
Tensor(amax_state_u, device=device), fxn=fxn)
grad_h_from_fp8 = grad_x_t
grad_weight_uop = grad_weight_partial_t.sum(axis=0).cast(dtypes.bfloat16).uop
if h_grad_u is not None:
h_grad_t = Tensor(h_grad_u, device=device).cast(dtypes.bfloat16)
grad_total = (grad_h_from_fp8 + h_grad_t) if grad_h_from_fp8 is not None else h_grad_t
else:
grad_total = grad_h_from_fp8
return grad_total.uop, grad_weight_uop
def _fused_bwd(gradient:UOp, kernel:UOp):
# NOTE: fwd inputs (fp8_out, x_normed_out, rrms_out, amax_out, x, weight, amax_state)
_, x_normed_u, rrms_u, _, x_u, weight_u, amax_state_u = kernel.src[1:]
grad_x, grad_w = _bwd_common(gradient, None, x_u, x_normed_u, rrms_u, weight_u, amax_state_u, kernel)
return (None, None, None, None, grad_x, grad_w, None)
def _fused_add_bwd(*args, **kwargs):
# Two invocation modes: 1 grad => positional; >1 grads => kwarg `call=`.
# Outputs: (fp8_out, h_out, x_normed_out, rrms_out, amax_buf). Both fp8 and h may be consumed
# downstream — TUPLE order in gradient.py preserves kernel-output slot order.
# Don't dispatch by dtype: matmul's bwd emits fp8 grad as bf16 (no explicit cast), so
# dtype-detection collapses both into h_grad and silently drops the rmsnorm-bwd path.
if 'call' in kwargs:
kernel, all_grads = kwargs['call'], list(args)
else:
gradient, kernel = args
all_grads = [gradient]
fp8_grad_u = h_grad_u = None
if len(all_grads) >= 2:
fp8_grad_u, h_grad_u = all_grads[0], all_grads[1]
elif len(all_grads) == 1:
g = all_grads[0]
if g.dtype == dtypes.bfloat16: h_grad_u = g
else: fp8_grad_u = g
_, _, x_normed_u, rrms_u, _, x_u, _, weight_u, amax_state_u = kernel.src[1:]
grad_h, grad_w = _bwd_common(fp8_grad_u, h_grad_u, x_u, x_normed_u, rrms_u, weight_u, amax_state_u, kernel)
return (None, None, None, None, None, grad_h, grad_h, grad_w, None)
def fused_rmsnorm_mul_quantize_fp8(x:Tensor, weight:Tensor, amax_state:Tensor, eps:float, fp8_dtype,
amax_out:Tensor) -> tuple[Tensor, Tensor, Tensor]:
# NOTE: rmsnorm(x) * weight -> fp8 + amax. Returns (fp8, x_normed, rrms).
# x_normed + rrms are saved for the rmsnorm backward (also recomputed here from x regs).
assert x.dtype == dtypes.bfloat16 and weight.dtype == dtypes.bfloat16
assert x.shape[-1] == weight.shape[-1], f"HIDDEN mismatch: x={x.shape}, weight={weight.shape}"
MBS, SEQ, HIDDEN = x.shape
axis = x.uop.axis if isinstance(x.device, tuple) else None
if isinstance(x.device, tuple): assert axis in (None, 0, 1), f"unsupported sharding axis={axis}"
fp8_out = alloc_like((MBS, SEQ, HIDDEN), fp8_dtype, x.device, axis)
x_normed_out = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis)
rrms_out = alloc_like((MBS, SEQ), dtypes.float32, x.device, axis)
fxn = functools.partial(_custom_fwd, dname=dname_of(x.device), eps_val=eps)
fp8_out, x_normed_out, rrms_out, amax_out, *_ = Tensor.custom_kernel(
fp8_out, x_normed_out, rrms_out, amax_out, x, weight, amax_state, fxn=fxn, grad_fxn=_fused_bwd)
return fp8_out, x_normed_out, rrms_out
def fused_add_rmsnorm_mul_quantize_fp8(x:Tensor, residual:Tensor, weight:Tensor, amax_state:Tensor,
eps:float, fp8_dtype, amax_out:Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]:
# NOTE: h = x + residual; y_normed = rmsnorm(h); fp8 = quantize(y_normed * weight).
# Returns (fp8, h, x_normed, rrms). h is also written so downstream can
# reuse it without recomputing x+residual — eliminates the separate residual-add kernel.
assert x.dtype == dtypes.bfloat16 and residual.dtype == dtypes.bfloat16 and weight.dtype == dtypes.bfloat16
assert x.shape == residual.shape
MBS, SEQ, HIDDEN = x.shape
axis = x.uop.axis if isinstance(x.device, tuple) else None
if isinstance(x.device, tuple): assert axis in (None, 0, 1), f"unsupported sharding axis={axis}"
fp8_out = alloc_like((MBS, SEQ, HIDDEN), fp8_dtype, x.device, axis)
h_out = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis)
x_normed_out = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis)
rrms_out = alloc_like((MBS, SEQ), dtypes.float32, x.device, axis)
fxn = functools.partial(_custom_fwd_add, dname=dname_of(x.device), eps_val=eps)
fp8_out, h_out, x_normed_out, rrms_out, amax_out, *_ = Tensor.custom_kernel(
fp8_out, h_out, x_normed_out, rrms_out, amax_out, x, residual, weight, amax_state,
fxn=fxn, grad_fxn=_fused_add_bwd)
return fp8_out, h_out, x_normed_out, rrms_out
tinygrad_repo/extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/fused_rmsnorm_mul_quantize_fp8.cpp
-155
View File
@@ -1,155 +0,0 @@
#include <hip/hip_runtime.h>
#include <hip/hip_bf16.h>
#include <hip/hip_fp8.h>
// Fuses the full pre-matmul preparation for a layer into a single HBM pass:
// y = rmsnorm(x) * weight (reduce-mean-square + rsqrt + per-elem mul)
// fp8 = fp8_sat(y * (FP8_MAX / amax_state))
// Also writes:
// rrms[row] — saved for the rmsnorm backward
// amax_out — scalar |y| via global atomic max
//
// Layout: one WG per row, ROWS_PER_WG rows per WG via grid-stride (ROWS = N_ELEMS / HIDDEN).
// Each thread handles HIDDEN / THREADS_PER_WG elements per row.
#ifndef N_ELEMS
#define N_ELEMS 67108864
#endif
#ifndef HIDDEN
#define HIDDEN 4096
#endif
#ifndef NUM_WG
#define NUM_WG 1024
#endif
#ifndef THREADS_PER_WG
#define THREADS_PER_WG 256
#endif
#ifndef EPS_LITERAL
#define EPS_LITERAL 1e-5f
#endif
#ifndef HAS_RESIDUAL
#define HAS_RESIDUAL 0
#endif
constexpr int VEC = 8;
constexpr float FP8_MAX = 448.0f;
static_assert(N_ELEMS % HIDDEN == 0, "N_ELEMS must be a multiple of HIDDEN");
static_assert(HIDDEN % (THREADS_PER_WG * VEC) == 0, "HIDDEN must be divisible by THREADS_PER_WG*VEC");
constexpr int ROWS = N_ELEMS / HIDDEN;
constexpr int ELEMS_PER_THREAD = HIDDEN / THREADS_PER_WG; // each thread sees this many elems per row
constexpr int VECS_PER_THREAD = ELEMS_PER_THREAD / VEC; // number of 8-wide vec loads
#if HAS_RESIDUAL
extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
fused_add_rmsnorm_mul_quantize_fp8(
__hip_fp8_storage_t* __restrict__ fp8_out, // fp8, ROWS*HIDDEN
__hip_bfloat16* __restrict__ h_out, // bf16, ROWS*HIDDEN — x + residual (saved for downstream)
__hip_bfloat16* __restrict__ x_normed_out, // bf16, ROWS*HIDDEN
float* __restrict__ rrms_out, // fp32, ROWS
float* __restrict__ amax_out, // fp32 scalar, initialized to 0 before launch
const __hip_bfloat16* __restrict__ x, // bf16, ROWS*HIDDEN
const __hip_bfloat16* __restrict__ residual, // bf16, ROWS*HIDDEN — added into x before rmsnorm
const __hip_bfloat16* __restrict__ weight, // bf16, HIDDEN
const float* __restrict__ amax_state) // fp32 scalar
{
#else
extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
fused_rmsnorm_mul_quantize_fp8(
__hip_fp8_storage_t* __restrict__ fp8_out, // fp8, ROWS*HIDDEN
__hip_bfloat16* __restrict__ x_normed_out, // bf16, ROWS*HIDDEN (saved for rmsnorm bwd)
float* __restrict__ rrms_out, // fp32, ROWS (fp32 to match rmsnorm_bwd.cpp expectation)
float* __restrict__ amax_out, // fp32 scalar, initialized to 0 before launch
const __hip_bfloat16* __restrict__ x, // bf16, ROWS*HIDDEN
const __hip_bfloat16* __restrict__ weight, // bf16, HIDDEN (per-hidden scale)
const float* __restrict__ amax_state) // fp32 scalar
{
#endif
__shared__ float sdata[THREADS_PER_WG];
const int tid = threadIdx.x;
const int wg = blockIdx.x;
const float scale = FP8_MAX / (static_cast<float>(*amax_state) + 1e-8f);
const float inv_hidden = 1.0f / static_cast<float>(HIDDEN);
float local_max = 0.0f;
// Grid-stride over rows. Each WG processes rows (wg, wg+NUM_WG, wg+2*NUM_WG, ...).
for (int row = wg; row < ROWS; row += NUM_WG) {
const int row_off = row * HIDDEN;
// Load row (+ residual if present) into registers.
float regs[ELEMS_PER_THREAD];
float sum_sq = 0.0f;
#pragma unroll
for (int v = 0; v < VECS_PER_THREAD; v++) {
const int h_base = tid * VEC + v * THREADS_PER_WG * VEC;
float4 raw = *reinterpret_cast<const float4*>(&x[row_off + h_base]);
const __hip_bfloat16 *xi = reinterpret_cast<const __hip_bfloat16*>(&raw);
#if HAS_RESIDUAL
float4 res_raw = *reinterpret_cast<const float4*>(&residual[row_off + h_base]);
const __hip_bfloat16 *ri = reinterpret_cast<const __hip_bfloat16*>(&res_raw);
__hip_bfloat16 h_buf[VEC];
#endif
#pragma unroll
for (int i = 0; i < VEC; i++) {
#if HAS_RESIDUAL
const float f = static_cast<float>(xi[i]) + static_cast<float>(ri[i]);
h_buf[i] = static_cast<__hip_bfloat16>(f);
#else
const float f = static_cast<float>(xi[i]);
#endif
regs[v * VEC + i] = f;
sum_sq += f * f;
}
#if HAS_RESIDUAL
*reinterpret_cast<float4*>(&h_out[row_off + h_base]) = *reinterpret_cast<float4*>(h_buf);
#endif
}
// LDS tree-reduce sum_sq across the WG.
sdata[tid] = sum_sq;
__syncthreads();
for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
if (tid < s) sdata[tid] = sdata[tid] + sdata[tid + s];
__syncthreads();
}
const float mean_sq = sdata[0] * inv_hidden;
const float rrms = 1.0f / sqrtf(mean_sq + EPS_LITERAL);
if (tid == 0) rrms_out[row] = rrms;
// Normalize, multiply by weight, quantize. Also write x_normed (for rmsnorm bwd).
#pragma unroll
for (int v = 0; v < VECS_PER_THREAD; v++) {
const int h_base = tid * VEC + v * THREADS_PER_WG * VEC;
float4 w_raw = *reinterpret_cast<const float4*>(&weight[h_base]);
const __hip_bfloat16 *wi = reinterpret_cast<const __hip_bfloat16*>(&w_raw);
__hip_fp8_storage_t out[VEC];
__hip_bfloat16 xn[VEC];
#pragma unroll
for (int i = 0; i < VEC; i++) {
const float x_normed = regs[v * VEC + i] * rrms;
xn[i] = static_cast<__hip_bfloat16>(x_normed);
const float y = x_normed * static_cast<float>(wi[i]);
local_max = fmaxf(local_max, fabsf(y));
const float scaled = fmaxf(-FP8_MAX, fminf(FP8_MAX, y * scale));
out[i] = __hip_cvt_float_to_fp8(scaled, __HIP_SATFINITE, __HIP_E4M3);
}
*reinterpret_cast<uint64_t*>(&fp8_out[row_off + h_base]) = *reinterpret_cast<uint64_t*>(out);
*reinterpret_cast<float4*>(&x_normed_out[row_off + h_base]) = *reinterpret_cast<float4*>(xn);
}
__syncthreads(); // before next row's sum_sq reduce reuses sdata
}
// Final per-WG amax reduce, then global atomic into the scalar.
sdata[tid] = local_max;
__syncthreads();
for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
if (tid < s) sdata[tid] = fmaxf(sdata[tid], sdata[tid + s]);
__syncthreads();
}
if (tid == 0 && sdata[0] > *amax_out) atomicMax(reinterpret_cast<int32_t*>(amax_out), __float_as_int(sdata[0]));
}
tinygrad_repo/extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/fused_rmsnorm_mul_quantize_fp8_bwd.cpp
-147
View File
@@ -1,147 +0,0 @@
#include <hip/hip_runtime.h>
#include <hip/hip_bf16.h>
// Full backward for fused_rmsnorm_mul_quantize_fp8.cpp. One HBM pass per row produces:
// grad_x (bf16) — gradient w.r.t. pre-rmsnorm x
// grad_weight_partial (fp32) — per-WG partial of the weight gradient, reduced later
//
// Input (all read):
// grad_fp8 (bf16) — upstream grad w.r.t. fp8_out (bf16-typed gradient value)
// x_normed (bf16) — saved from the fwd kernel, shape (ROWS, HIDDEN)
// rrms (fp32) — saved rrms per row
// weight (bf16) — per-HIDDEN rmsnorm weight
// amax_state (bf16) — delayed amax used to compute the fp8 scale in fwd
//
// Chain: y = x_normed * weight; fp8 = sat(y * scale). Through STE: grad_y = grad_fp8 * scale.
// grad_x_normed = grad_y * weight.
// grad_weight = sum_rows(grad_y * x_normed).
// grad_x = rrms * (grad_x_normed - x_normed * mean(grad_x_normed * x_normed, last_dim)).
#ifndef N_ELEMS
#define N_ELEMS 67108864
#endif
#ifndef HIDDEN
#define HIDDEN 4096
#endif
#ifndef NUM_WG
#define NUM_WG 1024
#endif
#ifndef THREADS_PER_WG
#define THREADS_PER_WG 256
#endif
constexpr int VEC = 8;
constexpr float FP8_MAX = 448.0f;
static_assert(N_ELEMS % HIDDEN == 0, "N_ELEMS must be a multiple of HIDDEN");
static_assert(HIDDEN % (THREADS_PER_WG * VEC) == 0, "HIDDEN must be divisible by THREADS_PER_WG*VEC");
constexpr int ROWS = N_ELEMS / HIDDEN;
constexpr int ELEMS_PER_THREAD = HIDDEN / THREADS_PER_WG;
constexpr int VECS_PER_THREAD = ELEMS_PER_THREAD / VEC;
extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
fused_rmsnorm_mul_quantize_fp8_bwd(
__hip_bfloat16* __restrict__ grad_x, // out: bf16, ROWS*HIDDEN
float* __restrict__ grad_weight_partial, // out: fp32, NUM_WG*HIDDEN
const __hip_bfloat16* __restrict__ grad_fp8, // in: bf16, ROWS*HIDDEN (grad of fp8_out)
const __hip_bfloat16* __restrict__ x_normed, // in: bf16, ROWS*HIDDEN
const float* __restrict__ rrms, // in: fp32, ROWS
const __hip_bfloat16* __restrict__ weight, // in: bf16, HIDDEN
const float* __restrict__ amax_state) // in: fp32 scalar
{
__shared__ float sdata[THREADS_PER_WG];
const int tid = threadIdx.x;
const int wg = blockIdx.x;
const float scale = FP8_MAX / (static_cast<float>(*amax_state) + 1e-8f);
const float inv_hidden = 1.0f / static_cast<float>(HIDDEN);
// Per-thread accumulator for grad_weight (across all rows this WG touches).
float gw_accum[ELEMS_PER_THREAD];
#pragma unroll
for (int i = 0; i < ELEMS_PER_THREAD; i++) gw_accum[i] = 0.0f;
// Preload weight into registers (same across rows). Use ELEMS_PER_THREAD entries.
float w_regs[ELEMS_PER_THREAD];
#pragma unroll
for (int v = 0; v < VECS_PER_THREAD; v++) {
const int h_base = tid * VEC + v * THREADS_PER_WG * VEC;
float4 w_raw = *reinterpret_cast<const float4*>(&weight[h_base]);
const __hip_bfloat16 *wi = reinterpret_cast<const __hip_bfloat16*>(&w_raw);
#pragma unroll
for (int i = 0; i < VEC; i++) w_regs[v * VEC + i] = static_cast<float>(wi[i]);
}
for (int row = wg; row < ROWS; row += NUM_WG) {
const int row_off = row * HIDDEN;
const float rrms_v = rrms[row];
// Load grad_fp8 and x_normed rows into registers, compute grad_y and grad_x_normed.
float g_y_regs[ELEMS_PER_THREAD];
float xn_regs[ELEMS_PER_THREAD];
float g_xn_regs[ELEMS_PER_THREAD]; // grad_x_normed
float local_dot = 0.0f; // sum(grad_x_normed * x_normed) for mean
#pragma unroll
for (int v = 0; v < VECS_PER_THREAD; v++) {
const int h_base = tid * VEC + v * THREADS_PER_WG * VEC;
float4 g_raw = *reinterpret_cast<const float4*>(&grad_fp8[row_off + h_base]);
float4 xn_raw = *reinterpret_cast<const float4*>(&x_normed[row_off + h_base]);
const __hip_bfloat16 *gi = reinterpret_cast<const __hip_bfloat16*>(&g_raw);
const __hip_bfloat16 *xni = reinterpret_cast<const __hip_bfloat16*>(&xn_raw);
#pragma unroll
for (int i = 0; i < VEC; i++) {
const int idx = v * VEC + i;
const float g_y = static_cast<float>(gi[i]) * scale;
const float xn = static_cast<float>(xni[i]);
g_y_regs[idx] = g_y;
xn_regs[idx] = xn;
g_xn_regs[idx] = g_y * w_regs[idx]; // grad_x_normed = grad_y * weight
gw_accum[idx] += g_y * xn; // grad_weight contrib
local_dot += g_xn_regs[idx] * xn; // for mean
}
}
// LDS reduce local_dot to sdata[0].
sdata[tid] = local_dot;
__syncthreads();
for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
if (tid < s) sdata[tid] = sdata[tid] + sdata[tid + s];
__syncthreads();
}
const float mean_term = sdata[0] * inv_hidden;
// Compute grad_x = rrms * (grad_x_normed - x_normed * mean_term) and write.
#pragma unroll
for (int v = 0; v < VECS_PER_THREAD; v++) {
const int h_base = tid * VEC + v * THREADS_PER_WG * VEC;
__hip_bfloat16 out[VEC];
#pragma unroll
for (int i = 0; i < VEC; i++) {
const int idx = v * VEC + i;
const float dx = rrms_v * (g_xn_regs[idx] - xn_regs[idx] * mean_term);
out[i] = static_cast<__hip_bfloat16>(dx);
}
*reinterpret_cast<float4*>(&grad_x[row_off + h_base]) = *reinterpret_cast<float4*>(out);
}
__syncthreads();
}
// Write this WG's grad_weight partial to HBM (fp32, NUM_WG x HIDDEN layout).
const int gw_row_off = wg * HIDDEN;
#pragma unroll
for (int v = 0; v < VECS_PER_THREAD; v++) {
const int h_base = tid * VEC + v * THREADS_PER_WG * VEC;
// Write 8 fp32 values with two float4 stores.
float4 out_lo, out_hi;
out_lo.x = gw_accum[v * VEC + 0]; out_lo.y = gw_accum[v * VEC + 1];
out_lo.z = gw_accum[v * VEC + 2]; out_lo.w = gw_accum[v * VEC + 3];
out_hi.x = gw_accum[v * VEC + 4]; out_hi.y = gw_accum[v * VEC + 5];
out_hi.z = gw_accum[v * VEC + 6]; out_hi.w = gw_accum[v * VEC + 7];
*reinterpret_cast<float4*>(&grad_weight_partial[gw_row_off + h_base + 0]) = out_lo;
*reinterpret_cast<float4*>(&grad_weight_partial[gw_row_off + h_base + 4]) = out_hi;
}
}
tinygrad_repo/extra/llama_kernels/fused_silu_mul_quantize_mxfp8/__init__.py
-104
View File
@@ -1,104 +0,0 @@
import functools
from tinygrad import Tensor, dtypes
from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
from extra.llama_kernels import FP8_MAX, THREADS_PER_WG, alloc_like
BLK = 32
PACK = 4
LOG2E = 1.4426950408889634
@functools.cache
def _custom_silu_mul_quantize_mxfp8(fp8_out:UOp, e8_out:UOp, si_out:UOp, x_w1:UOp, x_w3:UOp) -> UOp:
rows, K = x_w1.shape
scale_K = K // BLK
n_elems = rows * K
n_super = n_elems // (BLK * PACK)
sk4 = scale_K // PACK
assert n_super % THREADS_PER_WG == 0, f"{n_super=} must divide over {THREADS_PER_WG=}"
nwg = n_super // THREADS_PER_WG
x_w1, x_w3 = x_w1.reshape(n_elems), x_w3.reshape(n_elems)
fp8_out = fp8_out.reshape(n_elems)
e8_out = e8_out.reshape(rows * scale_K)
si_out = si_out.reshape(sk4 * rows)
wg = UOp.range(nwg, 0, AxisType.GLOBAL)
tid = UOp.range(THREADS_PER_WG, 1, AxisType.LOCAL)
sb = UOp.range(PACK, 2, AxisType.UNROLL)
lane = UOp.range(BLK, 3, AxisType.UNROLL)
super_idx = wg * THREADS_PER_WG + tid
idx = super_idx * (BLK * PACK) + sb * BLK + lane
w1 = x_w1[idx].cast(dtypes.float)
w3 = x_w3[idx].cast(dtypes.float)
sig = (1.0 + (w1 * -LOG2E).exp2()).reciprocal()
act = w1 * sig * w3
abs_a = (act < 0.0).where(-act, act)
blk_max = abs_a.reduce(lane, arg=Ops.MAX)
e8f = (blk_max.maximum(1e-38).log2().floor() + 127.0).maximum(0.0).minimum(254.0)
qscale = (127.0 - e8f).exp2()
scaled = (act * qscale).maximum(-FP8_MAX).minimum(FP8_MAX)
e8u8 = e8f.cast(dtypes.uint8)
fp8_store = fp8_out[idx].store(scaled.cast(fp8_out.dtype)).end(lane)
e8_store = e8_out.after(fp8_store)[super_idx * PACK + sb].store(e8u8)
packed = (e8u8.cast(dtypes.uint32) << (sb.cast(dtypes.uint32) * 8)).reduce(sb, arg=Ops.ADD)
row, col4 = super_idx // sk4, super_idx % sk4
si_store = si_out.after(e8_store.end(sb))[col4 * rows + row].store(packed)
return si_store.end(tid, wg).sink(arg=KernelInfo(f"silu_mul_quantize_mxfp8_{n_elems}", opts_to_apply=()))
@functools.cache
def _custom_silu_mul_bwd_mxfp8(gx1_out:UOp, gx3_out:UOp, x_w1:UOp, x_w3:UOp, grad_aq:UOp, e8:UOp) -> UOp:
rows, K = x_w1.shape
scale_K = K // BLK
n_elems = rows * K
VEC = 8
assert n_elems % (THREADS_PER_WG * VEC) == 0, f"{n_elems=} must divide {THREADS_PER_WG*VEC=}"
nwg = n_elems // (THREADS_PER_WG * VEC)
x_w1, x_w3, grad_aq = x_w1.reshape(n_elems), x_w3.reshape(n_elems), grad_aq.reshape(n_elems)
gx1_out, gx3_out, e8 = gx1_out.reshape(n_elems), gx3_out.reshape(n_elems), e8.reshape(rows * scale_K)
wg = UOp.range(nwg, 0, AxisType.GLOBAL)
tid = UOp.range(THREADS_PER_WG, 1, AxisType.LOCAL)
lane = UOp.range(VEC, 2, AxisType.UNROLL)
idx = (wg * THREADS_PER_WG + tid) * VEC + lane
e8v = e8[idx // BLK].cast(dtypes.float)
qscale = (127.0 - e8v).exp2()
ga = grad_aq[idx].cast(dtypes.float) * qscale
w1 = x_w1[idx].cast(dtypes.float)
w3 = x_w3[idx].cast(dtypes.float)
sig = (1.0 + (w1 * -LOG2E).exp2()).reciprocal()
s = w1 * sig
sprime = sig * (1.0 + w1 * (1.0 - sig))
gx1 = gx1_out[idx].store((ga * sprime * w3).cast(gx1_out.dtype))
gx3 = gx3_out.after(gx1)[idx].store((ga * s).cast(gx3_out.dtype))
return gx3.end(lane, tid, wg).sink(arg=KernelInfo(f"silu_mul_bwd_mxfp8_{n_elems}", opts_to_apply=()))
def _silu_mul_quantize_mxfp8_bwd(gradient:UOp, kernel:UOp):
_, e8_out, _, x_w1, x_w3 = kernel.src[1:]
device = x_w1.device
rows, K = x_w1.shape
axis = x_w1.axis if isinstance(device, tuple) else None
gx1 = alloc_like((rows, K), dtypes.bfloat16, device, axis)
gx3 = alloc_like((rows, K), dtypes.bfloat16, device, axis)
gx1, gx3, *_ = Tensor.custom_kernel(gx1, gx3, Tensor(x_w1, device=device), Tensor(x_w3, device=device),
Tensor(gradient, device=device).cast(dtypes.bfloat16), Tensor(e8_out.after(kernel), device=device),
fxn=_custom_silu_mul_bwd_mxfp8)
return (None, None, None, gx1.uop, gx3.uop)
def fused_silu_mul_quantize_mxfp8(x_w1:Tensor, x_w3:Tensor) -> tuple[Tensor, Tensor, Tensor]:
assert x_w1.shape == x_w3.shape, f"{x_w1.shape} != {x_w3.shape}"
assert x_w1.dtype == dtypes.bfloat16 and x_w3.dtype == dtypes.bfloat16
assert x_w1.ndim == 2, f"expected 2d, got {x_w1.shape}"
from extra.gemm.cdna_asm_gemm import FP8_DTYPE
rows, K = x_w1.shape
scale_K = K // BLK
axis = x_w1.uop.axis if isinstance(x_w1.device, tuple) else None
fp8_out = alloc_like((rows, K), FP8_DTYPE, x_w1.device, axis)
e8_out = alloc_like((rows, scale_K), dtypes.uint8, x_w1.device, axis)
si_out = alloc_like((scale_K // PACK, rows), dtypes.uint32, x_w1.device, None if axis is None else (1 if axis == 0 else 0))
fp8_out, e8_out, si_out, *_ = Tensor.custom_kernel(fp8_out, e8_out, si_out, x_w1, x_w3,
fxn=_custom_silu_mul_quantize_mxfp8, grad_fxn=_silu_mul_quantize_mxfp8_bwd)
return fp8_out, e8_out, si_out
tinygrad_repo/extra/llama_kernels/quantize_fp8_delayed/__init__.py
-96
View File
@@ -1,96 +0,0 @@
import functools
from tinygrad import Tensor, dtypes
from tinygrad.dtype import AddrSpace
from tinygrad.helpers import prod
from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, alloc_like
@functools.cache
def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state:UOp, device=None) -> UOp:
VEC = 8
n_elems = prod(x.shape)
assert n_elems % (NUM_WG * THREADS_PER_WG * VEC) == 0
x = x.reshape(n_elems)
fp8_out = fp8_out.reshape(n_elems)
wg = UOp.range(NUM_WG, 0, AxisType.GLOBAL)
tid = UOp.range(THREADS_PER_WG, 1, AxisType.LOCAL)
it = UOp.range((n_elems // VEC) // (NUM_WG * THREADS_PER_WG), 2, AxisType.WEAK)
lane = UOp.range(VEC, 3, AxisType.UNROLL)
idx = (((it * NUM_WG + wg) * THREADS_PER_WG + tid) * VEC) + lane
scale = FP8_MAX / (amax_state[0].cast(dtypes.float) + 1e-8)
x_f = x[idx].cast(dtypes.float)
abs_x = (x_f < 0.0).where(-x_f, x_f)
scaled = (x_f * scale).maximum(-FP8_MAX).minimum(FP8_MAX)
fp8_store = fp8_out[idx].store(scaled.cast(fp8_out.dtype)).end(lane)
lane_max = abs_x.reduce(lane, arg=Ops.MAX)
lmax = UOp.placeholder((1,), dtypes.float, slot=1, addrspace=AddrSpace.REG)
lmax_init = lmax.after(wg, tid)[0].store(0.0)
lmax_prev = lmax.after(lmax_init, it)[0]
lmax_store = lmax.after(fp8_store)[0].store(lmax_prev.maximum(lane_max))
lmax_val = lmax.after(lmax_store.end(it))[0]
lds = UOp.placeholder((THREADS_PER_WG,), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL)
lds = lds.after(lds[tid].store(lmax_val))
step = THREADS_PER_WG // 2
while step:
active = tid < step
other = lds[(tid + step).valid(active)].load()
lds = lds.after(lds[tid.valid(active)].store(lds[tid].maximum(other)))
step //= 2
device = device[0].split(":")[0] if isinstance(device, tuple) else device.split(":")[0]
if device in {"AMD", "NULL"}: atomic_arg = "if ({2} > {3}) __hip_atomic_fetch_max((int*){0}, {1}, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);"
else: raise NotImplementedError(f"no atomic max for device {device}")
amax_idx = amax_out.reshape((1,)).index(UOp.const(0))
max_val = lds[0].load()
atomic = UOp(Ops.CUSTOM, src=(amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=(atomic_arg, dtypes.void))
return atomic.end(tid, wg).sink(arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", opts_to_apply=()))
@functools.cache
def _custom_quantize_fp8_scalar(fp8_out:UOp, x:UOp, amax_state:UOp) -> UOp:
n_elems = prod(x.shape)
i = UOp.range(n_elems, 0)
x_f = x.reshape(n_elems)[i].cast(dtypes.float)
scale = FP8_MAX / (amax_state[0].cast(dtypes.float) + 1e-8)
store = fp8_out.reshape(n_elems)[i].store((x_f * scale).cast(fp8_out.dtype))
return store.end(i).sink(arg=KernelInfo(f"quantize_fp8_scalar_{n_elems}"))
def _quantize_fp8_delayed_bwd(gradient:UOp, kernel:UOp):
# NOTE: STE-equivalent backward — grad_x = grad_fp8 * scale, scale = FP8_MAX / amax_state.
# `gradient` is bf16 grad w.r.t. fp8 output (asm_gemm bwd already applied x_scale).
_, _, x, amax_state = kernel.src[1:]
device = x.device
scale = FP8_MAX / (Tensor(amax_state, device=device).float() + 1e-8)
grad_x = (Tensor(gradient, device=device).float() * scale).cast(dtypes.bfloat16)
return (None, None, grad_x.uop, None)
def quantize_fp8_delayed(x:Tensor, amax_state:Tensor, amax_out:Tensor, fp8_dtype=dtypes.fp8e4m3) -> tuple[Tensor, Tensor]:
# NOTE: one-pass bf16 -> fp8 quantize with delayed scaling.
# Fused kernel reads x once and writes fp8 + scalar amax via global atomic max.
assert x.dtype == dtypes.bfloat16, f"expected bf16, got {x.dtype}"
axis = x.uop.axis if isinstance(x.device, tuple) else None
fp8_out = alloc_like(x.shape, fp8_dtype, x.device, axis)
n_elems = prod(x.uop.shard_shape)
assert n_elems % NUM_WG == 0, f"{n_elems=} must divide over {NUM_WG=}"
fxn = functools.partial(_custom_quantize_fp8_with_amax, device=x.device)
fp8_out, amax_out, *_ = Tensor.custom_kernel(fp8_out, amax_out, x, amax_state,
fxn=fxn, grad_fxn=_quantize_fp8_delayed_bwd)
inv_scale = (amax_state.float() + 1e-8) / FP8_MAX
return fp8_out, inv_scale
def quantize_fp8_scalar(x:Tensor, amax_state:Tensor, fp8_dtype=dtypes.fp8e4m3) -> Tensor:
# NOTE: pure one-pass bf16 -> fp8 quantize with delayed scalar scale. No amax computation.
axis = x.uop.axis if isinstance(x.device, tuple) else None
fp8_out = alloc_like(x.shape, fp8_dtype, x.device, axis)
fxn = _custom_quantize_fp8_scalar
fp8_out, *_ = Tensor.custom_kernel(fp8_out, x, amax_state, fxn=fxn)
return fp8_out
tinygrad_repo/extra/llama_kernels/quantize_mxfp8_fused/__init__.py
-71
View File
@@ -1,71 +0,0 @@
import functools
from tinygrad import Tensor, dtypes
from tinygrad.helpers import prod
from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
from extra.llama_kernels import FP8_MAX, THREADS_PER_WG, alloc_like
BLK = 32
PACK = 4
@functools.cache
def _custom_quantize_mxfp8(fp8_out:UOp, e8_out:UOp, si_out:UOp, x:UOp) -> UOp:
rows, K = x.shape
scale_K = K // BLK
n_elems = rows * K
n_super = n_elems // (BLK * PACK)
sk4 = scale_K // PACK
assert n_super % THREADS_PER_WG == 0, f"{n_super=} must divide over {THREADS_PER_WG=}"
nwg = n_super // THREADS_PER_WG
x = x.reshape(n_elems)
fp8_out = fp8_out.reshape(n_elems)
e8_out = e8_out.reshape(rows * scale_K)
si_out = si_out.reshape(sk4 * rows)
wg = UOp.range(nwg, 0, AxisType.GLOBAL)
tid = UOp.range(THREADS_PER_WG, 1, AxisType.LOCAL)
sb = UOp.range(PACK, 2, AxisType.UNROLL)
lane = UOp.range(BLK, 3, AxisType.UNROLL)
super_idx = wg * THREADS_PER_WG + tid
idx = super_idx * (BLK * PACK) + sb * BLK + lane
x_f = x[idx].cast(dtypes.float)
abs_x = (x_f < 0.0).where(-x_f, x_f)
blk_max = abs_x.reduce(lane, arg=Ops.MAX)
e8f = (blk_max.maximum(1e-38).log2().floor() + 127.0).maximum(0.0).minimum(254.0)
qscale = (127.0 - e8f).exp2()
scaled = (x_f * qscale).maximum(-FP8_MAX).minimum(FP8_MAX)
e8u8 = e8f.cast(dtypes.uint8)
fp8_store = fp8_out[idx].store(scaled.cast(fp8_out.dtype)).end(lane)
e8_store = e8_out.after(fp8_store)[super_idx * PACK + sb].store(e8u8)
# pack the 4 e8 of this super-block into one uint32 (little-endian: byte sb), write transposed (sk4, row)
packed = (e8u8.cast(dtypes.uint32) << (sb.cast(dtypes.uint32) * 8)).reduce(sb, arg=Ops.ADD)
row, col4 = super_idx // sk4, super_idx % sk4
si_store = si_out.after(e8_store.end(sb))[col4 * rows + row].store(packed)
return si_store.end(tid, wg).sink(arg=KernelInfo(f"quantize_mxfp8_{n_elems}", opts_to_apply=()))
def _quantize_mxfp8_fused_bwd(gradient:UOp, kernel:UOp):
_, e8_out, _, x = kernel.src[1:]
device = x.device
rows, K = x.shape
scale_K = K // BLK
e8 = Tensor(e8_out, device=device).reshape(rows, scale_K)
qscale = (127.0 - e8.cast(dtypes.float32)).exp2().reshape(rows, scale_K, 1).expand(rows, scale_K, BLK).reshape(rows, K)
grad_x = (Tensor(gradient, device=device).float() * qscale).cast(dtypes.bfloat16)
return (None, None, None, grad_x.uop)
def quantize_mxfp8_fused(x:Tensor) -> tuple[Tensor, Tensor, Tensor]:
assert x.dtype == dtypes.bfloat16, f"expected bf16, got {x.dtype}"
assert x.ndim == 2, f"expected 2d (rows, K), got {x.shape}"
from extra.gemm.cdna_asm_gemm import FP8_DTYPE
rows, K = x.shape
scale_K = K // BLK
axis = x.uop.axis if isinstance(x.device, tuple) else None
fp8_out = alloc_like((rows, K), FP8_DTYPE, x.device, axis)
e8_out = alloc_like((rows, scale_K), dtypes.uint8, x.device, axis)
si_out = alloc_like((scale_K // PACK, rows), dtypes.uint32, x.device, None if axis is None else (1 if axis == 0 else 0))
fp8_out, e8_out, si_out, *_ = Tensor.custom_kernel(fp8_out, e8_out, si_out, x, fxn=_custom_quantize_mxfp8, grad_fxn=_quantize_mxfp8_fused_bwd)
return fp8_out, e8_out, si_out
tinygrad_repo/extra/mlx_driver/loopback.py
+4 -5
View File
@@ -3,7 +3,6 @@
import struct
from tinygrad.helpers import getenv, round_up
from tinygrad.device import Device, BufferSpec, Buffer
from tinygrad.dtype import dtypes
from tinygrad.runtime.support.system import PCIDevice
from tinygrad.runtime.support.memory import AddrSpace
from tinygrad.runtime.ops_amd import AMDComputeQueue
@@ -31,8 +30,8 @@ print(f"[init] loopback connect QP 0x{qp.qp_info['qpn']:x}")
qp.connect(qp.qp_info['qpn'], dev.mac, int.from_bytes(dev.local_gid, 'big'))
# allocate src/dst via AMD GPU allocator
buf_src = Buffer(gpu.device, BUF_SIZE, dtypes.uint8, options=BufferSpec(nolru=True), preallocate=True)
buf_dst = Buffer(gpu.device, BUF_SIZE, dtypes.uint8, options=BufferSpec(nolru=True), preallocate=True)
buf_src = Buffer(gpu.device, BUF_SIZE, options=BufferSpec(nolru=True), preallocate=True)
buf_dst = Buffer(gpu.device, BUF_SIZE, options=BufferSpec(nolru=True), preallocate=True)
bar_base = gpu.iface.pci_dev.bar_info(gpu.iface.vram_bar)[0]
src_paddr = buf_src.meta.mapping.paddrs[0][0] + bar_base
@@ -41,8 +40,8 @@ print(f"src paddr=0x{src_paddr:x} dst paddr=0x{dst_paddr:x}")
# fill src, zero dst
test_msg = b"Hello from loopback send/recv!"
buf_src.copy_from(Buffer("PYTHON", BUF_SIZE, dtypes.uint8, opaque=memoryview(bytearray(test_msg.ljust(BUF_SIZE, b'\x00')))))
buf_dst.copy_from(Buffer("PYTHON", BUF_SIZE, dtypes.uint8, opaque=memoryview(bytearray(BUF_SIZE))))
buf_src.copy_from(Buffer("PYTHON", BUF_SIZE, initial_value=bytes(test_msg.ljust(BUF_SIZE, b'\x00'))))
buf_dst.copy_from(Buffer("PYTHON", BUF_SIZE, initial_value=bytes(BUF_SIZE)))
gpu.synchronize()
# post recv WQE on RQ from CPU (scatter entry: byte_count, lkey, addr)
tinygrad_repo/extra/models/t5.py
+1 -2
View File
@@ -62,8 +62,7 @@ class T5LayerNorm:
# w/o mean and there is no bias. Additionally we want to make sure that the accumulation for
# half-precision inputs is done in fp32
variance = hidden_states.cast(dtypes.float32).pow(2).mean(-1, keepdim=True)
hidden_states = hidden_states * Tensor.rsqrt(variance + self.variance_epsilon)
hidden_states = hidden_states.cast(dtypes.float32).rmsnorm(eps=self.variance_epsilon)
# convert into half-precision if necessary
if self.weight.dtype in [dtypes.float16, dtypes.bfloat16]:

Some files were not shown because too many files have changed in this diff Show More