diff --git a/INFERENCE-PROFILE-GFX906.md b/INFERENCE-PROFILE-GFX906.md new file mode 100644 index 000000000..734f1c9c3 --- /dev/null +++ b/INFERENCE-PROFILE-GFX906.md @@ -0,0 +1,320 @@ +# Profilo delle fasi di inferenza su 6x AMD gfx906 + +## Configurazione misurata + +Le statistiche di questo documento provengono da un test eseguito direttamente +sulla configurazione GPU usata dai launcher gfx906: + +- modello: DeepSeek-V4-Flash IQ2XXS; +- backend: ROCm 7.2.4, architettura `gfx906` wave64; +- topologia: pipeline parallela PP6; +- coordinatore: MI50 32 GB, layer `0:12` e output head; +- worker: cinque GPU gfx906 da 16 GB; +- suddivisione worker: `13:18`, `19:24`, `25:30`, `31:36`, `37:42`; +- esperti completamente residenti in VRAM, senza streaming SSD; +- contesto configurato: 300.000 token; +- posizione del contesto durante la misura: circa 2.400 token; +- chunk di prefill: 64 token; +- finestra distribuita: 5 chunk; +- trasporto delle attivazioni: FP16; +- prompt: 2.418 token; +- generazione: 64 token. + +Il profiler e la telemetria erano abilitati soltanto durante questa misura. I +launcher normali continuano a tenere disattivato il logging della telemetria. + +## Ripartizione complessiva della richiesta + +| Macro-fase | Tempo | Percentuale sul totale | +|---|---:|---:| +| Prefill | 33,22 s | 86,7% | +| Generazione di 64 token | 5,11 s | 13,3% | +| Tokenizzazione e gestione API | circa 2 ms | <0,1% | +| **Totale richiesta** | **38,33 s** | **100,0%** | + +La percentuale complessiva tra prefill e generazione dipende dal rapporto tra +token di input e token generati. Le tabelle successive descrivono separatamente +le due fasi e sono quindi più utili per individuare i colli di bottiglia. + +## Prefill + +Il prefill usa una pipeline: mentre una GPU elabora un chunk, le altre possono +elaborare chunk precedenti. Di conseguenza, i tempi dei singoli stadi si +sovrappongono e non possono essere sommati per ottenere direttamente il tempo +wall-clock. + +La tabella seguente normalizza le percentuali rispetto al lavoro seriale +equivalente di un chunk completo da 64 token, circa 2.829 ms. Grazie alla +sovrapposizione della pipeline, a regime viene completato un chunk ogni circa +819 ms. + +| Stadio | Tempo medio per chunk | Percentuale del lavoro seriale | Durata rispetto al bottleneck | +|---|---:|---:|---:| +| Coordinatore, layer `0:12` e dispatch | circa 819 ms | 29,0% | 100,0% | +| Worker ROCR 0, layer `13:18` | 405 ms | 14,3% | 49,4% | +| Worker ROCR 1, layer `19:24` | 392 ms | 13,9% | 47,9% | +| Worker ROCR 3, layer `25:30` | 392 ms | 13,9% | 47,9% | +| Worker ROCR 4, layer `31:36` | 405 ms | 14,3% | 49,4% | +| Worker ROCR 5, layer `37:42` | 390 ms | 13,8% | 47,6% | +| Trasferimenti delle attivazioni tra worker | 25,7 ms | 0,9% | 3,1% | +| **Totale seriale equivalente** | **2.829 ms** | **100,0%** | — | + +### Tempo wall-clock del prefill + +| Componente | Tempo | Percentuale del prefill | +|---|---:|---:| +| Pipeline a regime | 30,70 s | 92,4% | +| Riempimento iniziale della pipeline | 2,52 s | 7,6% | +| **Totale prefill** | **33,22 s** | **100,0%** | + +### Interpretazione del prefill + +Il coordinatore è il collo di bottiglia principale. Esegue 13 layer, mentre +ogni worker ne esegue 6, e richiede circa 819 ms per stadio contro 390-405 ms +dei worker. I worker impiegano quindi soltanto il 48-49% circa del tempo dello +stadio più lento. + +I trasferimenti tra worker rappresentano circa lo 0,9% del lavoro seriale. La +rete locale non è il limite dominante del prefill; le opportunità principali +sono il riequilibrio dei layer e l'ottimizzazione del calcolo sul coordinatore. + +## Generazione autoregressiva di token + +Durante il decode i blocchi devono essere attraversati in sequenza per ogni +token. Le percentuali di questa tabella sono quindi direttamente additive e +sono normalizzate rispetto al tempo end-to-end medio di 79,86 ms per token. + +| Fase | Tempo medio per token | Percentuale del decode | +|---|---:|---:| +| Embedding e coordinatore, layer `0:12` | 20,99 ms | 26,28% | +| Worker ROCR 0, layer `13:18` | 10,76 ms | 13,47% | +| Worker ROCR 1, layer `19:24` | 11,21 ms | 14,04% | +| Worker ROCR 3, layer `25:30` | 11,08 ms | 13,88% | +| Worker ROCR 4, layer `31:36` | 12,00 ms | 15,03% | +| Worker ROCR 5, layer `37:42` | 11,13 ms | 13,93% | +| Output head | 1,50 ms | 1,87% | +| Conversioni FP16, TCP e scheduling distribuito | 1,11 ms | 1,39% | +| Sampling e gestione API | 0,09 ms | 0,11% | +| **Totale per token** | **79,86 ms** | **100,00%** | + +Il calcolo dei transformer layer occupa complessivamente circa il 96,63% del +tempo di generazione: + +- coordinatore: 26,28%; +- cinque worker: 70,35%; +- totale calcolo dei layer: 96,63%. + +Le componenti esterne al calcolo dei layer pesano complessivamente circa il +3,37%: + +- output head: 1,87%; +- trasporto, conversioni e scheduling: 1,39%; +- sampling e API: 0,11%. + +## Throughput + +| Misura | Risultato | +|---|---:| +| Prefill medio | 72,80 token/s | +| Decode con profiler e telemetria | 12,52 token/s | +| Decode normale senza telemetria | circa 14,33 token/s | +| Overhead osservato del profiling | circa 12,6% | + +Il profiling genera una riga diagnostica per ogni worker e per ogni token. I +valori assoluti del decode risultano pertanto inferiori al funzionamento +normale, mentre la distribuzione percentuale rimane utile per identificare gli +stadi dominanti. + +## Distanza dalla capacità teorica delle GPU + +Il fatto che circa il 96,63% del decode sia trascorso nei transformer layer non +significa che le unità di calcolo delle GPU siano sature. La voce `eval` +comprende anche letture HBM, dequantizzazione IQ2/Q2, routing degli esperti, +riduzioni, sincronizzazioni e tempi di lancio dei kernel. + +ROCm espone sei GPU Vega 20 `gfx906`, tutte con 60 compute unit wave64: + +| GPU | Quantità | CU per GPU | Clock massimo | FP16 teorico per GPU | +|---|---:|---:|---:|---:| +| Radeon VII | 2 | 60 | 1.801 MHz | circa 27,66 TFLOPS | +| Radeon Pro VII | 3 | 60 | 1.700 MHz | circa 26,11 TFLOPS | +| GPU da 32 GB usata come coordinatore | 1 | 60 | 1.730 MHz | circa 26,57 TFLOPS | +| **Totale** | **6** | **360** | — | **circa 160,24 TFLOPS** | + +Il picco FP32 aggregato è circa 80,12 TFLOPS. Ogni scheda dispone inoltre di +circa 1 TB/s di banda HBM2, per un massimo nominale aggregato vicino a 6 TB/s. +I riferimenti AMD sono disponibili per +[Radeon VII](https://www.amd.com/en/support/downloads/drivers.html/graphics/radeon-rx/radeon-rx-vega-series/amd-radeon-vii.html), +[Radeon Pro VII](https://www.amd.com/en/products/specifications/professional-graphics.html) +e [Instinct MI50](https://www.amd.com/en/support/downloads/drivers.html/accelerators/instinct/instinct-mi-series/instinct-mi50.html). + +### Lavoro attivo del modello MoE + +Il GGUF descrive 284,33 miliardi di parametri logici, ma il modello seleziona +soltanto 6 dei 256 esperti routed per layer. Il lavoro minimo attivo per token +può essere stimato così: + +| Componente | Parametri logici attivi per token | +|---|---:| +| Pesi non-routed sempre utilizzati | circa 7,30 miliardi | +| Sei esperti routed selezionati | circa 6,49 miliardi | +| **Totale attivo** | **circa 13,80 miliardi** | + +Contando una moltiplicazione e un'addizione per peso, i matvec rappresentano +circa 27,6 GFLOP utili per token. Attenzione, normalizzazioni e riduzioni +aggiungono altro lavoro, ma non cambiano l'ordine di grandezza del confronto. + +### Utilizzo equivalente del picco FP16 + +| Indicatore | Prefill | Decode | +|---|---:|---:| +| Throughput normale | 72,80 token/s | circa 14,33 token/s | +| Calcolo utile equivalente | circa 2,01 TFLOPS | circa 0,40 TFLOPS | +| Percentuale dei 160,24 TFLOPS aggregati | **circa 1,25%** | **circa 0,25%** | +| Picco disponibile considerando la topologia | circa 91 TFLOPS | circa 26,7 TFLOPS | +| Utilizzo rispetto al picco disponibile | **circa 2,2%** | **circa 1,5%** | +| Distanza dal picco disponibile | **circa 45 volte** | **circa 67 volte** | + +Nel prefill la pipeline permette alle GPU di sovrapporre il lavoro, ma lo +sbilanciamento tra i 13 layer del coordinatore e i 6 layer di ogni worker lascia +i worker occupati soltanto per circa il 48-49% del tempo dello stadio più lento. +La capacità teorica realmente disponibile, pesata per questi periodi di +attività, è quindi circa 91 TFLOPS invece dei 160,24 TFLOPS nominali. + +Nel decode di una singola sessione le slice devono essere attraversate in +sequenza. In ogni istante lavora principalmente una sola GPU, quindi il tetto +istantaneo pertinente è circa 26-27 TFLOPS, non la somma da 160,24 TFLOPS. La +percentuale aggregata dello 0,25% descrive l'utilizzo dell'intero sistema, ma +l'1,5% rispetto alla GPU attiva è il confronto più significativo. + +Queste percentuali sono espresse in FLOP utili del modello. Non misurano +direttamente quante istruzioni intere o FP16 esegue l'hardware durante la +dequantizzazione e non devono essere interpretate come un possibile speedup +automatico di 45-67 volte. + +### Utilizzo equivalente della banda HBM nel decode + +Dal contenuto del GGUF risultano circa 8,20 GiB di pesi non-routed e 72,56 GiB +di pesi routed. Se vengono letti soltanto i sei esperti selezionati, il minimo +di pesi compressi da attraversare per token è: + +```text +8,20 GiB + 72,56 GiB * 6 / 256 = circa 9,90 GiB/token +``` + +A 14,33 token/s questo equivale a circa 142 GiB/s. Poiché durante il decode +single-stream è attiva principalmente una GPU alla volta, il confronto utile è +con circa 954 GiB/s della singola scheda, non con i 6 TB/s aggregati. + +| Indicatore decode | Valore | +|---|---:| +| Banda minima utile stimata | circa 142 GiB/s | +| Banda teorica della GPU attiva | circa 954 GiB/s | +| Percentuale del picco HBM | **circa 15%** | +| Distanza dal limite HBM ideale | **circa 6,7 volte** | +| Roofline puramente teorico | circa 94-96 token/s | + +Il roofline da 94-96 token/s è un limite matematico non raggiungibile nella +pratica: presuppone il 100% della banda HBM e nessun costo per dequantizzazione, +routing, cache, attenzione, sincronizzazione o lancio dei kernel. + +### Stiamo sfruttando tutta la capacità computazionale? + +**No.** I dati disponibili mostrano che né il prefill né il decode sono vicini +alla capacità teorica delle GPU: + +- il prefill realizza circa il 2,2% dei FLOP utili rispetto alla capacità + teorica pesata per il tempo in cui gli stadi possono essere attivi; +- il decode realizza circa l'1,5% dei FLOP utili della GPU attiva; +- il decode utilizza circa il 15% della banda HBM teorica; +- la pipeline single-stream impedisce al decode di utilizzare + contemporaneamente tutte e sei le GPU; +- il coordinatore sbilanciato limita il parallelismo del prefill. + +Il modello MoE riduce molto il numero di esperti calcolati, ma non spiega da +solo l'efficienza osservata: rimangono circa 13,80 miliardi di parametri attivi +e circa 9,90 GiB di pesi compressi da leggere per token. I risultati indicano +quindi margine nei kernel gfx906, nella fusione delle operazioni, nel numero di +launch, nell'accesso agli esperti e nell'orchestrazione della pipeline. + +Per attribuire con precisione la perdita residua servono contatori hardware +ROCm, in particolare: + +- banda HBM effettivamente letta e scritta; +- occupazione delle compute unit e wave attive; +- utilizzo VALU/FMA e istruzioni intere di dequantizzazione; +- hit rate delle cache L1/L2; +- durata aggregata per famiglia di kernel MoE, attention e quantized matvec; +- tempo tra un kernel e il successivo. + +Senza questi contatori è possibile dimostrare che la capacità non è sfruttata +completamente, ma non separare ancora con esattezza quanta parte del limite +derivi da HBM, dequantizzazione, occupazione, kernel launch o dipendenze della +pipeline. + +## Colli di bottiglia osservati + +1. **Prefill sbilanciato sul coordinatore.** I 13 layer locali richiedono circa + il doppio del tempo delle slice da 6 layer assegnate ai worker. +2. **Decode dominato dal calcolo GPU.** Il 96,63% del tempo è trascorso nei + transformer layer; ottimizzare TCP o la finestra distribuita può produrre + soltanto miglioramenti limitati nel decode. +3. **Worker `31:36` leggermente più lento.** La media è 12,00 ms per token e + sono stati osservati picchi di circa 15,2 ms. Conviene verificare clock, + temperatura e power cap della GPU corrispondente a ROCR 4. +4. **Costo della rete contenuto.** Trasporto e orchestrazione pesano circa + l'1,39% nel decode; i trasferimenti tra worker valgono circa lo 0,9% del + lavoro seriale del prefill. + +## Limiti della misura + +Queste percentuali descrivono la modalità residente e una posizione di +contesto di circa 2.400 token. Con un contesto già riempito da centinaia di +migliaia di token, il costo dell'attenzione aumenta e la ripartizione può +cambiare. + +La modalità di `run-context.sh` non è rappresentata da queste tabelle. Quando +si abilita lo streaming SSD degli esperti, lettura, caricamento e gestione della +cache degli esperti diventano componenti importanti e richiedono un profiling +separato. + +## Ottimizzazione della ripartizione per il profilo velocità + +Il profilo originale usava 13 layer sul coordinatore e 6 layer su ciascun +worker. A `ctx=32768` è possibile mantenere sette layer residenti su ogni GPU da +16 GB e usare una ripartizione più bilanciata: + +```text +coordinatore: 0:7 +worker: 8:14, 15:21, 22:28, 29:35, 36:42 +``` + +Il confronto è stato eseguito con lo stesso prompt da 2.418 token, 64 token di +output, chunk 64, finestra 5 e telemetria attiva. + +| Misura | Split 13 + 6x5 | Split 8 + 7x5 | Variazione | +|---|---:|---:|---:| +| Prefill | 72,80 token/s | 109,28 token/s | **+50,1%** | +| Tempo prefill | 33,22 s | 22,13 s | **-33,4%** | +| Decode profilato | 12,52 token/s | 12,61 token/s | **+0,7%** | +| Tempo totale 2.418+64 | 38,33 s | 27,21 s | **-29,0%** | + +La modifica risolve il principale sbilanciamento del prefill, ma non accelera +in modo sostanziale il decode, che continua ad attraversare le sei slice in +sequenza. Per questo `run-speed.sh` usa lo split 8 + 7x5 a 32K, mentre +`run-context.sh` conserva lo split 13 + 6x5 per lasciare spazio alla KV cache da +1M token e allo streaming SSD. + +## Valutazione preliminare del tensor parallelism + +RCCL 2.27 è installato, ma l'inizializzazione TP2 sulle coppie PCIe non è +utilizzabile con l'attuale configurazione di boot. RCCL segnala l'assenza di +`iommu=pt` come possibile causa di instabilità o blocco e il test P2P si arresta +durante l'inizializzazione. Anche il fallback con P2P disabilitato non è +affidabile e ha terminato con `SIGSEGV`. + +Il tensor parallelism non viene quindi integrato nel percorso di produzione +finché il sistema non viene avviato con una configurazione IOMMU compatibile e +un microbenchmark RCCL TP2 non completa correttamente. Il PP6 bilanciato offre +nel frattempo un miglioramento verificato senza modificare kernel, formato dei +pesi o semantica del modello. diff --git a/Makefile b/Makefile index 9711dc1a4..dcdf800c7 100644 --- a/Makefile +++ b/Makefile @@ -33,6 +33,10 @@ CPU_CORE_OBJS = ds4_cpu.o ds4_distributed.o ds4_ssd.o CUDA_LDLIBS ?= -lm -Xcompiler -pthread -L$(CUDA_HOME)/targets/sbsa-linux/lib -L$(CUDA_HOME)/lib64 -lcudart -lcublas HIPCC ?= $(shell command -v hipcc 2>/dev/null || echo /opt/rocm/bin/hipcc) ROCM_ARCH ?= gfx1151 +ROCM_CPPFLAGS := +ifeq ($(ROCM_ARCH),gfx906) +ROCM_CPPFLAGS += -DDS4_GFX906 +endif ROCM_CFLAGS ?= -O3 -ffast-math -g -fno-finite-math-only -pthread -D__HIP_PLATFORM_AMD__ -Wno-unused-command-line-argument --offload-arch=$(ROCM_ARCH) ROCM_LDLIBS ?= -lm -pthread -lhipblas -lhipblaslt DS4_LINK ?= $(NVCC) $(NVCCFLAGS) @@ -40,7 +44,7 @@ DS4_LINK_LIBS ?= $(CUDA_LDLIBS) METAL_LDLIBS := $(LDLIBS) endif -.PHONY: all help clean test cpu cuda cuda-spark cuda-generic cuda-regression strix-halo rocm +.PHONY: all help clean test cpu cuda cuda-spark cuda-generic cuda-regression strix-halo rocm gfx906 rocm-regression ifeq ($(UNAME_S),Darwin) all: ds4 ds4-server ds4-bench ds4-eval ds4-agent @@ -85,7 +89,9 @@ help: @echo " make cuda-generic Build CUDA for a generic local CUDA GPU" @echo " make cuda CUDA_ARCH=sm_N Build CUDA with an explicit nvcc -arch value" @echo " make strix-halo Build ROCm for Strix Halo / gfx1151" - @echo " make rocm Alias for make strix-halo" + @echo " make rocm ROCM_ARCH=gfxN Build ROCm for an explicit AMD GPU architecture" + @echo " make gfx906 Build ROCm for Vega 20 / gfx906" + @echo " make rocm-regression ROCM_ARCH=gfxN Build and run ROCm GPU regressions" @echo " make cpu Build CPU-only ./ds4, ./ds4-server, ./ds4-bench, ./ds4-eval, and ./ds4-agent" @echo " make test Build and run tests" @echo " make clean Remove build outputs" @@ -104,14 +110,17 @@ cuda: fi $(MAKE) -B ds4 ds4-server ds4-bench ds4-eval ds4-agent CUDA_ARCH="$(CUDA_ARCH)" -strix-halo: +rocm: $(MAKE) -B ds4 ds4-server ds4-bench ds4-eval ds4-agent \ CORE_OBJS="ds4.o ds4_distributed.o ds4_ssd.o ds4_rocm.o" \ CFLAGS="$(CFLAGS) -DDS4_ROCM_BUILD" \ - DS4_LINK="$(HIPCC) $(ROCM_CFLAGS)" \ + DS4_LINK="$(HIPCC) $(ROCM_CFLAGS) $(ROCM_CPPFLAGS)" \ DS4_LINK_LIBS="$(ROCM_LDLIBS)" +strix-halo: + $(MAKE) rocm ROCM_ARCH=gfx1151 -rocm: strix-halo +gfx906: + $(MAKE) rocm ROCM_ARCH=gfx906 ds4: ds4_cli.o ds4_help.o linenoise.o $(CORE_OBJS) $(DS4_LINK) -o $@ $^ $(DS4_LINK_LIBS) @@ -137,6 +146,13 @@ cpu: ds4_cli_cpu.o ds4_server_cpu.o ds4_bench_cpu.o ds4_eval_cpu.o ds4_agent_cpu cuda-regression: tests/cuda_long_context_smoke ./tests/cuda_long_context_smoke + +rocm-regression: tests/rocm_long_context_smoke + ./tests/rocm_long_context_smoke +ifeq ($(ROCM_ARCH),gfx906) + $(MAKE) tests/gfx906_wmma_test ROCM_ARCH=gfx906 + ./tests/gfx906_wmma_test +endif endif ds4.o: ds4.c ds4.h ds4_ssd.h ds4_distributed.h ds4_gpu.h @@ -212,11 +228,17 @@ ds4_cuda.o: ds4_cuda.cu ds4_gpu.h ds4_iq2_tables_cuda.inc $(NVCC) $(NVCCFLAGS) -c -o $@ ds4_cuda.cu ds4_rocm.o: ds4_rocm.cu ds4_gpu.h ds4_iq2_tables_cuda.inc $(ROCM_SRCS) - $(HIPCC) $(ROCM_CFLAGS) -c -o $@ ds4_rocm.cu + $(HIPCC) $(ROCM_CFLAGS) $(ROCM_CPPFLAGS) -c -o $@ ds4_rocm.cu tests/cuda_long_context_smoke: tests/cuda_long_context_smoke.o ds4_cuda.o $(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS) +tests/rocm_long_context_smoke: tests/cuda_long_context_smoke.o ds4_rocm.o + $(HIPCC) $(ROCM_CFLAGS) $(ROCM_CPPFLAGS) -o $@ $^ $(ROCM_LDLIBS) + +tests/gfx906_wmma_test: tests/shim_test.cu rocm/ds4_rocm_wmma_gfx906.cuh rocm/ds4_rocm_q8.cuh rocm/ds4_rocm_router.cuh + $(HIPCC) $(ROCM_CFLAGS) $(ROCM_CPPFLAGS) -I. -o $@ tests/shim_test.cu + ds4_test: ds4_test.o ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS) ifeq ($(UNAME_S),Darwin) $(CC) $(CFLAGS) -o $@ ds4_test.o ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS) $(METAL_LDLIBS) @@ -241,4 +263,4 @@ q4k-dot-test: tests/test_q4k_dot.c ./tests/test_q4k_dot clean: - rm -f ds4 ds4-server ds4-bench ds4-eval ds4-agent ds4_cpu ds4_native ds4_server_test ds4_test ds4_agent_test tests/test_q4k_dot *.o tests/cuda_long_context_smoke tests/cuda_long_context_smoke.o + rm -f ds4 ds4-server ds4-bench ds4-eval ds4-agent ds4_cpu ds4_native ds4_server_test ds4_test ds4_agent_test tests/test_q4k_dot *.o tests/cuda_long_context_smoke tests/rocm_long_context_smoke tests/gfx906_wmma_test tests/cuda_long_context_smoke.o diff --git a/README.md b/README.md index 785695284..ab580f21c 100644 --- a/README.md +++ b/README.md @@ -181,10 +181,14 @@ Q4 requires the larger-memory machine class, so M3 Max Q4 numbers are `N/A`. The normal Metal path tries to make the model resident in GPU-addressable memory. This is the fastest path and should remain your default when the model -fits. When it does not fit, DwarfStar also has a Metal-only **SSD streaming** -capacity mode. In this mode the non-routed model weights stay resident, while -routed MoE experts are kept in an in-memory cache and loaded from the GGUF file -on cache misses. +fits. When it does not fit, DwarfStar also has an **SSD streaming** capacity +mode. Backend- and topology-specific streaming paths must be qualified before +production use. Distributed ROCm layer slices use the bounded selected-address +expert path; their cache must hold at least one layer's expert table +(`n_expert` entries), otherwise startup rejects the configuration instead of +silently mapping full routed layers. In this mode the non-routed model weights +stay resident, while routed MoE experts are kept in an in-memory cache and +loaded from the GGUF file on cache misses. Streaming is not as fast as fitting the full model in RAM. It still needs memory for non-routed weights, KV cache, graph scratch, activations, and the routed @@ -215,6 +219,34 @@ then uses the rest for routed experts. Leave the hot expert preload enabled for normal use; use `--ssd-streaming-cold` and `--ssd-streaming-preload-experts N` only for measurements. +For distributed ROCm, pass the same streaming options to the coordinator and +every worker. DeepSeek V4 Flash has 256 experts and about 6.75 MiB per complete +IQ2 expert, so `2GB` is a practical minimum that enables selected-address +prefill while leaving headroom for KV and scratch buffers: + +```sh +--ssd-streaming --ssd-streaming-cache-experts 2GB +``` + +Larger caches improve the hit rate and generation speed. The fully resident +path remains substantially faster when the model and requested context fit. + +The six-gfx906 PP6 example includes two self-contained profiles. Edit +`MODEL_PATH` inside the selected file and run it directly; no launch-time +environment variables are required. The speed profile uses a balanced 8 + 7x5 +layer split at 32K context, while the capacity profile keeps the conservative +13 + 6x5 split needed by its larger KV cache: + +```sh +./run-speed.sh # resident weights, 32K context, maximum throughput +./run-context.sh # 2GB expert cache per GPU, native 1M context +``` + +`run.sh` itself is not tied to that machine: `COORD_DEVICE`, `COORD_LAYERS`, +and the space-separated `WORKER_SPECS` (`DEVICE,START:END`) describe an +arbitrary distributed gfx906 route. The two profile files are concrete +six-GPU examples whose settings can be edited in place. + ### Practical SSD streaming examples On 64GB MacBooks, start with the 2-bit Flash GGUF and a moderate expert cache: diff --git a/ds4.c b/ds4.c index 640511eb0..7a1d32dad 100644 --- a/ds4.c +++ b/ds4.c @@ -2280,7 +2280,8 @@ static bool accelerator_cache_model_tensors(ds4_backend backend, const ds4_model *m, const uint64_t *span_offsets, const uint64_t *span_sizes, - uint32_t span_count) { + uint32_t span_count, + bool spans_preinstalled) { if (backend != DS4_BACKEND_CUDA) return true; if (!m || !m->map || m->size == 0) return false; #ifndef DS4_ROCM_BUILD @@ -2291,8 +2292,13 @@ static bool accelerator_cache_model_tensors(ds4_backend backend, const double t0 = now_sec(); uint64_t prepared = 0; - if (!accelerator_prepare_model_tensor_spans(m, span_offsets, span_sizes, span_count, &prepared)) { - return false; + if (spans_preinstalled) { + for (uint32_t i = 0; i < span_count; i++) prepared += span_sizes[i]; + } else { + if (!accelerator_prepare_model_tensor_spans(m, span_offsets, span_sizes, + span_count, &prepared)) { + return false; + } } if (!accelerator_cache_q8_tensors(m, span_offsets, span_sizes, span_count)) return false; const double t1 = now_sec(); @@ -2302,8 +2308,11 @@ static bool accelerator_cache_model_tensors(ds4_backend backend, const char *accelerator_name = "CUDA"; #endif fprintf(stderr, - "ds4: %s startup model preparation covered %.2f GiB of tensor spans in %.3fs\n", - accelerator_name, (double)prepared / 1073741824.0, t1 - t0); + "ds4: %s startup model preparation %s %.2f GiB of tensor spans in %.3fs\n", + accelerator_name, + spans_preinstalled ? "reused preinstalled" : "covered", + (double)prepared / 1073741824.0, + t1 - t0); return true; } #else @@ -2311,12 +2320,14 @@ static bool accelerator_cache_model_tensors(ds4_backend backend, const ds4_model *m, const uint64_t *span_offsets, const uint64_t *span_sizes, - uint32_t span_count) { + uint32_t span_count, + bool spans_preinstalled) { (void)backend; (void)m; (void)span_offsets; (void)span_sizes; (void)span_count; + (void)spans_preinstalled; return true; } #endif @@ -10359,6 +10370,12 @@ typedef struct { uint32_t spec_prefix1_n_comp[DS4_MAX_LAYER]; uint32_t spec_prefix1_n_index_comp[DS4_MAX_LAYER]; bool spec_capture_prefix1; + /* Distributed graph instances only own persistent KV state for their + * local transformer slice. Non-distributed graphs use the full layer + * range. Keeping this range in the graph also makes reset/validation use + * exactly the same ownership policy as allocation. */ + uint32_t cache_layer_start; + uint32_t cache_layer_end; uint32_t raw_cap; /* Maximum compressed-row capacity across layers. Shared work buffers use * this worst-case size because ratio-4 indexer layers can still reach it. */ @@ -10760,13 +10777,17 @@ static bool metal_graph_apply_directional_steering_ffn( return metal_graph_apply_directional_steering(g, x, il, rows, g ? g->directional_steering_ffn_scale : 0.0f); } -static uint64_t metal_graph_kv_cache_bytes_for_context(uint32_t ctx_size, uint32_t raw_cap) { - uint64_t bytes = (uint64_t)DS4_N_LAYER * +static uint64_t metal_graph_kv_cache_bytes_for_context( + uint32_t ctx_size, + uint32_t raw_cap, + uint32_t layer_start, + uint32_t layer_end) { + uint64_t bytes = (uint64_t)(layer_end - layer_start + 1u) * raw_cap * DS4_N_HEAD_DIM * sizeof(float); - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + for (uint32_t il = layer_start; il <= layer_end; il++) { const uint32_t ratio = ds4_layer_compress_ratio(il); if (ratio == 0) continue; const uint64_t comp_cap = (uint64_t)(ctx_size / ratio + 2u); @@ -10783,16 +10804,19 @@ static uint64_t metal_graph_context_bytes_for_kv_policy( uint32_t ctx_size, uint32_t raw_cap, uint32_t prefill_cap, + uint32_t layer_start, + uint32_t layer_end, uint64_t *kv_cache_bytes_out) { uint32_t min_ratio = UINT32_MAX; - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + for (uint32_t il = layer_start; il <= layer_end; il++) { const uint32_t ratio = ds4_layer_compress_ratio(il); if (ratio != 0 && ratio < min_ratio) min_ratio = ratio; } if (min_ratio == UINT32_MAX) min_ratio = ctx_size ? ctx_size : 1u; uint64_t comp_cap = (uint64_t)(ctx_size / min_ratio + 2u); if (comp_cap < 2u) comp_cap = 2u; - const uint64_t kv_cache_bytes = metal_graph_kv_cache_bytes_for_context(ctx_size, raw_cap); + const uint64_t kv_cache_bytes = metal_graph_kv_cache_bytes_for_context( + ctx_size, raw_cap, layer_start, layer_end); if (kv_cache_bytes_out) *kv_cache_bytes_out = kv_cache_bytes; uint64_t bytes = kv_cache_bytes + 2ull * comp_cap * prefill_cap * sizeof(float); @@ -10961,9 +10985,17 @@ static bool metal_graph_alloc_raw_cap( uint32_t raw_cap, uint32_t ctx_size, uint32_t prefill_cap, - bool enable_mtp) { + bool enable_mtp, + uint32_t cache_layer_start, + uint32_t cache_layer_end) { memset(g, 0, sizeof(*g)); + if (cache_layer_start > cache_layer_end || + cache_layer_end >= (uint32_t)DS4_N_LAYER) { + return false; + } g->mtp_enabled = enable_mtp; + g->cache_layer_start = cache_layer_start; + g->cache_layer_end = cache_layer_end; if (raw_cap == 0) raw_cap = 1; if (ctx_size == 0) ctx_size = raw_cap; if (prefill_cap == 0) prefill_cap = 1; @@ -10977,7 +11009,7 @@ static bool metal_graph_alloc_raw_cap( g->raw_window = raw_window; g->prefill_cap = prefill_cap; uint32_t min_ratio = UINT32_MAX; - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + for (uint32_t il = cache_layer_start; il <= cache_layer_end; il++) { const uint32_t ratio = ds4_layer_compress_ratio(il); if (ratio != 0 && ratio < min_ratio) min_ratio = ratio; } @@ -11014,7 +11046,9 @@ static bool metal_graph_alloc_raw_cap( const uint64_t pc = prefill_cap; uint64_t kv_cache_bytes = 0; const uint64_t context_bytes = - metal_graph_context_bytes_for_kv_policy(ctx_size, raw_cap, prefill_cap, &kv_cache_bytes); + metal_graph_context_bytes_for_kv_policy(ctx_size, raw_cap, prefill_cap, + cache_layer_start, cache_layer_end, + &kv_cache_bytes); const bool managed_kv_cache = ds4_gpu_should_use_managed_kv_cache(kv_cache_bytes, context_bytes) != 0; if (managed_kv_cache) { @@ -11054,7 +11088,7 @@ static bool metal_graph_alloc_raw_cap( g->kv_raw = ds4_gpu_tensor_alloc((uint64_t)DS4_N_HEAD_DIM * sizeof(float)); g->kv = ds4_gpu_tensor_alloc((uint64_t)DS4_N_HEAD_DIM * sizeof(float)); bool state_init_ok = true; - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + for (uint32_t il = cache_layer_start; il <= cache_layer_end; il++) { g->layer_raw_cache[il] = metal_graph_alloc_kv_cache_tensor( managed_kv_cache, (uint64_t)raw_cap * DS4_N_HEAD_DIM * sizeof(float)); @@ -11216,7 +11250,9 @@ static bool metal_graph_alloc_raw_cap( g->batch_routed_out = ds4_gpu_tensor_alloc(pc * DS4_N_EMBD * sizeof(float)); bool layer_cache_ok = true; - for (uint32_t il = 0; layer_cache_ok && il < DS4_N_LAYER; il++) { + for (uint32_t il = cache_layer_start; + layer_cache_ok && il <= cache_layer_end; + il++) { layer_cache_ok = g->layer_raw_cache[il] != NULL; const uint32_t ratio = ds4_layer_compress_ratio(il); if (layer_cache_ok && ratio != 0) { @@ -11292,7 +11328,9 @@ static bool metal_graph_alloc( ds4_gpu_graph *g, const ds4_weights *weights, const ds4_layer_weights *layer) { - return metal_graph_alloc_raw_cap(g, weights, layer, DS4_N_SWA, DS4_N_SWA, 1, false); + return metal_graph_alloc_raw_cap(g, weights, layer, + DS4_N_SWA, DS4_N_SWA, 1, false, + 0, DS4_N_LAYER - 1u); } static bool metal_graph_install_model_spans( @@ -11600,10 +11638,12 @@ static bool metal_graph_stream_prefill_batch_selected_addr_enabled( const ds4_gpu_graph *g, const ds4_weights *weights, uint32_t n_tokens) { + const ds4_layer_weights *shape_layer = + weights ? weights_first_bound_layer(weights) : NULL; if (!g || !g->ssd_streaming || g->quality || - !weights || + !shape_layer || n_tokens <= 1 || getenv("DS4_METAL_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") != NULL || getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") != NULL || @@ -11611,9 +11651,12 @@ static bool metal_graph_stream_prefill_batch_selected_addr_enabled( getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") != NULL || DS4_N_LAYER == 0 || DS4_N_EXPERT_USED != 6 || - weights->layer[0].ffn_gate_exps->type != DS4_TENSOR_IQ2_XXS || - weights->layer[0].ffn_up_exps->type != DS4_TENSOR_IQ2_XXS || - weights->layer[0].ffn_down_exps->type != DS4_TENSOR_Q2_K) { + !shape_layer->ffn_gate_exps || + !shape_layer->ffn_up_exps || + !shape_layer->ffn_down_exps || + shape_layer->ffn_gate_exps->type != DS4_TENSOR_IQ2_XXS || + shape_layer->ffn_up_exps->type != DS4_TENSOR_IQ2_XXS || + shape_layer->ffn_down_exps->type != DS4_TENSOR_Q2_K) { return false; } @@ -20272,7 +20315,7 @@ static bool metal_graph_reset_prefill_state(ds4_gpu_graph *g) { memset(g->layer_n_comp, 0, sizeof(g->layer_n_comp)); memset(g->layer_n_index_comp, 0, sizeof(g->layer_n_index_comp)); g->mtp_n_raw = 0; - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + for (uint32_t il = g->cache_layer_start; il <= g->cache_layer_end; il++) { const uint32_t ratio = ds4_layer_compress_ratio(il); if (ratio == 0) continue; const uint32_t coff = ratio == 4 ? 2u : 1u; @@ -21510,7 +21553,8 @@ static int metal_graph_prompt_logits_test( ds4_gpu_graph g; bool ok = metal_graph_alloc_raw_cap(&g, weights, &weights->layer[0], - raw_cap, (uint32_t)ctx_size, (uint32_t)n_test, false); + raw_cap, (uint32_t)ctx_size, (uint32_t)n_test, false, + 0, DS4_N_LAYER - 1u); if (!ok) { metal_graph_free(&g); fprintf(stderr, "ds4: failed to initialize Metal graph prompt test runtime\n"); @@ -22956,7 +23000,8 @@ static int generate_metal_graph_raw_swa( } ds4_gpu_graph g; bool ok = metal_graph_alloc_raw_cap(&g, weights, &weights->layer[0], - raw_cap, (uint32_t)ctx_size, prefill_cap, false); + raw_cap, (uint32_t)ctx_size, prefill_cap, false, + 0, DS4_N_LAYER - 1u); if (!ok) { fprintf(stderr, "ds4: failed to allocate GPU graph runtime\n"); return 1; @@ -25023,7 +25068,8 @@ int ds4_engine_collect_imatrix(ds4_engine *e, ds4_gpu_graph g; bool ok = metal_graph_alloc_raw_cap(&g, weights, &weights->layer[0], - raw_cap, (uint32_t)ctx_size, prefill_cap, false); + raw_cap, (uint32_t)ctx_size, prefill_cap, false, + 0, DS4_N_LAYER - 1u); if (!ok) { fprintf(stderr, "ds4: failed to allocate imatrix Metal graph runtime\n"); free(dataset); @@ -25728,6 +25774,30 @@ int ds4_engine_open(ds4_engine **out, const ds4_engine_options *opt) { *out = NULL; return 1; } +#ifdef DS4_ROCM_BUILD + /* Distributed ROCm prefill uses a selected-address table to keep + * routed experts out of the ordinary model-span cache. The table is + * enabled only when the resident expert budget can cover one complete + * layer; below that threshold a worker would fall back to mapping a + * full routed layer and defeat the bounded-cache memory policy. */ + if (e->ssd_streaming && + e->distributed.role != DS4_DISTRIBUTED_NONE && + e->ssd_streaming_cache_experts < DS4_N_EXPERT) { + uint64_t per_expert_bytes = 0; + (void)ds4_streaming_routed_expert_bytes(&e->weights, + &per_expert_bytes); + fprintf(stderr, + "ds4: distributed ROCm SSD streaming requires a cache of " + "at least %u experts (%.2f GiB for this model); got %u\n", + (uint32_t)DS4_N_EXPERT, + (double)((uint64_t)DS4_N_EXPERT * per_expert_bytes) / + 1073741824.0, + e->ssd_streaming_cache_experts); + ds4_engine_close(e); + *out = NULL; + return 1; + } +#endif ds4_gpu_set_streaming_expert_cache_budget(e->ssd_streaming_cache_experts); if (e->ssd_streaming) { /* @@ -25932,7 +26002,8 @@ int ds4_engine_open(ds4_engine **out, const ds4_engine_options *opt) { (void)ds4_gpu_set_model_fd_for_map(e->model.fd, e->model.map); if (!accelerator_cache_model_tensors(e->backend, &e->model, load_offsets, load_sizes, - load_span_count)) { + load_span_count, + e->ssd_streaming)) { fprintf(stderr, "ds4: %s failed to prepare optional model cache\n", ds4_backend_name(e->backend)); free(load_offsets); @@ -25948,7 +26019,8 @@ int ds4_engine_open(ds4_engine **out, const ds4_engine_options *opt) { if (e->mtp_ready) { (void)ds4_gpu_set_model_fd_for_map(e->mtp_model.fd, e->mtp_model.map); if (!accelerator_cache_model_tensors(e->backend, &e->mtp_model, - NULL, NULL, 0)) { + NULL, NULL, 0, + e->ssd_streaming)) { fprintf(stderr, "ds4: %s failed to prepare optional MTP model cache\n", ds4_backend_name(e->backend)); ds4_engine_close(e); @@ -26071,8 +26143,16 @@ int ds4_session_create(ds4_session **out, ds4_engine *e, int ctx_size) { free(s); return 1; } + uint32_t cache_layer_start = 0; + uint32_t cache_layer_end = DS4_N_LAYER - 1u; + if (e->distributed.role != DS4_DISTRIBUTED_NONE && + e->distributed.layers.set) { + cache_layer_start = e->distributed.layers.start; + cache_layer_end = e->distributed.layers.end; + } if (!metal_graph_alloc_raw_cap(&s->graph, &e->weights, shape_layer, - raw_cap, (uint32_t)ctx_size, s->prefill_cap, e->mtp_ready)) + raw_cap, (uint32_t)ctx_size, s->prefill_cap, + e->mtp_ready, cache_layer_start, cache_layer_end)) { free(s); return 1; @@ -26584,11 +26664,26 @@ int ds4_session_eval_layer_slice(ds4_session *s, ds4_gpu_tensor *last_hc = NULL; ds4_gpu_tensor *saved_cur = NULL; + /* Selected-address expert streaming is layer-number agnostic: cache keys + * and address tables carry the absolute layer id. Distributed workers + * therefore use the same bounded-cache path as a coordinator whose slice + * starts at zero, instead of mapping every routed tensor in their slice. */ const bool batch_selected_addr = g->ssd_streaming && - layer_start == 0 && (metal_graph_stream_prefill_batch_selected_addr_enabled(g, &e->weights, n_tokens) || metal_graph_cuda_stream_prefill_batch_selected_addr_enabled(g, &e->weights, n_tokens)); + if (g->ssd_streaming && getenv("DS4_ROCM_STREAM_CACHE_STATS") != NULL) { + fprintf(stderr, + "ds4: ROCm layer-slice streaming map mode=%s layers=%u:%u tokens=%u " + "cache_experts=%u quality=%d input_hc=%d\n", + batch_selected_addr ? "selected-address" : "full-layer", + layer_start, + layer_end, + n_tokens, + ds4_gpu_stream_expert_cache_configured_count(), + g->quality ? 1 : 0, + input_hc ? 1 : 0); + } if (g->ssd_streaming) { for (uint32_t il = layer_start; ok && il <= layer_end; il++) { g->streaming_static_decode_map_current = false; diff --git a/ds4_rocm.h b/ds4_rocm.h index 50481b8d5..f983e362c 100644 --- a/ds4_rocm.h +++ b/ds4_rocm.h @@ -4,8 +4,15 @@ #include #include #include +#if defined(DS4_GFX906) +// gfx906 (Vega 20) lacks MFMA matrix cores and rocWMMA refuses to compile for +// it (static_assert "Unsupported architecture"). Use our device-side shim that +// emulates the small rocWMMA subset ds4 uses via __shfl width=32 wave32 lanes. +#include "rocm/ds4_rocm_wmma_gfx906.cuh" +#else #include #include +#endif #define cudaError_t hipError_t #define cudaStream_t hipStream_t diff --git a/ds4_server.c b/ds4_server.c index 34a9d5084..3b59ef9a3 100644 --- a/ds4_server.c +++ b/ds4_server.c @@ -1175,13 +1175,273 @@ static bool parse_tool_calls_value(const char **p, tool_calls *calls) { return true; } +static void json_escape(buf *b, const char *s); + +typedef struct { + char *key; + char *value; +} json_canonical_member; + +static int json_canonical_member_cmp(const void *a, const void *b) { + const json_canonical_member *ma = a; + const json_canonical_member *mb = b; + int cmp = strcmp(ma->key, mb->key); + return cmp ? cmp : strcmp(ma->value, mb->value); +} + +static void json_canonical_members_free(json_canonical_member *members, int len) { + for (int i = 0; i < len; i++) { + free(members[i].key); + free(members[i].value); + } + free(members); +} + +static bool json_canonical_value(const char **p, buf *out, int depth, + bool preserve_object_order); + +static bool json_canonical_object(const char **p, buf *out, int depth, + bool preserve_order) { + if (depth >= JSON_MAX_NESTING) return false; + json_ws(p); + if (**p != '{') return false; + (*p)++; + + json_canonical_member *members = NULL; + int len = 0, cap = 0; + json_ws(p); + while (**p && **p != '}') { + char *key = NULL; + if (!json_string(p, &key)) goto bad; + json_ws(p); + if (**p != ':') { + free(key); + goto bad; + } + (*p)++; + + buf value = {0}; + const bool property_map = !strcmp(key, "properties"); + if (!json_canonical_value(p, &value, depth + 1, property_map)) { + free(key); + buf_free(&value); + goto bad; + } + for (int i = 0; i < len; i++) { + if (!strcmp(members[i].key, key)) { + free(key); + buf_free(&value); + goto bad; + } + } + if (len == cap) { + cap = cap ? cap * 2 : 8; + members = xrealloc(members, (size_t)cap * sizeof(members[0])); + } + members[len++] = (json_canonical_member){ + .key = key, + .value = buf_take(&value), + }; + + json_ws(p); + if (**p == ',') { + (*p)++; + json_ws(p); + continue; + } + if (**p != '}') goto bad; + } + if (**p != '}') goto bad; + (*p)++; + + if (!preserve_order && len > 1) + qsort(members, (size_t)len, sizeof(members[0]), + json_canonical_member_cmp); + buf_putc(out, '{'); + for (int i = 0; i < len; i++) { + if (i) buf_putc(out, ','); + json_escape(out, members[i].key); + buf_putc(out, ':'); + buf_puts(out, members[i].value); + } + buf_putc(out, '}'); + json_canonical_members_free(members, len); + return true; + +bad: + json_canonical_members_free(members, len); + return false; +} + +static bool json_canonical_array(const char **p, buf *out, int depth) { + if (depth >= JSON_MAX_NESTING) return false; + json_ws(p); + if (**p != '[') return false; + (*p)++; + buf_putc(out, '['); + bool wrote = false; + json_ws(p); + while (**p && **p != ']') { + if (wrote) buf_putc(out, ','); + if (!json_canonical_value(p, out, depth + 1, false)) return false; + wrote = true; + json_ws(p); + if (**p == ',') { + (*p)++; + json_ws(p); + continue; + } + if (**p != ']') return false; + } + if (**p != ']') return false; + (*p)++; + buf_putc(out, ']'); + return true; +} + +static bool json_canonical_value(const char **p, buf *out, int depth, + bool preserve_object_order) { + json_ws(p); + if (**p == '{') + return json_canonical_object(p, out, depth, preserve_object_order); + if (**p == '[') return json_canonical_array(p, out, depth); + if (**p == '"') { + char *s = NULL; + if (!json_string(p, &s)) return false; + json_escape(out, s); + free(s); + return true; + } + if (json_lit(p, "true")) { + buf_puts(out, "true"); + return true; + } + if (json_lit(p, "false")) { + buf_puts(out, "false"); + return true; + } + if (json_lit(p, "null")) { + buf_puts(out, "null"); + return true; + } + + const char *start = *p; + double number = 0.0; + if (!json_number(p, &number)) return false; + buf_append(out, start, (size_t)(*p - start)); + return true; +} + +static char *json_canonicalize(const char *json) { + const char *p = json ? json : "null"; + buf out = {0}; + if (!json_canonical_value(&p, &out, 0, false)) { + buf_free(&out); + return NULL; + } + json_ws(&p); + if (*p) { + buf_free(&out); + return NULL; + } + return buf_take(&out); +} + static void append_raw_json_line(buf *b, const char *json) { if (!json || !json[0]) return; if (b->len) buf_putc(b, '\n'); buf_puts(b, json); } -static void json_escape(buf *b, const char *s); +typedef struct { + char *name; + char *json; +} tool_schema_line; + +static char *tool_schema_json_name(const char *json) { + const char *p = json ? json : ""; + json_ws(&p); + if (*p != '{') return NULL; + p++; + json_ws(&p); + while (*p && *p != '}') { + char *key = NULL; + if (!json_string(&p, &key)) return NULL; + json_ws(&p); + if (*p != ':') { + free(key); + return NULL; + } + p++; + if (!strcmp(key, "name")) { + char *name = NULL; + free(key); + return json_string(&p, &name) ? name : NULL; + } + free(key); + if (!json_skip_value(&p)) return NULL; + json_ws(&p); + if (*p == ',') p++; + json_ws(&p); + } + return NULL; +} + +static int tool_schema_line_cmp(const void *a, const void *b) { + const tool_schema_line *sa = a; + const tool_schema_line *sb = b; + int cmp = strcmp(sa->name ? sa->name : "", sb->name ? sb->name : ""); + return cmp ? cmp : strcmp(sa->json, sb->json); +} + +static char *tool_schema_lines_sorted(const char *schemas) { + if (!schemas || !schemas[0]) return xstrdup(""); + char *copy = xstrdup(schemas); + tool_schema_line *lines = NULL; + int len = 0, cap = 0; + char *save = NULL; + for (char *line = strtok_r(copy, "\n", &save); + line; + line = strtok_r(NULL, "\n", &save)) + { + if (len == cap) { + cap = cap ? cap * 2 : 8; + lines = xrealloc(lines, (size_t)cap * sizeof(lines[0])); + } + lines[len++] = (tool_schema_line){ + .name = tool_schema_json_name(line), + .json = xstrdup(line), + }; + } + /* Duplicate or missing names are malformed/ambiguous tool sets. Keep + * their input order so canonicalization cannot silently change which + * duplicate schema supplies the argument order used for DSML replay. */ + bool sortable = true; + for (int i = 0; i < len && sortable; i++) { + if (!lines[i].name || !lines[i].name[0]) { + sortable = false; + break; + } + for (int j = 0; j < i; j++) { + if (!strcmp(lines[i].name, lines[j].name)) { + sortable = false; + break; + } + } + } + if (sortable && len > 1) + qsort(lines, (size_t)len, sizeof(lines[0]), tool_schema_line_cmp); + + buf out = {0}; + for (int i = 0; i < len; i++) { + append_raw_json_line(&out, lines[i].json); + free(lines[i].name); + free(lines[i].json); + } + free(lines); + free(copy); + return buf_take(&out); +} static char *openai_function_schema_from_tool(const char *raw) { const char *p = raw; @@ -1538,8 +1798,11 @@ static bool append_responses_namespace_tool_schemas(buf *schemas, char *schema = responses_namespace_function_schema_from_tool(tool_raw, name, &wire_name); if (schema) { - append_raw_json_line(schemas, schema); - tool_schema_orders_add_json_wire(orders, schema, name, wire_name, false); + char *canonical = json_canonicalize(schema); + const char *stable = canonical ? canonical : schema; + append_raw_json_line(schemas, stable); + tool_schema_orders_add_json_wire(orders, stable, name, wire_name, false); + free(canonical); appended = true; } free(schema); @@ -1579,17 +1842,26 @@ static bool parse_tools_value(const char **p, char **out, tool_schema_orders *or if (!json_raw_value(p, &raw)) goto bad; char *function = openai_function_schema_from_tool(raw); if (function) { - append_raw_json_line(&schemas, function); - tool_schema_orders_add_json(orders, function); + char *canonical = json_canonicalize(function); + const char *stable = canonical ? canonical : function; + append_raw_json_line(&schemas, stable); + tool_schema_orders_add_json(orders, stable); + free(canonical); } else if (!append_responses_namespace_tool_schemas(&schemas, orders, raw)) { char *special = responses_special_schema_from_tool(raw); if (special) { - append_raw_json_line(&schemas, special); - tool_schema_orders_add_json_wire(orders, special, + char *canonical = json_canonicalize(special); + const char *stable = canonical ? canonical : special; + append_raw_json_line(&schemas, stable); + tool_schema_orders_add_json_wire(orders, stable, NULL, NULL, true); + free(canonical); } else { - append_raw_json_line(&schemas, raw); - tool_schema_orders_add_json(orders, raw); + char *canonical = json_canonicalize(raw); + const char *stable = canonical ? canonical : raw; + append_raw_json_line(&schemas, stable); + tool_schema_orders_add_json(orders, stable); + free(canonical); } free(special); } @@ -1601,7 +1873,9 @@ static bool parse_tools_value(const char **p, char **out, tool_schema_orders *or } if (**p != ']') goto bad; (*p)++; - *out = buf_take(&schemas); + char *joined = buf_take(&schemas); + *out = tool_schema_lines_sorted(joined); + free(joined); return true; bad: buf_free(&schemas); @@ -11887,6 +12161,142 @@ static void test_tool_schema_order_from_openai_tools(void) { tool_schema_orders_free(&orders); } +static void test_tool_schema_list_order_is_canonical(void) { + const char *ab = + "[{\"type\":\"function\",\"function\":{\"name\":\"read\"," + "\"description\":\"read\",\"parameters\":{\"type\":\"object\"}}}," + "{\"type\":\"function\",\"function\":{\"name\":\"stat\"," + "\"description\":\"stat\",\"parameters\":{\"type\":\"object\"}}}]"; + const char *ba = + "[{\"type\":\"function\",\"function\":{\"name\":\"stat\"," + "\"description\":\"stat\",\"parameters\":{\"type\":\"object\"}}}," + "{\"type\":\"function\",\"function\":{\"name\":\"read\"," + "\"description\":\"read\",\"parameters\":{\"type\":\"object\"}}}]"; + const char *pa = ab; + const char *pb = ba; + char *sa = NULL, *sb = NULL; + tool_schema_orders oa = {0}, ob = {0}; + TEST_ASSERT(parse_tools_value(&pa, &sa, &oa)); + TEST_ASSERT(parse_tools_value(&pb, &sb, &ob)); + TEST_ASSERT(sa && sb && !strcmp(sa, sb)); + const char *read = sa ? strstr(sa, "\"name\":\"read\"") : NULL; + const char *stat = sa ? strstr(sa, "\"name\":\"stat\"") : NULL; + TEST_ASSERT(read && stat && read < stat); + free(sa); + free(sb); + tool_schema_orders_free(&oa); + tool_schema_orders_free(&ob); +} + +static void test_tool_schema_json_key_order_is_canonical(void) { + const char *normal = + "[{\"type\":\"function\",\"function\":{\"name\":\"read\"," + "\"description\":\"read one file\",\"parameters\":{" + "\"type\":\"object\",\"properties\":{\"path\":{\"type\":\"string\"}}}}}]"; + const char *rekeyed = + "[{\"function\":{\"parameters\":{\"properties\":{" + "\"path\":{\"type\":\"string\"}},\"type\":\"object\"}," + "\"description\":\"read one file\",\"name\":\"read\"}," + "\"type\":\"function\"}]"; + const char *pa = normal; + const char *pb = rekeyed; + char *sa = NULL, *sb = NULL; + tool_schema_orders oa = {0}, ob = {0}; + TEST_ASSERT(parse_tools_value(&pa, &sa, &oa)); + TEST_ASSERT(parse_tools_value(&pb, &sb, &ob)); + TEST_ASSERT(sa && sb && !strcmp(sa, sb)); + const tool_schema_order *order = tool_schema_orders_find(&ob, "read"); + TEST_ASSERT(order && order->len == 1 && !strcmp(order->prop[0], "path")); + free(sa); + free(sb); + tool_schema_orders_free(&oa); + tool_schema_orders_free(&ob); +} + +static void test_tool_schema_semantic_changes_remain_distinct(void) { + const char *one = + "[{\"type\":\"function\",\"function\":{\"name\":\"read\"," + "\"description\":\"read one file\",\"parameters\":{\"type\":\"object\"}}}]"; + const char *two = + "[{\"type\":\"function\",\"function\":{\"name\":\"read\"," + "\"description\":\"read two files\",\"parameters\":{\"type\":\"object\"}}}]"; + const char *pa = one; + const char *pb = two; + char *sa = NULL, *sb = NULL; + tool_schema_orders oa = {0}, ob = {0}; + TEST_ASSERT(parse_tools_value(&pa, &sa, &oa)); + TEST_ASSERT(parse_tools_value(&pb, &sb, &ob)); + TEST_ASSERT(sa && sb && strcmp(sa, sb)); + free(sa); + free(sb); + tool_schema_orders_free(&oa); + tool_schema_orders_free(&ob); +} + +static void test_tool_schema_property_order_is_preserved(void) { + const char *path_first = + "[{\"type\":\"function\",\"function\":{\"name\":\"read\"," + "\"parameters\":{\"type\":\"object\",\"properties\":{" + "\"path\":{\"type\":\"string\"},\"encoding\":{\"type\":\"string\"}}}}}]"; + const char *encoding_first = + "[{\"type\":\"function\",\"function\":{\"name\":\"read\"," + "\"parameters\":{\"type\":\"object\",\"properties\":{" + "\"encoding\":{\"type\":\"string\"},\"path\":{\"type\":\"string\"}}}}}]"; + const char *pa = path_first; + const char *pb = encoding_first; + char *sa = NULL, *sb = NULL; + tool_schema_orders oa = {0}, ob = {0}; + TEST_ASSERT(parse_tools_value(&pa, &sa, &oa)); + TEST_ASSERT(parse_tools_value(&pb, &sb, &ob)); + TEST_ASSERT(sa && sb && strcmp(sa, sb)); + const tool_schema_order *a = tool_schema_orders_find(&oa, "read"); + const tool_schema_order *b = tool_schema_orders_find(&ob, "read"); + TEST_ASSERT(a && a->len == 2 && !strcmp(a->prop[0], "path")); + TEST_ASSERT(b && b->len == 2 && !strcmp(b->prop[0], "encoding")); + free(sa); + free(sb); + tool_schema_orders_free(&oa); + tool_schema_orders_free(&ob); +} + +static void test_duplicate_tool_names_are_not_reordered(void) { + const char *ab = + "[{\"name\":\"same\",\"description\":\"a\"}," + "{\"name\":\"same\",\"description\":\"b\"}]"; + const char *ba = + "[{\"name\":\"same\",\"description\":\"b\"}," + "{\"name\":\"same\",\"description\":\"a\"}]"; + const char *pa = ab; + const char *pb = ba; + char *sa = NULL, *sb = NULL; + tool_schema_orders oa = {0}, ob = {0}; + TEST_ASSERT(parse_tools_value(&pa, &sa, &oa)); + TEST_ASSERT(parse_tools_value(&pb, &sb, &ob)); + TEST_ASSERT(sa && sb && strcmp(sa, sb)); + free(sa); + free(sb); + tool_schema_orders_free(&oa); + tool_schema_orders_free(&ob); +} + +static void test_duplicate_json_keys_are_not_canonicalized(void) { + const char *ab = + "[{\"name\":\"same\",\"description\":\"a\",\"description\":\"b\"}]"; + const char *ba = + "[{\"name\":\"same\",\"description\":\"b\",\"description\":\"a\"}]"; + const char *pa = ab; + const char *pb = ba; + char *sa = NULL, *sb = NULL; + tool_schema_orders oa = {0}, ob = {0}; + TEST_ASSERT(parse_tools_value(&pa, &sa, &oa)); + TEST_ASSERT(parse_tools_value(&pb, &sb, &ob)); + TEST_ASSERT(sa && sb && strcmp(sa, sb)); + free(sa); + free(sb); + tool_schema_orders_free(&oa); + tool_schema_orders_free(&ob); +} + static void test_tool_schema_order_from_responses_tool_search(void) { const char *json = "[{\"type\":\"tool_search\",\"execution\":\"client\"," @@ -15768,6 +16178,12 @@ static void ds4_server_unit_tests_run(void) { test_render_chat_prompt_text_renders_tools_before_system(); test_tool_schema_order_from_anthropic_schema(); test_tool_schema_order_from_openai_tools(); + test_tool_schema_list_order_is_canonical(); + test_tool_schema_json_key_order_is_canonical(); + test_tool_schema_semantic_changes_remain_distinct(); + test_tool_schema_property_order_is_preserved(); + test_duplicate_tool_names_are_not_reordered(); + test_duplicate_json_keys_are_not_canonicalized(); test_tool_schema_order_from_responses_tool_search(); test_responses_function_named_tool_search_stays_function_call(); test_responses_namespace_tool_schemas_restore_wire_namespace(); diff --git a/rocm/ds4_rocm_attention.cuh b/rocm/ds4_rocm_attention.cuh index 0d688326b..67f17a410 100644 --- a/rocm/ds4_rocm_attention.cuh +++ b/rocm/ds4_rocm_attention.cuh @@ -445,6 +445,28 @@ __device__ __forceinline__ static float attention_dot_f32_vec4_oldhip(const floa return s; } +__device__ __forceinline__ static float attention_dot_f32_vec4_coop_oldhip_w32( + const float *a, const float *b, uint32_t n, uint32_t lane) { + float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f; + const uint32_t n4 = n >> 2u; + const float4 *a4 = (const float4 *)a; + const float4 *b4 = (const float4 *)b; + for (uint32_t i = lane; i < n4; i += 32u) { + const float4 av = a4[i]; + const float4 bv = b4[i]; + s0 += av.x * bv.x; + s1 += av.y * bv.y; + s2 += av.z * bv.z; + s3 += av.w * bv.w; + } + s0 = attention_warp_sum_oldhip_w32(s0); + s1 = attention_warp_sum_oldhip_w32(s1); + s2 = attention_warp_sum_oldhip_w32(s2); + s3 = attention_warp_sum_oldhip_w32(s3); + return (s0 + s1) + (s2 + s3); +} + +template __global__ static void attention_decode_mixed_one_fast_oldhip_kernel( float *heads, const float *q, @@ -469,30 +491,63 @@ __global__ static void attention_decode_mixed_one_fast_oldhip_kernel( const float scale = rsqrtf((float)head_dim); float local_max = sinks[h]; - for (uint32_t r = tid; r < n_raw; r += blockDim.x) { - const uint32_t row = raw_cap ? ((raw_start + r) % raw_cap) : r; - const float *kv = raw_kv + (uint64_t)row * head_dim; - float s = use_vec4 ? attention_dot_f32_vec4_oldhip(qh, kv, head_dim) : 0.0f; - if (!use_vec4) { - for (uint32_t i = 0; i < head_dim; i++) s += qh[i] * kv[i]; + if (COOPERATIVE_DOT) { + const uint32_t lane = tid & 31u; + const uint32_t group = tid >> 5u; + const uint32_t groups = blockDim.x >> 5u; + for (uint32_t rr = group; rr < n_rows; rr += groups) { + const bool raw = rr < n_raw; + const uint32_t logical_row = raw ? rr : rr - n_raw; + const bool valid = raw || !(use_mask && comp_mask && + comp_mask[logical_row] <= -5.0e29f); + float s = -3.4e38f; + if (valid) { + const uint32_t row = raw && raw_cap + ? ((raw_start + logical_row) % raw_cap) : logical_row; + const float *kv = (raw ? raw_kv : comp_kv) + + (uint64_t)row * head_dim; + float dot = use_vec4 + ? attention_dot_f32_vec4_coop_oldhip_w32(qh, kv, head_dim, lane) + : 0.0f; + if (!use_vec4) { + for (uint32_t i = lane; i < head_dim; i += 32u) + dot += qh[i] * kv[i]; + dot = attention_warp_sum_oldhip_w32(dot); + } + s = dot * scale; + if (!raw && use_mask && comp_mask) s += comp_mask[logical_row]; + } + if (lane == 0u) { + scores[rr] = s; + local_max = fmaxf(local_max, s); + } } - s *= scale; - scores[r] = s; - local_max = fmaxf(local_max, s); - } - for (uint32_t c = tid; c < n_comp; c += blockDim.x) { - float s = -3.4e38f; - if (!(use_mask && comp_mask && comp_mask[c] <= -5.0e29f)) { - const float *kv = comp_kv + (uint64_t)c * head_dim; - float dot = use_vec4 ? attention_dot_f32_vec4_oldhip(qh, kv, head_dim) : 0.0f; + } else { + for (uint32_t r = tid; r < n_raw; r += blockDim.x) { + const uint32_t row = raw_cap ? ((raw_start + r) % raw_cap) : r; + const float *kv = raw_kv + (uint64_t)row * head_dim; + float s = use_vec4 ? attention_dot_f32_vec4_oldhip(qh, kv, head_dim) : 0.0f; if (!use_vec4) { - for (uint32_t i = 0; i < head_dim; i++) dot += qh[i] * kv[i]; + for (uint32_t i = 0; i < head_dim; i++) s += qh[i] * kv[i]; } - s = dot * scale; - if (use_mask && comp_mask) s += comp_mask[c]; + s *= scale; + scores[r] = s; + local_max = fmaxf(local_max, s); + } + for (uint32_t c = tid; c < n_comp; c += blockDim.x) { + float s = -3.4e38f; + if (!(use_mask && comp_mask && comp_mask[c] <= -5.0e29f)) { + const float *kv = comp_kv + (uint64_t)c * head_dim; + float dot = use_vec4 ? attention_dot_f32_vec4_oldhip(qh, kv, head_dim) : 0.0f; + if (!use_vec4) { + for (uint32_t i = 0; i < head_dim; i++) dot += qh[i] * kv[i]; + } + s = dot * scale; + if (use_mask && comp_mask) s += comp_mask[c]; + } + scores[n_raw + c] = s; + local_max = fmaxf(local_max, s); } - scores[n_raw + c] = s; - local_max = fmaxf(local_max, s); } const float max_score = attention_block_max_oldhip_w32(local_max); @@ -717,7 +772,7 @@ __global__ static void attention_decode_mixed_kernel( if (kvrow) { float dot = 0.0f; for (uint32_t d = qlane; d < head_dim; d += 8u) dot += qh[d] * kvrow[d]; - const uint32_t mask = 0xffu << (threadIdx.x & 24u); + const MASK_T mask = ds4_subgroup_mask(threadIdx.x, 8u); for (uint32_t off = 4u; off > 0u; off >>= 1u) { dot += __shfl_down_sync(static_cast(mask), dot, off, 8); } @@ -962,7 +1017,7 @@ __global__ static void attention_indexed_mixed_kernel( : comp_kv + (uint64_t)comp_rows[row - raw_count] * head_dim; float dot = 0.0f; for (uint32_t d = qlane; d < head_dim; d += 8u) dot += qh[d] * kvrow[d]; - const uint32_t mask = 0xffu << (threadIdx.x & 24u); + const MASK_T mask = ds4_subgroup_mask(threadIdx.x, 8u); for (uint32_t off = 4u; off > 0u; off >>= 1u) { dot += __shfl_down_sync(static_cast(mask), dot, off, 8); } @@ -1146,7 +1201,7 @@ __global__ static void attention_indexed_mixed_heads8_online_kernel( dot4_f32(q2, k2) + dot4_f32(q3, k3); score = warp_sum_f32(score) * scale; - score = __shfl_sync(FULL_WARP_MASK, score, 0); + score = __shfl_sync(FULL_WARP_MASK, score, 0, 32); const float new_m = fmaxf(max_s, score); const float old_scale = expf(max_s - new_m); @@ -1257,7 +1312,7 @@ __global__ static void attention_static_mixed_heads8_online_kernel( const float *score_row = scores + warp * 768u; for (uint32_t i = lane; i < n_score; i += 32u) max_s = fmaxf(max_s, score_row[i]); max_s = warp_max_f32(max_s); - max_s = __shfl_sync(FULL_WARP_MASK, max_s, 0); + max_s = __shfl_sync(FULL_WARP_MASK, max_s, 0, 32); } float den = 0.0f; if (valid_head) { @@ -1269,7 +1324,7 @@ __global__ static void attention_static_mixed_heads8_online_kernel( } den = warp_sum_f32(den); den += expf(sinks[head] - max_s); - den = __shfl_sync(FULL_WARP_MASK, den, 0); + den = __shfl_sync(FULL_WARP_MASK, den, 0, 32); } float4 o0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); @@ -1422,7 +1477,7 @@ __global__ static void attention_decode_mixed_heads8_online_kernel( dot4_f32(q2, k2) + dot4_f32(q3, k3); score = warp_sum_f32(score) * scale; - score = __shfl_sync(FULL_WARP_MASK, score, 0); + score = __shfl_sync(FULL_WARP_MASK, score, 0, 32); max_s = fmaxf(max_s, score); } @@ -1458,7 +1513,7 @@ __global__ static void attention_decode_mixed_heads8_online_kernel( dot4_f32(q2, k2) + dot4_f32(q3, k3); score = warp_sum_f32(score) * scale; - score = __shfl_sync(FULL_WARP_MASK, score, 0); + score = __shfl_sync(FULL_WARP_MASK, score, 0, 32); const float row_scale = expf(score - max_s); sum_s += row_scale; diff --git a/rocm/ds4_rocm_attention_launch.cuh b/rocm/ds4_rocm_attention_launch.cuh index b9b43d958..d3c15ad9f 100644 --- a/rocm/ds4_rocm_attention_launch.cuh +++ b/rocm/ds4_rocm_attention_launch.cuh @@ -60,7 +60,25 @@ extern "C" int ds4_gpu_attention_decode_heads_tensor( if (cfg->oldhip_attention_decode) { const uint32_t rows = n_raw + n_comp; const size_t shmem = (size_t)(rows ? rows : 1u) * sizeof(float); - attention_decode_mixed_one_fast_oldhip_kernel<<<(unsigned)n_head, 256, shmem>>>( +#if defined(DS4_GFX906) + /* Cooperative rows coalesce the score-dot loads and win for the + * short decode shapes. Once the cache grows, shuffle overhead + * outweighs that benefit; keep the scalar-per-row path for long + * contexts (8,228 rows measured 7.7% faster on Vega 20). */ + if (cfg->attention_cooperative_dot && rows <= 1024u) + attention_decode_mixed_one_fast_oldhip_kernel<<< + (unsigned)n_head, 256, shmem>>>( + (float *)heads->ptr, + (const float *)q->ptr, + (const float *)raw_kv->ptr, + n_comp ? (const float *)comp_kv->ptr : NULL, + use_mask ? (const float *)comp_mask->ptr : NULL, + sinks, n_raw, raw_cap, raw_start, n_comp, use_mask, + n_head, head_dim, (uint32_t)((head_dim & 3u) == 0u)); + else +#endif + attention_decode_mixed_one_fast_oldhip_kernel<<< + (unsigned)n_head, 256, shmem>>>( (float *)heads->ptr, (const float *)q->ptr, (const float *)raw_kv->ptr, diff --git a/rocm/ds4_rocm_common.cuh b/rocm/ds4_rocm_common.cuh index ebe786286..1cc4a1f3f 100644 --- a/rocm/ds4_rocm_common.cuh +++ b/rocm/ds4_rocm_common.cuh @@ -3,6 +3,25 @@ // Included from ds4_cuda.cu before more specialized modules; these helpers are // intentionally kept static in the single translation unit. +/* + * Return the mask for the power-of-two logical subgroup containing + * linear_lane. CUDA hardware uses 32-lane warps, while gfx9 uses 64-lane + * wavefronts and packs two logical DS4 wave32 groups into one wavefront. + * Keeping this calculation in one helper avoids accidentally addressing the + * wrong half of a wave64 in shuffle reductions. + */ +__device__ static MASK_T ds4_subgroup_mask(uint32_t linear_lane, uint32_t width) { +#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__) + const uint32_t wave_width = (uint32_t)__builtin_amdgcn_wavefrontsize(); +#else + const uint32_t wave_width = (uint32_t)warpSize; +#endif + if (width >= wave_width) return (MASK_T)FULL_WARP_MASK; + const MASK_T low = ((MASK_T)1u << width) - (MASK_T)1u; + const uint32_t wave_lane = linear_lane & (wave_width - 1u); + return low << (wave_lane & ~(width - 1u)); +} + __global__ static void fill_f32_kernel(float *x, uint64_t n, float v) { uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; if (i < n) x[i] = v; @@ -374,4 +393,3 @@ __device__ static float f16_bits_to_f32(uint16_t bits) { __device__ static float dot4_f32(float4 a, float4 b) { return a.x * b.x + a.y * b.y + a.z * b.z + a.w * b.w; } - diff --git a/rocm/ds4_rocm_matmul.cuh b/rocm/ds4_rocm_matmul.cuh index 0f9e06625..7e32f210f 100644 --- a/rocm/ds4_rocm_matmul.cuh +++ b/rocm/ds4_rocm_matmul.cuh @@ -275,6 +275,7 @@ static int cuda_matmul_q8_0_tensor_labeled(ds4_gpu_tensor *out, const void *mode } if (n_tok > 1) { #if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__) +#if !defined(DS4_GFX906) if (!g_quality_mode && (in_dim % 32u) == 0u && out_dim >= 1024u && n_tok >= 256u && @@ -292,6 +293,7 @@ static int cuda_matmul_q8_0_tensor_labeled(ds4_gpu_tensor *out, const void *mode blocks * 34u); return cuda_ok(cudaGetLastError(), "matmul_q8_0 f32 batch wmma 4w launch"); } +#endif #endif if ((in_dim & 31u) == 0u && out_dim <= UINT32_MAX && n_tok <= UINT32_MAX) { const uint32_t rows_per_block = 32u; @@ -731,7 +733,7 @@ extern "C" int ds4_gpu_matmul_f16_pair_tensor( if (!w0 || !w1) return 0; if (!g_quality_mode && !cuda_runtime_config()->graph_dump) { if (in_dim <= 8192u && in_dim * sizeof(float) <= 65536u) { - const uint32_t rows_per_block = 32u; + const uint32_t rows_per_block = cuda_runtime_config()->f16_pair_decode_rpb; matmul_f16_pair_f32_sharedx_warp_rows_w32_kernel<<< ((unsigned)out_dim + rows_per_block - 1u) / rows_per_block, rows_per_block * 32u, diff --git a/rocm/ds4_rocm_moe.cuh b/rocm/ds4_rocm_moe.cuh index 956b1fb22..6481dc688 100644 --- a/rocm/ds4_rocm_moe.cuh +++ b/rocm/ds4_rocm_moe.cuh @@ -4,8 +4,12 @@ // ds4_rocm_moe_launch.cuh so host policy/glue can keep using these static kernels directly. #if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__) +#if defined(DS4_GFX906) +#include "ds4_rocm_wmma_gfx906.cuh" +#else #include #endif +#endif __device__ static float dev_f16_to_f32(uint16_t v) { return __half2float(*reinterpret_cast(&v)); @@ -463,7 +467,7 @@ __device__ static void dev_dot_q2_K_q8_K_block8( } __device__ static float half_warp_sum_f32(float v, uint32_t lane16) { - uint32_t mask = 0xffffu << (threadIdx.x & 16u); + const MASK_T mask = ds4_subgroup_mask(threadIdx.x, 16u); for (int offset = 8; offset > 0; offset >>= 1) { v += __shfl_down_sync(static_cast(mask), v, offset, 16); } @@ -472,7 +476,7 @@ __device__ static float half_warp_sum_f32(float v, uint32_t lane16) { } __device__ static float quarter_warp_sum_f32(float v, uint32_t lane8) { - uint32_t mask = 0xffu << (threadIdx.x & 24u); + const MASK_T mask = ds4_subgroup_mask(threadIdx.x, 8u); for (int offset = 4; offset > 0; offset >>= 1) { v += __shfl_down_sync(static_cast(mask), v, offset, 8); } diff --git a/rocm/ds4_rocm_q8.cuh b/rocm/ds4_rocm_q8.cuh index 5b2423de3..50cf42080 100644 --- a/rocm/ds4_rocm_q8.cuh +++ b/rocm/ds4_rocm_q8.cuh @@ -4,8 +4,12 @@ // private/static while we gradually split the custom ROCm backend into modules. #if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__) +#if defined(DS4_GFX906) +#include "ds4_rocm_wmma_gfx906.cuh" +#else #include #endif +#endif __device__ __forceinline__ static int32_t load_i8x4_i32_aligned(const int8_t *p) { return *(const int32_t *)p; @@ -679,6 +683,7 @@ typedef float __attribute__((ext_vector_type(8))) ds4_q8_float8_t; * into LDS as f16, while each wave owns 16 output rows and computes four * 16-token WMMA columns. It is opt-in from host code because it only wins once * the token batch is large enough to amortize the bigger tile. */ +#if !defined(DS4_GFX906) __launch_bounds__(128, 2) __global__ static void matmul_q8_0_f32_batch_wmma_4w_kernel( float *out, @@ -780,6 +785,7 @@ __global__ static void matmul_q8_0_f32_batch_wmma_4w_kernel( } } } +#endif template __global__ static void matmul_q8_0_f32_batch_wmma_onthefly_kernel( diff --git a/rocm/ds4_rocm_router.cuh b/rocm/ds4_rocm_router.cuh index 8e04613ec..24fbf5f69 100644 --- a/rocm/ds4_rocm_router.cuh +++ b/rocm/ds4_rocm_router.cuh @@ -8,7 +8,7 @@ __device__ __forceinline__ static bool router_score_better(float av, uint32_t ai return av > bv || (av == bv && ai < bi); } -template +template __global__ static void router_select_warp_topk_kernel( int32_t *selected, float *weights, @@ -23,22 +23,24 @@ __global__ static void router_select_warp_topk_kernel( float expert_weight_scale, int has_bias, int hash_mode) { + static_assert(N_EXPERT % LOGICAL_WIDTH == 0u, + "router expert count must be divisible by the logical wave width"); const uint32_t lane = threadIdx.x; const uint32_t row_in_block = threadIdx.y; const uint32_t t = blockIdx.x * blockDim.y + row_in_block; - if (t >= n_tokens || lane >= 32u) return; + if (t >= n_tokens || lane >= LOGICAL_WIDTH) return; const float *log = logits + (uint64_t)t * N_EXPERT; float *prob = probs + (uint64_t)t * N_EXPERT; int32_t *sel = selected + (uint64_t)t * DS4_ROCM_N_EXPERT_USED; float *w = weights + (uint64_t)t * DS4_ROCM_N_EXPERT_USED; __shared__ float sprob[4][N_EXPERT]; - float local_prob[N_EXPERT / 32u]; - float local_score[N_EXPERT / 32u]; + float local_prob[N_EXPERT / LOGICAL_WIDTH]; + float local_score[N_EXPERT / LOGICAL_WIDTH]; #pragma unroll - for (uint32_t j = 0; j < N_EXPERT / 32u; j++) { - const uint32_t e = lane + j * 32u; + for (uint32_t j = 0; j < N_EXPERT / LOGICAL_WIDTH; j++) { + const uint32_t e = lane + j * LOGICAL_WIDTH; const float p = ds4_precise_sqrtf(softplus_dev(log[e])); local_prob[j] = p; local_score[j] = p + (has_bias ? bias[e] : 0.0f); @@ -76,8 +78,8 @@ __global__ static void router_select_warp_topk_kernel( float best_prob = 0.0f; uint32_t best_idx = UINT32_MAX; #pragma unroll - for (uint32_t j = 0; j < N_EXPERT / 32u; j++) { - const uint32_t e = lane + j * 32u; + for (uint32_t j = 0; j < N_EXPERT / LOGICAL_WIDTH; j++) { + const uint32_t e = lane + j * LOGICAL_WIDTH; const float s = local_score[j]; if (router_score_better(s, e, best_score, best_idx)) { best_score = s; @@ -86,10 +88,12 @@ __global__ static void router_select_warp_topk_kernel( } } #pragma unroll - for (uint32_t mask = 16u; mask > 0u; mask >>= 1u) { - const float other_score = __shfl_xor_sync(FULL_WARP_MASK, best_score, mask); - const float other_prob = __shfl_xor_sync(FULL_WARP_MASK, best_prob, mask); - const uint32_t other_idx = __shfl_xor_sync(FULL_WARP_MASK, best_idx, mask); + for (uint32_t mask = LOGICAL_WIDTH >> 1u; mask > 0u; mask >>= 1u) { + const uint32_t lane_abs = threadIdx.y * blockDim.x + threadIdx.x; + const MASK_T row_mask = ds4_subgroup_mask(lane_abs, LOGICAL_WIDTH); + const float other_score = __shfl_xor_sync(row_mask, best_score, mask); + const float other_prob = __shfl_xor_sync(row_mask, best_prob, mask); + const uint32_t other_idx = __shfl_xor_sync(row_mask, best_idx, mask); if (router_score_better(other_score, other_idx, best_score, best_idx)) { best_score = other_score; best_prob = other_prob; @@ -97,8 +101,8 @@ __global__ static void router_select_warp_topk_kernel( } } #pragma unroll - for (uint32_t j = 0; j < N_EXPERT / 32u; j++) { - const uint32_t e = lane + j * 32u; + for (uint32_t j = 0; j < N_EXPERT / LOGICAL_WIDTH; j++) { + const uint32_t e = lane + j * LOGICAL_WIDTH; if (e == best_idx) local_score[j] = -INFINITY; } if (lane == 0) { @@ -121,6 +125,7 @@ __global__ static void router_select_warp_topk_kernel( } } +#ifndef DS4_ROCM_ROUTER_KERNEL_ONLY extern "C" int ds4_gpu_router_select_tensor(ds4_gpu_tensor *selected, ds4_gpu_tensor *weights, ds4_gpu_tensor *probs, const void *model_map, uint64_t model_size, uint64_t bias_offset, uint64_t hash_offset, uint32_t hash_rows, uint32_t token, uint32_t n_expert, uint32_t n_expert_used, float expert_weight_scale, uint32_t n_expert_groups, uint32_t n_group_used, bool has_bias, bool hash_mode, const ds4_gpu_tensor *logits) { const uint32_t active_n_expert = n_expert != 0u ? n_expert : DS4_ROCM_N_EXPERT; const float active_scale = expert_weight_scale != 0.0f ? expert_weight_scale : DS4_ROCM_EXPERT_WEIGHT_SCALE; @@ -151,19 +156,38 @@ extern "C" int ds4_gpu_router_select_tensor(ds4_gpu_tensor *selected, ds4_gpu_te } } if (ok) { - dim3 block(32, 4, 1); - if (active_n_expert == DS4_ROCM_MAX_N_EXPERT) { - router_select_warp_topk_kernel<<<1, block>>>( - (int32_t *)selected->ptr, (float *)weights->ptr, (float *)probs->ptr, - bias, hash, (const float *)logits->ptr, NULL, tok, hash_rows, 1, - active_scale, has_bias && !hash_mode, hash_mode); - } else { - router_select_warp_topk_kernel<<<1, block>>>( - (int32_t *)selected->ptr, (float *)weights->ptr, (float *)probs->ptr, - bias, hash, (const float *)logits->ptr, NULL, tok, hash_rows, 1, - active_scale, has_bias && !hash_mode, hash_mode); +#if defined(DS4_GFX906) + if (cuda_runtime_config()->router_wave64) { + dim3 block(64, 1, 1); + if (active_n_expert == DS4_ROCM_MAX_N_EXPERT) { + router_select_warp_topk_kernel<<<1, block>>>( + (int32_t *)selected->ptr, (float *)weights->ptr, (float *)probs->ptr, + bias, hash, (const float *)logits->ptr, NULL, tok, hash_rows, 1, + active_scale, has_bias && !hash_mode, hash_mode); + } else { + router_select_warp_topk_kernel<<<1, block>>>( + (int32_t *)selected->ptr, (float *)weights->ptr, (float *)probs->ptr, + bias, hash, (const float *)logits->ptr, NULL, tok, hash_rows, 1, + active_scale, has_bias && !hash_mode, hash_mode); + } + ok = cuda_ok(cudaGetLastError(), "router_select wave64 launch"); + } else +#endif + { + dim3 block(32, 4, 1); + if (active_n_expert == DS4_ROCM_MAX_N_EXPERT) { + router_select_warp_topk_kernel<<<1, block>>>( + (int32_t *)selected->ptr, (float *)weights->ptr, (float *)probs->ptr, + bias, hash, (const float *)logits->ptr, NULL, tok, hash_rows, 1, + active_scale, has_bias && !hash_mode, hash_mode); + } else { + router_select_warp_topk_kernel<<<1, block>>>( + (int32_t *)selected->ptr, (float *)weights->ptr, (float *)probs->ptr, + bias, hash, (const float *)logits->ptr, NULL, tok, hash_rows, 1, + active_scale, has_bias && !hash_mode, hash_mode); + } + ok = cuda_ok(cudaGetLastError(), "router_select launch"); } - ok = cuda_ok(cudaGetLastError(), "router_select launch"); } return ok; } @@ -197,6 +221,23 @@ extern "C" int ds4_gpu_router_select_batch_tensor(ds4_gpu_tensor *selected, ds4_ if (!hash) return 0; } dim3 block(32, 4, 1); +#if defined(DS4_GFX906) + if (cuda_runtime_config()->router_wave64) { + block = dim3(64, 4, 1); + if (active_n_expert == DS4_ROCM_MAX_N_EXPERT) { + router_select_warp_topk_kernel<<<(n_tokens + 3u) / 4u, block>>>( + (int32_t *)selected->ptr, (float *)weights->ptr, (float *)probs->ptr, + bias, hash, (const float *)logits->ptr, (const int32_t *)tokens->ptr, + 0, hash_rows, n_tokens, active_scale, has_bias && !hash_mode, hash_mode); + } else { + router_select_warp_topk_kernel<<<(n_tokens + 3u) / 4u, block>>>( + (int32_t *)selected->ptr, (float *)weights->ptr, (float *)probs->ptr, + bias, hash, (const float *)logits->ptr, (const int32_t *)tokens->ptr, + 0, hash_rows, n_tokens, active_scale, has_bias && !hash_mode, hash_mode); + } + return cuda_ok(cudaGetLastError(), "router_select batch wave64 launch"); + } +#endif if (active_n_expert == DS4_ROCM_MAX_N_EXPERT) { router_select_warp_topk_kernel<<<(n_tokens + 3u) / 4u, block>>>( (int32_t *)selected->ptr, @@ -230,3 +271,4 @@ extern "C" int ds4_gpu_router_select_batch_tensor(ds4_gpu_tensor *selected, ds4_ } return cuda_ok(cudaGetLastError(), "router_select launch"); } +#endif diff --git a/rocm/ds4_rocm_runtime.cuh b/rocm/ds4_rocm_runtime.cuh index 3bd786f8e..b21a1312e 100644 --- a/rocm/ds4_rocm_runtime.cuh +++ b/rocm/ds4_rocm_runtime.cuh @@ -969,8 +969,35 @@ static int cuda_stream_resident_evict_one( return 1; } -static uint64_t cuda_stream_resident_free_reserve_bytes(void) { - return 16ull * 1024ull * 1024ull * 1024ull; +static uint64_t cuda_stream_resident_free_reserve_bytes(uint64_t total_bytes) { + const char *env = getenv("DS4_ROCM_STREAM_FREE_RESERVE_MB"); + if (env && env[0]) { + char *end = NULL; + const unsigned long long mb = strtoull(env, &end, 10); + if (end != env && mb >= 256u && mb <= 16384u) { + return (uint64_t)mb * 1048576ull; + } + } + if (total_bytes == 0u) { + size_t free_b = 0, total_b = 0; + if (cudaMemGetInfo(&free_b, &total_b) == cudaSuccess) { + (void)free_b; + total_bytes = (uint64_t)total_b; + } else { + (void)cudaGetLastError(); + } + } + + /* Expert-cache capacity scales with the card. A fixed 16 GiB reserve made + * streaming impossible on 16 GiB discrete GPUs. Keep at least 4 GiB for + * KV/scratch plus the roughly 1.7 GiB transient full-layer prefill span; + * larger cards retain one eighth of VRAM, capped at the old 16 GiB. */ + const uint64_t min_reserve = 4ull * 1073741824ull; + const uint64_t max_reserve = 16ull * 1073741824ull; + uint64_t reserve = total_bytes / 8u; + if (reserve < min_reserve) reserve = min_reserve; + if (reserve > max_reserve) reserve = max_reserve; + return reserve; } static int cuda_stream_resident_make_room( @@ -986,9 +1013,9 @@ static int cuda_stream_resident_make_room( size_t free_b = 0; size_t total_b = 0; - const uint64_t reserve = cuda_stream_resident_free_reserve_bytes(); while (cudaMemGetInfo(&free_b, &total_b) == cudaSuccess) { - (void)total_b; + const uint64_t reserve = + cuda_stream_resident_free_reserve_bytes((uint64_t)total_b); if ((uint64_t)free_b >= reserve && bytes <= (uint64_t)free_b - reserve) { return 1; @@ -1028,7 +1055,7 @@ static int cuda_stream_resident_alloc( (double)bytes / 1048576.0, layer, expert, - (double)cuda_stream_resident_free_reserve_bytes() / 1073741824.0); + (double)cuda_stream_resident_free_reserve_bytes(0) / 1073741824.0); return -1; } @@ -3515,6 +3542,15 @@ static uint32_t cuda_rows_per_block_or_default(uint32_t v, uint32_t def) { return (v == 1u || v == 2u || v == 4u || v == 8u || v == 16u || v == 32u) ? v : def; } +static uint32_t cuda_rows_per_block_env(const char *name, uint32_t def) { + const char *value = getenv(name); + if (!value || !value[0]) return def; + char *end = NULL; + const unsigned long parsed = strtoul(value, &end, 10); + if (end == value || *end != '\0' || parsed > UINT32_MAX) return def; + return cuda_rows_per_block_or_default((uint32_t)parsed, def); +} + struct ds4_rocm_runtime_config { int initialized; int q8_prequant_decode; @@ -3523,7 +3559,10 @@ struct ds4_rocm_runtime_config { int attention_output_cublas_all; int shared_down_cublas; int graph_dump; + int router_wave64; + int attention_cooperative_dot; uint32_t q8_decode_rpb; + uint32_t f16_pair_decode_rpb; uint32_t q8_hc_decode_rpb; uint32_t attn_out_low_decode_rpb; uint32_t moe_decode_rpb; @@ -3540,7 +3579,31 @@ static const ds4_rocm_runtime_config *cuda_runtime_config(void) { g_rocm_cfg.attention_output_cublas_all = !g_quality_mode; g_rocm_cfg.shared_down_cublas = !g_quality_mode; g_rocm_cfg.graph_dump = cuda_env_present(getenv("DS4_METAL_GRAPH_DUMP_PREFIX")); - g_rocm_cfg.q8_decode_rpb = g_quality_mode ? 8u : 1u; +#if defined(DS4_GFX906) + g_rocm_cfg.router_wave64 = + !cuda_env_present(getenv("DS4_ROCM_DISABLE_ROUTER_WAVE64")); + g_rocm_cfg.attention_cooperative_dot = + !cuda_env_present(getenv("DS4_ROCM_DISABLE_ATTENTION_COOP_DOT")); +#else + g_rocm_cfg.router_wave64 = 0; + g_rocm_cfg.attention_cooperative_dot = 0; +#endif + uint32_t q8_decode_default = g_quality_mode ? 8u : 1u; +#if defined(DS4_GFX906) + /* Pack two independent software-wave32 rows into each hardware wave64. + * Row arithmetic and reduction order remain unchanged. */ + if (!g_quality_mode) q8_decode_default = 2u; +#endif + g_rocm_cfg.q8_decode_rpb = cuda_rows_per_block_env( + "DS4_ROCM_Q8_DECODE_RPB", q8_decode_default); + uint32_t f16_pair_default = 32u; +#if defined(DS4_GFX906) + /* A 1024-thread block exposes too few workgroups for the 60 CUs in + * Vega 20 at the 256/512/1024-row decode shapes. */ + f16_pair_default = 8u; +#endif + g_rocm_cfg.f16_pair_decode_rpb = cuda_rows_per_block_env( + "DS4_ROCM_F16_PAIR_DECODE_RPB", f16_pair_default); g_rocm_cfg.q8_hc_decode_rpb = g_quality_mode ? 8u : 16u; g_rocm_cfg.attn_out_low_decode_rpb = g_quality_mode ? 8u : 32u; g_rocm_cfg.moe_decode_rpb = g_quality_mode ? 8u : 1u; @@ -3556,7 +3619,7 @@ static uint64_t cuda_q8_f16_cache_limit_bytes(void) { static uint64_t cuda_q8_f16_cache_reserve_bytes(uint64_t total_bytes) { if (g_ssd_streaming_mode) { - return cuda_stream_resident_free_reserve_bytes(); + return cuda_stream_resident_free_reserve_bytes(total_bytes); } if (total_bytes >= 112ull * 1024ull * 1024ull * 1024ull) { return 512ull * 1048576ull; @@ -4132,7 +4195,21 @@ static uint64_t cuda_model_cache_limit_bytes(void) { } static uint64_t cuda_model_arena_chunk_bytes(uint64_t need) { - uint64_t bytes = 1792ull * 1048576ull; + uint64_t mb = 1792; + const char *env = getenv("DS4_ROCM_WEIGHT_ARENA_CHUNK_MB"); + if (!env || !env[0]) env = getenv("DS4_CUDA_WEIGHT_ARENA_CHUNK_MB"); + if (env && env[0]) { + char *end = NULL; + unsigned long long v = strtoull(env, &end, 10); + if (end != env && v > 0) mb = (uint64_t)v; + } + if (mb < 256) mb = 256; + if (mb > 8192) mb = 8192; + uint64_t bytes = mb * 1048576ull; + if (need > bytes / 2u) { + const uint64_t align = 64ull * 1048576ull; + return (need + align - 1u) & ~(align - 1u); + } if (bytes < need) { const uint64_t align = 256ull * 1048576ull; bytes = (need + align - 1u) & ~(align - 1u); @@ -4253,6 +4330,15 @@ static const char *cuda_model_range_ptr_from_fd( g_model_ranges.push_back({model_map, offset, bytes, dev, NULL, NULL, 0, 0, 1}); g_model_range_by_offset[offset] = g_model_ranges.size() - 1u; g_model_range_bytes += bytes; + if (cuda_stream_cache_stats_on()) { + fprintf(stderr, + DS4_GPU_LOG_PREFIX "model arena map %-18s off=%.3f GiB bytes=%.2f MiB total=%.2f GiB ranges=%zu\n", + what ? what : "weights", + (double)offset / 1073741824.0, + (double)bytes / 1048576.0, + (double)g_model_range_bytes / 1073741824.0, + g_model_ranges.size()); + } cuda_model_load_progress_note(g_model_range_bytes); return (const char *)dev; } diff --git a/rocm/ds4_rocm_wmma_gfx906.cuh b/rocm/ds4_rocm_wmma_gfx906.cuh new file mode 100644 index 000000000..e59067a6d --- /dev/null +++ b/rocm/ds4_rocm_wmma_gfx906.cuh @@ -0,0 +1,196 @@ +// Device-side compatibility implementation for the rocWMMA subset used by ds4 +// on AMD gfx906 GPUs (Vega 20: Radeon VII / Pro VII / MI50). These devices do +// not have the MFMA matrix instructions introduced with gfx908, and the +// official rocWMMA headers reject the architecture. +// +// The ROCm backend uses software-wave32 kernel layouts. This header follows the +// same convention on wave64 hardware: lane = tid & 31 and shuffles use an +// explicit width of 32. +// +// Implemented API subset: +// - fragment +// - load_matrix_sync(frag, ptr, ldm) // default row_major +// - load_matrix_sync(frag, ptr, ldm, layout) // explicit layout +// - fill_fragment(frag, val) +// - mma_sync(d, a, b, c) // d = a*b + c (in-place ok) +// - store_matrix_sync(ptr, frag, ldm, layout) +// +// The only supported tile is M=N=K=16, half x half -> float. The internal +// fragment layout does not need to match rocWMMA: 256 row-major elements are +// distributed over 32 logical lanes, eight elements per lane. +// lane l owns tile[l/2][(l%2)*8 .. (l%2)*8+7] +// +// mma_sync su 32 lane via __shfl width=32: +// - A[lane/2][k] k=0..15: 8 owned + 8 shfl da lane^1 (8 shfl half) +// - B[k][(l%2)*8+c] k=0..15,c=0..7: shfl da lane 2*k+(l%2) (128 shfl half) +// - 8 FMA per (c,k) → 128 FMA/lane per mma_sync +// This prioritizes compatibility over performance. Packed shuffles or a native +// wave64 kernel can replace it later without changing call sites. +#pragma once + +#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__) +#include +#include + +namespace rocwmma { + +// Matrix-kind tags. +struct matrix_a {}; +struct matrix_b {}; +struct accumulator {}; + +// Fragment-layout tags. +struct row_major {}; +struct col_major {}; + +// Store layout, matching the names used by rocWMMA call sites. +enum layout_t { mem_row_major, mem_col_major, mem_stride }; + +namespace detail { + +__device__ __forceinline__ int wmma_lane() { return static_cast(threadIdx.x) & 31; } + +} // namespace detail + +// Eight elements per lane (256 tile elements / 32 logical lanes). Layout +// defaults to row_major, as it does for ds4's rocWMMA accumulator fragments. +template +struct fragment { + static_assert(M == 16 && N == 16 && K == 16, + "ds4 gfx906 shim supports only 16x16x16 tiles"); + T x[8]; +}; + +// ===================== load_matrix_sync ===================== +// matrix_a, row_major: A[i][j] is ptr + i*ldm + j. +template +__device__ __forceinline__ void +load_matrix_sync(fragment &frag, + const half *ptr, uint32_t ldm) { + const int l = detail::wmma_lane(); + const int i = l >> 1; + const int j0 = (l & 1) * 8; + const half *p = ptr + (uint64_t)i * ldm + j0; + #pragma unroll + for (int c = 0; c < 8; ++c) frag.x[c] = p[c]; +} + +// matrix_a, col_major: A[i][j] is ptr + j*ldm + i. +template +__device__ __forceinline__ void +load_matrix_sync(fragment &frag, + const half *ptr, uint32_t ldm) { + const int l = detail::wmma_lane(); + const int i = l >> 1; + const int j0 = (l & 1) * 8; + #pragma unroll + for (int c = 0; c < 8; ++c) frag.x[c] = ptr[(uint64_t)(j0 + c) * ldm + i]; +} + +// matrix_b, row_major: B[k][j] is ptr + k*ldm + j. +template +__device__ __forceinline__ void +load_matrix_sync(fragment &frag, + const half *ptr, uint32_t ldm) { + const int l = detail::wmma_lane(); + const int k = l >> 1; + const int j0 = (l & 1) * 8; + const half *p = ptr + (uint64_t)k * ldm + j0; + #pragma unroll + for (int c = 0; c < 8; ++c) frag.x[c] = p[c]; +} + +// matrix_b, col_major: B[k][j] is ptr + j*ldm + k. +template +__device__ __forceinline__ void +load_matrix_sync(fragment &frag, + const half *ptr, uint32_t ldm) { + const int l = detail::wmma_lane(); + const int k = l >> 1; + const int j0 = (l & 1) * 8; + #pragma unroll + for (int c = 0; c < 8; ++c) frag.x[c] = ptr[(uint64_t)(j0 + c) * ldm + k]; +} + +// ===================== fill_fragment ===================== +template +__device__ __forceinline__ void +fill_fragment(fragment &frag, float v) { + #pragma unroll + for (int c = 0; c < 8; ++c) frag.x[c] = v; +} +template +__device__ __forceinline__ void +fill_fragment(fragment &frag, float v) { + #pragma unroll + for (int c = 0; c < 8; ++c) frag.x[c] = v; +} + +// ===================== mma_sync ===================== +// d = a * b + c (tile 16x16: d[i][j] = sum_k a[i][k]*b[k][j]) +// d/c are float accumulators; a/b are half fragments. In-place d == c is safe +// because c is copied to lane-local accumulators before d is written. +template +__device__ __forceinline__ void +mma_sync(fragment &d, + const fragment &a, + const fragment &b, + const fragment &cfrag) { + const int l = detail::wmma_lane(); + const int j0 = (l & 1) * 8; + + // Lane-local accumulator, allowing d to alias cfrag. + float acc[8]; + #pragma unroll + for (int c = 0; c < 8; ++c) acc[c] = cfrag.x[c]; + + // Gather a full A row: eight owned values plus eight from the peer lane. + half a_row[16]; + #pragma unroll + for (int c = 0; c < 8; ++c) a_row[j0 + c] = a.x[c]; + const int peer = l ^ 1; + #pragma unroll + for (int c = 0; c < 8; ++c) + a_row[(1 - (l & 1)) * 8 + c] = __shfl(a.x[c], peer, 32); + + // Gather B[k][j0..j0+7] from lane 2*k+(l&1), then accumulate. + #pragma unroll + for (int k = 0; k < 16; ++k) { + const int src = (k << 1) + (l & 1); + const float af = __half2float(a_row[k]); + #pragma unroll + for (int c = 0; c < 8; ++c) { + const half bh = __shfl(b.x[c], src, 32); + acc[c] = fmaf(af, __half2float(bh), acc[c]); + } + } + + // Write the destination fragment. + #pragma unroll + for (int c = 0; c < 8; ++c) d.x[c] = acc[c]; +} + +// ===================== store_matrix_sync ===================== +template +__device__ __forceinline__ void +store_matrix_sync(float *ptr, + const fragment &frag, + uint32_t ldm, layout_t layout) { + const int l = detail::wmma_lane(); + const int i = l >> 1; + const int j0 = (l & 1) * 8; + if (layout == mem_col_major) { + // C[i][j] is ptr + j*ldm + i. + #pragma unroll + for (int c = 0; c < 8; ++c) + ptr[(uint64_t)(j0 + c) * ldm + i] = frag.x[c]; + } else { // Treat mem_stride as row-major for this restricted API. + float *p = ptr + (uint64_t)i * ldm + j0; + #pragma unroll + for (int c = 0; c < 8; ++c) p[c] = frag.x[c]; + } +} + +} // namespace rocwmma + +#endif // __HIP_PLATFORM_AMD__ || __HIPCC__ diff --git a/run-context.sh b/run-context.sh new file mode 100755 index 000000000..2ba522cbf --- /dev/null +++ b/run-context.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash +# Six-gfx906 capacity profile: bounded SSD expert cache and native 1M context. +# Edit MODEL_PATH here; no launch-time environment variables are required. +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +export MODEL_PATH="${SCRIPT_DIR}/ds4flash.gguf" +export CTX=1000000 +export MAX_TOKENS=1000000 +export PREFILL_CHUNK=64 +export DIST_WINDOW=5 +export WORKER_START_DELAY=5 +export HTTP_HOST=0.0.0.0 +export HTTP_PORT=8080 +export DIST_HOST=127.0.0.1 +export DIST_PORT=19000 +export DS4_PROFILE=0 +export DS4_COORD_SERIALIZE=0 +export SSD_STREAMING=1 +export SSD_STREAMING_CACHE_EXPERTS=2GB +export SSD_STREAMING_PRELOAD_EXPERTS= +export SSD_STREAMING_COLD=0 +export COORD_DEVICE=2 +export COORD_LAYERS=0:12 +export WORKER_SPECS="0,13:18 1,19:24 3,25:30 4,31:36 5,37:42" + +exec "${SCRIPT_DIR}/run.sh" diff --git a/run-speed.sh b/run-speed.sh new file mode 100755 index 000000000..75eb511d8 --- /dev/null +++ b/run-speed.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +# Six-gfx906 profile for maximum throughput: all model weights stay resident. +# Edit MODEL_PATH here; no launch-time environment variables are required. +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +export MODEL_PATH="${SCRIPT_DIR}/ds4flash.gguf" +# Seven resident layers fit the 16GB workers at 32K context. This balanced +# split is about 50% faster in prefill than the capacity-oriented 13 + 6x5. +export CTX=32768 +export MAX_TOKENS=32768 +export PREFILL_CHUNK=64 +export DIST_WINDOW=5 +export WORKER_START_DELAY=5 +export HTTP_HOST=0.0.0.0 +export HTTP_PORT=8080 +export DIST_HOST=127.0.0.1 +export DIST_PORT=19000 +export DS4_PROFILE=0 +export DS4_COORD_SERIALIZE=0 +export SSD_STREAMING=0 +export SSD_STREAMING_CACHE_EXPERTS= +export SSD_STREAMING_PRELOAD_EXPERTS= +export SSD_STREAMING_COLD=0 +export COORD_DEVICE=2 +export COORD_LAYERS=0:7 +export WORKER_SPECS="0,8:14 1,15:21 3,22:28 4,29:35 5,36:42" + +exec "${SCRIPT_DIR}/run.sh" diff --git a/run.sh b/run.sh new file mode 100755 index 000000000..1aae333d1 --- /dev/null +++ b/run.sh @@ -0,0 +1,238 @@ +#!/usr/bin/env bash +# +# run.sh — configurable distributed DeepSeek-V4-Flash launcher for AMD gfx906. +# The supplied run-speed.sh and run-context.sh are validated six-GPU examples. +# +# Adapted from the 4x RTX 3090 recipe: +# https://github.com/Forge-the-Kingdom/inference-serving-recipes/blob/main/recipes/dwarfstar/deepseek-v4-iq2-4x3090-gpu-resident.md +# +# Key changes vs the CUDA recipe: +# - Backend flag is --rocm (not --cuda) in a DS4_ROCM_BUILD. +# - Per-process GPU pinning uses ROCR_VISIBLE_DEVICES (not CUDA_VISIBLE_DEVICES). +# - Coordinator and worker device indexes are explicit below because ROCR +# indexes need not match the rocm-smi card numbers. +# +# Prerequisites: +# - ds4 and ds4-server built for your ROCm arch: make rocm ROCM_ARCH= +# gfx906 builds use the in-tree wave64-safe rocWMMA compatibility shim. +# - Enough free VRAM on every device selected by the chosen profile. +# - The Flash IQ2XXS GGUF downloaded, e.g.: +# ./download_model.sh q2-imatrix +# which links ./ds4flash.gguf to the file below. +# +# Usage: +# 1. Set MODEL_PATH and the device/layer map below for your system. +# 2. ./run.sh # starts workers then the coordinator +# 3. In another shell: curl -s http://127.0.0.1:8080/v1/models +# Ctrl+C the coordinator; workers exit when the route drops. +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +# ============================================================================= +# Override MODEL_PATH when the GGUF is not available through ./ds4flash.gguf. +MODEL_PATH="${MODEL_PATH:-${SCRIPT_DIR}/ds4flash.gguf}" +# ============================================================================= + +if [ -z "${MODEL_PATH}" ]; then + echo "run.sh: MODEL_PATH is empty. Edit this script and set MODEL_PATH to the" >&2 + echo " DeepSeek-V4-Flash IQ2XXS GGUF, e.g. /path/to/model.gguf" >&2 + exit 1 +fi +if [ ! -f "${MODEL_PATH}" ]; then + echo "run.sh: model file not found: ${MODEL_PATH}" >&2 + exit 1 +fi + +# --- binaries (built with: make rocm ROCM_ARCH=) --- +DS4_SERVER="${SCRIPT_DIR}/ds4-server" # coordinator +DS4_WORKER="${SCRIPT_DIR}/ds4" # worker +for b in "${DS4_SERVER}" "${DS4_WORKER}"; do + if [ ! -x "${b}" ]; then + echo "run.sh: missing executable: ${b}" >&2 + echo " build with: make rocm ROCM_ARCH=" >&2 + exit 1 + fi +done + +# --- distributed pipeline tuning (from the 4x3090 recipe) --- +DIST_HOST="${DIST_HOST:-127.0.0.1}" +DIST_PORT="${DIST_PORT:-19000}" +CTX="${CTX:-28672}" # tested safe ceiling with the GUI active +MAX_TOKENS="${MAX_TOKENS:-${CTX}}" +PREFILL_CHUNK="${PREFILL_CHUNK:-64}" # small chunks keep all 6 cards busy +DIST_WINDOW="${DIST_WINDOW:-5}" # larger windows can OOM 16 GB workers +WORKER_START_DELAY="${WORKER_START_DELAY:-20}" +DS4_PROFILE="${DS4_PROFILE:-0}" +DS4_COORD_SERIALIZE="${DS4_COORD_SERIALIZE:-0}" +ROCPROF_COORD_OUTPUT_DIR="${ROCPROF_COORD_OUTPUT_DIR:-}" +ROCPROF_WORKER_DEVICE="${ROCPROF_WORKER_DEVICE:-}" +ROCPROF_WORKER_OUTPUT_DIR="${ROCPROF_WORKER_OUTPUT_DIR:-}" +TRACE_FILE="${TRACE_FILE:-}" +TRACE_ARGS=() +if [ -n "${TRACE_FILE}" ]; then + TRACE_ARGS+=(--trace "${TRACE_FILE}") +fi + +# --- optional SSD-backed routed-expert cache --- +# Keep disabled for the fastest fully resident path. Enabling it leaves +# non-routed weights resident and loads routed MoE experts into a bounded cache +# from the GGUF on demand. The cache budget is per process/GPU. +SSD_STREAMING="${SSD_STREAMING:-0}" +SSD_STREAMING_CACHE_EXPERTS="${SSD_STREAMING_CACHE_EXPERTS:-}" +SSD_STREAMING_PRELOAD_EXPERTS="${SSD_STREAMING_PRELOAD_EXPERTS:-}" +SSD_STREAMING_COLD="${SSD_STREAMING_COLD:-0}" +SSD_ARGS=() +if [ "${SSD_STREAMING}" != "0" ]; then + SSD_ARGS+=(--ssd-streaming) + if [ -n "${SSD_STREAMING_CACHE_EXPERTS}" ]; then + SSD_ARGS+=(--ssd-streaming-cache-experts "${SSD_STREAMING_CACHE_EXPERTS}") + fi + if [ -n "${SSD_STREAMING_PRELOAD_EXPERTS}" ]; then + SSD_ARGS+=(--ssd-streaming-preload-experts "${SSD_STREAMING_PRELOAD_EXPERTS}") + fi + if [ "${SSD_STREAMING_COLD}" != "0" ]; then + SSD_ARGS+=(--ssd-streaming-cold) + fi +fi + +# --- HTTP API --- +HTTP_HOST="${HTTP_HOST:-0.0.0.0}" +HTTP_PORT="${HTTP_PORT:-8080}" + +# --- MTP speculative decoding (draft model, P1 speedup) --- +# Disabled: the draft model does not fit reliably beside the resident profile +# and did not improve quality-adjusted throughput on this PP6 topology. +MTP_PATH="${MTP_PATH:-${SCRIPT_DIR}/DeepSeek-V4-Flash-MTP-Q4K-Q8_0-F32.gguf}" +MTP_DRAFT="${MTP_DRAFT:-2}" +MTP_ARGS=() + +# --- layer split: all 43 layers (0..42), coordinator owns output head --- +# The context profile keeps the conservative 13 + 6x5 split. The speed +# profile uses 8 + 7x5 at a smaller context: equal-size GPU stages remove the +# coordinator prefill bottleneck while still fitting the 16GB workers. +# +# A worker spec is DEVICE,LAYER_START:LAYER_END. Keeping the route in data +# instead of hard-coding five workers makes this shared launcher usable for +# other gfx906 counts and memory layouts. The two example profiles below set +# the exact six-GPU maps validated on the development machine. +COORD_DEVICE="${COORD_DEVICE:-2}" +COORD_LAYERS="${COORD_LAYERS:-0:12}" +WORKER_SPECS="${WORKER_SPECS:-0,13:18 1,19:24 3,25:30 4,31:36 5,37:42}" + +WORKER_PIDS=() + +cleanup() { + local pid + for pid in "${WORKER_PIDS[@]}"; do + kill "${pid}" 2>/dev/null || true + done + pkill -f "ds4.*--role worker.*${DIST_PORT}" 2>/dev/null || true +} +trap cleanup EXIT INT TERM +cleanup + +start_worker() { # [extra ds4 arguments...] + local dev="$1" layers="$2" + local arg ssd_worker=0 + local -a worker_env=(env + -u DS4_DIST_DECODE_PROFILE + DS4_ROCM_WEIGHT_ARENA_CHUNK_MB=256 + DS4_LOCK_FILE="/tmp/ds4-worker-${dev}.lock" + ROCR_VISIBLE_DEVICES="${dev}") + local -a worker_prefix=() + shift 2 + for arg in "$@"; do + if [ "${arg}" = "--ssd-streaming" ]; then + ssd_worker=1 + break + fi + done + if [ "${DS4_PROFILE}" != "0" ] && [ "${ssd_worker}" != "0" ]; then + worker_env+=(DS4_ROCM_STREAM_CACHE_STATS=1) + fi + if [ -n "${ROCPROF_WORKER_DEVICE}" ] && + [ "${ROCPROF_WORKER_DEVICE}" = "${dev}" ]; then + if [ -z "${ROCPROF_WORKER_OUTPUT_DIR}" ]; then + echo "run.sh: ROCPROF_WORKER_OUTPUT_DIR is required for profiling" >&2 + exit 1 + fi + worker_prefix=(rocprofv3 --kernel-trace --stats --output-format csv + --output-directory "${ROCPROF_WORKER_OUTPUT_DIR}" --) + echo "run.sh: profiling worker ROCR ${dev} into ${ROCPROF_WORKER_OUTPUT_DIR}" + fi + "${worker_env[@]}" \ + "${worker_prefix[@]}" \ + "${DS4_WORKER}" -m "${MODEL_PATH}" --rocm \ + --prefill-chunk "${PREFILL_CHUNK}" \ + --ctx "${CTX}" \ + --role worker --layers "${layers}" \ + --coordinator "${DIST_HOST}" "${DIST_PORT}" \ + "$@" & + WORKER_PIDS+=("$!") + echo "run.sh: worker started on ROCR device ${dev}, layers ${layers} (pid $!)" +} + +check_last_worker() { + local pid="${WORKER_PIDS[${#WORKER_PIDS[@]}-1]}" + sleep "${WORKER_START_DELAY}" + if ! kill -0 "${pid}" 2>/dev/null; then + echo "run.sh: worker ${pid} exited during startup" >&2 + wait "${pid}" || true + exit 1 + fi +} + +if [ "${SSD_STREAMING}" != "0" ]; then + echo "run.sh: starting SSD expert-streaming profile, ctx=${CTX}, chunk=${PREFILL_CHUNK}, window=${DIST_WINDOW}" +else + echo "run.sh: starting GPU-resident profile, ctx=${CTX}, chunk=${PREFILL_CHUNK}, window=${DIST_WINDOW}" +fi +for worker_spec in ${WORKER_SPECS}; do + IFS=, read -r worker_dev worker_layers worker_extra <<< "${worker_spec}" + if [[ ! "${worker_dev}" =~ ^[0-9]+$ ]] || + [[ ! "${worker_layers}" =~ ^[0-9]+:[0-9]+$ ]] || + [ -n "${worker_extra}" ]; then + echo "run.sh: invalid worker spec '${worker_spec}'; expected DEVICE,START:END" >&2 + exit 1 + fi + start_worker "${worker_dev}" "${worker_layers}" "${SSD_ARGS[@]}" + check_last_worker +done + +# Let workers fail visibly before the coordinator commits its allocation. +sleep 2 + +echo "run.sh: starting coordinator on ROCR device ${COORD_DEVICE}, layers ${COORD_LAYERS} plus local output head" +COORD_ENV=(env + -u DS4_DIST_DECODE_PROFILE + DS4_ROCM_WEIGHT_ARENA_CHUNK_MB=256 + DS4_LOCK_FILE="/tmp/ds4-coordinator-${COORD_DEVICE}.lock" + ROCR_VISIBLE_DEVICES="${COORD_DEVICE}") +if [ "${DS4_COORD_SERIALIZE}" != "0" ]; then + COORD_ENV+=(AMD_SERIALIZE_KERNEL="${DS4_COORD_SERIALIZE}") +fi +COORD_PREFIX=() +if [ -n "${ROCPROF_COORD_OUTPUT_DIR}" ]; then + COORD_PREFIX=(rocprofv3 --kernel-trace --stats --output-format csv + --output-directory "${ROCPROF_COORD_OUTPUT_DIR}" --) + echo "run.sh: profiling coordinator into ${ROCPROF_COORD_OUTPUT_DIR}" +fi + +"${COORD_ENV[@]}" \ +"${COORD_PREFIX[@]}" \ +"${DS4_SERVER}" -m "${MODEL_PATH}" --rocm \ + --prefill-chunk "${PREFILL_CHUNK}" \ + --ctx "${CTX}" \ + --tokens "${MAX_TOKENS}" \ + --host "${HTTP_HOST}" --port "${HTTP_PORT}" \ + "${TRACE_ARGS[@]}" \ + --role coordinator --layers "${COORD_LAYERS}" \ + "${MTP_ARGS[@]}" \ + "${SSD_ARGS[@]}" \ + --listen "${DIST_HOST}" "${DIST_PORT}" \ + --dist-prefill-chunk "${PREFILL_CHUNK}" \ + --dist-prefill-window "${DIST_WINDOW}" \ + --dist-activation-bits 16 & +COORD_PID=$! +wait "${COORD_PID}" diff --git a/tests/cuda_long_context_smoke.c b/tests/cuda_long_context_smoke.c index eac1d75bc..651b551c5 100644 --- a/tests/cuda_long_context_smoke.c +++ b/tests/cuda_long_context_smoke.c @@ -1,6 +1,7 @@ #include "ds4_gpu.h" #include +#include #include #include #include @@ -40,7 +41,7 @@ static int check_large_topk(void) { double elapsed = 0.0; if (scores && selected && ds4_gpu_tensor_write(scores, 0, scores_host, score_count * sizeof(float))) { - /* Exclude one-time CUDA module/kernel setup from the throughput guard. */ + /* Exclude one-time GPU module/kernel setup from the throughput guard. */ if (!ds4_gpu_indexer_topk_tensor(selected, scores, n_comp, n_tokens, top_k) || !ds4_gpu_synchronize()) { rc = 1; @@ -69,8 +70,9 @@ static int check_large_topk(void) { } } if (rc == 0) { - const double max_seconds = getenv_seconds("DS4_CUDA_TOPK_REGRESSION_SEC", 2.0); - fprintf(stderr, "cuda-regression: top-k n_comp=%u n_tokens=%u elapsed=%.3fs\n", + const double max_seconds = getenv_seconds("DS4_GPU_TOPK_REGRESSION_SEC", + getenv_seconds("DS4_CUDA_TOPK_REGRESSION_SEC", 2.0)); + fprintf(stderr, "gpu-regression: top-k n_comp=%u n_tokens=%u elapsed=%.3fs\n", n_comp, n_tokens, elapsed); if (elapsed > max_seconds) { fprintf(stderr, "top-k regression: %.3fs exceeds %.3fs\n", elapsed, max_seconds); @@ -156,11 +158,123 @@ static int check_decode_attention_overflow_path(void) { return rc; } +static int check_decode_attention_ring_reference(void) { + const uint32_t n_head = 8; + const uint32_t head_dim = 512; + const uint32_t n_raw = 257; + const uint32_t raw_cap = 300; + const uint32_t raw_start = 270; + const uint32_t n_comp = 511; + const uint64_t q_count = (uint64_t)n_head * head_dim; + const uint64_t raw_count = (uint64_t)raw_cap * head_dim; + const uint64_t comp_count = (uint64_t)n_comp * head_dim; + float *sinks = (float *)malloc(n_head * sizeof(float)); + float *q_host = (float *)malloc((size_t)q_count * sizeof(float)); + float *raw_host = (float *)malloc((size_t)raw_count * sizeof(float)); + float *comp_host = (float *)malloc((size_t)comp_count * sizeof(float)); + float *mask_host = (float *)malloc(n_comp * sizeof(float)); + float *heads_host = (float *)malloc((size_t)q_count * sizeof(float)); + float *reference = (float *)malloc((size_t)q_count * sizeof(float)); + float *scores = (float *)malloc((n_raw + n_comp) * sizeof(float)); + if (!sinks || !q_host || !raw_host || !comp_host || !mask_host || + !heads_host || !reference || !scores) return 1; + + for (uint32_t h = 0; h < n_head; h++) sinks[h] = -0.25f + (float)h * 0.01f; + for (uint64_t i = 0; i < q_count; i++) + q_host[i] = (float)((int)((i * 13u + 5u) % 101u) - 50) / 700.0f; + for (uint64_t i = 0; i < raw_count; i++) + raw_host[i] = (float)((int)((i * 17u + 3u) % 127u) - 63) / 500.0f; + for (uint64_t i = 0; i < comp_count; i++) + comp_host[i] = (float)((int)((i * 19u + 7u) % 131u) - 65) / 550.0f; + for (uint32_t c = 0; c < n_comp; c++) + mask_host[c] = (c % 17u == 0u) ? -1.0e30f : -(float)(c % 7u) * 0.015f; + + const float scale = 1.0f / sqrtf((float)head_dim); + for (uint32_t h = 0; h < n_head; h++) { + const float *qh = q_host + (uint64_t)h * head_dim; + float max_score = sinks[h]; + for (uint32_t r = 0; r < n_raw; r++) { + const uint32_t row = (raw_start + r) % raw_cap; + const float *kv = raw_host + (uint64_t)row * head_dim; + float dot = 0.0f; + for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kv[d]; + scores[r] = dot * scale; + if (scores[r] > max_score) max_score = scores[r]; + } + for (uint32_t c = 0; c < n_comp; c++) { + float s = -3.4e38f; + if (mask_host[c] > -5.0e29f) { + const float *kv = comp_host + (uint64_t)c * head_dim; + float dot = 0.0f; + for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kv[d]; + s = dot * scale + mask_host[c]; + } + scores[n_raw + c] = s; + if (s > max_score) max_score = s; + } + float denom = expf(sinks[h] - max_score); + for (uint32_t r = 0; r < n_raw + n_comp; r++) { + scores[r] = expf(scores[r] - max_score); + denom += scores[r]; + } + for (uint32_t d = 0; d < head_dim; d++) { + float acc = 0.0f; + for (uint32_t r = 0; r < n_raw; r++) { + const uint32_t row = (raw_start + r) % raw_cap; + acc += scores[r] * raw_host[(uint64_t)row * head_dim + d]; + } + for (uint32_t c = 0; c < n_comp; c++) + acc += scores[n_raw + c] * comp_host[(uint64_t)c * head_dim + d]; + reference[(uint64_t)h * head_dim + d] = acc / denom; + } + } + + ds4_gpu_tensor *heads = ds4_gpu_tensor_alloc(q_count * sizeof(float)); + ds4_gpu_tensor *q = ds4_gpu_tensor_alloc(q_count * sizeof(float)); + ds4_gpu_tensor *raw = ds4_gpu_tensor_alloc(raw_count * sizeof(float)); + ds4_gpu_tensor *comp = ds4_gpu_tensor_alloc(comp_count * sizeof(float)); + ds4_gpu_tensor *mask = ds4_gpu_tensor_alloc((uint64_t)n_comp * sizeof(float)); + int rc = 1; + if (heads && q && raw && comp && mask && + ds4_gpu_tensor_write(q, 0, q_host, q_count * sizeof(float)) && + ds4_gpu_tensor_write(raw, 0, raw_host, raw_count * sizeof(float)) && + ds4_gpu_tensor_write(comp, 0, comp_host, comp_count * sizeof(float)) && + ds4_gpu_tensor_write(mask, 0, mask_host, (uint64_t)n_comp * sizeof(float)) && + ds4_gpu_attention_decode_heads_tensor(heads, sinks, n_head * sizeof(float), 0, + q, raw, n_raw, raw_cap, raw_start, + comp, 0, n_comp, mask, 1, + n_head, head_dim) && + ds4_gpu_synchronize() && + ds4_gpu_tensor_read(heads, 0, heads_host, q_count * sizeof(float))) { + float max_abs = 0.0f; + float max_rel = 0.0f; + for (uint64_t i = 0; i < q_count; i++) { + const float abs_err = fabsf(heads_host[i] - reference[i]); + const float rel_err = abs_err / fmaxf(1.0e-4f, fabsf(reference[i])); + if (abs_err > max_abs) max_abs = abs_err; + if (rel_err > max_rel) max_rel = rel_err; + } + fprintf(stderr, "gpu-regression: attention ring reference max_abs=%g max_rel=%g\n", + (double)max_abs, (double)max_rel); + rc = (max_abs <= 2.0e-5f && max_rel <= 2.0e-3f) ? 0 : 1; + } + + ds4_gpu_tensor_free(mask); + ds4_gpu_tensor_free(comp); + ds4_gpu_tensor_free(raw); + ds4_gpu_tensor_free(q); + ds4_gpu_tensor_free(heads); + free(scores); free(reference); free(heads_host); free(mask_host); + free(comp_host); free(raw_host); free(q_host); free(sinks); + return rc; +} + int main(void) { if (!ds4_gpu_init()) return 1; int rc = check_large_topk(); if (check_decode_attention_overflow_path() != 0) rc = 1; + if (check_decode_attention_ring_reference() != 0) rc = 1; ds4_gpu_cleanup(); - if (rc == 0) puts("cuda long-context regression: OK"); + if (rc == 0) puts("GPU long-context regression: OK"); return rc; } diff --git a/tests/shim_test.cu b/tests/shim_test.cu new file mode 100644 index 000000000..8d5cb7dae --- /dev/null +++ b/tests/shim_test.cu @@ -0,0 +1,340 @@ +// Validates ds4_rocm_wmma_gfx906.cuh on real gfx906 hardware. +// Mimics the ds4 MoE tile kernel pattern: blockDim=256, wave=tid>>5 (8 software +// wave32 groups on wave64 hardware), rocwmma fragment load/mma/store. +// Verifies: launch safety (no HSA exception) + numerics vs CPU reference. +#include "ds4_rocm.h" +#include +#include +#include +#include +#include + +#define FULL_WARP_MASK 0xFFFFFFFFFFFFFFFFULL +#define MASK_T uint64_t +#define DS4_ROCM_UNUSED __attribute__((unused)) +#include "rocm/ds4_rocm_common.cuh" +#include "rocm/ds4_rocm_q8.cuh" + +enum { + DS4_ROCM_N_EXPERT = 256u, + DS4_ROCM_MAX_N_EXPERT = 384u, + DS4_ROCM_N_EXPERT_USED = 6u +}; +#define DS4_ROCM_ROUTER_KERNEL_ONLY +#include "rocm/ds4_rocm_router.cuh" + +#define MTILES 8 +#define NB 4 // 4 B fragments per wave like the MoE gate_up_mid kernel (bg0,bu0,bg1,bu1) + +static bool hip_ok(hipError_t err, const char *what) { + if (err == hipSuccess) return true; + std::fprintf(stderr, "%s failed: %s\n", what, hipGetErrorString(err)); + return false; +} + +__global__ void shim_mma_kernel(const half *A, const half *B, float *C, int n_waves) { + const uint32_t tid = threadIdx.x; + const uint32_t wave = tid >> 5u; + if (wave >= (uint32_t)n_waves) return; + + rocwmma::fragment a; + rocwmma::fragment b0, b1, b2, b3; + rocwmma::fragment c0, c1, c2, c3; + + rocwmma::fill_fragment(c0, 0.0f); + rocwmma::fill_fragment(c1, 0.0f); + rocwmma::fill_fragment(c2, 0.0f); + rocwmma::fill_fragment(c3, 0.0f); + + // 2 K-step iterations to exercise accumulation (like k0 loop in ds4). + for (uint32_t k0 = 0; k0 < 2u; k0++) { + rocwmma::load_matrix_sync(a, A + (size_t)wave * 512 + k0 * 256, 16); + rocwmma::load_matrix_sync(b0, B + ((size_t)wave * NB + 0) * 512 + k0 * 256, 16); + rocwmma::load_matrix_sync(b1, B + ((size_t)wave * NB + 1) * 512 + k0 * 256, 16); + rocwmma::load_matrix_sync(b2, B + ((size_t)wave * NB + 2) * 512 + k0 * 256, 16); + rocwmma::load_matrix_sync(b3, B + ((size_t)wave * NB + 3) * 512 + k0 * 256, 16); + rocwmma::mma_sync(c0, a, b0, c0); + rocwmma::mma_sync(c1, a, b1, c1); + rocwmma::mma_sync(c2, a, b2, c2); + rocwmma::mma_sync(c3, a, b3, c3); + } + + rocwmma::store_matrix_sync(C + ((size_t)wave * NB + 0) * 256, c0, 16, rocwmma::mem_row_major); + rocwmma::store_matrix_sync(C + ((size_t)wave * NB + 1) * 256, c1, 16, rocwmma::mem_row_major); + rocwmma::store_matrix_sync(C + ((size_t)wave * NB + 2) * 256, c2, 16, rocwmma::mem_row_major); + rocwmma::store_matrix_sync(C + ((size_t)wave * NB + 3) * 256, c3, 16, rocwmma::mem_row_major); +} + +// CPU reference: C[i][j] = sum over 2 K-tiles of sum_k A[i][k]*B[k][j] +static void ref_gemm(const std::vector &A, const std::vector &B, + std::vector &C, int wave, int nb) { + for (int i = 0; i < 16; i++) + for (int j = 0; j < 16; j++) { + float acc = 0.0f; + for (int kt = 0; kt < 2; kt++) + for (int k = 0; k < 16; k++) + acc += A[(size_t)wave * 512 + kt * 256 + i * 16 + k] * + B[((size_t)wave * NB + nb) * 512 + kt * 256 + k * 16 + j]; + C[((size_t)wave * NB + nb) * 256 + i * 16 + j] = acc; + } +} + +static bool test_q8_wave64_row_packing() { + const uint64_t in_dim = 4096; + const uint64_t out_dim = 4097; + const uint64_t blocks = in_dim / 32; + const size_t weight_bytes = (size_t)out_dim * blocks * 34u; + + std::vector hw(weight_bytes); + std::vector hxq(blocks * 32u); + std::vector hxscale(blocks); + std::vector h1(out_dim), h2(out_dim); + for (uint64_t row = 0; row < out_dim; row++) { + for (uint64_t block = 0; block < blocks; block++) { + const half scale = __float2half(0.0025f * (float)(1u + ((row + block) % 13u))); + unsigned char *dst = hw.data() + ((size_t)row * blocks + block) * 34u; + std::memcpy(dst, &scale, sizeof(scale)); + for (uint32_t i = 0; i < 32u; i++) + dst[2u + i] = (unsigned char)(int8_t)(((row * 3u + block * 5u + i * 7u) % 255u) - 127); + } + } + for (size_t i = 0; i < hxq.size(); i++) hxq[i] = (int8_t)(((i * 11u) % 255u) - 127); + for (size_t i = 0; i < hxscale.size(); i++) hxscale[i] = 0.001f * (float)(1u + (i % 17u)); + + unsigned char *dw = nullptr; + int8_t *dxq = nullptr; + float *dxscale = nullptr, *d1 = nullptr, *d2 = nullptr; + const bool allocated = + hip_ok(hipMalloc(&dw, hw.size()), "q8 malloc weights") && + hip_ok(hipMalloc(&dxq, hxq.size()), "q8 malloc activation") && + hip_ok(hipMalloc(&dxscale, hxscale.size() * sizeof(float)), "q8 malloc scales") && + hip_ok(hipMalloc(&d1, h1.size() * sizeof(float)), "q8 malloc rpb1") && + hip_ok(hipMalloc(&d2, h2.size() * sizeof(float)), "q8 malloc rpb2"); + if (!allocated) { + (void)hipFree(d2); (void)hipFree(d1); (void)hipFree(dxscale); + (void)hipFree(dxq); (void)hipFree(dw); + return false; + } + bool ok = + hip_ok(hipMemcpy(dw, hw.data(), hw.size(), hipMemcpyHostToDevice), "q8 copy weights") && + hip_ok(hipMemcpy(dxq, hxq.data(), hxq.size(), hipMemcpyHostToDevice), "q8 copy activation") && + hip_ok(hipMemcpy(dxscale, hxscale.data(), hxscale.size() * sizeof(float), hipMemcpyHostToDevice), + "q8 copy scales"); + if (ok) { + matmul_q8_0_preq_rows_w32_kernel<<<(unsigned)out_dim, 32>>>( + d1, dw, dxq, dxscale, in_dim, out_dim, blocks, 1u, 1); + matmul_q8_0_preq_rows_w32_kernel<<<((unsigned)out_dim + 1u) / 2u, 64>>>( + d2, dw, dxq, dxscale, in_dim, out_dim, blocks, 2u, 1); + ok = hip_ok(hipDeviceSynchronize(), "q8 packed rows launch") && + hip_ok(hipMemcpy(h1.data(), d1, h1.size() * sizeof(float), hipMemcpyDeviceToHost), + "q8 copy rpb1") && + hip_ok(hipMemcpy(h2.data(), d2, h2.size() * sizeof(float), hipMemcpyDeviceToHost), + "q8 copy rpb2"); + } + size_t mismatches = 0; + if (ok) { + for (size_t i = 0; i < h1.size(); i++) + if (std::memcmp(&h1[i], &h2[i], sizeof(float)) != 0) mismatches++; + } + std::printf("Q8 WAVE64 ROW PACKING %s (%zu rows, %zu bit mismatches)\n", + ok && mismatches == 0 ? "PASS" : "FAIL", h1.size(), mismatches); + (void)hipFree(d2); (void)hipFree(d1); (void)hipFree(dxscale); + (void)hipFree(dxq); (void)hipFree(dw); + return ok && mismatches == 0; +} + +static bool test_f16_pair_workgroup_sizing() { + const uint32_t in_dim = 1024; + const uint32_t out_dim = 513; + std::vector hw0((size_t)in_dim * out_dim), hw1(hw0.size()); + std::vector hx(in_dim), href0(out_dim), href1(out_dim), h0(out_dim), h1(out_dim); + for (size_t i = 0; i < hw0.size(); i++) { + hw0[i] = __float2half((float)((int)(i % 31u) - 15) / 32.0f); + hw1[i] = __float2half((float)((int)(i % 29u) - 14) / 32.0f); + } + for (uint32_t i = 0; i < in_dim; i++) hx[i] = (float)((int)(i % 23u) - 11) / 16.0f; + + half *dw0 = nullptr, *dw1 = nullptr; + float *dx = nullptr, *do0 = nullptr, *do1 = nullptr; + const bool allocated = + hip_ok(hipMalloc(&dw0, hw0.size() * sizeof(half)), "f16 pair malloc w0") && + hip_ok(hipMalloc(&dw1, hw1.size() * sizeof(half)), "f16 pair malloc w1") && + hip_ok(hipMalloc(&dx, hx.size() * sizeof(float)), "f16 pair malloc x") && + hip_ok(hipMalloc(&do0, href0.size() * sizeof(float)), "f16 pair malloc out0") && + hip_ok(hipMalloc(&do1, href1.size() * sizeof(float)), "f16 pair malloc out1"); + if (!allocated) { + (void)hipFree(do1); (void)hipFree(do0); (void)hipFree(dx); + (void)hipFree(dw1); (void)hipFree(dw0); + return false; + } + bool ok = + hip_ok(hipMemcpy(dw0, hw0.data(), hw0.size() * sizeof(half), hipMemcpyHostToDevice), + "f16 pair copy w0") && + hip_ok(hipMemcpy(dw1, hw1.data(), hw1.size() * sizeof(half), hipMemcpyHostToDevice), + "f16 pair copy w1") && + hip_ok(hipMemcpy(dx, hx.data(), hx.size() * sizeof(float), hipMemcpyHostToDevice), + "f16 pair copy x"); + if (ok) { + matmul_f16_pair_f32_sharedx_warp_rows_w32_kernel<<<(out_dim + 31u) / 32u, + 1024u, (size_t)in_dim * sizeof(float)>>>(do0, do1, dw0, dw1, dx, in_dim, out_dim); + ok = hip_ok(hipDeviceSynchronize(), "f16 pair rpb32 launch") && + hip_ok(hipMemcpy(href0.data(), do0, href0.size() * sizeof(float), hipMemcpyDeviceToHost), + "f16 pair copy reference 0") && + hip_ok(hipMemcpy(href1.data(), do1, href1.size() * sizeof(float), hipMemcpyDeviceToHost), + "f16 pair copy reference 1"); + } + if (ok) { + matmul_f16_pair_f32_sharedx_warp_rows_w32_kernel<<<(out_dim + 7u) / 8u, + 256u, (size_t)in_dim * sizeof(float)>>>(do0, do1, dw0, dw1, dx, in_dim, out_dim); + ok = hip_ok(hipDeviceSynchronize(), "f16 pair rpb8 launch") && + hip_ok(hipMemcpy(h0.data(), do0, h0.size() * sizeof(float), hipMemcpyDeviceToHost), + "f16 pair copy rpb8 0") && + hip_ok(hipMemcpy(h1.data(), do1, h1.size() * sizeof(float), hipMemcpyDeviceToHost), + "f16 pair copy rpb8 1"); + } + size_t mismatches = 0; + if (ok) { + for (size_t i = 0; i < h0.size(); i++) + if (std::memcmp(&h0[i], &href0[i], sizeof(float)) != 0 || + std::memcmp(&h1[i], &href1[i], sizeof(float)) != 0) mismatches++; + } + std::printf("F16 PAIR WORKGROUP SIZING %s (%zu rows, %zu bit mismatches)\n", + ok && mismatches == 0 ? "PASS" : "FAIL", h0.size(), mismatches); + (void)hipFree(do1); (void)hipFree(do0); (void)hipFree(dx); + (void)hipFree(dw1); (void)hipFree(dw0); + return ok && mismatches == 0; +} + +template +static bool test_router_wave64() { + std::vector hlogits(N_EXPERT), hbias(N_EXPERT); + std::vector href_sel(DS4_ROCM_N_EXPERT_USED), h_sel(DS4_ROCM_N_EXPERT_USED); + std::vector href_weights(DS4_ROCM_N_EXPERT_USED), h_weights(DS4_ROCM_N_EXPERT_USED); + std::vector href_probs(N_EXPERT), h_probs(N_EXPERT); + for (uint32_t i = 0; i < N_EXPERT; i++) { + hlogits[i] = (float)((int)((i * 17u) % 101u) - 50) / 9.0f; + hbias[i] = (float)((int)((i * 7u) % 31u) - 15) / 100.0f; + } + /* Deliberate equal scores exercise the expert-index tie break. */ + hlogits[3] = hlogits[131] = 4.0f; + hbias[3] = hbias[131] = 0.25f; + + float *dlogits = nullptr, *dbias = nullptr, *dweights = nullptr, *dprobs = nullptr; + int32_t *dselected = nullptr; + const bool allocated = + hip_ok(hipMalloc(&dlogits, hlogits.size() * sizeof(float)), "router malloc logits") && + hip_ok(hipMalloc(&dbias, hbias.size() * sizeof(float)), "router malloc bias") && + hip_ok(hipMalloc(&dselected, h_sel.size() * sizeof(int32_t)), "router malloc selected") && + hip_ok(hipMalloc(&dweights, h_weights.size() * sizeof(float)), "router malloc weights") && + hip_ok(hipMalloc(&dprobs, h_probs.size() * sizeof(float)), "router malloc probs"); + if (!allocated) { + (void)hipFree(dprobs); (void)hipFree(dweights); (void)hipFree(dselected); + (void)hipFree(dbias); (void)hipFree(dlogits); + return false; + } + bool ok = + hip_ok(hipMemcpy(dlogits, hlogits.data(), hlogits.size() * sizeof(float), hipMemcpyHostToDevice), + "router copy logits") && + hip_ok(hipMemcpy(dbias, hbias.data(), hbias.size() * sizeof(float), hipMemcpyHostToDevice), + "router copy bias"); + if (ok) { + router_select_warp_topk_kernel<<<1, dim3(32, 4, 1)>>>( + dselected, dweights, dprobs, dbias, nullptr, dlogits, nullptr, 0, + 0, 1, 1.5f, 1, 0); + ok = hip_ok(hipDeviceSynchronize(), "router wave32 launch") && + hip_ok(hipMemcpy(href_sel.data(), dselected, href_sel.size() * sizeof(int32_t), + hipMemcpyDeviceToHost), "router copy selected reference") && + hip_ok(hipMemcpy(href_weights.data(), dweights, href_weights.size() * sizeof(float), + hipMemcpyDeviceToHost), "router copy weights reference") && + hip_ok(hipMemcpy(href_probs.data(), dprobs, href_probs.size() * sizeof(float), + hipMemcpyDeviceToHost), "router copy probs reference"); + } + if (ok) { + router_select_warp_topk_kernel<<<1, dim3(64, 1, 1)>>>( + dselected, dweights, dprobs, dbias, nullptr, dlogits, nullptr, 0, + 0, 1, 1.5f, 1, 0); + ok = hip_ok(hipDeviceSynchronize(), "router wave64 launch") && + hip_ok(hipMemcpy(h_sel.data(), dselected, h_sel.size() * sizeof(int32_t), + hipMemcpyDeviceToHost), "router copy selected wave64") && + hip_ok(hipMemcpy(h_weights.data(), dweights, h_weights.size() * sizeof(float), + hipMemcpyDeviceToHost), "router copy weights wave64") && + hip_ok(hipMemcpy(h_probs.data(), dprobs, h_probs.size() * sizeof(float), + hipMemcpyDeviceToHost), "router copy probs wave64"); + } + size_t mismatches = 0; + if (ok) { + for (size_t i = 0; i < h_sel.size(); i++) { + mismatches += h_sel[i] != href_sel[i]; + mismatches += std::memcmp(&h_weights[i], &href_weights[i], sizeof(float)) != 0; + } + for (size_t i = 0; i < h_probs.size(); i++) + mismatches += std::memcmp(&h_probs[i], &href_probs[i], sizeof(float)) != 0; + } + std::printf("ROUTER WAVE64 %s (%u experts, %zu bit mismatches)\n", + ok && mismatches == 0 ? "PASS" : "FAIL", N_EXPERT, mismatches); + (void)hipFree(dprobs); (void)hipFree(dweights); (void)hipFree(dselected); + (void)hipFree(dbias); (void)hipFree(dlogits); + return ok && mismatches == 0; +} + +int main() { + const int waves = MTILES; + std::vector hA(waves * 512), hB(waves * NB * 512); + std::vector hC_ref(waves * NB * 256, 0.0f); + std::vector hC(waves * NB * 256, -1.0f); + + srand(42); + for (auto &v : hA) v = ((rand() % 17) - 8) / 8.0f; + for (auto &v : hB) v = ((rand() % 13) - 6) / 8.0f; + + for (int w = 0; w < waves; w++) + for (int nb = 0; nb < NB; nb++) ref_gemm(hA, hB, hC_ref, w, nb); + + half *dA = nullptr, *dB = nullptr; + float *dC = nullptr; + std::vector hAh(hA.size()), hBh(hB.size()); + for (size_t i = 0; i < hA.size(); i++) hAh[i] = __float2half(hA[i]); + for (size_t i = 0; i < hB.size(); i++) hBh[i] = __float2half(hB[i]); + + if (!hip_ok(hipMalloc(&dA, hAh.size() * sizeof(half)), "malloc A") || + !hip_ok(hipMalloc(&dB, hBh.size() * sizeof(half)), "malloc B") || + !hip_ok(hipMalloc(&dC, hC.size() * sizeof(float)), "malloc C") || + !hip_ok(hipMemcpy(dA, hAh.data(), hAh.size() * sizeof(half), hipMemcpyHostToDevice), "copy A") || + !hip_ok(hipMemcpy(dB, hBh.data(), hBh.size() * sizeof(half), hipMemcpyHostToDevice), "copy B") || + !hip_ok(hipMemset(dC, 0, hC.size() * sizeof(float)), "clear C")) { + (void)hipFree(dC); + (void)hipFree(dB); + (void)hipFree(dA); + return 2; + } + + shim_mma_kernel<<<1, 256>>>(dA, dB, dC, waves); + hipError_t err = hipDeviceSynchronize(); + if (err != hipSuccess) { + printf("SHIM TEST LAUNCH FAILURE: %s\n", hipGetErrorString(err)); + return 3; + } + if (!hip_ok(hipMemcpy(hC.data(), dC, hC.size() * sizeof(float), hipMemcpyDeviceToHost), "copy C")) { + (void)hipFree(dC); + (void)hipFree(dB); + (void)hipFree(dA); + return 3; + } + + double max_err = 0.0; + for (size_t i = 0; i < hC.size(); i++) { + double e = fabs((double)hC[i] - (double)hC_ref[i]); + if (e > max_err) max_err = e; + } + printf("max_err=%g (256 tiles checked, %zu values)\n", max_err, hC.size()); + (void)hipFree(dC); + (void)hipFree(dB); + (void)hipFree(dA); + if (max_err > 1e-3) { printf("SHIM TEST NUMERIC MISMATCH\n"); return 4; } + if (!test_q8_wave64_row_packing()) return 5; + if (!test_f16_pair_workgroup_sizing()) return 6; + if (!test_router_wave64()) return 7; + if (!test_router_wave64()) return 8; + printf("SHIM TEST PASS\n"); + return 0; +}