{"id":29992,"date":"2026-08-06T10:06:51","date_gmt":"2026-08-06T02:06:51","guid":{"rendered":"https:\/\/aif.amtbbs.org\/?p=29992"},"modified":"2026-08-06T10:06:51","modified_gmt":"2026-08-06T02:06:51","slug":"ai-infra-%e8%bf%9b%e9%98%b6%ef%bc%9a%e5%a6%82%e4%bd%95%e8%ae%a9%e5%a4%a7%e6%a8%a1%e5%9e%8b%e8%be%93%e5%87%ba%e7%a1%ae%e5%ae%9a%e7%9a%84%e7%bb%93%e6%9e%9c","status":"publish","type":"post","link":"https:\/\/aif.amtbbs.org\/index.php\/2026\/08\/06\/29992\/","title":{"rendered":"AI Infra \u8fdb\u9636\uff1a\u5982\u4f55\u8ba9\u5927\u6a21\u578b\u8f93\u51fa\u786e\u5b9a\u7684\u7ed3\u679c"},"content":{"rendered":"<div class=\"article-desc\">\u5927\u6a21\u578b\u5728\u63a8\u7406\u8fc7\u7a0b\u4e2d\u5bfc\u81f4Batch Variance\u7684\u6839\u672c\u539f\u56e0\uff0c\u5728\u4e8e\u5e95\u5c42\u7b97\u5b50\u4e3a\u6700\u5927\u5316\u786c\u4ef6\u8d44\u6e90\u5229\u7528\u7387\uff0c\u52a8\u6001\u8c03\u6574\u4e86\u89c4\u7ea6\u8f74\uff08\u5982 GEMM \u4e2d\u7684 Split-K \u6216 Attention \u4e2d\u7684 KV \u8f74\uff09\u7684\u5207\u5206\u7b56\u7565\uff0c\u8fdb\u800c\u6539\u53d8\u4e86\u6d6e\u70b9\u6570\u7d2f\u52a0\u6811\u7684\u62d3\u6251\u7ed3\u6784\u3002<\/div>\n<div id=\"postspictures\" class=\"article-content\">\n<div id=\"container\" class=\"container am-engine\" data-v-01a18e2f=\"\" data-element=\"root\">\n<p>\u4f5c\u8005 | binnnliu<\/p>\n<p>\u4f60\u6709\u6ca1\u6709\u53d1\u73b0\u8ddf\u5927\u6a21\u578b\u5bf9\u8bdd\uff0c\u540c\u6837\u7684\u63d0\u793a\u8bcd\u6bcf\u6b21\u7ed3\u679c\u90fd\u4e0d\u4e00\u6837\u3002 \u8fd9\u7b26\u5408\u6211\u4eec\u5bf9\u4e8e\u201c\u5927\u6a21\u578b\u7684\u672c\u8d28\u5c31\u662f\u4e2a\u731c\u8bcd\u5668\u201d\u7684\u4e00\u8d2f\u8ba4\u77e5\u3002<\/p>\n<p><img data-dominant-color=\"6195e6\" data-has-transparency=\"false\" style=\"--dominant-color: #6195e6;\" loading=\"lazy\" decoding=\"async\" class=\"not-transparent alignnone size-full wp-image-29993\" src=\"https:\/\/aiforumimage.oss-cn-shanghai.aliyuncs.com\/wp-content\/uploads\/2026\/08\/b777a2c41560a7116d90599537eb6836288690.jpg\" width=\"1248\" height=\"696\" srcset=\"https:\/\/aiforumimage.oss-cn-shanghai.aliyuncs.com\/wp-content\/uploads\/2026\/08\/b777a2c41560a7116d90599537eb6836288690.jpg 1248w, https:\/\/aiforumimage.oss-cn-shanghai.aliyuncs.com\/wp-content\/uploads\/2026\/08\/b777a2c41560a7116d90599537eb6836288690-300x167.jpg 300w, https:\/\/aiforumimage.oss-cn-shanghai.aliyuncs.com\/wp-content\/uploads\/2026\/08\/b777a2c41560a7116d90599537eb6836288690-1024x571.jpg 1024w, https:\/\/aiforumimage.oss-cn-shanghai.aliyuncs.com\/wp-content\/uploads\/2026\/08\/b777a2c41560a7116d90599537eb6836288690-150x84.jpg 150w, https:\/\/aiforumimage.oss-cn-shanghai.aliyuncs.com\/wp-content\/uploads\/2026\/08\/b777a2c41560a7116d90599537eb6836288690-768x428.jpg 768w\" sizes=\"auto, (max-width: 1248px) 100vw, 1248px\" \/><\/p>\n<p>\u7136\u800c\uff0c\u53ef\u91cd\u590d\u6027\u662f\u79d1\u5b66\u8fdb\u6b65\u7684\u57fa\u77f3\u3002\u56e0\u6b64\u8ba9\u5927\u6a21\u578b\u8f93\u51fa\u5b8c\u5168\u786e\u5b9a\u7684\u7ed3\u679c\u662f\u4e00\u4e2a\u975e\u5e38\u503c\u5f97\u7814\u7a76\u7684\u95ee\u9898\uff0c\u7279\u522b\u662f\u5f3a\u5316\u5b66\u4e60\u9700\u8981\u786e\u5b9a\u6027\u7684 Rollout\uff0c\u6765\u4fdd\u8bc1\u5b9e\u9a8c\u7684\u53ef\u590d\u73b0\u6027\u548c\u8bad\u7ec3\u8fc7\u7a0b\u7684\u7a33\u5b9a\u6027\u3002<\/p>\n<blockquote><p>&nbsp;<\/p>\n<p>Reproducibility is a bedrock of scientific progress.<\/p>\n<p>&nbsp;<\/p><\/blockquote>\n<p>\u5176\u5b9e\u5f88\u591a\u540c\u5b66\u7b2c\u4e00\u65f6\u95f4\u4f1a\u60f3\u5230\uff1a \u540c\u6837\u7684\u786c\u4ef6\uff0c\u540c\u6837\u7684\u63d0\u793a\u8bcd\uff1a<\/p>\n<ul data-id=\"u738a58b-44aTuomf\">\n<li data-id=\"ld70c578-nettIP00\">\u662f\u4e0d\u662f\u628a Temperature \u8bbe\u7f6e\u4e3a 0\uff0c\u5173\u6389\u968f\u673a\u91c7\u6837\u5c31\u597d\u4e86\uff1f<\/li>\n<li data-id=\"ld70c578-HGCPSJd7\">\u5982\u679c\u4e1a\u52a1\u9700\u8981 Temperature &gt; 0\uff0c\u90a3\u628a\u968f\u673a\u6570\u79cd\u5b50\uff08Seed\uff09\u9501\u6b7b\uff0c\u662f\u4e0d\u662f\u4e5f\u80fd\u4fdd\u8bc1\u6bcf\u6b21\u8f93\u51fa\u4e00\u6837\uff1f<\/li>\n<\/ul>\n<p>\u6211\u5df2\u7ecf\u6d4b\u8bd5\u8fc7\u4e86\uff0c\u73af\u5883\uff1avllm serve Qwen\/Qwen3-8B\u00a0\u7b54\u6848\uff1a \u662f\u4e5f\u4e0d\u662f\u3002<\/p>\n<ul data-id=\"u738a58b-wH1lGCKC\">\n<li data-id=\"ld70c578-nW2WRPgs\">\u662f\uff08Run-to-run \u786e\u5b9a\u6027\uff09\uff1a \u540c\u4e00\u65f6\u95f4\uff0c\u6ca1\u6709\u5176\u4ed6\u8bf7\u6c42\u65f6\uff0c\u591a\u6b21\u540c\u6837\u7684\u63d0\u793a\u8bcd\u8bf7\u6c42\uff0c\u56de\u7b54\u662f\u4e00\u6837\u7684\uff1b<\/li>\n<li data-id=\"ld70c578-1MQZXWNy\">\u4e0d\u662f\uff08Batch Invariance \/ \u6279\u6b21\u4e0d\u53d8\u6027\uff09\uff1a \u540c\u4e00\u65f6\u95f4\uff0c\u6709\u5176\u4ed6\u8bf7\u6c42\u65f6\uff0c\u591a\u6b21\u540c\u6837\u7684\u63d0\u793a\u8bcd\u8bf7\u6c42\uff0c\u56de\u7b54\u662f\u4e0d\u4e00\u6837\u7684\uff1b<\/li>\n<\/ul>\n<p>\u4ec0\u4e48\uff01\u96be\u9053\u4e0d\u540c\u7684\u8bf7\u6c42\u95f4\u4f1a\u76f8\u4e92\u5f71\u54cd\uff1f\u4e0d\u53ef\u80fd\uff0c\u4e0d\u540c\u8bf7\u6c42\u95f4\u7684\u6ce8\u610f\u529b\u673a\u5236\u662f\u4e25\u683c\u7269\u7406\u9694\u79bb\uff0c\u4e0d\u5b58\u5728\u4e0a\u4e0b\u6587\u6c61\u67d3\u7684\u95ee\u9898\u3002<\/p>\n<p>\u5176\u5b9e\u8fd9\u4e00\u5207\u7684\u6839\u6e90\u662f\u63a8\u7406\u5f15\u64ce\u5e95\u5c42\u7684\u52a8\u6001\u7ec4\u6279\u4e0e\u7b97\u5b50\u8c03\u5ea6\u7b56\u7565\u5f15\u53d1\u4e86\u6d6e\u70b9\u52a0\u6cd5\u7684\u987a\u5e8f\u53d8\u5316\uff0c\u800c\u6d6e\u70b9\u52a0\u6cd5\u4e0d\u6ee1\u8db3\u7ed3\u5408\u5f8b\uff1a(a + b) + c \u2260 a + (b + c)\u3002<\/p>\n<p>\u4ece\u7cfb\u7edf\u8bbe\u8ba1\u7684\u89d2\u5ea6\u6765\u770b\uff0c\u5bfc\u81f4\u7ed3\u679c\u6ce2\u52a8\u7684\u6839\u672c\u539f\u56e0\uff0c\u4e0d\u5728\u4e8e\u4fe1\u606f\u5e72\u6270\uff0c\u800c\u5728\u4e8e\u7cfb\u7edf\u4e3a\u4e86\u63a9\u76d6\u8bbf\u5b58\u5ef6\u8fdf\uff0c\u5728\u5e95\u5c42\u89e6\u53d1\u4e86\u975e\u786e\u5b9a\u6027\u7684\u786c\u4ef6\u7ea7\u5e76\u884c\u4f18\u5316\uff08Non-deterministic Hardware-level Parallel Optimizations\uff09\u3002<\/p>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/42243cb339515a435bd512183195d03f3a4597.webp\" data-type=\"block\" \/><\/p>\n<p>\u76f8\u5bf9\u4e8e\u5b9a\u70b9\u6570\uff0c\u6d6e\u70b9\u6570\u7684\u672c\u8d28\u662f\u79d1\u5b66\u8ba1\u6570\u6cd5\uff0c\u4e3b\u8981\u662f\u4e3a\u4e86\u5728\u6709\u9650\u7684\u4f4d\u6570\u5185\uff08\u5982\u4ec5\u4ec5 16 bit \u6216 32 bit\uff09\uff0c\u5b9e\u73b0\u52a8\u6001\u7684\u6570\u503c\u8303\u56f4\u4e0e\u7edd\u5bf9\u7cbe\u5ea6\u95f4\u7684trade-off\u3002\u6307\u6570\u8d8a\u5927\u80fd\u8868\u793a\u7684\u6570\u503c\u8303\u56f4\u4e5f\u5c31\u8d8a\u5927\uff0c\u540c\u65f6\u56e0\u4e3a\u5c3e\u6570\u4f4d\u6570\u957f\u5ea6\u56fa\u5b9a\uff0c\u6307\u6570\u8d8a\u5927\u80fd\u8868\u793a\u7684\u7edd\u5bf9\u7cbe\u5ea6\u4e5f\u5c31\u8d8a\u4f4e\u3002\u7edd\u5bf9\u7cbe\u5ea6\uff08\u5373\u6b65\u957f\uff09\u7684\u8ba1\u7b97\u516c\u5f0f\u4e3a\uff1aULP(x) = 2E\u00a0&#8211;\u00a0M\uff0c\u5176\u4e2d\u00a0E\u00a0\u662f\u8be5\u6d6e\u70b9\u6570\u7684\u771f\u5b9e\u6307\u6570\uff0cM\u00a0\u662f\u5c3e\u6570\u4f4d\u6570\u3002<\/p>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/56afdd502d5c1dc8e5a447fd0029f8bff23b43.webp\" data-type=\"block\" \/><\/p>\n<p>\u6839\u636e IEEE 754 \u6807\u51c6\uff0cFP16 \u5305\u542b 1bit\u7b26\u53f7\u4f4d\u30015bit\u6307\u6570\u4f4d\u548c 10bit\u5c3e\u6570\u4f4d\u3002\u5f53\u6570\u503c\u4e3a 2048\uff08\u5373\u00a0211\uff09\u65f6\uff0c\u5176\u89c4\u683c\u5316\u8868\u793a\u4e3a\u00a01.00000000002\u00a0\u00d7 211\u3002\u6b64\u65f6\uff0c\u76f8\u90bb\u53ef\u8868\u793a\u6570\u4e4b\u95f4\u7684\u6b65\u957f\uff08ULP\uff09\u589e\u81f3 2\uff0c\u8fd9\u610f\u5473\u7740 FP16 \u7684\u4e0b\u4e00\u4e2a\u53ef\u8868\u793a\u6570\u4e3a 2050\uff0c\u65e0\u6cd5\u7cbe\u786e\u8868\u793a 2049\u3002\u5f53\u5e95\u5c42\u786c\u4ef6\u6267\u884c\u00a02048 + 1\u00a0\u65f6\uff0c\u8fd0\u7b97\u5668\u5185\u90e8\u4f1a\u501f\u52a9\u66f4\u5bbd\u7684 GRS \u6269\u5c55\u4f4d\uff08\u4fdd\u62a4\u4f4d\u3001\u820d\u5165\u4f4d\u3001\u7c98\u6ede\u4f4d\uff09\u5f97\u51fa\u7cbe\u786e\u7ed3\u679c 2049\u3002\u4f46\u5f53\u7ed3\u679c\u9700\u8981\u5199\u56de 10 \u4f4d\u5c3e\u6570\u65f6\uff0c\u7531\u4e8e 2049 \u6070\u597d\u4f4d\u4e8e 2048 \u4e0e 2050 \u7684\u6b63\u4e2d\u95f4\uff0c\u7cfb\u7edf\u89e6\u53d1\u4e86\u5411\u5076\u6570\u820d\u5165\uff08Round to Nearest, ties to Even\uff09\u7684\u89c4\u5219\uff0c\u6700\u7ec8\u7ed3\u679c\u88ab\u5f3a\u884c\u820d\u5165\u56de 2048\u3002<\/p>\n<p>\u8fd9\u91cc\u9700\u8981\u8bf4\u660e\u662f\uff1a<\/p>\n<p>\u8fd9\u4e2a\u4f8b\u5b50\u53ea\u662f\u4e3a\u4e86\u8bf4\u660e\u201c\u6d6e\u70b9\u52a0\u6cd5\u4e0d\u6ee1\u8db3\u7ed3\u5408\u5f8b\u201d\u3002\u5728\u5b9e\u9645 \u63a8\u7406\u8fc7\u7a0b\u4e2d\uff0c GEMM \u662f FP16\/BF16 \u8f93\u5165\u3001FP32 \u7d2f\u52a0\uff1bbatch variance \u4f9d\u7136\u5b58\u5728\uff0c\u6839\u56e0\u662f reduction topology \u6539\u53d8\uff0c\u800c\u4e0d\u53ea\u662f FP16 \u5b58\u50a8\u7cbe\u5ea6\u4f4e\u3002<\/p>\n<table class=\"data-table\" data-transient-attributes=\"class\" data-width=\"676.997px\">\n<colgroup data-id=\"c7104f7d-9J0PXdOU\">\n<col span=\"1\" width=\"112.83\" data-id=\"cd89ecb0-cBXYVI97\" \/>\n<col span=\"1\" width=\"112.83\" data-id=\"cd89ecb0-0WiJ0ijl\" \/>\n<col span=\"1\" width=\"112.83\" data-id=\"cd89ecb0-5hSEUFVR\" \/>\n<col span=\"1\" width=\"112.83\" data-id=\"cd89ecb0-iRO7oEke\" \/>\n<col span=\"1\" width=\"112.83\" data-id=\"cd89ecb0-UBXSWXW0\" \/>\n<col span=\"1\" width=\"112.847\" data-id=\"cd89ecb0-KLI8YQHP\" \/><\/colgroup>\n<tbody data-id=\"t6d5e859-13Xl7kD2\">\n<tr data-id=\"t31e458f-GBY7DZ2L\">\n<td data-id=\"t6267798-4mDENmRg\" data-transient-attributes=\"table-cell-selection\">\u6570\u503c<\/td>\n<td data-id=\"t6267798-LHFg7B3a\" data-transient-attributes=\"table-cell-selection\">\u89c4\u683c\u5316\u8868\u793a<\/td>\n<td data-id=\"t6267798-e6XamMPJ\" data-transient-attributes=\"table-cell-selection\">\u771f\u5b9e\u6307\u6570<\/td>\n<td data-id=\"t6267798-H84fU5hZ\" data-transient-attributes=\"table-cell-selection\">\u5b9e\u9645\u5b58\u50a8\u6307\u6570<\/td>\n<td data-id=\"t6267798-OScOAS38\" data-transient-attributes=\"table-cell-selection\">\u7269\u7406\u5c3e\u6570\u4f4d (10 bit)<\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-9lPlP5rL\" data-transient-attributes=\"table-cell-selection\">\u7d2f\u52a0\u62c6\u89e3<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-RSGkm3c8\">\n<td data-id=\"t6267798-KlMYHDnf\" data-transient-attributes=\"table-cell-selection\">0.5<\/td>\n<td data-id=\"t6267798-RKiRqSro\" data-transient-attributes=\"table-cell-selection\">1.00000000002\u00a0\u00d7 2-1<\/td>\n<td data-id=\"t6267798-i8sAeo3B\" data-transient-attributes=\"table-cell-selection\">-1<\/td>\n<td data-id=\"t6267798-YgAZeiFP\" data-transient-attributes=\"table-cell-selection\">14 (<code>01110<\/code>)<\/td>\n<td data-id=\"t6267798-F9TOwHq8\" data-transient-attributes=\"table-cell-selection\"><code>0000000000<\/code><\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-hswQjjw3\" data-transient-attributes=\"table-cell-selection\">2-1<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-e4oTdnNa\">\n<td data-id=\"t6267798-Xr5nvVqe\" data-transient-attributes=\"table-cell-selection\">1<\/td>\n<td data-id=\"t6267798-wfcD9qGo\" data-transient-attributes=\"table-cell-selection\">1.00000000002\u00a0\u00d7 20<\/td>\n<td data-id=\"t6267798-IfJqzHic\" data-transient-attributes=\"table-cell-selection\">0<\/td>\n<td data-id=\"t6267798-V39VGPcq\" data-transient-attributes=\"table-cell-selection\">15 (<code>01111<\/code>)<\/td>\n<td data-id=\"t6267798-LnqmVEnX\" data-transient-attributes=\"table-cell-selection\"><code>0000000000<\/code><\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-HOlrCgYv\" data-transient-attributes=\"table-cell-selection\">20<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-cSpHTbm4\">\n<td data-id=\"t6267798-607csJ1U\" data-transient-attributes=\"table-cell-selection\">2<\/td>\n<td data-id=\"t6267798-H38h08XE\" data-transient-attributes=\"table-cell-selection\">1.00000000002\u00a0\u00d7 21<\/td>\n<td data-id=\"t6267798-EZFQnTj8\" data-transient-attributes=\"table-cell-selection\">1<\/td>\n<td data-id=\"t6267798-mxGTvrIS\" data-transient-attributes=\"table-cell-selection\">16 (<code>10000<\/code>)<\/td>\n<td data-id=\"t6267798-wmVAEdVb\" data-transient-attributes=\"table-cell-selection\"><code>0000000000<\/code><\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-26c7gIw6\" data-transient-attributes=\"table-cell-selection\">21<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-8kAYVAQD\">\n<td data-id=\"t6267798-s2LArvOX\" data-transient-attributes=\"table-cell-selection\">2046<\/td>\n<td data-id=\"t6267798-EvzCod6L\" data-transient-attributes=\"table-cell-selection\">1.11111111102\u00a0\u00d7 210<\/td>\n<td data-id=\"t6267798-nzt73oGm\" data-transient-attributes=\"table-cell-selection\">10<\/td>\n<td data-id=\"t6267798-FBirbIbR\" data-transient-attributes=\"table-cell-selection\">25 (<code>11001<\/code>)<\/td>\n<td data-id=\"t6267798-9Cls4lgw\" data-transient-attributes=\"table-cell-selection\"><code>1111111110<\/code><\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-HIpeweCU\" data-transient-attributes=\"table-cell-selection\">210\u00a0+ 29\u00a0+ \u22ef + 21<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-kdbSbMRJ\">\n<td data-id=\"t6267798-Lcaen4jY\" data-transient-attributes=\"table-cell-selection\">2047<\/td>\n<td data-id=\"t6267798-CPja4YM7\" data-transient-attributes=\"table-cell-selection\">1.11111111112\u00a0\u00d7 210<\/td>\n<td data-id=\"t6267798-EDpL2bX1\" data-transient-attributes=\"table-cell-selection\">10<\/td>\n<td data-id=\"t6267798-O6py2ECx\" data-transient-attributes=\"table-cell-selection\">25 (<code>11001<\/code>)<\/td>\n<td data-id=\"t6267798-rTKGYYIQ\" data-transient-attributes=\"table-cell-selection\"><code>1111111111<\/code><\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-e8iv7f9H\" data-transient-attributes=\"table-cell-selection\">210\u00a0+ 29\u00a0+ \u22ef + 20<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-U4542ekd\">\n<td data-id=\"t6267798-baN13mCt\" data-transient-attributes=\"table-cell-selection\">2048<\/td>\n<td data-id=\"t6267798-jdCqCLXv\" data-transient-attributes=\"table-cell-selection\">1.00000000002\u00a0\u00d7 211<\/td>\n<td data-id=\"t6267798-6udliU9W\" data-transient-attributes=\"table-cell-selection\">11<\/td>\n<td data-id=\"t6267798-ZN6BmBjD\" data-transient-attributes=\"table-cell-selection\">26 (<code>11010<\/code>)<\/td>\n<td data-id=\"t6267798-bjAN8OYa\" data-transient-attributes=\"table-cell-selection\"><code>0000000000<\/code><\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-CNg49abw\" data-transient-attributes=\"table-cell-selection\">211<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-UCACJadi\">\n<td class=\"table-last-column\" data-id=\"t6267798-Zq05E7iv\" data-transient-attributes=\"table-cell-selection\">2050<\/td>\n<td class=\"table-last-column\" data-id=\"t6267798-IJWJes4t\" data-transient-attributes=\"table-cell-selection\">1.00000000012\u00a0\u00d7 211<\/td>\n<td class=\"table-last-column\" data-id=\"t6267798-u6YlG8Mj\" data-transient-attributes=\"table-cell-selection\">11<\/td>\n<td class=\"table-last-column\" data-id=\"t6267798-7NiByicM\" data-transient-attributes=\"table-cell-selection\">26 (<code>11010<\/code>)<\/td>\n<td class=\"table-last-column\" data-id=\"t6267798-7PzSrj80\" data-transient-attributes=\"table-cell-selection\"><code>0000000001<\/code><\/td>\n<td class=\"table-last-column table-last-row\" data-id=\"t6267798-4tHuOxDt\" data-transient-attributes=\"table-cell-selection\">211\u00a0+ 21<\/td>\n<\/tr>\n<\/tbody>\n<\/table>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/31b82e442019e20ecf096779893433fa34eb05.webp\" data-type=\"block\" \/><\/p>\n<p>\u90a3\u4e48\u95ee\u9898\u53c8\u6765\u4e86\uff0c\u90fd\u6709\u54ea\u4e9b\u64cd\u4f5c\u4f1a\u6539\u53d8\u6d6e\u70b9\u52a0\u6cd5\u987a\u5e8f\u5462\uff1f\u8fd9\u4e2a\u6211\u4eec\u8981\u4ece\u77e9\u9635\u8fd0\u7b97\u8bf4\u8d77\u3002<\/p>\n<h3>\u4e00\u3001GEMM\u7684Batch Invariance<\/h3>\n<p>GEMM\u64cd\u4f5c\u4e3a\u5565\u4f1a\u6539\u53d8\u6d6e\u70b9\u52a0\u6cd5\u987a\u5e8f\u5462\uff1f\u8fd9\u8981\u4eceGEMM\u7684\u4f18\u5316\u7b56\u7565\u8bf4\u8d77\u3002<\/p>\n<h4>1. GEMM\u7684\u4f18\u5316\u7b56\u7565<\/h4>\n<p>\u8981\u7406\u89e3\u8fd9\u4e2a\u95ee\u9898\uff0c\u5fc5\u987b\u8981\u5148\u7406\u89e3\u5185\u5b58\u5899\u4e0b GPU GEMM\u7b97\u5b50\u7684\u6f14\u8fdb\u8def\u7ebf\u3002\u5f88\u591a\u4eba\u53ef\u80fd\u4f1a\u95ee\uff1a\u4e0d\u5c31\u662f\u4e00\u4e2a\u57fa\u7840\u7684\u77e9\u9635\u4e58\u6cd5\u5417\uff0c\u600e\u4e48\u641e\u5f97\u8fd9\u4e48\u590d\u6742\uff1f\u7b54\u6848\u662f\uff1a\u6781\u5176\u590d\u6742\uff0c\u751a\u81f3\u6574\u4e2a AI Infra \u9886\u57df\u90fd\u5728\u56f4\u7ed5\u5b83\u75af\u72c2\u5377\u7ec6\u8282\u3002\u00a0\u6838\u5fc3\u539f\u56e0\u5728\u4e8e\uff1a\u73b0\u4ee3 GPU \u7684\u8ba1\u7b97\u5355\u5143\uff08Tensor Core\uff09\u7b97\u529b\u589e\u957f\u592a\u731b\uff0c\u8fdc\u8fdc\u7529\u5f00\u4e86\u663e\u5b58\u5e26\u5bbd\u7684\u589e\u901f\u3002\u8fd9\u5c31\u5bfc\u81f4 GPU \u5927\u90e8\u5206\u65f6\u95f4\u90fd\u5728\u7b49\u6570\u636e\u642c\u8fd0\uff0c\u8ba1\u7b97\u5355\u5143\u7531\u4e8e\u672a\u80fd\u53ca\u65f6\u83b7\u53d6\u8f93\u5165\u6570\u636e\uff0c\u5927\u91cf\u65f6\u95f4\u5904\u4e8e\u7a7a\u95f2\u7b49\u5f85\u72b6\u6001\uff08\u5373 Compute Bound \u9000\u5316\u4e3a Memory Bound\uff09\u3002\u8fd9\u5c31\u662f\u6240\u8c13\u7684\u5185\u5b58\u5899\uff08Memory Wall\uff09\u3002<\/p>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/c5e25a402517f1da0a051162bad8ee9e991bd6.webp\" data-type=\"block\" \/><\/p>\n<div>\n<div class=\"hljs-cto\">\n<div class=\"hljs-cto\"><button class=\"copy_btn disable\" data-clipboard-target=\"#code_id_0\">\u590d\u5236<\/button><\/p>\n<div class=\"code-toolbar\">\n<pre class=\"has-pre-numbering language-javascript\" tabindex=\"0\"><code class=\"language-javascript\"><span class=\"token keyword\">for<\/span> i from <span class=\"token number\">0<\/span> to <span class=\"token constant\">M<\/span><span class=\"token operator\">:<\/span>          <span class=\"token comment\">\/\/ \u5916\u5c42\uff1a\u904d\u5386 A \u7684\u884c (\u603b\u5171 M \u884c)<\/span>\r\n    <span class=\"token keyword\">for<\/span> j from <span class=\"token number\">0<\/span> to <span class=\"token constant\">N<\/span><span class=\"token operator\">:<\/span>      <span class=\"token comment\">\/\/ \u4e2d\u5c42\uff1a\u904d\u5386 B \u7684\u5217 (\u603b\u5171 N \u5217)<\/span>\r\n        <span class=\"token keyword\">for<\/span> k from <span class=\"token number\">0<\/span> to <span class=\"token constant\">K<\/span><span class=\"token operator\">:<\/span>  <span class=\"token comment\">\/\/ \u5185\u5c42\uff1a\u8ba1\u7b97 A\u7b2ci\u884c \u548c B\u7b2cj\u5217 \u7684\u70b9\u79ef (\u957f\u5ea6\u4e3a K)<\/span>\r\n            <span class=\"token constant\">C<\/span><span class=\"token punctuation\">[<\/span>i<span class=\"token punctuation\">,<\/span>j<span class=\"token punctuation\">]<\/span> <span class=\"token operator\">+=<\/span> <span class=\"token constant\">A<\/span><span class=\"token punctuation\">[<\/span>i<span class=\"token punctuation\">,<\/span>k<span class=\"token punctuation\">]<\/span> <span class=\"token operator\">*<\/span> <span class=\"token constant\">B<\/span><span class=\"token punctuation\">[<\/span>k<span class=\"token punctuation\">,<\/span>j<span class=\"token punctuation\">]<\/span><\/code><\/pre>\n<ul id=\"code_id_0\" class=\"pre-numbering\">\n<li>1.<\/li>\n<li>2.<\/li>\n<li>3.<\/li>\n<li>4.<\/li>\n<\/ul>\n<div class=\"toolbar\"><\/div>\n<\/div>\n<\/div>\n<\/div>\n<\/div>\n<h4>2. Tiling &#8211; \u5206\u5757\u77e9\u9635\u4e58\u6cd5<\/h4>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/82a28a443e23929383802233f46c41442dda4b.webp\" data-type=\"block\" \/><\/p>\n<p>\u5982\u679c\u5355\u7eaf\u6309\u7167\u4e0a\u8ff0\u516c\u5f0f\u9010\u4e2a\u5143\u7d20\u53bb\u8ba1\u7b97\uff0c\u6bcf\u6b21\u4e58\u52a0\u8fd0\u7b97\u90fd\u9700\u8981\u4ece\u6781\u5176\u7f13\u6162\u7684\u5168\u5c40\u663e\u5b58\uff08HBM\uff09\u4e2d\u8bfb\u53d6\u6570\u636e\u3002\u7531\u4e8e\u5b8c\u5168\u6ca1\u6709\u5229\u7528\u5230\u7247\u4e0a\u9ad8\u901f\u7f13\u5b58\uff08SRAM \/ \u5171\u4eab\u5185\u5b58\uff09\u548c\u5bc4\u5b58\u5668\uff08Registers\uff09\uff0c\u8fd9\u4f1a\u5bfc\u81f4\u6d77\u91cf\u7684 HBM \u91cd\u590d\u8bfb\u64cd\u4f5c\uff0c\u4f7f\u5f97\u7b97\u529b\u65e0\u6cd5\u7684\u5145\u5206\u5229\u7528\u3002<\/p>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/b23dbfb21c82e258f88527f984979670db5b46.webp\" data-type=\"block\" \/><\/p>\n<p>\u4e3a\u4e86\u6253\u7834\u5185\u5b58\u5e26\u5bbd\u7684\u74f6\u9888\uff0c\u4e1a\u754c\u6f14\u5316\u51fa\u4e86\u00a0IO-Aware\uff08I\/O \u611f\u77e5\uff09\u00a0\u7684\u5206\u5757\u7b97\u6cd5\u3002\u5176\u6838\u5fc3\u601d\u60f3\u662f\uff1a\u5316\u6574\u4e3a\u96f6\uff0c\u5c06\u5927\u77e9\u9635\u5207\u5206\u6210\u9002\u5408\u653e\u5165\u7f13\u5b58\u7684\u5c0f\u5757\u3002\u00a0\u8c03\u5ea6\u65f6\uff0c\u5c06\u6bcf\u4e2aCTA\uff08Cooperative Thread Array\/ Thread Block\uff09\u4e0e\u7ed3\u679c\u77e9\u9635\u00a0C\u00a0\u7684\u4e00\u4e2a\u7279\u5b9a\u5c0f\u5757\u00a0Cij\u00a0\u7ed1\u5b9a\uff0c\u8ba9\u8be5 CTA \u5168\u6743\u8d1f\u8d23\u8fd9\u4e2a\u5c0f\u5757\u7684\u8ba1\u7b97\u3002<\/p>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/695b0dc33007b229718077197fbcae018c83cb.webp\" data-type=\"block\" \/><\/p>\n<p>\u4e00\u4e2a CTA \u8d1f\u8d23\u8ba1\u7b97\u00a0Cij\u00a0\u7684\u5b8c\u6574\u6570\u636e\u6d41\u8f6c\u903b\u8f91\u5982\u4e0b\uff1a<\/p>\n<ul data-id=\"u738a58b-N4zAzEd4\">\n<li data-id=\"ld70c578-bbmLV0GC\">\u521d\u59cb\u5316\u7d2f\u52a0\u5668\uff1a\u5206\u914d\u5bc4\u5b58\u5668\uff0c\u4e13\u95e8\u7528\u4e8e\u5b58\u653e\u76ee\u6807\u5206\u5757\u00a0Cij\u00a0\u7684\u4e2d\u95f4\u7d2f\u52a0\u7ed3\u679c\u3002<\/li>\n<li data-id=\"ld70c578-nJXpC1d4\">\u6cbf\u7740K \u7ef4\u5ea6\u89c4\u7ea6\uff1a\u6cbf\u7740\u516c\u5171\u7ef4\u5ea6\u00a0K\u00a0\u6b65\u8fdb\u626b\u63cf\u3002\u5728\u6bcf\u4e00\u4e2a\u6b65\u957f\uff08Step\u00a0p\uff09\u5185\u6267\u884c\u4ee5\u4e0b\u6d41\u6c34\u7ebf\uff1a<\/li>\n<li data-id=\"ld70c578-AznIUqMp\">Load (HBM\u00a0\u2192\u00a0SRAM)\uff1a\u5c06\u5f53\u524d\u8ba1\u7b97\u6240\u9700\u7684\u5b50\u5757\u00a0Aip\u00a0\u548c\u00a0Bpj\u00a0\u4ece\u7f13\u6162\u7684\u5168\u5c40\u663e\u5b58\uff0c\u4e00\u6b21\u6027\u6279\u91cf\u642c\u8fd0\u5230\u901f\u5ea6\u66f4\u5feb\u7684\u5171\u4eab\u5185\u5b58\uff08SRAM\uff09\u4e2d\u3002<\/li>\n<li data-id=\"ld70c578-WKmo9AdH\">Compute (Warp GEMM)\uff1a\u5404\u4e2a Warp \u4ece SRAM \u4e2d\u5c06\u6570\u636e\u63d0\u53d6\u81f3\u5bc4\u5b58\u5668\uff0c\u4ea4\u7531 Tensor Core \u6267\u884c\u6781\u901f\u7684\u77e9\u9635\u4e58\u6cd5\u3002<\/li>\n<li data-id=\"ld70c578-SVjjjMqi\">Accumulate\uff1a\u5c06\u672c\u8f6e\u7b97\u51fa\u7684\u4e58\u79ef\u7ed3\u679c\uff0c\u5c31\u5730\u4e0e\u5bc4\u5b58\u5668\u4e2d\u7684\u7d2f\u52a0\u5668\u76f8\u52a0\u3002\u4e0d\u5199\u56de\u663e\u5b58\u3002<\/li>\n<\/ul>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/7769c04792ed0cec3a694945e796e95f9f1605.webp\" data-type=\"block\" \/><\/p>\n<ul data-id=\"u738a58b-Vtmt98CW\">\n<li data-id=\"ld70c578-v3Cd4ILu\">\u6536\u5c3e\u4e0e\u5199\u56de\uff08Epilogue\uff09\uff1a\u5f53 K \u7ef4\u5ea6\u7684\u5faa\u73af\u5168\u90e8\u8dd1\u5b8c\uff0c\u5bc4\u5b58\u5668\u4e2d\u7684\u7d2f\u52a0\u5668\u4fbf\u5f97\u5230\u4e86\u00a0Cij\u00a0\u7684\u6700\u7ec8\u7cbe\u786e\u503c\u3002\u6b64\u65f6\uff0c\u518d\u5c06\u5176\u7edf\u4e00\u5199\u56de\u5168\u5c40\u663e\u5b58\uff08HBM\uff09\u3002<\/li>\n<\/ul>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/464365a3543fb57ecb2534f467c5093c06136e.webp\" data-type=\"block\" \/><\/p>\n<h4>3. Split-K<\/h4>\n<p>\u57fa\u4e8e\u6211\u4eec\u521a\u521a\u68b3\u7406\u7684\u903b\u8f91\uff1a\u6bcf\u4e2a CTA \u8d1f\u8d23\u8f93\u51fa\u77e9\u9635\u00a0C\u00a0\u7684\u4e00\u4e2a Tile\uff0c\u5e76\u5728\u5185\u90e8\u6cbf\u7740\u00a0K\u00a0\u7ef4\u5ea6\u4e32\u884c\u8dd1\u5faa\u73af\u3002<\/p>\n<p>\u4f46\u73b0\u5b9e\u4e2d\u7ecf\u5e38\u4f1a\u9047\u5230\u4e00\u79cd\u6781\u7aef\u60c5\u51b5\uff1a\u5047\u8bbe\u00a0K\u00a0\u7684\u7ef4\u5ea6\u6781\u5176\u5e9e\u5927\uff0c\u800c\u8f93\u51fa\u77e9\u9635\u00a0C\u00a0\uff08\u5373\u00a0M\u00a0\u548c\u00a0N\u00a0\u7ef4\u5ea6\uff09\u975e\u5e38\u5c0f\uff0c\u5c0f\u5230\u53ea\u80fd\u5207\u51fa 4 \u4e2a Tile\u3002\u6b64\u65f6\uff0cGPU \u786c\u4ef6\u8c03\u5ea6\u5668\u53ea\u4f1a\u62c9\u8d77 4 \u4e2a CTA \u53bb\u5e72\u6d3b\u3002\u8981\u77e5\u9053\uff0c\u4e00\u5757 NVIDIA H100 \u62e5\u6709 132 \u4e2a SM\uff0c\u8fd9\u5c31\u610f\u5473\u7740\u6709 128 \u4e2a SM \u5904\u4e8e\u5b8c\u5168\u7a7a\u95f2\u7684\u72b6\u6001\uff0c\u90fd\u5728\u7b49\u5f85\u90a3 4 \u4e2a CTA \u5728\u6781\u5176\u6f2b\u957f\u7684\u00a0K\u00a0\u7ef4\u5ea6\u4e0a\u82e6\u54c8\u54c8\u5730\u8dd1 Loop\uff0cGPU \u7684\u7b97\u529b\u4e5f\u88ab\u6781\u5927\u7684\u6d6a\u8d39\u4e86\u3002<\/p>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/29dbe0220ed0d06674453870018f9f216df800.webp\" data-type=\"block\" \/><\/p>\n<p>\u65e2\u7136\u5728\u7a7a\u95f4\u7ef4\u5ea6\uff08M,\u00a0N\uff09\u4e0a\u5207\u4e0d\u51fa\u66f4\u591a\u7684\u4efb\u52a1\uff0c\u6bd4\u5982\u5728 LLM \u63a8\u7406\u7684 Decoding \u9636\u6bb5\uff0cBatch Size \u6781\u5c0f\uff0cM=1\uff0c\u5982\u679c\u5f3a\u884c\u53ea\u6309\u7a7a\u95f4\u5207\u5206\uff0c\u4f1a\u5bfc\u81f4 GPU \u4e0a\u5927\u91cf\u6838\u5fc3\u65e0\u6d3b\u53ef\u5e72\uff0c\u90a3\u80fd\u4e0d\u80fd\u5728K\u7ef4\u5ea6\u62c6\u5206\uff0c\u5206\u914d\u7ed9\u4e0d\u540c\u7684CTA\u8ba1\u7b97\u5462\uff1f\u2014\u2014\u8fd9\u5c31\u662f\u6240\u8c13\u7684\u00a0Split-K \u4f18\u5316\u7b56\u7565\u3002\u5b83\u7684\u6838\u5fc3\u601d\u60f3\u662f\uff1a\u6253\u7834\u5355\u4e2a CTA \u72ec\u81ea\u8ba1\u7b97\u6574\u4e2a\u00a0K\u00a0\u7ef4\u5ea6\u7684\u903b\u8f91\uff0c\u628a\u00a0K\u00a0\u7ef4\u5ea6\u5207\u5206\u6210\u591a\u6bb5\uff08Split_K\uff09\uff0c\u8ba9\u591a\u4e2a\u4e0d\u540c\u7684 CTA \u540c\u65f6\u5904\u7406\u540c\u4e00\u4e2a\u00a0Cij\u00a0\u5757\u5728\u4e0d\u540c\u00a0K\u00a0\u7247\u6bb5\u4e0a\u7684\u5c40\u90e8\u4e58\u79ef\u7d2f\u52a0\u3002<\/p>\n<p>\uff08\u7279\u522b\u6ce8\u610f\uff1a\u5176\u5b9e\u5728\u666e\u901a\u7684 Tiling \u8ba1\u7b97\u65f6\uff0c\u4e5f\u662f\u6cbf\u7740\u00a0K\u00a0\u7ef4\u5ea6\u5206\u6bb5\u8ba1\u7b97\u7684\u3002\u4f46\u4e24\u8005\u6709\u672c\u8d28\u533a\u522b\uff1aBLOCK_K\u00a0\u662f\u65f6\u95f4\u4e0a\u7684\u4e32\u884c\uff0c\u540c\u4e00\u4e2a CTA \u6bcf\u6b21\u4ece\u5168\u5c40\u663e\u5b58\u642c\u8fd0\u00a0BLOCK_K\u00a0\u5927\u5c0f\u7684\u6570\u636e\uff0c\u5728\u5bc4\u5b58\u5668\u4e2d\u6309\u56fa\u5b9a\u987a\u5e8f\u7d2f\u52a0\uff1b\u800c\u00a0SPLIT_K\u00a0\u662f\u7a7a\u95f4\u4e0a\u7684\u5e76\u884c\uff0c\u5f3a\u884c\u628a\u4efb\u52a1\u5206\u53d1\u7ed9\u7269\u7406\u4e0a\u72ec\u7acb\u7684\u591a\u4e2a CTA \u540c\u65f6\u8ba1\u7b97\u3002\uff09<\/p>\n<p>\u5982\u56fe\u6240\u793a\uff0c\u5047\u8bbe\u8bbe\u7f6e\u00a0Split-K = 2\uff0c\u4e0d\u540c\u989c\u8272\u7684 Tile \u5206\u522b\u7531\u4e24\u4e2a\u72ec\u7acb\u7684 CTA \u5e76\u884c\u5904\u7406\uff0c\u7b97\u529b\u5229\u7528\u7387\u77ac\u95f4\u7ffb\u500d\u3002\u6700\u7ec8\u901a\u8fc7atomic_add\u5c06\u7ed3\u679c\u7d2f\u52a0\u5230C\u7684\u540c\u4e00\u4f4d\u7f6e\u3002 \u8fd9\u91cc\u9700\u8981\u6ce8\u610f\u7684\u662f\uff1aSPLIT_K \u592a\u5927, atomic_add \u7ade\u4e89\u589e\u52a0,\u6536\u76ca\u4f1a\u9012\u51cf\u3002<\/p>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/24f943563d6c7d4e62f08142905e5f5bb51a45.webp\" data-type=\"block\" \/><\/p>\n<p>\u5f53\u7136\u9664\u4e86atomic_add\u7684\u65b9\u5f0f\uff0c\u8fd8\u6709\u4e00\u79cd\u5b9e\u73b0\u65b9\u5f0f\uff1a Workspace Reduction\uff0c\u5177\u4f53\u5de5\u4f5c\u6d41\u7a0b\uff1a<\/p>\n<ul data-id=\"u738a58b-6d8zOQDg\">\n<li data-id=\"ld70c578-8XUb6oKY\">\u5206\u914d Workspace\uff08\u663e\u5b58\u5206\u914d\uff09\uff1a\u5728HBM\u4e2d\u989d\u5916\u5206\u914d\u4e00\u5757\u5927\u5c0f\u4e3a\u00a0[split_k, M, N]\u00a0\u7684\u4e34\u65f6\u5185\u5b58\u7f13\u51b2\u533a\uff08Workspace\uff09\u3002<\/li>\n<li data-id=\"ld70c578-tsBUkqJz\">\u6267\u884c Partial GEMM Kernel: \u542f\u52a8\u4e3b\u77e9\u9635\u4e58\u6cd5 Kernel\u3002\u6bcf\u4e2a Thread Block \u8d1f\u8d23\u8ba1\u7b97 K \u7ef4\u5ea6\u4e0a\u7684\u4e00\u4e2a\u5207\u7247\u3002\u8ba1\u7b97\u5b8c\u6210\u540e\uff0c\u4e0d\u4f7f\u7528\u539f\u5b50\u52a0\uff0c\u800c\u662f\u5c06\u5c40\u90e8\u7ed3\u679c\u76f4\u63a5\u5199\u5165\u5230 Workspace \u4e2d\u5c5e\u4e8e\u81ea\u5df1\u7684\u90a3\u4e2a\u5207\u7247\u4f4d\u7f6e\uff08\u5373\u00a0workspace[k_idx, m, n]\uff09\u3002<\/li>\n<li data-id=\"ld70c578-a3cr6qLJ\">\u6267\u884c Reduction Kernel\uff1a\u4e3b Kernel \u6267\u884c\u5b8c\u6bd5\u5e76\u540c\u6b65\u540e\uff0c\u542f\u52a8\u7b2c\u4e8c\u4e2a\u72ec\u7acb\u7684\u8f7b\u91cf\u7ea7\u89c4\u7ea6 Kernel\u3002\u8fd9\u4e2a Kernel \u8d1f\u8d23\u6cbf\u7740 split_k\u7ef4\u5ea6\uff0c\u5c06 Workspace \u4e2d\u7684\u5c40\u90e8\u7ed3\u679c\u76f8\u52a0\uff0c\u5e76\u5c06\u6700\u7ec8\u7684\u6c42\u548c\u7ed3\u679c\u5199\u5165\u5230\u76ee\u6807\u8f93\u51fa\u77e9\u9635 C \u4e2d\u3002<\/li>\n<\/ul>\n<h4>4. GROUP_M &#8211; Swizzle L2 Cache<\/h4>\n<p>\u901a\u8fc7\u7ea6\u675f CTA\u5728\u77e9\u9635C\u4e2d\u7684\u8c03\u5ea6\u987a\u5e8f\uff0c\u907f\u514d\u4e86\u7531\u4e8e\u8de8\u5ea6\u8fc7\u5927\u7684\u79bb\u6563\u5185\u5b58\u8bbf\u95ee\u800c\u5bfc\u81f4\u7684 L2 Cache \u6296\u52a8\uff08Cache Thrashing\uff09\u548c\u9891\u7e41\u7684\u6570\u636e\u9a71\u9010\u3002\u5b83\u5728\u903b\u8f91\u4e0a\u5c06\u591a\u4e2a\u72ec\u7acb\u7684 CTA \u91cd\u65b0\u7ec4\u5408\u6210\u4e00\u4e2a\u7ef4\u5ea6\u4e3a\u00a0GROUP_SIZE_M \u00d7 N\u00a0\u7684\u5b8f\u89c2\u8c03\u5ea6\u77e9\u9635\u3002\u5728\u8fd9\u4e2a\u88ab\u9650\u5b9a\u7684\u8fde\u7eed\u6267\u884c\u533a\u95f4\u5185\uff0c\u7531\u4e8e\u88ab\u8c03\u5ea6\u7684 CTA \u96c6\u4e2d\u5904\u7406\u7a7a\u95f4\u4e0a\u76f8\u90bb\u7684\u8f93\u51fa\u5757\uff0c\u5b83\u4eec\u80fd\u591f\u5171\u4eab\u5df2\u52a0\u8f7d\u81f3 L2 Cache \u4e2d\u7684\u77e9\u9635 A \u548c B \u7684\u6570\u636e\uff0c\u4ece\u800c\u6700\u5927\u5316\u6570\u636e\u7684\u65f6\u95f4\u5c40\u90e8\u6027\u4e0e L2 Cache \u7684\u590d\u7528\u7387\u3002\u8fd9\u4f55\u5c1d\u4e0d\u662f\u53e6\u5916\u4e00\u79cd\u7ef4\u5ea6\u7684Tiling\u5462\u00a0? \u53ea\u662f\u4e3a\u4e86\u590d\u7528L2\u3002\u800c\u4e4b\u524d\u6211\u4eec\u63d0\u5230\u7684Tiling\uff0c\u662f\u4e3a\u4e86\u590d\u7528SRAM\u548c\u5bc4\u5b58\u5668\u3002\u66f4\u8fdb\u4e00\u6b65Tensor Parallel(Column\/Row Parallel)\u5176\u5b9e\u4e5f\u662f\u66f4\u9ad8\u7ef4\u5ea6\u7684Tiling\u5462\uff5e<\/p>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/1330f9397bf822d1f9f650431e5c4a5f4ff5d7.webp\" data-type=\"block\" \/><\/p>\n<h4>5. \u7b97\u5b50\u8c03\u4f18<\/h4>\n<p>\u524d\u6587\u63a2\u8ba8\u4e86 Tiling\u3001Split-K\u3001Swizzle \u7b49\u65e8\u5728\u7f13\u89e3\u8bbf\u5b58\u74f6\u9888\u7684\u6838\u5fc3\u4f18\u5316\u7b56\u7565\u3002\u4f46\u5728\u5b9e\u9645\u7684 GPU \u786c\u4ef6\u6267\u884c\u5c42\u9762\uff0c\u9488\u5bf9\u4e0d\u540c\u89c4\u6a21\u7684\u8f93\u5165\u5f20\u91cf\uff08Tensor Shape\uff09\uff0c\u7cfb\u7edf\u5fc5\u987b\u786e\u5b9a\u5177\u4f53\u7684\u5e76\u884c\u5316\u5207\u5206\u914d\u7f6e\uff1a\u5373\u77e9\u9635\u5206\u5757\u7684\u5177\u4f53\u7ef4\u5ea6\uff08BLOCK_M \/ BLOCK_N \/ BLOCK_K\uff09\u4ee5\u53ca\u00a0K\u00a0\u7ef4\u5ea6\u7684\u5207\u5206\u6bb5\u6570\uff08SPLIT_K\uff09\u3002\u8fd9\u4e9b\u5207\u5206\u53c2\u6570\u7684\u9009\u62e9\uff0c\u4e0d\u4ec5\u76f4\u63a5\u51b3\u5b9a\u4e86\u7b97\u5b50\u7684\u8bbf\u5b58\u6548\u7387\u4e0e\u786c\u4ef6\u5229\u7528\u7387\uff08Occupancy\uff09\uff0c\u66f4\u5173\u952e\u7684\u662f\u2014\u2014\u5206\u5757\u53c2\u6570\u7684\u52a8\u6001\u53d8\u5316\u4f1a\u91cd\u5851\u5e95\u5c42\u6d6e\u70b9\u6570\u7d2f\u52a0\u7684\u5f52\u7ea6\u62d3\u6251\uff08Reduction Tree\uff09\uff0c\u8fd9\u662f\u5bfc\u81f4\u5927\u6a21\u578b\u63a8\u7406\u5728\u4e0d\u540c\u6279\u6b21\u4e0b\u5931\u53bbBatch Invariance\u7684\u6839\u672c\u539f\u56e0\u3002<\/p>\n<p>\u4e3a\u4e86\u89e3\u6790\u63a8\u7406\u5f15\u64ce\u5e95\u5c42\u4e3a\u4f55\u4f1a\u6839\u636e\u8f93\u5165\u7279\u5f81\u52a8\u6001\u53d8\u66f4\u8fd9\u4e9b\u8c03\u5ea6\u53c2\u6570\uff0c\u6211\u4eec\u9700\u8981\u5148\u56de\u987e GPU \u7b97\u5b50\u7684\u7f16\u7a0b\u8303\u5f0f\u6f14\u8fdb\uff0c\u5e76\u6df1\u5165\u63a2\u8ba8\u73b0\u4ee3\u7f16\u8bd1\u5668\u5f15\u5165\u7684\u81ea\u52a8\u8c03\u4f18\uff08AutoTune\uff09\u673a\u5236\u3002<\/p>\n<p><strong>(1) CUDA VS Triton<\/strong><\/p>\n<p>\u5728\u4e4b\u524d\u7684\u300aAI Infra\u5165\u95e8\uff1aGPU\u662f\u5982\u4f55\u5de5\u4f5c\u7684\u300b\u4e2d\uff0c\u6211\u4eec\u63a2\u8ba8\u4e86 CUDA \u7f16\u7a0b\u6a21\u578b\u4e0e GPU \u786c\u4ef6\u6267\u884c\u6a21\u578b\uff1a Grid\u5b9a\u4e49\u4e86\u5185\u90e8Thread Block\u7684\u7ec4\u7ec7\u5f62\u5f0f(gridDim)\uff1bThread Block\u5b9a\u4e49\u4e86\u5185\u90e8thread\u7684\u7ec4\u7ec7\u5f62\u5f0f(blockDim)\u3002<\/p>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/b357938656ee030f8f79471b747de98f670017.webp\" data-type=\"block\" \/><\/p>\n<p>\u7136\u800c\uff0c\u539f\u751f CUDA \u7f16\u7a0b\u7684\u95e8\u69db\u4f9d\u7136\u6781\u9ad8\uff0c\u5f00\u53d1\u8005\u9700\u8981\u624b\u52a8\u5b9e\u73b0\u5404\u79cd\u5e95\u5c42\u4e14\u7e41\u7410\u7684\u4f18\u5316\u903b\u8f91\u3002OpenAI Triton \u7684\u51fa\u73b0\uff0c\u5f7b\u5e95\u6539\u53d8\u4e86\u8fd9\u4e00\u73b0\u72b6\u3002\u5927\u5e45\u964d\u4f4e\u4e86\u7f16\u5199\u9ad8\u6027\u80fd GPU \u7b97\u5b50\uff08Kernel\uff09\u7684\u95e8\u69db\uff0c\u8ba9\u5f00\u53d1\u8005\u80fd\u4ee5\u63a5\u8fd1 Python \u7684\u751f\u4ea7\u529b\uff0c\u5199\u51fa\u63a5\u8fd1 CUDA C++ \u4e13\u5bb6\u7ea7\u6027\u80fd\u7684\u4ee3\u7801\u3002\u501f\u52a9\u5185\u6838\u878d\u5408\uff08Kernel Fusion\uff09\u3001IO-Aware \u7b49\u5e95\u5c42\u4f18\u5316\u6280\u672f\uff0cTriton \u4e0d\u4ec5\u4fdd\u8bc1\u4e86\u6781\u81f4\u6027\u80fd\uff0c\u8fd8\u517c\u987e\u4e86\u6781\u4f73\u7684\u786c\u4ef6\u65e0\u5173\u6027\u3002<\/p>\n<p>\u4e24\u8005\u6838\u5fc3\u601d\u60f3\u7684\u5dee\u5f02\u5728\u4e8e\u7f16\u7a0b\u8303\u5f0f\uff1a<\/p>\n<ul data-id=\"u738a58b-FZqLq0Y0\">\n<li data-id=\"ld70c578-dC72Fdnl\">CUDA\uff08\u4ee5\u7ebf\u7a0b\u4e3a\u4e2d\u5fc3\uff09\uff1a\u5f00\u53d1\u8005\u9700\u8981\u663e\u5f0f\u5b9a\u4e49\u5e95\u5c42\u7f51\u683c\uff0c\u4f8b\u5982 &lt;&lt;&gt;&gt;\uff0c\u7531\u5e95\u5c42\u751f\u6210 128 \u4e2a\u7ebf\u7a0b\uff0c\u786c\u4ef6\u518d\u5c06\u5176\u5212\u5206\u4e3a 4 \u4e2a Warp \u8c03\u5ea6\u6267\u884c\u3002<\/li>\n<li data-id=\"ld70c578-o1Juu4Hc\">Triton\uff08\u4ee5Tiling\/Block \u4e3a\u4e2d\u5fc3\u4e0e\u81ea\u52a8\u8c03\u4f18\uff09\uff1a\u5f00\u53d1\u8005\u53ea\u9700\u5173\u6ce8\u6570\u636e\u5206\u5757\uff0c\u5e76\u53ef\u901a\u8fc7 @triton.autotune \u88c5\u9970\u5668\u63d0\u4f9b\u4e00\u4e2a\u914d\u7f6e\u641c\u7d22\u7a7a\u95f4\uff08\u5982\u4e0d\u540c\u7684\u5206\u5757\u5927\u5c0f\u3001\u6d4b\u8bd5 4 \u4e2a\u6216 8 \u4e2a Warp\uff09\u3002Triton \u7f16\u8bd1\u5668\u4e0d\u4ec5\u4f1a\u5728\u5e95\u5c42\u81ea\u52a8\u5c06\u9700\u6c42\uff08\u5982 4 \u4e2a Warp\uff09\u7cbe\u51c6\u6620\u5c04\u4e3a 4\u00d732=128 \u4e2a\u7ebf\u7a0b\u4ee5\u5c4f\u853d\u786c\u4ef6\u7ec6\u8282\uff08\u5373&lt;&lt;&gt;&gt; \uff09\uff0c\u66f4\u4f1a\u5728\u8fd0\u884c\u65f6\u81ea\u52a8\u8fdb\u884c\u57fa\u51c6\u6d4b\u8bd5\uff0c\u667a\u80fd\u9009\u51fa\u5f53\u524d\u786c\u4ef6\u4e0b\u7684\u6700\u4f18\u53c2\u6570\u7ec4\u5408\u3002<\/li>\n<\/ul>\n<p>\u5728 PyTorch 2.0 \u53ca\u66f4\u9ad8\u7248\u672c\u4e2d\uff0c\u5f15\u5165\u4e86OpenAI Triton\u4f5c\u4e3a\u7f16\u8bd1\u5668\uff1a<\/p>\n<ul data-id=\"u738a58b-NPS46ZVg\">\n<li data-id=\"ld70c578-nDiLh0KA\">\u524d\u7aef\u5206\u6790\uff1a\u5f53\u7528\u6237\u8c03\u7528 torch.compile(model) \u65f6\uff0cPyTorch \u7684\u524d\u7aef\uff08\u5982 TorchDynamo\uff09\u4f1a\u6355\u83b7\u8ba1\u7b97\u56fe\u3002<\/li>\n<li data-id=\"ld70c578-c6hhaX8X\">\u540e\u7aef\u4f18\u5316\u4e0e\u4ee3\u7801\u751f\u6210\uff1a\u9ed8\u8ba4\u540e\u7aef TorchInductor \u4f1a\u5206\u6790\u56fe\u4e2d\u53ef\u4ee5\u88ab Kernel Fusion \u7684\u64cd\u4f5c\uff0c\u751f\u6210 OpenAI Triton \u4ee3\u7801\u3002<\/li>\n<li data-id=\"ld70c578-hkdjPuO0\">JIT\u7f16\u8bd1\u4e0e\u6267\u884c\uff1a\u6700\u540e\uff0c\u7531 Triton \u7684 JIT (Just-In-Time) \u7f16\u8bd1\u5668\u63a5\u7ba1\u8fd9\u4e9b\u4ee3\u7801\uff0c\u5c06\u5176\u7f16\u8bd1\u6210\u4e00\u4e2a\u4e3a\u7279\u5b9a GPU \u786c\u4ef6\u9ad8\u5ea6\u4f18\u5316\u7684\u3001\u5355\u4e00\u7684\u878d\u5408\u5185\u6838\uff08Fused Kernel\uff09\uff0c\u6700\u7ec8\u9ad8\u6548\u6267\u884c\u3002<\/li>\n<\/ul>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/b2a00e64871d7831dc20901939994e1b0d93cc.webp\" data-type=\"block\" \/><\/p>\n<p>\u5728\u8fd9\u4e00\u751f\u6001\u4f53\u7cfb\u4e0b\uff0cPyTorch\uff08\u501f\u52a9 Inductor + Triton\uff09\u81ea\u52a8\u63a5\u7ba1\u4e86\u8fc7\u53bb\u9700\u8981 CUDA \u4e13\u5bb6\u8017\u8d39\u5927\u91cf\u7cbe\u529b\u624d\u80fd\u5b8c\u6210\u7684\u6027\u80fd\u8c03\u4f18\u5de5\u4f5c\u3002\u5c3d\u7ba1\u81ea\u52a8\u8c03\u4f18\u7684 Inductor+Triton \u5728\u7edd\u5927\u591a\u6570\u573a\u666f\u4e0b\u90fd\u80fd\u5e26\u6765\u5de8\u5927\u7684\u6027\u80fd\u98de\u8dc3\uff0c\u4f46\u5b83\u672c\u8d28\u4e0a\u4ecd\u662f\u4e00\u4e2a\u901a\u7528\u578b\u7684\u81ea\u52a8\u5316\u65b9\u6848\u3002\u5f53\u9762\u5bf9\u6781\u5176\u5173\u952e\u3001\u5bf9\u6027\u80fd\u538b\u69a8\u5230\u6781\u81f4\u7684\u7279\u6b8a\u7b97\u5b50\u65f6\uff0c\u7eaf\u624b\u5de5\u7684\u5e95\u5c42\u4f18\u5316\u4f9d\u7136\u4e0d\u53ef\u6216\u7f3a\u3002<\/p>\n<p>\u5f53\u7136\uff0c\u65e0\u8bba\u662f\u8ba9\u7f16\u8bd1\u5668\u81ea\u52a8\u8c03\u4f18\uff0c\u8fd8\u662f\u6211\u4eec\u624b\u52a8\u8c03\u4f18\uff0c\u90fd\u5fc5\u987b\u7a7f\u900f\u62bd\u8c61\uff0c\u7406\u89e3 GPU \u5e95\u5c42\u6700\u6838\u5fc3\u7684\u7269\u7406\u8fb9\u754c\u4e0e\u8c03\u4f18\u903b\u8f91\u3002<\/p>\n<p>\u5728\u786c\u4ef6\u6267\u884c\u5c42\uff0c\u591a\u4e2a Thread Block \u53ef\u4ee5\u5171\u4eab\u540c\u4e00\u4e2a\u6d41\u591a\u5904\u7406\u5668\uff08SM\uff09\u3002\u53ea\u8981 SM \u7684\u7269\u7406\u8d44\u6e90\uff08\u5bc4\u5b58\u5668\u3001\u5171\u4eab\u5185\u5b58\u7b49\uff09\u8fd8\u6ca1\u88ab\u5360\u6ee1\uff0c\u786c\u4ef6\u8c03\u5ea6\u5668\u5c31\u4f1a\u628a\u66f4\u591a\u7684 Block \u585e\u8fdb\u540c\u4e00\u4e2a SM \u91cc\u5e76\u53d1\u6267\u884c\u3002\u7136\u800c\uff0c\u6bcf\u4e2a SM \u90fd\u6709\u786c\u6027\u7684\u5e76\u53d1\u89c4\u683c\u4e0a\u9650\uff08\u4f8b\u5982\u6700\u5927\u7ebf\u7a0b\u6570 2048\uff0c\u6700\u5927 Block \u6570 32\uff09\u3002<\/p>\n<p>&nbsp;<\/p>\n<p>\u540c\u4e00\u4e2aThread Block \u5185\u7684 Thread \u53ef\u4ee5\u901a\u8fc7\u5171\u4eab\u5185\u5b58\uff08Shared Memory\uff09\u8fdb\u884c\u6570\u636e\u4ea4\u6362\uff0c\u5e76\u4e14\u53ef\u4ee5\u8fdb\u884c\u540c\u6b65\uff08__syncthreads()\uff09\u3002<\/p>\n<p>&nbsp;<\/p>\n<p><strong>(2) \u5bc4\u5b58\u5668\u9650\u5236<\/strong><\/p>\n<p>A100\/H100 \u6bcf\u4e2a SM \u7684\u7269\u7406\u5bc4\u5b58\u5668\u6587\u4ef6\u5927\u5c0f\u56fa\u5b9a\u4e3a 65536 \u4e2a 32-bit \u5bc4\u5b58\u5668\uff08256 KB\uff09\u3002\u6bcf\u4e2a\u786c\u4ef6\u7ebf\u7a0b\u6700\u591a\u53ea\u80fd\u5206\u914d 255 \u4e2a\u5bc4\u5b58\u5668\u3002 \u800c\u5355\u4e2a\u7ebf\u7a0b\u9700\u8981\u7684\u5bc4\u5b58\u5668\u6570\u7531\u4e24\u90e8\u5206\u7ec4\u6210\uff1a<\/p>\n<ul data-id=\"u738a58b-YXCX2NbI\">\n<li data-id=\"ld70c578-PtoEy79D\">\u57fa\u7840\u5f00\u9500\uff08Base Overhead\uff09\uff1a\u6bcf\u4e2a\u7ebf\u7a0b\u79c1\u6709\u7684\uff0c\u7528\u6765\u5b58\u5185\u5b58\u6307\u9488\u3001\u5faa\u73af\u8ba1\u6570\u5668\u3001TMA \u72b6\u6001\u7b49\uff1b<\/li>\n<li data-id=\"ld70c578-fnYiLFro\">\u7d2f\u52a0\u5668\u5206\u644a\uff08Accumulator Share\uff09\uff1a\u5047\u8bbe\u8f93\u51fa\u5757\u5927\u5c0f\uff0c\u5373Accumulator \u4e3a BLOCK_M \u00d7 BLOCK_N\uff0c\u8fd9\u5bc4\u5b58\u5668\u662f\u88ab\u6240\u6709\u7ebf\u7a0b\u5e73\u644a\u7684\u3002\u6bcf\u4e2a\u8fdb\u7a0b\u7684\u5bc4\u5b58\u5668\u7528\u91cf\u4e3a\uff1aBLOCK_M \u00d7 BLOCK_N \/ NUM_WARPS \/ 32\u3002\u5982\u679cTiling\u53d8\u5927\uff0c\u5fc5\u987b\u540c\u6b65\u589e\u52a0num_warps,\u5426\u5219\u4f1a\u5bfc\u81f4Register Spilling\uff1b<\/li>\n<\/ul>\n<p>\u5173\u8054\u516c\u5f0f\uff1a(BLOCK_M * BLOCK_N) \/ (num_warps * 32) &lt;= 255 \u901a\u5e38\u9700\u63a7\u5236\u5728 128 \u5de6\u53f3, \u5982\u679c\u8d85\u8fc7255\uff0cnum_warps\u00a0\u5fc5\u987b\u8c03\u5927\uff0c\u6216\u8005\u00a0BLOCK\u00a0\u5fc5\u987b\u8c03\u5c0f\u3002<\/p>\n<p><strong>(3) \u6267\u884c\u8303\u5f0f\u4e0e\u5171\u4eab\u5185\u5b58\u9650\u5236<\/strong><\/p>\n<p>\u5728\u4f20\u7edf\u7684 CUDA \u4f18\u5316\u5b9e\u8df5\u4e2d\uff0c\u63d0\u5347 Occupancy\uff08\u5360\u7528\u7387\uff09\u4ee5\u9690\u85cf\u8bbf\u5b58\u5ef6\u8fdf\u662f\u4e00\u9879\u6838\u5fc3\u539f\u5219\u3002\u5373\u7ebf\u7a0b\u7ea7\u5e76\u884c\uff08TLP\uff09\u9690\u85cf\u5ef6\u8fdf\uff1aSM \u4e0a\u9a7b\u7559\u7684\u6d3b\u8dc3 Warp \u6570\u91cf\u8d8a\u591a\uff08\u5373 Occupancy \u8d8a\u9ad8\uff09\uff0c\u786c\u4ef6\u8c03\u5ea6\u5668\uff08Warp Scheduler\uff09\u5728\u5f53\u524d Warp \u56e0\u8bbf\u5b58\u800c\u963b\u585e\uff08Stall\uff09\u65f6\uff0c\u5c31\u8d8a\u5bb9\u6613\u627e\u5230\u5176\u4ed6\u5904\u4e8e\u5c31\u7eea\u72b6\u6001\u7684 Warp \u8fdb\u884c\u5207\u6362\u3002\u901a\u8fc7\u8fd9\u79cd\u96f6\u5f00\u9500\u7684\u4e0a\u4e0b\u6587\u5207\u6362\uff0c\u7cfb\u7edf\u4f7f\u5f97\u8ba1\u7b97\u6307\u4ee4\u4e0e\u5185\u5b58\u8bbf\u95ee\u5728\u65f6\u95f4\u4e0a\u76f8\u4e92\u91cd\u53e0\uff0c\u4ece\u800c\u6709\u6548\u9690\u85cf\u4e86\u5168\u5c40\u5185\u5b58\u7684\u7269\u7406\u5ef6\u8fdf\u3002<\/p>\n<p>\u7136\u800c\uff0c\u5728 Ampere \/ Hopper \u67b6\u6784\u4e2d\uff0cTensor Core \u63d0\u4f9b\u4e86\u6781\u9ad8\u7684\u6d6e\u70b9\u541e\u5410\u91cf\uff0c\u4f46 HBM \u7684\u7269\u7406\u8bbf\u5b58\u5ef6\u8fdf\u5e76\u672a\u540c\u6bd4\u4f8b\u7f29\u51cf\u3002\u5728\u8fd9\u79cd\u8ba1\u7b97\u80fd\u529b\u8fdc\u8d85\u8bbf\u5b58\u5e26\u5bbd\u7684\u80cc\u666f\u4e0b\uff0c\u9ad8 Occupancy \u7b56\u7565\u7684\u8fb9\u9645\u6548\u76ca\u6025\u5267\u8870\u51cf\uff1a\u5373\u4f7f SM \u4e0a\u9a7b\u7559\u4e86\u5927\u91cf Warp\uff0c\u5b83\u4eec\u4e5f\u4f1a\u8fc5\u901f\u8017\u5c3d\u5f53\u524d\u8ba1\u7b97\u4efb\u52a1\uff0c\u5e76\u96c6\u4f53\u89e6\u53d1\u8bbf\u5b58\u8bf7\u6c42\uff0c\u5bfc\u81f4\u6240\u6709 Warp \u540c\u65f6\u9677\u5165\u963b\u585e\u3002\u6b64\u65f6\uff0c\u5355\u7eaf\u4f9d\u9760\u7ebf\u7a0b\u7ea7\u7684\u4e0a\u4e0b\u6587\u5207\u6362\u5df2\u65e0\u6cd5\u63a9\u76d6\u5185\u5b58\u7ea7\u522b\u7684\u7269\u7406\u5ef6\u8fdf\u3002<\/p>\n<p>\u4e3a\u7ef4\u6301 Tensor Core \u7684\u9ad8\u541e\u5410\u7387\uff0c\u73b0\u4ee3 GPU \u5f15\u5165\u4e86\u786c\u4ef6\u7ea7\u5f02\u6b65\u5185\u5b58\u62f7\u8d1d\u673a\u5236\uff08\u5982 Ampere \u67b6\u6784\u7684 cp.async \u548c Hopper \u67b6\u6784\u7684 TMA \u5f15\u64ce\uff09\u3002num_stages \u7684\u6838\u5fc3\u601d\u60f3\u662f\u5229\u7528\u591a\u7ea7\u7f13\u51b2\uff08Multi-Buffering\uff09\u5b9e\u73b0\u8ba1\u7b97\u6307\u4ee4\u4e0e\u6570\u636e\u642c\u8fd0\u6307\u4ee4\u7684\u5f02\u6b65\u5e76\u53d1\u3002\u5373GEMM\u7684\u4f18\u5316\u65b9\u5411\u4ece\u9ad8\u5e76\u53d1\u63a9\u76d6\u8bbf\u5b58\u5ef6\u8fdf\u8f6c\u5411\u901a\u8fc7\u00a0Asynchronous Pipelining (\u5f02\u6b65\u6d41\u6c34\u7ebf)\u00a0\u6765\u63a9\u76d6\u5ef6\u8fdf\u3002<\/p>\n<p>\u82e5 num_stages \u914d\u7f6e\u8fc7\u5c0f\uff08\u5982\u9ed8\u8ba4\u503c 2\uff09\uff0c\u8ba1\u7b97\u8fc7\u7a0b\u6613\u56e0\u7b49\u5f85\u6570\u636e\u5c31\u7eea\u800c\u4ea7\u751f\u6c14\u6ce1\uff08Pipeline Bubble\uff09\uff0c\u5bfc\u81f4\u603b\u4f53\u541e\u5410\u91cf\u53d7\u9650\u3002\u82e5\u914d\u7f6e\u8fc7\u5927\uff0c\u9664\u5bfc\u81f4 SRAM \u6ea2\u51fa\u89e6\u53d1\u7f16\u8bd1\u5931\u8d25\uff08Out of Shared Memory\uff09\u5916\uff0c\u7ba1\u7406\u591a\u7ea7\u6d41\u6c34\u7ebf\u72b6\u6001\u7684\u6307\u9488\u8fd8\u4f1a\u6d88\u8017\u989d\u5916\u7684\u5bc4\u5b58\u5668\u8d44\u6e90\uff0c\u9020\u6210\u6027\u80fd\u5012\u9000\u3002\u5728\u4e0d\u540c\u786c\u4ef6\u67b6\u6784\uff08\u5982 SRAM \u5bb9\u91cf\u66f4\u5927\u7684 H100\uff09\u4e0a\u8fdb\u884c\u9ad8\u9636\u8c03\u4f18\u65f6\uff0c\u6838\u5fc3\u75db\u70b9\u5728\u4e8e\u5bfb\u627e\u80fd\u591f\u6700\u5927\u5316\u91cd\u53e0\u7387\u3001\u4e14\u4e0d\u89e6\u53d1\u8d44\u6e90\u6ea2\u51fa\u7684\u6700\u4f73\u8fb9\u754c\u914d\u7f6e\u53c2\u6570\u3002<\/p>\n<p>(BLOCK_M * BLOCK_K + BLOCK_N * BLOCK_K) * \u5b57\u8282\u6570 * num_stages &lt;= SMEM\u7269\u7406\u4e0a\u9650 (\u5982 H100 \u662f 228KB)<\/p>\n<p><strong>(4) \u8c03\u4f18Trade-Off<\/strong><\/p>\n<p>\u4ee5 Llama-3-8B\u00a0Q\/K\/V Linear Proj &#8211; Fused QKV\u4e3a\u4f8b\uff0chidden_size\u4e3a4096\uff0c qkv_proj_size\u4e3a6144\uff0c\u5982\u4e0b\uff1a<\/p>\n<table class=\"data-table\" data-transient-attributes=\"class\" data-width=\"803.993px\">\n<colgroup data-id=\"c7104f7d-Kdj4XYNn\">\n<col span=\"1\" width=\"133.993\" data-id=\"cd89ecb0-VOdFeLYS\" \/>\n<col span=\"1\" width=\"133.993\" data-id=\"cd89ecb0-UB3d5fq3\" \/>\n<col span=\"1\" width=\"133.993\" data-id=\"cd89ecb0-gCTsZBYn\" \/>\n<col span=\"1\" width=\"133.993\" data-id=\"cd89ecb0-g1Mo7kMF\" \/>\n<col span=\"1\" width=\"133.993\" data-id=\"cd89ecb0-khIcGBOF\" \/>\n<col span=\"1\" width=\"134.01\" data-id=\"cd89ecb0-SJiDtEGH\" \/><\/colgroup>\n<tbody data-id=\"t6d5e859-QeDdkCDO\">\n<tr data-id=\"t31e458f-oeIka57q\">\n<td data-id=\"t6267798-8AXAZ1hi\" data-transient-attributes=\"table-cell-selection\">\u76ee\u6807\u77e9\u9635<\/td>\n<td data-id=\"t6267798-uPMkdGXm\" data-transient-attributes=\"table-cell-selection\">\u8ba1\u7b97\u516c\u5f0f<\/td>\n<td data-id=\"t6267798-KB953SFG\" data-transient-attributes=\"table-cell-selection\">\u8f93\u5165 X \u7ef4\u5ea6<\/td>\n<td data-id=\"t6267798-fvqKfPub\" data-transient-attributes=\"table-cell-selection\">\u6743\u91cd\u77e9\u9635\u00a0Wqkv\u00a0\u7ef4\u5ea6<\/td>\n<td data-id=\"t6267798-aWzs7Dz6\" data-transient-attributes=\"table-cell-selection\">\u8f93\u51fa\u7ed3\u679c\u7ef4\u5ea6<\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-dckPMros\" data-transient-attributes=\"table-cell-selection\">\u5907\u6ce8<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-ef8n7QVf\">\n<td class=\"table-last-column\" data-id=\"t6267798-Wm41sK8h\" data-transient-attributes=\"table-cell-selection\">Fused QKV<\/td>\n<td class=\"table-last-column\" data-id=\"t6267798-b3G4MDMk\" data-transient-attributes=\"table-cell-selection\">QKVout\u00a0=\u00a0X\u00a0\u00b7\u00a0Wqkv<\/td>\n<td class=\"table-last-column\" data-id=\"t6267798-KvhIn3Ln\" data-transient-attributes=\"table-cell-selection\"><code>[num_sched_tokens, hidden_size]<\/code><\/td>\n<td class=\"table-last-column\" data-id=\"t6267798-NwslpgS3\" data-transient-attributes=\"table-cell-selection\"><code>[hidden_size,qkv_proj_size]<\/code><\/td>\n<td class=\"table-last-column\" data-id=\"t6267798-INOlcCm6\" data-transient-attributes=\"table-cell-selection\"><code>[num_sched_tokens, qkv_proj_size]<\/code><\/td>\n<td class=\"table-last-column table-last-row\" data-id=\"t6267798-N5qjRYrv\" data-transient-attributes=\"table-cell-selection\">\u6743\u91cd\u6309\u5217\u62fc\u63a5\uff0c\u6267\u884c\u5355\u6b21\u5bbd\u77e9\u9635 GEMM<br \/>\n4096 (Q) + 1024 (K) + 1024 (V) = 6144<\/td>\n<\/tr>\n<\/tbody>\n<\/table>\n<p>\u5373\uff1aM=num_sched_tokens,N=6144,K=4096\uff0c\u5176\u4e2dBLOCK_K\u4e00\u822c\u8bbe\u7f6e\u4e3a 16\u300132 \u6216 64\u3002<\/p>\n<p>BLOCK_M \/ BLOCK_N\u7684\u8bbe\u7f6e\uff1a<\/p>\n<ul data-id=\"u738a58b-gyFBLy9S\">\n<li data-id=\"ld70c578-YGufPz21\">M\u5f88\u5927\u65f6\uff08M&gt;4096)\uff0c\u539f\u5219\u4e0a\u6211\u4eec\u5e0c\u671b\u63d0\u5347\u8ba1\u7b97\u5f3a\u5ea6\uff0cBLOCK_M \u00d7 BLOCK_N \u5c3d\u53ef\u80fd\u7684\u5927\uff0c\u6781\u5927\u5730\u590d\u7528SMEM\uff0c\u6b64\u65f6\u9700\u8981\u5bfb\u627e num_warps \u7684\u5e73\u8861\u70b9\uff0cnum_warps\u592a\u5c0f\u4f1a\u5bfc\u81f4\u5e72\u6d3b\u7684\u7ebf\u7a0b\u5c11\uff0c\u5355\u4e2a\u7ebf\u7a0b\u5206\u644a\u5230\u7684\u7d2f\u52a0\u5668\u5bc4\u5b58\u5668\u8fc7\u591a\uff0c\u8fdb\u800c\u51fb\u7a7f\u5355\u7ebf\u7a0b\u6700\u591a 255 \u4e2a\u5bc4\u5b58\u5668\u7684\u7269\u7406\u4e0a\u9650\uff0c\u89e6\u53d1Register Spilling\uff1b\u800cnum_warps\u592a\u5927\u4f1a\u5bfc\u81f4\u57fa\u7840\u5f00\u9500\u5360\u6bd4\u5927\uff0c\u4e00\u4e2aCTA\u5360\u7528\u5927\u91cfSM\u8d44\u6e90\uff0c\u4ece\u800c\u5bfc\u81f4GPU Occupancy \u964d\u4f4e\u3002\u9700\u8981\u4e00\u4e2a\u5408\u9002\u7684num_warps\u6765\u786e\u4fdd BLOCK_M \u00d7 BLOCK_N \u8db3\u591f\u5927\u3002<\/li>\n<li data-id=\"ld70c578-yjVR1GJR\">M\u5f88\u5c0f\u65f6 \uff08M&lt;64\uff09\uff0c\u7f29\u5c0f BLOCK_M \u8d34\u5408 M\uff0c\u4e3a\u4e86\u5145\u5206\u5229\u7528 SM \u7684\u7b97\u529b\uff0c\u53ef\u4ee5 1. \u9002\u5ea6\u7f29\u5c0f BLOCK_N\u4ece\u800c\u4ea7\u751f\u66f4\u591a\u7684Tile\u5206\u53d1\u7ed9\u4e0d\u540c\u7684 SM\uff1b 2. \u5f00\u542f Split-K\uff0c\u5728 K\u7ef4\u5ea6\u62c9\u8d77\u66f4\u591a CTA \u5e76\u884c\u8ba1\u7b97\u3002<\/li>\n<\/ul>\n<table class=\"data-table\" data-transient-attributes=\"class\" data-width=\"850px\">\n<colgroup data-id=\"c7104f7d-WIYbUhCZ\">\n<col span=\"1\" width=\"169.236\" data-id=\"cd89ecb0-u5HjmQ3w\" \/>\n<col span=\"1\" width=\"169.236\" data-id=\"cd89ecb0-1Zrcg1O0\" \/>\n<col span=\"1\" width=\"256.219\" data-id=\"cd89ecb0-HHkDlITl\" \/>\n<col span=\"1\" width=\"255.288\" data-id=\"cd89ecb0-nQihU0dV\" \/><\/colgroup>\n<tbody data-id=\"t6d5e859-K6ZHC2WH\">\n<tr data-id=\"t31e458f-0henA1N3\">\n<td data-id=\"t6267798-QFhLORAo\" data-transient-attributes=\"table-cell-selection\">\u6307\u6807<\/td>\n<td data-id=\"t6267798-E3Dtmyoe\" data-transient-attributes=\"table-cell-selection\">\u4f18\u5316\u7684\u786c\u4ef6\u76ee\u6807<\/td>\n<td data-id=\"t6267798-qzyn9JLs\" data-transient-attributes=\"table-cell-selection\">\u6838\u5fc3\u4f5c\u7528\u4e0e\u5b9a\u4e49<\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-lzEfhg6j\" data-transient-attributes=\"table-cell-selection\">Trade-Off<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-YTKcufnm\">\n<td data-id=\"t6267798-IqPCrKiL\" data-transient-attributes=\"table-cell-selection\">BLOCK_M \/ BLOCK_N<\/td>\n<td data-id=\"t6267798-eyLTy5Vv\" data-transient-attributes=\"table-cell-selection\">\u5171\u4eab\u5185\u5b58 (SMEM)<br \/>\n\u4e00\u7ea7\u7f13\u5b58 (L1)<\/td>\n<td data-id=\"t6267798-HvaHpJir\" data-transient-attributes=\"table-cell-selection\">\u77e9\u9635\u884c\u4e0e\u5217\u65b9\u5411\u7684\u5b50\u5757\u5927\u5c0f\uff08\u5982 128, 256\uff09\uff0c\u51b3\u5b9a\u4e86\u6bcf\u6b21\u8bfb\u5165\u5171\u4eab\u5185\u5b58\u7684\u6570\u636e\u9762\u79ef\u3002<\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-3xHgDqwG\" data-transient-attributes=\"table-cell-selection\">\u5982\u679c\u592a\u5927\uff0c\u4f1a\u5bfc\u81f4\u6bcf\u4e2a\u7ebf\u7a0b\u6d88\u8017\u8fc7\u591a\u5bc4\u5b58\u5668\uff0c\u5f15\u53d1 Register Spilling\uff08\u5bc4\u5b58\u5668\u6ea2\u51fa\u5230\u6781\u6162\u7684 Local Memory\uff09\uff0c\u5bfc\u81f4\u6027\u80fd\u96ea\u5d29\u3002<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-fXmW7QBj\">\n<td data-id=\"t6267798-3owYYIBL\" data-transient-attributes=\"table-cell-selection\">BLOCK_K<\/td>\n<td data-id=\"t6267798-sXYUKrVr\" data-transient-attributes=\"table-cell-selection\">\u5171\u4eab\u5185\u5b58 (SMEM)<\/td>\n<td data-id=\"t6267798-repZOyrj\" data-transient-attributes=\"table-cell-selection\">\u5728\u7ef4\u5ea6\u00a0K\u00a0\u4e0a\u6bcf\u6b21\u7d2f\u52a0\u7684\u957f\u5ea6\u3002<\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-gTSzIrU9\" data-transient-attributes=\"table-cell-selection\">\u540c\u4e0a\uff0c\u901a\u5e38\u8bbe\u7f6e\u4e3a 32 \u6216 64\u3002\u9700\u4e0e M\/N\/num_stages \u914d\u5408\u8ba1\u7b97\u603b SMEM \u5360\u7528\u3002<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-aWpbmWFs\">\n<td data-id=\"t6267798-gSVWgFOw\" data-transient-attributes=\"table-cell-selection\">num_warps<\/td>\n<td data-id=\"t6267798-OD8iQBtz\" data-transient-attributes=\"table-cell-selection\">\u5bc4\u5b58\u5668 (Registers)<br \/>\n\u5e76\u53d1\u5ea6 (Occupancy)<\/td>\n<td data-id=\"t6267798-8i7ke3cq\" data-transient-attributes=\"table-cell-selection\">\u6bcf\u4e2a Thread Block \u5206\u914d\u7684 Warp \u6570\u91cf\uff081 Warp = 32 \u7ebf\u7a0b\uff09\u3002\u672c\u8d28\u662f\u5206\u6bcd\uff0c\u7528\u6765\u7a00\u91ca\u6bcf\u4e2a\u7ebf\u7a0b\u7684\u8ba1\u7b97\u91cf\u3002<\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-bxMZDr35\" data-transient-attributes=\"table-cell-selection\">\u592a\u5c0f\uff1a\u7d2f\u52a0\u5668\u6491\u7206\u5355\u7ebf\u7a0b 255 \u4e2a\u5bc4\u5b58\u5668\u4e0a\u9650\uff0c\u53d1\u751f\u81f4\u547d\u7684 Register Spilling\u3002<br \/>\n\u592a\u5927\uff1a\u5355\u4e2a Block \u5360\u7528\u8fc7\u591a Warp\uff0c\u5bfc\u81f4\u4e00\u4e2a SM \u91cc\u88c5\u4e0d\u4e0b\u51e0\u4e2a Block\uff0c\u5e76\u53d1\u5ea6\uff08Occupancy\uff09\u66b4\u8dcc\u3002<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-QOfRwrWT\">\n<td data-id=\"t6267798-PLmaQhJx\" data-transient-attributes=\"table-cell-selection\">num_stages<\/td>\n<td data-id=\"t6267798-tH3IpkqL\" data-transient-attributes=\"table-cell-selection\">\u5168\u5c40\u663e\u5b58 (HBM)<br \/>\n\u5ef6\u8fdf\u63a9\u76d6<\/td>\n<td data-id=\"t6267798-ADd4myLV\" data-transient-attributes=\"table-cell-selection\">\u8f6f\u4ef6\u6d41\u6c34\u7ebf\u7ea7\u6570\u3002\u5f00\u8f9f\u591a\u4efd\u7f13\u5b58\uff0c\u8ba9\u8ba1\u7b97\u5355\u5143\u5728\u7b97\u5f53\u524d\u5757\u65f6\uff0c\u540e\u53f0\u5f02\u6b65\u53bb\u53d6\u540e\u9762\u7684\u5757\u3002<\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-baC5zcKI\" data-transient-attributes=\"table-cell-selection\">\u592a\u5927\u4f1a\u5bfc\u81f4 SMEM \u6ea2\u51fa<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-BMKYDszV\">\n<td data-id=\"t6267798-M9WBAtNe\" data-transient-attributes=\"table-cell-selection\">GROUP_M\u5373Swizzle<\/td>\n<td data-id=\"t6267798-irS9IJsI\" data-transient-attributes=\"table-cell-selection\">\u4e8c\u7ea7\u7f13\u5b58 (L2 Cache)<\/td>\n<td data-id=\"t6267798-dqEkZOuV\" data-transient-attributes=\"table-cell-selection\">\u5c06\u51e0\u4e2a\u8fde\u7eed\u7684\u884c\u6253\u5305\u6210\u4e00\u7ec4\uff0c\u6253\u7834\u9ed8\u8ba4\u7684\u9010\u884c\u626b\u63cf\u3002<\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-KgzgNmZG\" data-transient-attributes=\"table-cell-selection\">\u5171\u4eab\u5df2\u52a0\u8f7d\u81f3 L2 Cache \u4e2d\u7684\u77e9\u9635 A \u548c B \u7684\u6570\u636e\uff0c\u4ece\u800c\u6700\u5927\u5316\u6570\u636e\u7684\u65f6\u95f4\u5c40\u90e8\u6027\u4e0e L2 Cache \u7684\u590d\u7528\u7387<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-X1JbKAtt\">\n<td class=\"table-last-column\" data-id=\"t6267798-r9xZBOcg\" data-transient-attributes=\"table-cell-selection\">SPLIT_K<\/td>\n<td class=\"table-last-column\" data-id=\"t6267798-gmMlH91Q\" data-transient-attributes=\"table-cell-selection\">\u6d41\u591a\u5904\u7406\u5668 (SM)<br \/>\n\u7b97\u529b\u5229\u7528\u7387<\/td>\n<td class=\"table-last-column\" data-id=\"t6267798-3VhK8ydJ\" data-transient-attributes=\"table-cell-selection\">\u5c06\u6781\u957f\u7684\u00a0K\u00a0\u7ef4\u5ea6\u5207\u5206\u7ed9\u4e0d\u540c\u7684 Block \u540c\u65f6\u7b97\uff0c\u6700\u540e\u518d\u505a\u539f\u5b50\u52a0\u6cd5\uff08Atomic Add\uff09\u5408\u5e76\u3002<\/td>\n<td class=\"table-last-column table-last-row\" data-id=\"t6267798-xBhem0PU\" data-transient-attributes=\"table-cell-selection\">\u4ec5\u5728\u00a0M\u00a0\u548c\u00a0N\u00a0\u6781\u5c0f\u3001K\u00a0\u6781\u5927\u7684\u60c5\u51b5\uff08\u6bd4\u5982Flash-Decoding\uff09\u5f00\u542f\uff0c\u80fd\u63d0\u9ad8SM\u5229\u7528\u7387\u3002<\/td>\n<\/tr>\n<\/tbody>\n<\/table>\n<h4>6. Batch Invariance\u7684\u6765\u6e90<\/h4>\n<p>\u524d\u6587\u63a2\u8ba8\u4e86\u8bf8\u591a\u4f18\u5316\u7b56\u7565\uff0c\u90a3\u4e48\u5177\u4f53\u662f\u54ea\u4e9b\u53c2\u6570\u7684\u53d8\u52a8\u5bfc\u81f4\u4e86\u6d6e\u70b9\u52a0\u6cd5\u987a\u5e8f\u7684\u6539\u53d8\uff1f<\/p>\n<p>\u4ece\u6570\u5b66\u672c\u8d28\u6765\u770b\uff0c\u77e9\u9635\u4e58\u6cd5\u00a0C\u00a0=\u00a0A\u00a0\u00d7\u00a0B\uff08\u5176\u4e2d\u00a0A\u00a0\u4e3a\u00a0[M, K]\uff0cB\u00a0\u4e3a\u00a0[K, N]\uff0cC\u4e3a\u00a0[M, N]\uff09\u4e2d\u4efb\u610f\u5143\u7d20\u00a0Ci,j\u00a0\u7684\u8ba1\u7b97\u903b\u8f91\u662f\u6052\u5b9a\u7684\u2014\u2014\u5373\u53d6\u00a0A\u00a0\u7684\u7b2c\u00a0i\u00a0\u884c\u4e0e\u00a0B\u00a0\u7684\u7b2c\u00a0j\u00a0\u5217\u8fdb\u884c\u70b9\u79ef\uff1a<\/p>\n<blockquote><p>Ci,j\u00a0= \u03a3k=0K-1\u00a0Ai,k\u00a0\u00d7\u00a0Bk,j<\/p><\/blockquote>\n<p>\u5728 Triton \/ CUDA \u7684\u5e95\u5c42\u4f18\u5316\u4e2d\uff0c\u9488\u5bf9\u5404\u9879\u8c03\u5ea6\u53c2\u6570\u7684\u52a8\u6001\u8c03\u6574\uff0c\u5176\u5bf9\u786e\u5b9a\u6027\u7684\u5f71\u54cd\u6709\u7740\u672c\u8d28\u533a\u522b\uff1a<\/p>\n<ul data-id=\"u738a58b-Z9QOvuBk\">\n<li data-id=\"ld70c578-13IWRMIo\">BLOCK_M\u00a0\u4e0e\u00a0BLOCK_N\uff08\u65e0\u5f71\u54cd\uff09\uff1a\u5206\u5757\u903b\u8f91\u4e3a\u52a0\u8f7d Tile_A\u00a0[BLOCK_M, BLOCK_K]\u00a0\u4e0e Tile_B\u00a0[BLOCK_K, BLOCK_N]\uff0c\u5e76\u8ba1\u7b97\u8f93\u51fa Tile_C\u00a0[BLOCK_M, BLOCK_N]\u3002\u8fd9\u4ec5\u4ec5\u662f\u5728\u505a\u7a7a\u95f4\u7ef4\u5ea6\u7684\u4efb\u52a1\u6620\u5c04\uff0c\u51b3\u5b9a\u4e86\u54ea\u4e9b\u5143\u7d20\u88ab\u6253\u5305\u5728\u4e00\u8d77\u72ec\u7acb\u8ba1\u7b97\uff0c\u4ee5\u53ca\u7531\u54ea\u4e2a\u5177\u4f53\u7684 CTA (Thread Block) \u8d1f\u8d23\u8ba1\u7b97\u54ea\u4e00\u5757\u533a\u57df\u3002\u5bf9\u4e8e\u8f93\u51fa\u77e9\u9635\u00a0C\u00a0\u4e2d\u7684\u7279\u5b9a\u5143\u7d20\u00a0Ci,j\u00a0\u800c\u8a00\uff0c\u65e0\u8bba\u5b83\u88ab\u5206\u914d\u7ed9\u54ea\u4e2a CTA\u3001\u4e0e\u54ea\u4e9b\u76f8\u90bb\u5143\u7d20\u4e00\u8d77\u88ab\u8ba1\u7b97\uff0c\u5176\u5e95\u5c42\u7684\u00a0K\u00a0\u7ef4\u5ea6\u70b9\u79ef\u903b\u8f91\u5e76\u672a\u6539\u53d8\u3002\u56e0\u6b64\uff0c\u6539\u53d8\u8fd9\u4e24\u4e2a\u53c2\u6570\u4e0d\u4f1a\u5f71\u54cd\u6d6e\u70b9\u52a0\u6cd5\u7684\u987a\u5e8f\u3002<\/li>\n<li data-id=\"ld70c578-y1klIXAe\">GROUP_M \/ Swizzle\uff08\u65e0\u5f71\u54cd\uff09\uff1aSwizzle \u672c\u8d28\u4e0a\u662f\u6539\u53d8\u4e86\u591a\u4e2a CTA \u5728\u7f51\u683c\uff08Grid\uff09\u7ea7\u522b\u7684\u8c03\u5ea6\u987a\u5e8f\uff0c\u4ee5\u6b64\u6765\u63d0\u9ad8 L2 Cache \u7684\u547d\u4e2d\u7387\u3002\u5b83\u51b3\u5b9a\u7684\u662f\u5148\u7b97\u54ea\u4e00\u4e2a\u7a7a\u95f4\u5757\uff0c\u540e\u7b97\u54ea\u4e00\u4e2a\uff0c\u5b8c\u5168\u6ca1\u6709\u5e72\u6d89\u67d0\u4e00\u4e2a\u7279\u5b9a\u5206\u5757\u5185\u90e8\u7684\u4e58\u79ef\u7d2f\u52a0\u8fc7\u7a0b\u3002\u56e0\u6b64\uff0cSwizzle \u540c\u6837\u4e0d\u4f1a\u7834\u574fBatch Invariance\u3002<\/li>\n<li data-id=\"ld70c578-RiOnEqqH\">BLOCK_K\uff08\u5f15\u5165\u786e\u5b9a\u6027\u8bef\u5dee\uff09\uff1a\u5b9a\u4e49\u4e86\u5728\u89c4\u7ea6\u7ef4\u5ea6\uff08Reduction Dimension\uff09\u4e0a\u7684\u6b65\u957f\u3002\u6539\u53d8\u00a0BLOCK_K\u00a0\u4f1a\u76f4\u63a5\u6539\u53d8\u5355\u6b21\u5faa\u73af\u4e2d\u52a0\u8f7d\u5230\u5bc4\u5b58\u5668\u4e2d\u7684\u6570\u636e\u91cf\uff0c\u8fdb\u800c\u6539\u53d8 Tensor Core \u5185\u90e8 MMA\u6307\u4ee4\u7684\u7d2f\u52a0\u6811\u62d3\u6251\u3002\u8fd9\u79cd\u6539\u53d8\u4f1a\u5bfc\u81f4\u6d6e\u70b9\u52a0\u6cd5\u987a\u5e8f\u7684\u53d8\u5316\uff08\u5373\u4e0d\u540c\u00a0BLOCK_K\u00a0\u4ea7\u751f\u4e0d\u540c\u7ed3\u679c\uff0c\u4f46\u540c\u4e00\u00a0BLOCK_K\u00a0\u7ed3\u679c\u6052\u5b9a\uff09\u3002<\/li>\n<li data-id=\"ld70c578-x4IBsNg9\">SPLIT_K\uff08\u5f15\u5165\u975e\u786e\u5b9a\u6027\u8bef\u5dee\uff09\uff1a\u5c06\u00a0K\u00a0\u7ef4\u5ea6\u5f3a\u884c\u5207\u5206\u7ed9\u591a\u4e2a CTA \u5e76\u53d1\u6267\u884c\u3002\u5982\u679c\u91c7\u7528Atomic Add\u8fdb\u884c\u7ed3\u679c\u5408\u5e76\uff0c\u7531\u4e8e GPU \u786c\u4ef6\u8c03\u5ea6\u7ebf\u7a0b\u5757\u7684\u5148\u540e\u987a\u5e8f\u662f\u5b8c\u5168\u968f\u673a\u7684\uff0c\u52a0\u6cd5\u987a\u5e8f\u662f\u5b8c\u5168\u968f\u673a\u4e0d\u53ef\u63a7\u7684\u3002\u5373\u4fbf\u91c7\u7528Workspace Reduction\uff0cSPLIT_K\u00a0\u6bb5\u6570\u7684\u52a8\u6001\u53d8\u5316\u540c\u6837\u4f1a\u6539\u53d8\u7d2f\u52a0\u6811\u62d3\u6251\u3002<\/li>\n<li data-id=\"ld70c578-rSkTxnbU\">num_warps\u00a0\u548c\u00a0num_stages\u00a0\u7684\u672c\u8d28\u662f\u786c\u4ef6\u8d44\u6e90\u5206\u914d\u4e0e\u6d41\u6c34\u7ebf\u8c03\u5ea6\uff08\u51b3\u5b9a\u5206\u914d\u591a\u5c11\u4e2acta\u53bb\u7b97\u3001\u5f00\u8f9f\u51e0\u4efd\u7f13\u5b58\u6c60\u505a\u5f02\u6b65\u642c\u8fd0\uff09\u3002\u5e76\u4e0d\u6539\u53d8\u5b8f\u89c2\u7684\u77e9\u9635\u70b9\u79ef\u8ba1\u7b97\u903b\u8f91\u548cReduction Tree\u62d3\u6251\u3002<\/li>\n<\/ul>\n<p>\u7efc\u4e0a\u6240\u8ff0\uff0c\u63a8\u7406\u5f15\u64ce\u5728\u6027\u80fd\u4e0e\u786e\u5b9a\u6027\u4e4b\u95f4\u5b58\u5728\u7740\u56fa\u6709\u7684\u67b6\u6784\u51b2\u7a81\uff0c\u4e3a\u4e86\u8ffd\u6c42\u6781\u81f4\u7684\u8bbf\u5b58\u590d\u7528\uff08\u52a8\u6001\u8c03\u6574\u00a0BLOCK_K\uff09\u4e0e\u5e76\u63d0\u5347\u6d41\u591a\u5904\u7406\u5668\uff08SM\uff09\u7684\u5e76\u53d1\u5229\u7528\u7387\uff08\u52a8\u6001\u5f00\u542f\u00a0SPLIT_K\uff09\uff0c\u5e95\u5c42\u7684\u542f\u53d1\u5f0f\u8c03\u5ea6\u7b56\u7565\u4e0d\u53ef\u907f\u514d\u5730\u6539\u53d8\u4e86\u6d6e\u70b9\u8fd0\u7b97\u7684Reduction Tree\u62d3\u6251\uff0c\u8fd9\u6b63\u662f\u5927\u6a21\u578b\u5728\u52a8\u6001 Batch \u4e0b\u4e22\u5931Batch Invariance\u7684\u6839\u672c\u539f\u56e0\u3002\u90a3\u4e48\u5728\u771f\u5b9e\u7684\u63a8\u7406\u573a\u666f\u4e2d\uff0c\u8be5\u5982\u4f55\u4fee\u8865\u8fd9\u4e2a\u673a\u5236\uff1f\u63a5\u4e0b\u6765\u6211\u4eec\u6df1\u5165 vLLM \u7684\u5b9e\u73b0\uff0c\u770b\u770b\u5b83\u662f\u5982\u4f55\u5728\u5e95\u5c42\u786c\u4ef6\u7279\u6027\u548c\u4e0a\u5c42\u8c03\u5ea6\u903b\u8f91\u4e4b\u95f4\u505a Trade-off \u7684\u3002<\/p>\n<h4>7. vLLM\u4e2dGEMM\u7684Batch Invariance\u652f\u6301<\/h4>\n<p>\u600e\u4e48\u89e3\u51b3\u5462\uff1f\u76f4\u89c9\u4e0a\uff0c\u53ea\u9700\u8981\u5728\u63a8\u7406\u5f15\u64ce\u4e2d\u5168\u5c40\u7981\u7528\u00a0Split-K\u00a0\u5e76\u9501\u6b7b\u00a0BLOCK_K\u00a0\u5373\u53ef\u3002\u4f46\u73b0\u5b9e\u7684\u5de5\u7a0b\u5b9e\u73b0\u8fdc\u6bd4\u8fd9\u590d\u6742\u3002\u73b0\u4ee3\u63a8\u7406\u5f15\u64ce\uff08\u5982 vLLM\uff09\u5e95\u5c42\u7684 GEMM \u64cd\u4f5c\u7684\u4e0b\u53d1\u8def\u5f84\u6d89\u53ca\u591a\u5c42\u62bd\u8c61\u4e0e\u591a\u79cd\u540e\u7aef\u3002\u5b9e\u73b0\u4e25\u683c\u7684 Batch Invariance \u5e76\u975e\u5355\u4e00\u914d\u7f6e\u7684\u4fee\u6539\uff0c\u800c\u662f\u9700\u8981\u9488\u5bf9\u7279\u5b9a\u7684\u8f6f\u786c\u4ef6\u8fd0\u884c\u73af\u5883\uff0c\u8fdb\u884c\u591a\u7ef4\u5ea6\u7684\u6267\u884c\u8def\u5f84\u8def\u7531\u4e0e\u53c2\u6570\u7ea6\u675f\u3002\u5177\u4f53\u7684GEMM\u6267\u884c\u8def\u5f84\u51b3\u4e8e\u4ee5\u4e0b\u4e09\u4e2a\u7ef4\u5ea6\u7684\u7ec4\u5408\uff1a<\/p>\n<ul data-id=\"u738a58b-EtnHCc4e\">\n<li data-id=\"ld70c578-QWDhxWpL\">\u786c\u4ef6\u67b6\u6784\uff08Architecture\uff09\uff1aSM80 (Ampere) \u4e0e SM90\/SM100 (Hopper\/Blackwell) \u7684\u5e95\u5c42 GEMM \u6267\u884c\u8303\u5f0f\u622a\u7136\u4e0d\u540c\u2014\u2014\u524d\u8005\u9760 warp \u7ea7\u5e76\u53d1\u63a9\u76d6\u8bbf\u5b58\u5ef6\u8fdf\uff0c\u540e\u8005\u9760 TMA + WGMMA \u7684\u5f02\u6b65\u6d41\u6c34\u7ebf\u3002\u6267\u884c\u8303\u5f0f\u4e0d\u540c\uff0c\u5b9e\u73b0Batch Invariance\u7684\u903b\u8f91\u4e5f\u4f1a\u4e0d\u540c\u3002<\/li>\n<li data-id=\"ld70c578-0OyPeweZ\">\u6570\u636e\u7cbe\u5ea6\uff08Data Type\uff09\uff1abf16\/fp16 \u7531\u751f\u6001\u6700\u6210\u719f\u7684 cuBLASLt \u627f\u63a5\uff1b\u800c fp8\/fp4 \u7b49\u4f4e\u7cbe\u5ea6\u4f1a\u8f6c\u5411 CUTLASS\uff08\u4e43\u81f3 DeepGEMM \u7b49\uff09\u4e13\u7528 kernel\u3002\u4e0d\u540c\u540e\u7aef\u63a7\u5236\u903b\u8f91\u4e5f\u4e0d\u4e00\u81f4\u3002<\/li>\n<li data-id=\"ld70c578-w9iHKPDW\">\u7b97\u5b50\u5165\u53e3\uff08API Backend\uff09\uff1ann.Linear\u00a0\u8d70\u7684\u662f vLLM \u81ea\u5df1\u7684 dispatch\uff0c\u53ef\u76f4\u63a5\u8def\u7531\u5230\u5b9a\u5236\u9ad8\u6027\u80fd\u5185\u6838\uff1b\u800c\u88f8\u00a0torch.mm \/ bmm\u00a0\u8d70 PyTorch dispatcher \u2192 cuBLAS\/cuBLASLt\uff0c\u8c03\u5ea6\u6743\u5728\u6846\u67b6\u548c\u95ed\u6e90\u5e93\u624b\u91cc\u3002<\/li>\n<\/ul>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/68ef397716ec79368a1989c00217267b62af99.webp\" data-type=\"block\" \/><\/p>\n<p><strong>(1) SM8x vs SM90\/SM100<\/strong><\/p>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/34f7e7b418b87a4fbd5466f3e41eaadc129b9c.webp\" data-type=\"block\" \/><\/p>\n<p>SM80: warp\u7ef4\u5ea6\u7684tensor core(mma.sync)\u8c03\u5ea6 : \u4f18\u5316\u903b\u8f91\u662f\u63d0\u9ad8\u5360\u7528\u7387\uff0cwarp\u540c\u65f6\u8d1f\u8d23\u642c\u8fd0\u548c\u8ba1\u7b97\u3002<\/p>\n<div>\n<div class=\"hljs-cto\">\n<div class=\"hljs-cto\"><button class=\"copy_btn disable\" data-clipboard-target=\"#code_id_1\">\u590d\u5236<\/button><\/p>\n<div class=\"code-toolbar\">\n<pre class=\"has-pre-numbering language-javascript\" tabindex=\"0\"><code class=\"language-javascript\"><span class=\"token number\">4096<\/span> <span class=\"token operator\">*<\/span> <span class=\"token number\">4096<\/span> <span class=\"token operator\">-<\/span><span class=\"token operator\">&gt;<\/span><span class=\"token constant\">CTA<\/span> Tile\uff08\u5982<span class=\"token number\">128<\/span><span class=\"token operator\">*<\/span><span class=\"token number\">128<\/span>\uff09  <span class=\"token operator\">-<\/span><span class=\"token operator\">&gt;<\/span> Warp Tile\uff08<span class=\"token number\">32<\/span> <span class=\"token operator\">*<\/span> <span class=\"token number\">64<\/span>\uff09 <span class=\"token operator\">-<\/span><span class=\"token operator\">&gt;<\/span> mma <span class=\"token function\">\u6307\u4ee4\u7ea7<\/span><span class=\"token punctuation\">(<\/span><span class=\"token number\">16<\/span> <span class=\"token operator\">*<\/span> <span class=\"token number\">8<\/span> <span class=\"token operator\">*<\/span> <span class=\"token number\">16<\/span><span class=\"token punctuation\">)<\/span><\/code><\/pre>\n<ul id=\"code_id_1\" class=\"pre-numbering\">\n<li>1.<\/li>\n<\/ul>\n<div class=\"toolbar\"><\/div>\n<\/div>\n<\/div>\n<\/div>\n<\/div>\n<p>SM90\/100 TMA \u786c\u4ef6\u5355\u5143\u7684\u5f15\u5165\u548c WGMMA \u6307\u4ee4\uff08Warp Group 128 \u7ebf\u7a0b\uff09\u4fc3\u6210\u4e86\u8303\u5f0f\u8f6c\u79fb\u901a\u8fc7Asynchronous Pipelining (\u5f02\u6b65\u6d41\u6c34\u7ebf)\u6765\u63a9\u76d6\u5ef6\u8fdf\u3002<\/p>\n<div>\n<div class=\"hljs-cto\">\n<div class=\"hljs-cto\"><button class=\"copy_btn disable\" data-clipboard-target=\"#code_id_2\">\u590d\u5236<\/button><\/p>\n<div class=\"code-toolbar\">\n<pre class=\"has-pre-numbering language-javascript\" tabindex=\"0\"><code class=\"language-javascript\"><span class=\"token number\">4096<\/span> <span class=\"token operator\">*<\/span> <span class=\"token number\">4096<\/span> <span class=\"token operator\">-<\/span><span class=\"token operator\">&gt;<\/span> <span class=\"token constant\">CTA<\/span> <span class=\"token function\">Tile<\/span> <span class=\"token punctuation\">(<\/span>\u5982 <span class=\"token number\">256<\/span> <span class=\"token operator\">*<\/span><span class=\"token number\">128<\/span><span class=\"token punctuation\">)<\/span> <span class=\"token operator\">-<\/span><span class=\"token operator\">&gt;<\/span> Warp Group <span class=\"token function\">Tile<\/span> <span class=\"token punctuation\">(<\/span><span class=\"token number\">64<\/span> <span class=\"token operator\">*<\/span><span class=\"token number\">128<\/span><span class=\"token punctuation\">)<\/span> <span class=\"token operator\">-<\/span><span class=\"token operator\">&gt;<\/span> wgmma <span class=\"token function\">\u6307\u4ee4\u7ea7<\/span> <span class=\"token punctuation\">(<\/span><span class=\"token number\">64<\/span> <span class=\"token operator\">*<\/span><span class=\"token constant\">N<\/span> <span class=\"token operator\">*<\/span> <span class=\"token number\">16<\/span><span class=\"token punctuation\">)<\/span><\/code><\/pre>\n<ul id=\"code_id_2\" class=\"pre-numbering\">\n<li>1.<\/li>\n<\/ul>\n<div class=\"toolbar\"><\/div>\n<\/div>\n<\/div>\n<\/div>\n<\/div>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/c1d27a481f670356d719130a714d1461d284f1.webp\" data-type=\"block\" \/><\/p>\n<p>\u5bf9\u4e8e\u5927\u77e9\u9635\uff0c\u7531\u4e8e CTA \u5206\u914d\u5230\u7684\u603b\u4f53\u7a7a\u95f4\u8db3\u591f\u5927\uff0c\u6bcf\u4e2a Warp \u90fd\u80fd\u5206\u5230\u4e00\u5757\u8db3\u591f\u5927\u7684\u4e13\u5c5e\u7a7a\u95f4\uff08\u4f8b\u5982 32&#215;64\uff09\u3002\u7136\u540e\u76f4\u63a5\u628a\u5927 Tile \u5207\u6210\u5c0f Tile \u5206\u7ed9 Warp\uff0c\u8fd9\u6837\u591a\u4e2a Warp \u4e4b\u95f4\u4ea4\u66ff\u8bbf\u5b58\u548c\u8ba1\u7b97\u4ece\u800c\u63a9\u76d6\u5185\u5b58\u5ef6\u8fdf\u3002\u5927\u5bb6\u90fd\u5404\u81ea\u5728 K\u7ef4\u5ea6\u72ec\u7acb\u3001\u4e32\u884c\u6267\u884c\u5230\u5e95\uff0c\u4e92\u4e0d\u5e72\u6d89\uff0crun-to-run\u662f\u6ca1\u95ee\u9898\u7684\u3002<\/p>\n<p>\u4f46\u662fbatch-to-batch\u60c5\u51b5\u4e0bBLOCK_K\u7684\u53d1\u751f\u53d8\u5316\uff0c\u56e0\u6b64\u4f1a\u5f15\u53d1\u6d6e\u70b9\u52a0\u6cd5\u987a\u5e8f\u7684\u6ce2\u52a8\uff1f<\/p>\n<p>\u5b9e\u9645\u4e0a\u5462\uff0c\u4e5f\u4e0d\u4e00\u5b9a\uff0c\u8fd9\u53d6\u51b3\u4e8ewarp level\u7684K\u662f\u4e0d\u662f\u6052\u5b9a\u7684\u3002\u6bd4\u5982\u542f\u53d1\u5f0f\u7684auto tune\u53ef\u80fd\u5bfc\u81f4BLOCK_K\u4ece32\u53d8\u621064\uff0c\u4f46\u662f\u5b9e\u9645\u7684\u5e95\u5c42\u6267\u884c\u4ee5K=16\u4e3a\u6b65\u957f\uff0c\u4f9d\u6b21\u5904\u7406K\u7ef4\u5ea6\u7684tile\uff0c\u56e0\u6b64\u5c31\u4ece\u5e95\u5c42\u6d88\u9664\u4e86Batch Variance\u3002\u90a3\u662f\u4e0d\u662f\u5c31\u662f\u8bf4\u53ef\u4ee5\u5ffd\u7565BLOCK_K\u7684\u5f71\u54cd\uff0c\u53ea\u5173\u6ce8Split-K\u5c31\u884c\u4e86\u3002\u8bb2\u9053\u7406\uff0c\u8fd9\u4e2a\u5728\u5927\u77e9\u9635\u4e0a\u5927\u6982\u7387\u662f\u6210\u7acb\u7684\u3002<\/p>\n<p>\u5bf9\u4e8e\u5c0f\u77e9\u9635\uff0c\u6bd4\u598232*32\u5206\u914d\u5230\u4e00\u4e2aCTA\uff08K=4096\uff09\uff0c\u90a3\u6bcf\u4e2awarp\u53ef\u80fd\u53ea\u80fd\u5206\u914d\u5f88\u5c0f\u7684Tile\uff0c\u4ece\u800c\u6ca1\u6709\u8db3\u591f\u591a\u7684Tile\uff0c\u5373M\/N\u5e76\u884c\u5ea6\u4e0d\u8db3-&gt;\u6d3b\u7684warp\u5c11,\u5360\u7528\u7387\u4f4e-&gt; \u8fd9\u79cd\u60c5\u51b5\u4e0b\u5728K\u7ef4\u5ea6\u4e32\u884c\u6267\u884c\u5c31\u65e0\u6cd5\u5145\u5206\u5229\u7528\u9690\u85cf\u5ef6\u8fdf\u7684\u7279\u6027\u3002\u56e0\u6b64\uff0c\u7cfb\u7edf\u88ab\u8feb\u8ba9\u6240\u6709 Warp \u91cd\u53e0\u5728\u8fd9\u540c\u4e00\u5757 32 * 32 \u7684\u77e9\u9635\u4e0a\uff0c\u5f3a\u884c\u5728\u65f6\u95f4\uff08K\uff09\u7ef4\u5ea6\u4e0a\u8fdb\u884c\u5207\u5206\uff08Warp-level K-Slicing\uff09\u3002 \u591a\u4e2a Warp \u5404\u7b97\u4e00\u6bb5 K\uff0c\u6700\u540e\u5728\u5171\u4eab\u5185\u5b58\u91cc\u901a\u8fc7Reduction Tree\u5408\u5e76\u3002\u4e00\u65e6 Batch \u53c2\u6570\u53d8\u52a8\u5bfc\u81f4 K \u5207\u5206\u6bb5\u6570\u6539\u53d8\uff0c\u89c4\u7ea6\u6811\u5f62\u72b6\u5c31\u4f1a\u6539\u53d8\uff0c\u4ece\u800c\u5f15\u53d1\u4e86\u6d6e\u70b9\u52a0\u6cd5\u987a\u5e8f\u7684\u6ce2\u52a8\u3002<\/p>\n<p>&nbsp;<\/p>\n<p>\u5f53\u7136\u5bf9\u4e8e\u5c0f\u77e9\u9635\uff0cGEMM \u5e93\u6b64\u65f6\u53ef\u80fd\u9009\u62e9\u66f4\u5c0f\u7684 tile\u3001GEMV\/SIMT\u3001persistent kernel\u3001warp-level K-slicing\u3001CTA-level Split-K \u6216 Stream-K \u7b49\u4e0d\u540c\u7b97\u6cd5\u3002<\/p>\n<p>Warp-level K-slicing \u4f1a\u8ba9\u591a\u4e2a warp \u5206\u522b\u8ba1\u7b97\u540c\u4e00\u8f93\u51fa tile \u7684\u4e0d\u540c K \u5206\u7247\uff0c\u518d\u5728 CTA \u5185\u5408\u5e76 partial results\uff1b\u5b83\u53ef\u4ee5\u589e\u52a0\u5355 CTA \u5185\u7684\u6709\u6548 warp \u6570\uff0c\u4f46\u4e0d\u80fd\u589e\u52a0 grid \u4e2d\u7684 CTA \u6570\u91cf\u3002\u4e3a\u4e86\u63d0\u5347SM\u7684\u5229\u7528\u7387\uff0c\u901a\u5e38\u9700\u8981 CTA-level Split-K\u3002<\/p>\n<p>&nbsp;<\/p>\n<p>\u4e0d\u8fc7\u5230\u4e86SM90+,\u964d\u4f4e\u4e86\u5bf9\u4f20\u7edf\u9ad8 occupancy \u9690\u85cf\u8bbf\u5b58\u5ef6\u8fdf\u7684\u4f9d\u8d56\uff0c\u6807\u51c6 WGMMA mainloop \u4e2d\uff0c\u4e00\u4e2awarpgroup \u5171\u540c\u7ef4\u62a4\u5206\u5e03\u5f0f accumulator\uff0c\u5e76\u4f9d\u6b21\u5904\u7406 K tile\uff0c\u4e0d\u9700\u8981\u751f\u6210\u591a\u4e2a\u72ec\u7acb warp partial result \u540e\u518d\u505a CTA \u5185 reduction\u3002<\/p>\n<p>\u6240\u4ee5\u6211\u4eec\u53ef\u4ee5\u770b\u5230\u4ee3\u7801\u6ce8\u91ca\uff1a&#8221;Hopper (SM90) and Blackwell (SM100): the only source of batch variance is split-k&#8221;\uff0c \u5373\u5728SM90 Hopper \/ SM100 Blackwell\u4e0a\uff0c\u5f53\u524d\u53d7\u652f\u6301\u548c\u6d4b\u8bd5\u7684 FP16\/BF16 PyTorch\/cuBLASLt \u8def\u5f84\uff0c\u5728\u7981\u6b62 Split-K batch invariant\u3002<\/p>\n<p>\u8fd9\u91cc\u9700\u8981\u7279\u522b\u8bf4\u660e\u7684\u662f\uff0cSM90\/SM100\u9501\u5b9aSplit-K\u53ef\u4ee5\u83b7\u53d6batch invariance\uff0cSM8x\u9700\u8981\u540c\u65f6\u9501\u5b9aBLOCK-K\u548cSplit-K\u624d\u80fd\u83b7\u53d6batch invariance\u90fd\u662f\u57fa\u4e8e\u6d4b\u8bd5\u7684\uff0c\u56e0\u4e3a cuBLAS \/ cuBLASLt \u662f\u95ed\u6e90\u7684\u3002<\/p>\n<p>\u5f53\u7136\u4e5f\u6709\u5f88\u591a\u4eba\u8bef\u8ba4\u4e3aSplit-k\u591aworkspace\u65f6\uff0c\u89c4\u7ea6\u6811\u662f\u6309\u987a\u5e8f\u89c4\u7ea6\u7684\uff0c\u53ef\u4ee5\u4fdd\u8bc1\u7ed3\u679c\u4e0d\u53d8\uff0c\u4f46\u662f\u8fd9\u4e2a\u53ea\u662frun-to-run\u7ef4\u5ea6\u7684\uff0c\u4e00\u65e6K\u53d8\u5316\uff0c\u7ed3\u679c\u8fd8\u662f\u4f1a\u53d8\u5316\u3002<\/p>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/e37bb6b09ca1ec314f8279abefdf03f70e93aa.webp\" data-type=\"block\" \/><\/p>\n<p>\u8bf4\u4e86\u8fd9\u4e48\u591a\uff0c\u5176\u5b9e\u5462\uff0c\u6700\u7ec8\u6d6e\u70b9\u6570\u52a0\u6cd5\u7684\u987a\u5e8f\uff0c\u4e0d\u4ec5\u53d7mainloop \u7684 BLOCK_K staging \u51b3\u5b9a\u7684\uff08CTA Tiling \uff09\uff0c\u4e5f\u53d7\u66f4\u4f4e\u7ef4\u5ea6\u7684Warp Tiling\/Warp Group Tiling\uff08\u4ee5\u53ca\u786c\u4ef6\u539f\u5b50\u6307\u4ee4\u7c92\u5ea6\uff09\u51b3\u5b9a\u3002\u4f46\u662f\u5462\uff0c\u8981\u89e3\u51b3\u8fd9\u4e2a\u95ee\u9898\u5374\u8981\u4eceCTA Tiling\u89d2\u5ea6\u6765\u3002<\/p>\n<p><strong>(2) SM80\u4e0aGEMM\u7684Batch Invariance\u652f\u6301<\/strong><\/p>\n<p>\u5bf9\u4e8eSM80\u8981\u540c\u65f6\u5173\u6ce8SPLIT_K\u548cBLOCK_K\u3002\u5176\u5b9e\u5bf9\u4e8eLinear\u5c42\u7684GEMM\u65e0\u8bbaSM\u662f\u54ea\u4e2a\u7248\u672c\/\u65e0\u8bba\u662feager\u8fd8\u662fcompile mode\uff0c\u90fd\u4f1a\u8d70\u5230linear_batch_invariant -&gt; matmul_persistent\u3002<\/p>\n<div>\n<div class=\"hljs-cto\">\n<div class=\"hljs-cto\"><button class=\"copy_btn disable\" data-clipboard-target=\"#code_id_3\">\u590d\u5236<\/button><\/p>\n<div class=\"code-toolbar\">\n<pre class=\"has-pre-numbering language-javascript\" tabindex=\"0\"><code class=\"language-javascript\"># <span class=\"token keyword\">class<\/span> <span class=\"token class-name\">UnquantizedLinearMethod<\/span><span class=\"token punctuation\">(<\/span>LinearMethodBase<span class=\"token punctuation\">)<\/span><span class=\"token operator\">:<\/span>\r\ndef <span class=\"token function\">apply<\/span><span class=\"token punctuation\">(<\/span>\r\n    self<span class=\"token punctuation\">,<\/span>\r\n    <span class=\"token literal-property property\">layer<\/span><span class=\"token operator\">:<\/span> torch<span class=\"token punctuation\">.<\/span>nn<span class=\"token punctuation\">.<\/span>Module<span class=\"token punctuation\">,<\/span>\r\n    <span class=\"token literal-property property\">x<\/span><span class=\"token operator\">:<\/span> torch<span class=\"token punctuation\">.<\/span>Tensor<span class=\"token punctuation\">,<\/span>\r\n    <span class=\"token literal-property property\">bias<\/span><span class=\"token operator\">:<\/span> torch<span class=\"token punctuation\">.<\/span>Tensor <span class=\"token operator\">|<\/span> None <span class=\"token operator\">=<\/span> None<span class=\"token punctuation\">,<\/span>\r\n<span class=\"token punctuation\">)<\/span> <span class=\"token operator\">-<\/span><span class=\"token operator\">&gt;<\/span> torch<span class=\"token punctuation\">.<\/span>Tensor<span class=\"token operator\">:<\/span>\r\n    <span class=\"token keyword\">if<\/span> envs<span class=\"token punctuation\">.<\/span><span class=\"token constant\">VLLM_BATCH_INVARIANT<\/span> and current_platform<span class=\"token punctuation\">.<\/span><span class=\"token function\">is_cuda_alike<\/span><span class=\"token punctuation\">(<\/span><span class=\"token punctuation\">)<\/span><span class=\"token operator\">:<\/span>\r\n        <span class=\"token keyword\">return<\/span> <span class=\"token function\">linear_batch_invariant<\/span><span class=\"token punctuation\">(<\/span>x<span class=\"token punctuation\">,<\/span> layer<span class=\"token punctuation\">.<\/span>weight<span class=\"token punctuation\">,<\/span> bias<span class=\"token punctuation\">)<\/span>\r\n    <span class=\"token keyword\">return<\/span> <span class=\"token function\">dispatch_unquantized_gemm<\/span><span class=\"token punctuation\">(<\/span><span class=\"token punctuation\">)<\/span><span class=\"token punctuation\">(<\/span>layer<span class=\"token punctuation\">,<\/span> x<span class=\"token punctuation\">,<\/span> layer<span class=\"token punctuation\">.<\/span>weight<span class=\"token punctuation\">,<\/span> bias<span class=\"token punctuation\">)<\/span><\/code><\/pre>\n<ul id=\"code_id_3\" class=\"pre-numbering\">\n<li>1.<\/li>\n<li>2.<\/li>\n<li>3.<\/li>\n<li>4.<\/li>\n<li>5.<\/li>\n<li>6.<\/li>\n<li>7.<\/li>\n<li>8.<\/li>\n<li>9.<\/li>\n<li>10.<\/li>\n<\/ul>\n<div class=\"toolbar\"><\/div>\n<\/div>\n<\/div>\n<\/div>\n<\/div>\n<p>\u800c\u5bf9\u4e8e\u975e\u7ebf\u6027\u5c42\u7684\u8ba1\u7b97\uff0c\u6bd4\u5982\u4ee3\u7801\u4e2d\u76f4\u63a5\u8c03\u7528torch.mm,\u5219\u8fd8\u9700\u8981\u989d\u5916\u8003\u8651\u3002\u7531\u4e8eSM80\u5728cuBLASLt\u4e0a\u4e0d\u652f\u6301\u7981\u7528BLOCK_K,\u53ea\u80fd\u66f4\u6362Triton\u5b9e\u73b0\uff0c\u6700\u7ec8\u4e5f\u4f1a\u8d70\u5230matmul_persistent\u3002\u5176\u4e2dmatmul_persistent\u662f\u4e00\u4e2apersistent kernel \u5373griddim=num_sm\uff0c\u56fa\u5b9a K \u987a\u5e8f\u3001\u56fa\u5b9a tile\u3001\u4e0d\u5207\u00a0SPLIT_K\u3002<\/p>\n<div>\n<div class=\"hljs-cto\">\n<div class=\"hljs-cto\"><button class=\"copy_btn disable\" data-clipboard-target=\"#code_id_4\">\u590d\u5236<\/button><\/p>\n<div class=\"code-toolbar\">\n<pre class=\"has-pre-numbering language-javascript\" tabindex=\"0\"><code class=\"language-javascript\"><span class=\"token keyword\">if<\/span> current_platform<span class=\"token punctuation\">.<\/span><span class=\"token function\">is_device_capability_family<\/span><span class=\"token punctuation\">(<\/span><span class=\"token number\">80<\/span><span class=\"token punctuation\">)<\/span><span class=\"token operator\">:<\/span>\r\n    # <span class=\"token constant\">SM80<\/span> <span class=\"token punctuation\">(<\/span>Ampere<span class=\"token punctuation\">)<\/span> cannot rely on cuBLASLt<span class=\"token operator\">-<\/span>only determinism<span class=\"token punctuation\">;<\/span> install the\r\n    # triton persistent matmul overrides <span class=\"token keyword\">for<\/span> mm<span class=\"token operator\">\/<\/span>addmm<span class=\"token operator\">\/<\/span>matmul<span class=\"token operator\">\/<\/span>linear<span class=\"token punctuation\">.<\/span>\r\n    _batch_invariant_LIB<span class=\"token punctuation\">.<\/span><span class=\"token function\">impl<\/span><span class=\"token punctuation\">(<\/span><span class=\"token string\">\"aten::mm\"<\/span><span class=\"token punctuation\">,<\/span> mm_batch_invariant<span class=\"token punctuation\">,<\/span> <span class=\"token string\">\"CUDA\"<\/span><span class=\"token punctuation\">)<\/span>\r\n    _batch_invariant_LIB<span class=\"token punctuation\">.<\/span><span class=\"token function\">impl<\/span><span class=\"token punctuation\">(<\/span><span class=\"token string\">\"aten::addmm\"<\/span><span class=\"token punctuation\">,<\/span> addmm_batch_invariant<span class=\"token punctuation\">,<\/span> <span class=\"token string\">\"CUDA\"<\/span><span class=\"token punctuation\">)<\/span>\r\n    _batch_invariant_LIB<span class=\"token punctuation\">.<\/span><span class=\"token function\">impl<\/span><span class=\"token punctuation\">(<\/span><span class=\"token string\">\"aten::matmul\"<\/span><span class=\"token punctuation\">,<\/span> matmul_batch_invariant<span class=\"token punctuation\">,<\/span> <span class=\"token string\">\"CUDA\"<\/span><span class=\"token punctuation\">)<\/span>\r\n    _batch_invariant_LIB<span class=\"token punctuation\">.<\/span><span class=\"token function\">impl<\/span><span class=\"token punctuation\">(<\/span><span class=\"token string\">\"aten::linear\"<\/span><span class=\"token punctuation\">,<\/span> linear_batch_invariant<span class=\"token punctuation\">,<\/span> <span class=\"token string\">\"CUDA\"<\/span><span class=\"token punctuation\">)<\/span><\/code><\/pre>\n<ul id=\"code_id_4\" class=\"pre-numbering\">\n<li>1.<\/li>\n<li>2.<\/li>\n<li>3.<\/li>\n<li>4.<\/li>\n<li>5.<\/li>\n<li>6.<\/li>\n<li>7.<\/li>\n<\/ul>\n<div class=\"toolbar\"><\/div>\n<\/div>\n<\/div>\n<\/div>\n<\/div>\n<p>\u7279\u522b\u6ce8\u610f\uff1acompile mode\u4e0baten override\u4e0d\u751f\u6548\uff0c\u6240\u4ee5\u4f1a\u8d70\u5230\u539f\u59cb\u7684aten::mm\u3002<\/p>\n<p><strong>(3) matmul_persistent \u662f\u600e\u4e48\u505a\u5230\u786e\u5b9a\u6027\u7684<\/strong><\/p>\n<p>\u5bf9\u4e8ebf16\u548cfp16\uff0c\u6240\u6709SM\u7248\u672c\u7684SM\u7684 Linear\uff0c\u65e0\u8bba\u662feager\u8fd8\u662fcompile mode\uff0c\u6700\u7ec8\u90fd\u843d\u5230\u00a0matmul_persistent\u00a0\u8fd9\u4e2a Triton kernel\u3002\u5f53\u7136eager mode\u4e0b\uff0c SM80\u7684\u88f8 mm\u7684\u4e5f\u4f1a\u8d70\u5230matmul_persistent\u3002\uff08compile mode\u4e0baten override\u4e0d\u751f\u6548\uff0c\u6240\u4ee5\u4f1a\u8d70\u5230\u539f\u59cb\u7684aten::mm\uff09<\/p>\n<p>matmul_persistent\u00a0\u7684\u5b9e\u73b0\u903b\u8f91\u5f88\u7b80\u5355\uff0c\u76f8\u5bf9\u4e8e\u666e\u901a GEMM \u9760\u00a0@triton.autotune\u00a0\/ cuBLASLt \u542f\u53d1\u5f0f\u6309 shape \u6311 config\uff0cmatmul_persistent\u00a0\u76f4\u63a5\u56fa\u5b9a config\uff0c\u5f7b\u5e95\u5173\u6389 autotune\u3002 \u6bcf\u79cd dtype \u7528\u4e00\u5957\u786c\u7f16\u7801\u7684\u5206\u5757\u53c2\u6570\uff0c\u7279\u522b\u662fBLOCK_K\uff0c\u4e0e\u8f93\u5165\u7684 M\/N\/K \u65e0\u5173\uff1a<\/p>\n<div>\n<div class=\"hljs-cto\">\n<div class=\"hljs-cto\"><button class=\"copy_btn disable\" data-clipboard-target=\"#code_id_5\">\u590d\u5236<\/button><\/p>\n<div class=\"code-toolbar\">\n<pre class=\"has-pre-numbering language-javascript\" tabindex=\"0\"><code class=\"language-javascript\">configs <span class=\"token operator\">=<\/span> <span class=\"token punctuation\">{<\/span>\r\n    torch<span class=\"token punctuation\">.<\/span>bfloat16<span class=\"token operator\">:<\/span> <span class=\"token punctuation\">{<\/span><span class=\"token string-property property\">\"BLOCK_SIZE_M\"<\/span><span class=\"token operator\">:<\/span> <span class=\"token number\">128<\/span><span class=\"token punctuation\">,<\/span> <span class=\"token string-property property\">\"BLOCK_SIZE_N\"<\/span><span class=\"token operator\">:<\/span> <span class=\"token number\">128<\/span><span class=\"token punctuation\">,<\/span> <span class=\"token string-property property\">\"BLOCK_SIZE_K\"<\/span><span class=\"token operator\">:<\/span> <span class=\"token number\">64<\/span><span class=\"token punctuation\">,<\/span>\r\n                     <span class=\"token string-property property\">\"GROUP_SIZE_M\"<\/span><span class=\"token operator\">:<\/span> <span class=\"token number\">8<\/span><span class=\"token punctuation\">,<\/span> <span class=\"token string-property property\">\"num_stages\"<\/span><span class=\"token operator\">:<\/span> <span class=\"token number\">3<\/span><span class=\"token punctuation\">,<\/span> <span class=\"token string-property property\">\"num_warps\"<\/span><span class=\"token operator\">:<\/span> <span class=\"token number\">8<\/span><span class=\"token punctuation\">}<\/span><span class=\"token punctuation\">,<\/span>\r\n    # fp16 <span class=\"token operator\">\/<\/span> fp32 \u5404\u81ea\u4e00\u5957\uff0c\u540c\u6837\u662f\u5e38\u91cf\r\n<span class=\"token punctuation\">}<\/span><\/code><\/pre>\n<ul id=\"code_id_5\" class=\"pre-numbering\">\n<li>1.<\/li>\n<li>2.<\/li>\n<li>3.<\/li>\n<li>4.<\/li>\n<li>5.<\/li>\n<\/ul>\n<div class=\"toolbar\"><\/div>\n<\/div>\n<\/div>\n<\/div>\n<\/div>\n<p>&nbsp;<\/p>\n<p>\u8fd9\u91cc\u8981\u6f84\u6e05\u4e00\u70b9\uff1aBLOCK_M\/N\/K\u3001SPLIT_K\u00a0\u5bf9 Triton autotune \u6765\u8bf4\u90fd\u53ea\u662f\u666e\u901a constexpr\uff0c\u800c split-k \u672c\u8eab\u662f\u4e00\u79cd\u9700\u8981\u4f5c\u8005\u624b\u5199\u8fdb kernel \u7684\u5b9e\u73b0\u6a21\u5f0f\uff08\u591a\u4e00\u7ef4 grid + atomic\/reduction\uff09\uff0c\u5e76\u975e autotune \u7684\u81ea\u52a8\u4f18\u5316\u7b56\u7565\u3002matmul_persistent\u00a0\u6ca1\u6709\u5b9e\u73b0 split-k\uff0c\u6240\u4ee5\u5b83\u65e2\u4e0d\u53d7 autotune \u5f71\u54cd\uff0c\u4e5f\u5929\u7136\u4e0d\u5b58\u5728 Split-K \u4e71\u5e8f\u2014\u2014\u771f\u6b63\u88ab\u5199\u6b7b\u7684\u5c31\u53ea\u6709\u00a0BLOCK_K\u00a0\u8fd9\u4e00\u4e2a\u53d8\u91cf\u3002<\/p>\n<p>\u53e6\u5916\u51fd\u6570\u540d\u91cc\u7684persistent\u5176\u5b9e\u8ddfbatch invariance\u65e0\u5173\uff0c\u662f\u8bf4grid \u53ea\u5f00\u00a0min(NUM_SMS, tile \u6570)\u3001\u6bcf\u4e2a program \u5e38\u9a7b SM \u5faa\u73af\u5403\u591a\u4e2a tile\uff09\uff0c\u662f\u4e3a\u4e86\u907f\u514d\u9891\u7e41 launch \u5f00\u9500\uff0ctile\u2192CTA \u7684\u6620\u5c04\u672c\u5c31\u7531\u56fa\u5b9a config \u51b3\u5b9a\uff0c\u4e0e\u786e\u5b9a\u6027\u65e0\u5173\u3002<\/p>\n<p>&nbsp;<\/p>\n<p><strong>(4) SM90\/SM100\u4e0aGEMM\u7684Batch Invariance\u652f\u6301<\/strong><\/p>\n<p>\u5bf9\u4e8e\uff0cSM90\/SM100\u53ea\u9700\u8981\u5173\u6ce8Split-K\uff0c\u5373\u7981\u6b62SPLIT_K\u62c6\u5206\uff1b<\/p>\n<blockquote><p>bf16\/fp16<\/p><\/blockquote>\n<p>\u4e0a\u9762\u5df2\u7ecf\u63d0\u5230\u4e86\u5bf9\u4e8eLinear\u5c42\u7684GEMM\u65e0\u8bbaSM\u662f\u54ea\u4e2a\u7248\u672c\uff0c\u90fd\u4f1a\u8d70\u5230linear_batch_invariant\uff0c\u6700\u7ec8\u8d70\u5230matmul_persistent\u3002<\/p>\n<p>\u5bf9\u4e8e\u76f4\u63a5\u8c03\u7528torch.mm\u7684\u90e8\u5206\uff0c\u5f53\u524d\u7684\u903b\u8f91\u662f\u4ecd\u7136\u8d70cuBLASLt\uff0c\u53ea\u662f\u4f9d\u8d56\u8bbe\u7f6e\u786e\u4fddSplit-K=1\u3002<\/p>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/67bc8a050f14fe481376646b813de165b5c06e.webp\" data-type=\"block\" \/><\/p>\n<p>\u8bbe\u7f6e\u4e86preferred_blas_library(backend=&#8221;cublaslt&#8221;)\u4e4b\u540e\uff0cGEMM\u4f1a\u4f18\u5148\u8d70cuBLASLt\uff0c\u901a\u8fc7\u8bbe\u7f6ereduction mask\uff0c\u786e\u4fdd\u4e86bf16\/fp16 Split-K=1\u3002CUBLASLT_MATMUL_PREF_REDUCTION_SCHEME_MASK = CUBLASLT_REDUCTION_SCHEME_NONE<\/p>\n<div>\n<div class=\"hljs-cto\">\n<div class=\"hljs-cto\"><button class=\"copy_btn disable\" data-clipboard-target=\"#code_id_6\">\u590d\u5236<\/button><\/p>\n<div class=\"code-toolbar\">\n<pre class=\"has-pre-numbering language-javascript\" tabindex=\"0\"><code class=\"language-javascript\">reduced_precision_val <span class=\"token operator\">=<\/span> <span class=\"token punctuation\">(<\/span>\r\n    <span class=\"token punctuation\">(<\/span>False<span class=\"token punctuation\">,<\/span> False<span class=\"token punctuation\">)<\/span> <span class=\"token keyword\">if<\/span> <span class=\"token function\">is_torch_equal_or_newer<\/span><span class=\"token punctuation\">(<\/span><span class=\"token string\">\"2.10.0\"<\/span><span class=\"token punctuation\">)<\/span> <span class=\"token keyword\">else<\/span> False\r\n<span class=\"token punctuation\">)<\/span>\r\ntorch<span class=\"token punctuation\">.<\/span>backends<span class=\"token punctuation\">.<\/span>cuda<span class=\"token punctuation\">.<\/span>matmul<span class=\"token punctuation\">.<\/span>allow_fp16_reduced_precision_reduction <span class=\"token operator\">=<\/span> <span class=\"token punctuation\">(<\/span>\r\n    reduced_precision_val\r\n<span class=\"token punctuation\">)<\/span>\r\ntorch<span class=\"token punctuation\">.<\/span>backends<span class=\"token punctuation\">.<\/span>cuda<span class=\"token punctuation\">.<\/span>matmul<span class=\"token punctuation\">.<\/span>allow_bf16_reduced_precision_reduction <span class=\"token operator\">=<\/span> <span class=\"token punctuation\">(<\/span>\r\n    reduced_precision_val\r\n<span class=\"token punctuation\">)<\/span>\r\ntorch<span class=\"token punctuation\">.<\/span>backends<span class=\"token punctuation\">.<\/span>cuda<span class=\"token punctuation\">.<\/span><span class=\"token function\">preferred_blas_library<\/span><span class=\"token punctuation\">(<\/span>backend<span class=\"token operator\">=<\/span><span class=\"token string\">\"cublaslt\"<\/span><span class=\"token punctuation\">)<\/span><\/code><\/pre>\n<ul id=\"code_id_6\" class=\"pre-numbering\">\n<li>1.<\/li>\n<li>2.<\/li>\n<li>3.<\/li>\n<li>4.<\/li>\n<li>5.<\/li>\n<li>6.<\/li>\n<li>7.<\/li>\n<li>8.<\/li>\n<li>9.<\/li>\n<li>10.<\/li>\n<\/ul>\n<div class=\"toolbar\"><\/div>\n<\/div>\n<\/div>\n<\/div>\n<\/div>\n<p>\u4f20\u7edf cuBLAS \u540e\u7aef\u6ca1\u6709 reduction-scheme mask \u8fd9\u4e2a\u80fd\u529b,PyTorch \u91cc\u4e00\u65e6\u8bbe\u4e86 allow_splitk=False \u53c8\u6ca1\u5207\u5230 cuBLASLt,\u4f1a\u76f4\u63a5 TORCH_CHECK \u62a5\u9519:TORCH_CHECK(reduction != DisallowReducedPrecisionDisallowSplitK, &#8220;&#8230;allow_splitk=False) requires the cuBLASLt backend&#8221;);<\/p>\n<p>fp8<\/p>\n<p>fp8 \u7684\u5165\u53e3\u662f\u91cf\u5316\u7ebf\u6027\u5c42\u91cc\u7684\u00a0cutlass_scaled_mm\uff0c\u5e95\u5c42\u6709 CUTLASS \u4e0e cuBLASLt\u00a0scaled_gemm\u00a0\u4e24\u6761\u8def\u3002<\/p>\n<p>\u9009\u4e2d CUTLASS \u2192 \u8c03\u5b83\uff08csrc if(batch_invariant) \u9501 config \u4e0d\u968f M<\/p>\n<p>cuBLASLt scaled_gemm \u2192 \u7531\u4e8eAPI\u8bbe\u7f6emask\u4ec5\u652f\u6301bf16\/fp16\uff0c\u5bf9\u4e8efp8\u53ea\u80fd\u53e6\u5916\u60f3\u529e\u6cd5\u3002\u600e\u4e48\u529e\u5462\uff1f\u628aworkspace\u7684\u9650\u5236\u5230\u5c3d\u91cf\u5c0f\uff0c\u800csplit-k(COMPUTE\/OUTPUT) \u4e0e atomic(INPLACE) reduction \u90fd\u9700 workspace\uff0c\u6700\u7ec8\u53ef\u9009\u7b56\u7565\u53ea\u5269\u4e0d\u9700 workspace \u7684 NONE\uff0c\u5373 split-k=1\uff1b\u4f46 scaled_gemm \u6ca1\u8bbe\u7f6e reduction-mask\uff0c\u65e0 API \u7ea7\u4fdd\u8bc1\u3002<\/p>\n<p>os.environ[&#8220;CUBLASLT_WORKSPACE_SIZE&#8221;] = &#8220;1&#8221;<\/p>\n<p>FlashInfer \u2192 DeepGEMM<\/p>\n<p>fp4<\/p>\n<p>\u5982\u679c\u5f00\u542f\u4e86batch invariant\uff0c\u65e0\u8bba M \u591a\u5927\uff0c\u90fd\u56fa\u5b9a\u4f7f\u7528\u540c\u4e00\u4e2a\u914d\u7f6e\u00a0sm100_fp4_config_default\u3002\u53e6\u5916\u5c31\u662f\u9009\u62e9PersistentScheduler\uff0c\u786e\u4fddK\u7ef4\u5ea6\u4e0d\u62c6\u5206\u3002<\/p>\n<p>CUTLASS 3.x TileScheduler\u7684\u5b9a\u4e49\uff1a<\/p>\n<table class=\"data-table\" data-transient-attributes=\"class\" data-width=\"676.997px\">\n<colgroup data-id=\"c7104f7d-K9eAZ0pU\">\n<col span=\"1\" width=\"225.66\" data-id=\"cd89ecb0-gow3o23x\" \/>\n<col span=\"1\" width=\"225.66\" data-id=\"cd89ecb0-padlBMtd\" \/>\n<col span=\"1\" width=\"225.677\" data-id=\"cd89ecb0-yTlm5lDU\" \/><\/colgroup>\n<tbody data-id=\"t6d5e859-47dSk4eB\">\n<tr data-id=\"t31e458f-o7vK1BVO\">\n<td data-id=\"t6267798-7aPDXSWE\" data-transient-attributes=\"table-cell-selection\">\u7c7b\u578b<\/td>\n<td data-id=\"t6267798-zJsuc3hT\" data-transient-attributes=\"table-cell-selection\">\u8c03\u5ea6\u65b9\u5f0f<\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-pozeDlam\" data-transient-attributes=\"table-cell-selection\">K \u7ef4\u5ea6\u62c6\u5206<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-GxIvo0UJ\">\n<td data-id=\"t6267798-Km6Yhxh5\" data-transient-attributes=\"table-cell-selection\"><code>cutlass::gemm::PersistentScheduler<\/code><\/td>\n<td data-id=\"t6267798-j0ZSlTK6\" data-transient-attributes=\"table-cell-selection\">\u6301\u4e45\u5316<\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-bsqNbnIL\" data-transient-attributes=\"table-cell-selection\">data-parallel\uff0c\u4e0d\u62c6 K<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-YEi5oIAj\">\n<td class=\"table-last-column\" data-id=\"t6267798-nsc9NHx6\" data-transient-attributes=\"table-cell-selection\"><code>cutlass::gemm::StreamKScheduler<\/code><\/td>\n<td class=\"table-last-column\" data-id=\"t6267798-YMk6qi0u\" data-transient-attributes=\"table-cell-selection\">\u6301\u4e45\u5316<\/td>\n<td class=\"table-last-column table-last-row\" data-id=\"t6267798-dTboJuqm\" data-transient-attributes=\"table-cell-selection\">\u62c6 K\uff08stream-K\/split-K\uff09<\/td>\n<\/tr>\n<\/tbody>\n<\/table>\n<div>\n<div class=\"hljs-cto\">\n<div class=\"hljs-cto\"><button class=\"copy_btn disable\" data-clipboard-target=\"#code_id_7\">\u590d\u5236<\/button><\/p>\n<div class=\"code-toolbar\">\n<pre class=\"has-pre-numbering language-javascript\" tabindex=\"0\"><code class=\"language-javascript\"><span class=\"token comment\">\/\/ Dispatch function to select appropriate config based on M<\/span>\r\ntemplate <span class=\"token operator\">&lt;<\/span>typename OutType<span class=\"token operator\">&gt;<\/span>\r\n<span class=\"token keyword\">void<\/span> <span class=\"token function\">cutlass_fp4_gemm_dispatch<\/span><span class=\"token punctuation\">(<\/span><span class=\"token parameter\"><span class=\"token literal-property property\">torch<\/span><span class=\"token operator\">:<\/span><span class=\"token operator\">:<\/span>stable<span class=\"token operator\">:<\/span><span class=\"token operator\">:<\/span>Tensor<span class=\"token operator\">&amp;<\/span> <span class=\"token constant\">D<\/span><span class=\"token punctuation\">,<\/span>\r\n                               <span class=\"token literal-property property\">torch<\/span><span class=\"token operator\">:<\/span><span class=\"token operator\">:<\/span>stable<span class=\"token operator\">:<\/span><span class=\"token operator\">:<\/span>Tensor <span class=\"token keyword\">const<\/span><span class=\"token operator\">&amp;<\/span> <span class=\"token constant\">A<\/span><span class=\"token punctuation\">,<\/span>\r\n                               <span class=\"token literal-property property\">torch<\/span><span class=\"token operator\">:<\/span><span class=\"token operator\">:<\/span>stable<span class=\"token operator\">:<\/span><span class=\"token operator\">:<\/span>Tensor <span class=\"token keyword\">const<\/span><span class=\"token operator\">&amp;<\/span> <span class=\"token constant\">B<\/span><span class=\"token punctuation\">,<\/span>\r\n                               <span class=\"token literal-property property\">torch<\/span><span class=\"token operator\">:<\/span><span class=\"token operator\">:<\/span>stable<span class=\"token operator\">:<\/span><span class=\"token operator\">:<\/span>Tensor <span class=\"token keyword\">const<\/span><span class=\"token operator\">&amp;<\/span> A_sf<span class=\"token punctuation\">,<\/span>\r\n                               <span class=\"token literal-property property\">torch<\/span><span class=\"token operator\">:<\/span><span class=\"token operator\">:<\/span>stable<span class=\"token operator\">:<\/span><span class=\"token operator\">:<\/span>Tensor <span class=\"token keyword\">const<\/span><span class=\"token operator\">&amp;<\/span> B_sf<span class=\"token punctuation\">,<\/span>\r\n                               <span class=\"token literal-property property\">torch<\/span><span class=\"token operator\">:<\/span><span class=\"token operator\">:<\/span>stable<span class=\"token operator\">:<\/span><span class=\"token operator\">:<\/span>Tensor <span class=\"token keyword\">const<\/span><span class=\"token operator\">&amp;<\/span> alpha<span class=\"token punctuation\">,<\/span> int64_t m<span class=\"token punctuation\">,<\/span>\r\n                               int64_t n<span class=\"token punctuation\">,<\/span> int64_t k<span class=\"token punctuation\">,<\/span> cudaStream_t stream<\/span><span class=\"token punctuation\">)<\/span> <span class=\"token punctuation\">{<\/span>\r\n  <span class=\"token keyword\">if<\/span> <span class=\"token punctuation\">(<\/span>vllm<span class=\"token operator\">:<\/span><span class=\"token operator\">:<\/span><span class=\"token function\">vllm_is_batch_invariant<\/span><span class=\"token punctuation\">(<\/span><span class=\"token punctuation\">)<\/span><span class=\"token punctuation\">)<\/span> <span class=\"token punctuation\">{<\/span>\r\n    using BiGemm <span class=\"token operator\">=<\/span> Fp4GemmSm100<span class=\"token operator\">&lt;<\/span>sm100_fp4_config_default<span class=\"token punctuation\">,<\/span> OutType<span class=\"token operator\">&gt;<\/span><span class=\"token punctuation\">;<\/span>\r\n    <span class=\"token function\">static_assert<\/span><span class=\"token punctuation\">(<\/span>\r\n        <span class=\"token literal-property property\">cute<\/span><span class=\"token operator\">:<\/span><span class=\"token operator\">:<\/span>is_same_v<span class=\"token operator\">&lt;<\/span>typename BiGemm<span class=\"token operator\">:<\/span><span class=\"token operator\">:<\/span>TileScheduler<span class=\"token punctuation\">,<\/span>\r\n                        <span class=\"token literal-property property\">cutlass<\/span><span class=\"token operator\">:<\/span><span class=\"token operator\">:<\/span>gemm<span class=\"token operator\">:<\/span><span class=\"token operator\">:<\/span>PersistentScheduler<span class=\"token operator\">&gt;<\/span><span class=\"token punctuation\">,<\/span>\r\n        <span class=\"token string\">\"batch_invariant requires a persistent tile scheduler; stream-K or \"<\/span>\r\n        <span class=\"token string\">\"split-K would break numerical invariance\"<\/span><span class=\"token punctuation\">)<\/span><span class=\"token punctuation\">;<\/span>\r\n    runGemm<span class=\"token operator\">&lt;<\/span>BiGemm<span class=\"token operator\">&gt;<\/span><span class=\"token punctuation\">(<\/span><span class=\"token constant\">D<\/span><span class=\"token punctuation\">,<\/span> <span class=\"token constant\">A<\/span><span class=\"token punctuation\">,<\/span> <span class=\"token constant\">B<\/span><span class=\"token punctuation\">,<\/span> A_sf<span class=\"token punctuation\">,<\/span> B_sf<span class=\"token punctuation\">,<\/span> alpha<span class=\"token punctuation\">,<\/span> m<span class=\"token punctuation\">,<\/span> n<span class=\"token punctuation\">,<\/span> k<span class=\"token punctuation\">,<\/span> stream<span class=\"token punctuation\">)<\/span><span class=\"token punctuation\">;<\/span>\r\n    <span class=\"token keyword\">return<\/span><span class=\"token punctuation\">;<\/span>\r\n  <span class=\"token punctuation\">}<\/span>\r\n\r\n  uint32_t <span class=\"token keyword\">const<\/span> mp2 <span class=\"token operator\">=<\/span> std<span class=\"token operator\">:<\/span><span class=\"token operator\">:<\/span><span class=\"token function\">max<\/span><span class=\"token punctuation\">(<\/span>static_cast<span class=\"token operator\">&lt;<\/span>uint32_t<span class=\"token operator\">&gt;<\/span><span class=\"token punctuation\">(<\/span><span class=\"token number\">16<\/span><span class=\"token punctuation\">)<\/span><span class=\"token punctuation\">,<\/span> <span class=\"token function\">next_pow_2<\/span><span class=\"token punctuation\">(<\/span>m<span class=\"token punctuation\">)<\/span><span class=\"token punctuation\">)<\/span><span class=\"token punctuation\">;<\/span>\r\n\r\n  <span class=\"token keyword\">if<\/span> <span class=\"token punctuation\">(<\/span>mp2 <span class=\"token operator\">&lt;=<\/span> <span class=\"token number\">16<\/span><span class=\"token punctuation\">)<\/span> <span class=\"token punctuation\">{<\/span>\r\n    <span class=\"token comment\">\/\/ m in [1, 16]<\/span>\r\n    runGemm<span class=\"token operator\">&lt;<\/span>Fp4GemmSm100<span class=\"token operator\">&lt;<\/span>sm100_fp4_config_M16<span class=\"token punctuation\">,<\/span> OutType<span class=\"token operator\">&gt;&gt;<\/span><span class=\"token punctuation\">(<\/span>\r\n        <span class=\"token constant\">D<\/span><span class=\"token punctuation\">,<\/span> <span class=\"token constant\">A<\/span><span class=\"token punctuation\">,<\/span> <span class=\"token constant\">B<\/span><span class=\"token punctuation\">,<\/span> A_sf<span class=\"token punctuation\">,<\/span> B_sf<span class=\"token punctuation\">,<\/span> alpha<span class=\"token punctuation\">,<\/span> m<span class=\"token punctuation\">,<\/span> n<span class=\"token punctuation\">,<\/span> k<span class=\"token punctuation\">,<\/span> stream<span class=\"token punctuation\">)<\/span><span class=\"token punctuation\">;<\/span>\r\n  <span class=\"token punctuation\">}<\/span> <span class=\"token keyword\">else<\/span> <span class=\"token keyword\">if<\/span> <span class=\"token punctuation\">(<\/span>mp2 <span class=\"token operator\">&lt;=<\/span> <span class=\"token number\">256<\/span><span class=\"token punctuation\">)<\/span> <span class=\"token punctuation\">{<\/span>\r\n    <span class=\"token comment\">\/\/ m in (16, 256]<\/span>\r\n    runGemm<span class=\"token operator\">&lt;<\/span>Fp4GemmSm100<span class=\"token operator\">&lt;<\/span>sm100_fp4_config_M256<span class=\"token punctuation\">,<\/span> OutType<span class=\"token operator\">&gt;&gt;<\/span><span class=\"token punctuation\">(<\/span>\r\n        <span class=\"token constant\">D<\/span><span class=\"token punctuation\">,<\/span> <span class=\"token constant\">A<\/span><span class=\"token punctuation\">,<\/span> <span class=\"token constant\">B<\/span><span class=\"token punctuation\">,<\/span> A_sf<span class=\"token punctuation\">,<\/span> B_sf<span class=\"token punctuation\">,<\/span> alpha<span class=\"token punctuation\">,<\/span> m<span class=\"token punctuation\">,<\/span> n<span class=\"token punctuation\">,<\/span> k<span class=\"token punctuation\">,<\/span> stream<span class=\"token punctuation\">)<\/span><span class=\"token punctuation\">;<\/span>\r\n  <span class=\"token punctuation\">}<\/span> <span class=\"token keyword\">else<\/span> <span class=\"token punctuation\">{<\/span>\r\n    <span class=\"token comment\">\/\/ m in (256, inf)<\/span>\r\n    runGemm<span class=\"token operator\">&lt;<\/span>Fp4GemmSm100<span class=\"token operator\">&lt;<\/span>sm100_fp4_config_default<span class=\"token punctuation\">,<\/span> OutType<span class=\"token operator\">&gt;&gt;<\/span><span class=\"token punctuation\">(<\/span>\r\n        <span class=\"token constant\">D<\/span><span class=\"token punctuation\">,<\/span> <span class=\"token constant\">A<\/span><span class=\"token punctuation\">,<\/span> <span class=\"token constant\">B<\/span><span class=\"token punctuation\">,<\/span> A_sf<span class=\"token punctuation\">,<\/span> B_sf<span class=\"token punctuation\">,<\/span> alpha<span class=\"token punctuation\">,<\/span> m<span class=\"token punctuation\">,<\/span> n<span class=\"token punctuation\">,<\/span> k<span class=\"token punctuation\">,<\/span> stream<span class=\"token punctuation\">)<\/span><span class=\"token punctuation\">;<\/span>\r\n  <span class=\"token punctuation\">}<\/span>\r\n<span class=\"token punctuation\">}<\/span><\/code><\/pre>\n<ul id=\"code_id_7\" class=\"pre-numbering\">\n<li>1.<\/li>\n<li>2.<\/li>\n<li>3.<\/li>\n<li>4.<\/li>\n<li>5.<\/li>\n<li>6.<\/li>\n<li>7.<\/li>\n<li>8.<\/li>\n<li>9.<\/li>\n<li>10.<\/li>\n<li>11.<\/li>\n<li>12.<\/li>\n<li>13.<\/li>\n<li>14.<\/li>\n<li>15.<\/li>\n<li>16.<\/li>\n<li>17.<\/li>\n<li>18.<\/li>\n<li>19.<\/li>\n<li>20.<\/li>\n<li>21.<\/li>\n<li>22.<\/li>\n<li>23.<\/li>\n<li>24.<\/li>\n<li>25.<\/li>\n<li>26.<\/li>\n<li>27.<\/li>\n<li>28.<\/li>\n<li>29.<\/li>\n<li>30.<\/li>\n<li>31.<\/li>\n<li>32.<\/li>\n<li>33.<\/li>\n<li>34.<\/li>\n<li>35.<\/li>\n<li>36.<\/li>\n<\/ul>\n<div class=\"toolbar\"><\/div>\n<\/div>\n<\/div>\n<\/div>\n<\/div>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/47c891b6160d36ca4505459a34357d084c5de4.webp\" data-type=\"block\" \/><\/p>\n<h3>\u4e09\u3001RMSNorm\u7684Batch Invariance<\/h3>\n<h4>1. RMSNorm &#8211; \u5747\u65b9\u6839\u5f52\u4e00\u5316<\/h4>\n<p>\u8fd9\u4e2a\u7b97\u5b50\u7684\u8f93\u5165\u662f[num_tokens, hidden_size]\uff0c\u7136\u540e\u9700\u8981\u5bf9\u4e8e\u6bcf\u4e2atoken\u7684hidden\u505a\u5747\u65b9\u6839\u5f52\u4e00\u5316\uff0c\u5747\u65b9\u6839\u5f52\u4e00\u5316\u7684\u516c\u5f0f\u5982\u4e0b\uff1a<\/p>\n<blockquote><p>y\u00a0= (x)\/(RMS(x)) \u2299 \u03b3\u00a0\uff0c\u5176\u4e2d\u5747\u65b9\u6839\uff08RMS\uff09\u7684\u8ba1\u7b97\uff1aRMS(x) = \u221a((1)\/(d) \u03a3i=1d\u00a0xi2\u00a0+ \u03b5)<\/p><\/blockquote>\n<p>\u8981\u7406\u89e3\u5177\u4f53\u7684\u6267\u884c\u8fc7\u7a0b\u548c\u4f18\u5316\u903b\u8f91\u8981\u5bf9\u4e8eGPU\u7684\u4f53\u7cfb\u67b6\u6784\u6709\u4e00\u5b9a\u7684\u4e86\u89e3\uff0c\u53ef\u4ee5\u53c2\u8003\u672c\u4eba\u7684\u300aAI Infra\u5165\u95e8\uff1aGPU\u662f\u5982\u4f55\u5de5\u4f5c\u7684\u300b\uff0c\u8fd9\u91cc\u5c31\u4e0d\u5728\u8d58\u8ff0\u4e86\u3002<\/p>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/a7318249761a7806cf9278fb1041de5abe54db.webp\" data-type=\"block\" \/><\/p>\n<p>\u7531\u4e8eRMSNORM\u9700\u8981\u505atoken\u7ef4\u5ea6\u7684\u5747\u65b9\u6839\u5f52\u4e00\u5316\uff0c\u8fd9\u91cc\u663e\u800c\u6613\u89c1\uff0c\u8981\u786e\u4fdd\u6bcf\u4e2atoken\u7684hidden\u5206\u914d\u5230\u4e00\u4e2athread block\uff0c\u8fd9\u6837\u53ef\u4ee5\u907f\u514d\u8de8thread block\u7684\u540c\u6b65\u3002<\/p>\n<div>\n<div class=\"hljs-cto\">\n<div class=\"hljs-cto\"><button class=\"copy_btn disable\" data-clipboard-target=\"#code_id_8\">\u590d\u5236<\/button><\/p>\n<div class=\"code-toolbar\">\n<pre class=\"has-pre-numbering language-javascript\" tabindex=\"0\"><code class=\"language-javascript\">dim3 <span class=\"token function\">grid<\/span><span class=\"token punctuation\">(<\/span>num_tokens<span class=\"token punctuation\">)<\/span><span class=\"token punctuation\">;<\/span><\/code><\/pre>\n<ul id=\"code_id_8\" class=\"pre-numbering\">\n<li>1.<\/li>\n<\/ul>\n<div class=\"toolbar\"><\/div>\n<\/div>\n<\/div>\n<\/div>\n<\/div>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/f29b9fe75c5091e016148413d1f1d2dc45e933.webp\" data-type=\"block\" \/><\/p>\n<p>\u90a3\u600e\u4e48\u5b9a\u4e49blockDim\u5462\uff1f\u9996\u5148\u8981\u660e\u767dblockDim\u610f\u5473\u7740\u4ec0\u4e48, blockDim\u672c\u8d28\u662f\u5206\u914d\u7ed9\u8fd9\u4e2ablock\u7684\u5e76\u53d1thread\u3002\u8fd9\u4e2a\u4e1c\u897f\u592a\u5927\u7684\u8bdd\uff0cblock\u5185\u7684\u89c4\u7ea6\u548c\u540c\u6b65\u7684\u6210\u672c\u90fd\u4f1a\u53d8\u9ad8\uff1b\u592a\u4f4e\u5462\uff0c\u6bcf\u4e2aSM\u53c8\u6709max block\u4e0a\u9650\uff0c\u5bfc\u81f4GPU Occupancy\u4f4e\u3002<\/p>\n<div>\n<div class=\"hljs-cto\">\n<div class=\"hljs-cto\"><button class=\"copy_btn disable\" data-clipboard-target=\"#code_id_9\">\u590d\u5236<\/button><\/p>\n<div class=\"code-toolbar\">\n<pre class=\"has-pre-numbering language-javascript\" tabindex=\"0\"><code class=\"language-javascript\"><span class=\"token comment\">\/\/ For large num_tokens, use smaller blocks to increase SM concurrency.<\/span>\r\n<span class=\"token keyword\">const<\/span> int max_block_size <span class=\"token operator\">=<\/span> <span class=\"token punctuation\">(<\/span>num_tokens <span class=\"token operator\">&lt;<\/span> <span class=\"token number\">256<\/span><span class=\"token punctuation\">)<\/span> <span class=\"token operator\">?<\/span> <span class=\"token number\">1024<\/span> <span class=\"token operator\">:<\/span> <span class=\"token number\">256<\/span><span class=\"token punctuation\">;<\/span>\r\n<span class=\"token keyword\">const<\/span> int block_size <span class=\"token operator\">=<\/span> std<span class=\"token operator\">:<\/span><span class=\"token operator\">:<\/span><span class=\"token function\">min<\/span><span class=\"token punctuation\">(<\/span>hidden_size <span class=\"token operator\">\/<\/span> calculated_vec_size<span class=\"token punctuation\">,<\/span> max_block_size<span class=\"token punctuation\">)<\/span><span class=\"token punctuation\">;<\/span>\r\ndim3 <span class=\"token function\">block<\/span><span class=\"token punctuation\">(<\/span>block_size<span class=\"token punctuation\">)<\/span><span class=\"token punctuation\">;<\/span><\/code><\/pre>\n<ul id=\"code_id_9\" class=\"pre-numbering\">\n<li>1.<\/li>\n<li>2.<\/li>\n<li>3.<\/li>\n<li>4.<\/li>\n<\/ul>\n<div class=\"toolbar\"><\/div>\n<\/div>\n<\/div>\n<\/div>\n<\/div>\n<p>\u4e00\u822c\u60c5\u51b5\u4e0b\uff0c\u4efb\u52a1\u8db3\u591f\u505a\uff08\u5373num_tokens\u591f\u65f6\uff09\uff0c\u65e0\u8bbablock_size\u5927\u5c0f\uff0c\u90fd\u80fd\u6709\u5f88\u9ad8\u7684\u5360\u7528\u7387\u3002\u8fd9\u65f6\u5019\u00a0block_size\u00a0\u4e00\u822c\u8bbe\u7f6e\u4e3a 256 \uff0c\u53ef\u4ee5\u907f\u5f00\u5927 Block \u5e26\u6765\u7684\u00a0__syncthreads()\u00a0\u540c\u6b65\u7684\u95ee\u9898\uff0c\u5b9e\u73b0\u5404\u4e2a\u5c0f Block \u4e4b\u95f4\u7684\u9ad8\u6548\u5e76\u53d1\u3002\uff08\u4e5f\u6709\u4e9b\u573a\u666f\u5c0fBlock\u4f1a\u5bfc\u81f4grid\u5c42\u9762\u7684\u89c4\u7ea6\u53d8\u591a\uff0c\u4f46\u662f\u7531\u4e8eRMSNORM\u662ftoken\u7ef4\u5ea6\u7684\u8ba1\u7b97\uff0c\u53ea\u6709block\u5c42\u9762\u7684\u89c4\u7ea6\uff0c\u56e0\u6b64\u6ca1\u6709\u8fd9\u4e2a\u95ee\u9898\uff09<\/p>\n<p>\u800c\u5f53num_tokens\u5f88\u5c11\u65f6\uff0c\u60c5\u51b5\u5c31\u53d8\u4e86\u3002\u6bd4\u5982\u53ea\u67092\u4e2atoken\uff0c\u53732\u4e2athread block\uff0c\u90a3\u9876\u591a\u4e5f\u5c31\u5360\u75282\u4e2aSM\uff0c\u8fd9\u65f6\u5019\u5b8c\u5168\u4e0d\u7528\u8003\u8651\u5168\u5c40Occupancy\u7684\u95ee\u9898\u4e86\u3002\u8fd9\u65f6\u5019\u8981\u8003\u8651\u7684\u5c31\u662f\u5c3d\u53ef\u80fd\u5145\u5206\u5229\u7528SM\u7684\u5e76\u53d1\u7b97\u529b\uff0c\u5373\u8db3\u591f\u591a\u7684\u7ebf\u7a0b\/warp\uff0c\u8ba9stride loop\u6b21\u6570\u53d8\u5c11\uff0c\u56e0\u6b64\u8bbe\u7f6e\u4e00\u4e2a\u8f83\u5927\u7684block_size\u5373\u53ef\u3002<\/p>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/89b95d9422dba6a923194132fd037c548aa489.webp\" data-type=\"block\" \/><\/p>\n<p>\u7136\u800cblock_size\u5219\u4f1a\u5f71\u54cd1. \u6bcf\u4e2a\u7ebf\u7a0b\u672c\u5730\u7d2f\u52a0\u7684\u5143\u7d20\u5206\u7ec4\uff1b2. BlockReduce \u6811\u5f62\u5f52\u7ea6\u7684\u5f62\u72b6\uff08\u5c42\u6570\u3001warp \u5206\u7ec4\uff09\u3002\u4ece\u800c\u5bfc\u81f4\u6d6e\u70b9\u6570\u52a0\u6cd5\u987a\u5e8f\u53d1\u751f\u53d8\u5316\u3002<\/p>\n<p>\u90a3\u600e\u4e48\u907f\u514d\u5462\uff1f<\/p>\n<h4>2. vLLM\u4e2dRMSNorm\u7684Batch Invariance\u652f\u6301<\/h4>\n<p>\u4ee5\u4e0a\u53ef\u4ee5\u770b\u5230,rms_norm \u7684\u7b97\u6cd5\u4e0e\u4f18\u5316\u672c\u8eab\u90fd\u5f88\u7b80\u5355\u3002\u4f46\u5b83\u7684 dispatch \u903b\u8f91\u5176\u5b9e\u76f8\u5f53\u590d\u6742:\u4ec5\u4ec5\u662f\u786e\u5b9a&#8221;\u5728\u591a\u540e\u7aef\u3001\u591a\u5e73\u53f0\u3001eager\/\u7f16\u8bd1\u7b49\u591a\u79cd\u7ec4\u5408\u4e0b\u8be5\u8d70\u54ea\u4e2a\u7b97\u5b50\u5b9e\u73b0&#8221;,\u5c31\u8981\u82b1\u8d39\u5f88\u5927\u529b\u6c14\u3002\u8fd9\u91cc\u6d89\u53cavLLM IR\u3001CustomOp\u3001Torch Inductor 3\u5c42dispatcher,\u8fd9\u91cc\u5c31\u4e0d\u6df1\u5165\u4e86\uff0c\u540e\u7eed\u6df1\u5165\u6587\u7ae0\u518d\u6df1\u5165\u89e3\u8bfb\u4e0b\u3002<\/p>\n<table class=\"data-table\" data-transient-attributes=\"class\" data-width=\"676.997px\">\n<colgroup data-id=\"c7104f7d-PV6Ip5iG\">\n<col span=\"1\" width=\"112.83\" data-id=\"cd89ecb0-9S2tnaS5\" \/>\n<col span=\"1\" width=\"112.83\" data-id=\"cd89ecb0-7753Qmqz\" \/>\n<col span=\"1\" width=\"112.83\" data-id=\"cd89ecb0-Hk9apL4y\" \/>\n<col span=\"1\" width=\"112.83\" data-id=\"cd89ecb0-FdcIUH6H\" \/>\n<col span=\"1\" width=\"112.83\" data-id=\"cd89ecb0-3GBfocYa\" \/>\n<col span=\"1\" width=\"112.847\" data-id=\"cd89ecb0-DanpNt4c\" \/><\/colgroup>\n<tbody data-id=\"t6d5e859-nPUBBGZ7\">\n<tr data-id=\"t31e458f-56v6FyUf\">\n<td data-id=\"t6267798-Fze9HArY\" data-transient-attributes=\"table-cell-selection\">Mode<\/td>\n<td data-id=\"t6267798-CqJ9YhHx\" data-transient-attributes=\"table-cell-selection\">residual<\/td>\n<td data-id=\"t6267798-Bz6gDIkx\" data-transient-attributes=\"table-cell-selection\">Batch Invariance<\/td>\n<td data-id=\"t6267798-jJU3fE1M\" data-transient-attributes=\"table-cell-selection\">dtype \/ quant<\/td>\n<td data-id=\"t6267798-J7guRjWo\" data-transient-attributes=\"table-cell-selection\">Path<\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-Gp6FjdKq\" data-transient-attributes=\"table-cell-selection\">Terminal implementation<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-cCqqPzSd\">\n<td data-id=\"t6267798-qLFoEk9d\" data-transient-attributes=\"table-cell-selection\">eager<\/td>\n<td data-id=\"t6267798-DI8y1W3m\" data-transient-attributes=\"table-cell-selection\">none<\/td>\n<td data-id=\"t6267798-7zDkaKUO\" data-transient-attributes=\"table-cell-selection\">off<\/td>\n<td data-id=\"t6267798-raAA6EKP\" data-transient-attributes=\"table-cell-selection\">bf16\/fp16<\/td>\n<td data-id=\"t6267798-U1Ke5Ijl\" data-transient-attributes=\"table-cell-selection\"><code>forward_cuda<\/code><\/p>\n<p>\u2192<code>forward_native<\/code>\u2192IR<code>[vllm_c, native]<\/code><\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-HVHcHI1W\" data-transient-attributes=\"table-cell-selection\">CUDA<\/p>\n<p><code>rms_norm<\/code>\u00a0\u2192\u00a0<code>rms_norm_kernel&lt;scalar_t, vec_size, tensor_rank, has_weight&gt;<\/code><\/td>\n<\/tr>\n<tr data-id=\"t31e458f-nAnXjnEX\">\n<td data-id=\"t6267798-ScuubOoo\" data-transient-attributes=\"table-cell-selection\">eager<\/td>\n<td data-id=\"t6267798-WpWYNCMx\" data-transient-attributes=\"table-cell-selection\">yes<\/td>\n<td data-id=\"t6267798-nt5vrNeQ\" data-transient-attributes=\"table-cell-selection\">off<\/td>\n<td data-id=\"t6267798-7LJRvLk3\" data-transient-attributes=\"table-cell-selection\">bf16\/fp16<\/td>\n<td data-id=\"t6267798-nsyLtwAp\" data-transient-attributes=\"table-cell-selection\"><code>forward_cuda<\/code><\/p>\n<p>\u2192<code>forward_native<\/code>\u2192IR<code>[vllm_c, native]<\/code><\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-DBlgTj8H\" data-transient-attributes=\"table-cell-selection\">CUDA<\/p>\n<p><code>fused_add_rms_norm<\/code>\u00a0\u2192\u00a0<code>fused_add_rms_norm_kernel&lt;scalar_t, width, has_weight&gt;<\/code><\/td>\n<\/tr>\n<tr data-id=\"t31e458f-tAuE09Dt\">\n<td data-id=\"t6267798-lzrgbns3\" data-transient-attributes=\"table-cell-selection\">eager<\/td>\n<td data-id=\"t6267798-S3yQXUtB\" data-transient-attributes=\"table-cell-selection\">none<\/td>\n<td data-id=\"t6267798-hJAJ4qgK\" data-transient-attributes=\"table-cell-selection\">on<\/td>\n<td data-id=\"t6267798-Uh1BCWHl\" data-transient-attributes=\"table-cell-selection\">bf16\/fp16<\/td>\n<td data-id=\"t6267798-9wbW9ozy\" data-transient-attributes=\"table-cell-selection\"><code>forward_cuda<\/code><\/p>\n<p>\u2192<code>rms_norm_batch_invariant<\/code><\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-2nZbU6T2\" data-transient-attributes=\"table-cell-selection\">Triton<\/p>\n<p><code>_rms_norm_kernel<\/code><\/td>\n<\/tr>\n<tr data-id=\"t31e458f-138c6EVU\">\n<td data-id=\"t6267798-rlv7ds4J\" data-transient-attributes=\"table-cell-selection\">eager<\/td>\n<td data-id=\"t6267798-Zjf0jUnj\" data-transient-attributes=\"table-cell-selection\">yes<\/td>\n<td data-id=\"t6267798-SctlAJia\" data-transient-attributes=\"table-cell-selection\">on<\/td>\n<td data-id=\"t6267798-tVv4WueE\" data-transient-attributes=\"table-cell-selection\">bf16\/fp16<\/td>\n<td data-id=\"t6267798-dHVCAkXt\" data-transient-attributes=\"table-cell-selection\"><code>forward_cuda<\/code><\/p>\n<p>\u2192<code>rms_norm_batch_invariant<\/code>\u2192<code>ops.fused_add_rms_norm<\/code><\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-RYKZsoy1\" data-transient-attributes=\"table-cell-selection\">CUDA<\/p>\n<p><code>fused_add_rms_norm<\/code>\u00a0\u2192\u00a0<code>fused_add_rms_norm_kernel&lt;scalar_t, width, has_weight&gt;<\/code>; block\u00a0pinned \u2192 1024<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-QDTKMccI\">\n<td data-id=\"t6267798-bjRePzDy\" data-transient-attributes=\"table-cell-selection\">compile<\/td>\n<td data-id=\"t6267798-EHmXs1sy\" data-transient-attributes=\"table-cell-selection\">none<\/td>\n<td data-id=\"t6267798-xdIVqeIE\" data-transient-attributes=\"table-cell-selection\">off \/ on<\/td>\n<td data-id=\"t6267798-YpzSuaW3\" data-transient-attributes=\"table-cell-selection\">bf16\/fp16<\/td>\n<td data-id=\"t6267798-b9bsjIoM\" data-transient-attributes=\"table-cell-selection\"><code>forward_native<\/code><\/p>\n<p>\u2192IR<code>[native]<\/code><\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-lf3XdKBu\" data-transient-attributes=\"table-cell-selection\">native aten \u2192 Inductor Triton<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-b0K7iYvS\">\n<td data-id=\"t6267798-1MMLWzWo\" data-transient-attributes=\"table-cell-selection\">compile<\/td>\n<td data-id=\"t6267798-XxrCnF6a\" data-transient-attributes=\"table-cell-selection\">yes<\/td>\n<td data-id=\"t6267798-Sd4VI4aw\" data-transient-attributes=\"table-cell-selection\">off \/ on<\/td>\n<td data-id=\"t6267798-VqNHHUIf\" data-transient-attributes=\"table-cell-selection\">bf16\/fp16<\/td>\n<td data-id=\"t6267798-15VRwwAc\" data-transient-attributes=\"table-cell-selection\"><code>forward_native<\/code><\/p>\n<p>\u2192IR<code>[native]<\/code><\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-bmgeEqz1\" data-transient-attributes=\"table-cell-selection\">native aten (add+norm) \u2192 Inductor Triton<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-9KDjPiS3\">\n<td data-id=\"t6267798-DhcT64Fl\" data-transient-attributes=\"table-cell-selection\">compile<\/td>\n<td data-id=\"t6267798-VVg5ekVX\" data-transient-attributes=\"table-cell-selection\">none \/ yes<\/td>\n<td data-id=\"t6267798-1sFHGb2r\" data-transient-attributes=\"table-cell-selection\">off<\/td>\n<td data-id=\"t6267798-dlY3QMtt\" data-transient-attributes=\"table-cell-selection\">fp8 static per-tensor<\/td>\n<td data-id=\"t6267798-QJ3jrQ6X\" data-transient-attributes=\"table-cell-selection\">RMSNorm(+add) +\u00a0<code>static_scaled_fp8_quant<\/code>\u00a0\u2192 fused<\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-rFI9lpd7\" data-transient-attributes=\"table-cell-selection\">CUDA<\/p>\n<p><code>rms_norm_static_fp8_quant<\/code>\u00a0\/\u00a0<code>fused_add_rms_norm_static_fp8_quant<\/code>; block\u00a0<code>(num_tokens &lt; 256) ? 1024 : 256<\/code><\/td>\n<\/tr>\n<tr data-id=\"t31e458f-fl0DxSE9\">\n<td data-id=\"t6267798-7nQmvG2P\" data-transient-attributes=\"table-cell-selection\">compile<\/td>\n<td data-id=\"t6267798-JqrJCbgw\" data-transient-attributes=\"table-cell-selection\">none \/ yes<\/td>\n<td data-id=\"t6267798-pEN55UEn\" data-transient-attributes=\"table-cell-selection\">off<\/td>\n<td data-id=\"t6267798-9rsXd8T4\" data-transient-attributes=\"table-cell-selection\">fp8 dynamic per-token<\/td>\n<td data-id=\"t6267798-nPmAWfni\" data-transient-attributes=\"table-cell-selection\">RMSNorm(+add) +\u00a0<code>dynamic_per_token_..._quant<\/code>\u00a0\u2192 fused<\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-PRaESJ0q\" data-transient-attributes=\"table-cell-selection\">CUDA<\/p>\n<p><code>rms_norm_dynamic_per_token_quant<\/code>; block fixed 1024<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-FoePaUTI\">\n<td data-id=\"t6267798-rxO4MRHU\" data-transient-attributes=\"table-cell-selection\">compile<\/td>\n<td data-id=\"t6267798-2gYQHBik\" data-transient-attributes=\"table-cell-selection\">none \/ yes<\/td>\n<td data-id=\"t6267798-lWRL1vfi\" data-transient-attributes=\"table-cell-selection\">off<\/td>\n<td data-id=\"t6267798-7n4yPj19\" data-transient-attributes=\"table-cell-selection\">fp8 block g128\/64<\/td>\n<td data-id=\"t6267798-RVYrU2Td\" data-transient-attributes=\"table-cell-selection\">RMSNorm(+add) +\u00a0<code>per_token_group_fp8_quant<\/code>\u00a0\u2192 fused<\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-3O69eFMh\" data-transient-attributes=\"table-cell-selection\">CUDA<\/p>\n<p><code>rms_norm_per_block_quant<\/code>\u00a0\u2192\u00a0<code>rms_norm_per_block_quant_dispatch<\/code>; block\u00a0<code>(num_tokens &lt;= 256) ? 512 : 256<\/code><\/td>\n<\/tr>\n<tr data-id=\"t31e458f-RlBlcxxj\">\n<td data-id=\"t6267798-1pIzj1FY\" data-transient-attributes=\"table-cell-selection\">compile<\/td>\n<td data-id=\"t6267798-0tsxOU9a\" data-transient-attributes=\"table-cell-selection\">yes<\/td>\n<td data-id=\"t6267798-I8PzY78Z\" data-transient-attributes=\"table-cell-selection\">on<\/td>\n<td data-id=\"t6267798-ke9ol1LV\" data-transient-attributes=\"table-cell-selection\">fp8 static per-tensor<\/td>\n<td data-id=\"t6267798-F00LNYGh\" data-transient-attributes=\"table-cell-selection\"><code>fused_add_rms_norm<\/code><\/p>\n<p>node + quant \u2192 fusion matches<\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-iyK4vDFe\" data-transient-attributes=\"table-cell-selection\">CUDA<\/p>\n<p><code>fused_add_rms_norm_static_fp8_quant<\/code>; block\u00a0pinned \u2192 1024<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-HSjY5aO2\">\n<td data-id=\"t6267798-EGw9pWC9\" data-transient-attributes=\"table-cell-selection\">compile<\/td>\n<td data-id=\"t6267798-dTAfis16\" data-transient-attributes=\"table-cell-selection\">yes<\/td>\n<td data-id=\"t6267798-ipGxqCnu\" data-transient-attributes=\"table-cell-selection\">on<\/td>\n<td data-id=\"t6267798-9IUOKhy7\" data-transient-attributes=\"table-cell-selection\">fp8 dynamic per-token<\/td>\n<td data-id=\"t6267798-UMycmYmq\" data-transient-attributes=\"table-cell-selection\"><code>fused_add_rms_norm<\/code><\/p>\n<p>node + quant \u2192 fusion matches<\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-UUuIOBDG\" data-transient-attributes=\"table-cell-selection\">CUDA<\/p>\n<p><code>rms_norm_dynamic_per_token_quant<\/code>; block fixed 1024 (BI-safe by construction)<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-XZWm5TeN\">\n<td data-id=\"t6267798-T61xp6jn\" data-transient-attributes=\"table-cell-selection\">compile<\/td>\n<td data-id=\"t6267798-nzDb2pYd\" data-transient-attributes=\"table-cell-selection\">yes<\/td>\n<td data-id=\"t6267798-mOBOVXLN\" data-transient-attributes=\"table-cell-selection\">on<\/td>\n<td data-id=\"t6267798-IxGM55Dk\" data-transient-attributes=\"table-cell-selection\">fp8 block g128\/64<\/td>\n<td data-id=\"t6267798-QZRCUyGh\" data-transient-attributes=\"table-cell-selection\"><code>fused_add_rms_norm<\/code><\/p>\n<p>node + quant \u2192 fusion matches<\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-ChrKoWEC\" data-transient-attributes=\"table-cell-selection\">CUDA<\/p>\n<p><code>rms_norm_per_block_quant<\/code>; block\u00a0pinned \u2192 512<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-SokqDeJR\">\n<td class=\"table-last-column\" data-id=\"t6267798-zue6BnxO\" data-transient-attributes=\"table-cell-selection\">compile<\/td>\n<td class=\"table-last-column\" data-id=\"t6267798-aNMRsyWT\" data-transient-attributes=\"table-cell-selection\">none<\/td>\n<td class=\"table-last-column\" data-id=\"t6267798-CGWiQiVE\" data-transient-attributes=\"table-cell-selection\">on<\/td>\n<td class=\"table-last-column\" data-id=\"t6267798-ZlwIAPGY\" data-transient-attributes=\"table-cell-selection\">fp8 (any)<\/td>\n<td class=\"table-last-column\" data-id=\"t6267798-jxr7snvi\" data-transient-attributes=\"table-cell-selection\"><code>rms_norm_batch_invariant<\/code><\/p>\n<p>\u2192 Triton; quant stays separate<\/td>\n<td class=\"table-last-column table-last-row\" data-id=\"t6267798-gYHn27h6\" data-transient-attributes=\"table-cell-selection\">Triton<\/p>\n<p><code>_rms_norm_kernel<\/code>\u00a0+ separate quant op \u2014\u00a0not fused<\/td>\n<\/tr>\n<\/tbody>\n<\/table>\n<p>\u4e00\u4e2a\u6709\u610f\u601d\u7684\u5bf9\u7167: eager mode\uff0c\u6709residual\u7684\u60c5\u51b5\u4e0b\uff0c\u65e0\u8bba\u662f\u5426\u5f00\u542fBatch Invariance\u90fd\u4f1a\u8d70\u5230CUDA\u5b9e\u73b0torch.ops._C.fused_add_rms_norm\u3002\u800c\u65e0 residual \u7684\u666e\u901a rms_norm,\u5f00\u542fBatch Invariance\u65f6\u624d\u5fc5\u987b\u6362\u6210\u4e13\u95e8\u7684 Triton kernel\u00a0_rms_norm_kernel\u3002<\/p>\n<p>\u5f53\u7136\u4e5f\u53ef\u4ee5\u901a\u8fc7\u5982\u4e0b\u6307\u4ee4\u8c03\u6574IR OP\u4f18\u5148\u7ea7\uff1a<\/p>\n<div>\n<div class=\"hljs-cto\">\n<div class=\"hljs-cto\"><button class=\"copy_btn disable\" data-clipboard-target=\"#code_id_10\">\u590d\u5236<\/button><\/p>\n<div class=\"code-toolbar\">\n<pre class=\"has-pre-numbering language-javascript\" tabindex=\"0\"><code class=\"language-javascript\">vllm serve meta<span class=\"token operator\">-<\/span>llama<span class=\"token operator\">\/<\/span>Llama<span class=\"token operator\">-<\/span><span class=\"token number\">3.2<\/span><span class=\"token operator\">-<\/span>1B \\\r\n  <span class=\"token operator\">--<\/span>ir<span class=\"token operator\">-<\/span>op<span class=\"token operator\">-<\/span>priority<span class=\"token punctuation\">.<\/span>rms_norm<span class=\"token operator\">=<\/span>vllm_c \\\r\n  <span class=\"token operator\">--<\/span>ir<span class=\"token operator\">-<\/span>op<span class=\"token operator\">-<\/span>priority<span class=\"token punctuation\">.<\/span>fused_add_rms_norm<span class=\"token operator\">=<\/span>vllm_c<\/code><\/pre>\n<ul id=\"code_id_10\" class=\"pre-numbering\">\n<li>1.<\/li>\n<li>2.<\/li>\n<li>3.<\/li>\n<\/ul>\n<div class=\"toolbar\"><\/div>\n<\/div>\n<\/div>\n<\/div>\n<\/div>\n<p><strong>(1) eager mode\u4e0bRMSNorm\u7684Batch Invariance<\/strong><\/p>\n<p>eager mode \u4e0b\u5f00\u542fBatch Invariance \u65f6,\u65e0\u8bba\u662f\u6709residual\u7684CUDA\u00a0torch.ops._C.fused_add_rms_norm(blockDim=min(hidden,1024),\u9501\u4f4f max_block_size \u4e0a\u9650\u4e3a 1024)\u8fd8\u662f\u65e0 residual \u7684Triton\u00a0_rms_norm_kernel( BLOCK_SIZE=1024),\u90fd\u901a\u8fc7\u8ba9 reduction \u5bbd\u5ea6\u4e0e num_tokens \u89e3\u8026\u6765\u56fa\u5b9a\u5f52\u7ea6\u7ed3\u6784,\u4ece\u800c\u4fdd\u8bc1\u6d6e\u70b9\u52a0\u6cd5\u987a\u5e8f\u4e0d\u53d8\u3001\u8de8 batch bit-exact\u3002<\/p>\n<p>\u5982\u4e0b\u4ee5\u00a0_rms_norm_kernel\u4e3a\u4f8b\uff1a<\/p>\n<div>\n<div class=\"hljs-cto\">\n<div class=\"hljs-cto\"><button class=\"copy_btn disable\" data-clipboard-target=\"#code_id_11\">\u590d\u5236<\/button><\/p>\n<div class=\"code-toolbar\">\n<pre class=\"has-pre-numbering language-javascript\" tabindex=\"0\"><code class=\"language-javascript\">def <span class=\"token function\">forward_cuda<\/span><span class=\"token punctuation\">(<\/span>\r\n    self<span class=\"token punctuation\">,<\/span>\r\n    <span class=\"token literal-property property\">x<\/span><span class=\"token operator\">:<\/span> torch<span class=\"token punctuation\">.<\/span>Tensor<span class=\"token punctuation\">,<\/span>\r\n    <span class=\"token literal-property property\">residual<\/span><span class=\"token operator\">:<\/span> torch<span class=\"token punctuation\">.<\/span>Tensor <span class=\"token operator\">|<\/span> None <span class=\"token operator\">=<\/span> None<span class=\"token punctuation\">,<\/span>\r\n<span class=\"token punctuation\">)<\/span> <span class=\"token operator\">-<\/span><span class=\"token operator\">&gt;<\/span> torch<span class=\"token punctuation\">.<\/span>Tensor <span class=\"token operator\">|<\/span> tuple<span class=\"token punctuation\">[<\/span>torch<span class=\"token punctuation\">.<\/span>Tensor<span class=\"token punctuation\">,<\/span> torch<span class=\"token punctuation\">.<\/span>Tensor<span class=\"token punctuation\">]<\/span><span class=\"token operator\">:<\/span>\r\n    <span class=\"token keyword\">if<\/span> envs<span class=\"token punctuation\">.<\/span><span class=\"token constant\">VLLM_BATCH_INVARIANT<\/span><span class=\"token operator\">:<\/span>\r\n        assert self<span class=\"token punctuation\">.<\/span>variance_size_override is None<span class=\"token punctuation\">,<\/span> <span class=\"token punctuation\">(<\/span>\r\n            <span class=\"token string\">\"Batch invariance is not supported for variance_size_override\"<\/span>\r\n        <span class=\"token punctuation\">)<\/span>\r\n        <span class=\"token keyword\">return<\/span> <span class=\"token function\">rms_norm_batch_invariant<\/span><span class=\"token punctuation\">(<\/span>\r\n            x<span class=\"token punctuation\">,<\/span>\r\n            self<span class=\"token punctuation\">.<\/span>weight<span class=\"token punctuation\">.<\/span>data<span class=\"token punctuation\">,<\/span>\r\n            self<span class=\"token punctuation\">.<\/span>variance_epsilon<span class=\"token punctuation\">,<\/span>\r\n            residual<span class=\"token operator\">=<\/span>residual<span class=\"token punctuation\">,<\/span>\r\n        <span class=\"token punctuation\">)<\/span>\r\n\r\n    <span class=\"token keyword\">return<\/span> self<span class=\"token punctuation\">.<\/span><span class=\"token function\">forward_native<\/span><span class=\"token punctuation\">(<\/span>x<span class=\"token punctuation\">,<\/span> residual<span class=\"token punctuation\">)<\/span><\/code><\/pre>\n<ul id=\"code_id_11\" class=\"pre-numbering\">\n<li>1.<\/li>\n<li>2.<\/li>\n<li>3.<\/li>\n<li>4.<\/li>\n<li>5.<\/li>\n<li>6.<\/li>\n<li>7.<\/li>\n<li>8.<\/li>\n<li>9.<\/li>\n<li>10.<\/li>\n<li>11.<\/li>\n<li>12.<\/li>\n<li>13.<\/li>\n<li>14.<\/li>\n<li>15.<\/li>\n<li>16.<\/li>\n<li>17.<\/li>\n<\/ul>\n<div class=\"toolbar\"><\/div>\n<\/div>\n<\/div>\n<\/div>\n<\/div>\n<p>vLLM\u4e2d_rms_norm_kernel\u662f\u7528Triton\u5b9e\u73b0\u7684\uff0c\u903b\u8f91\u5f88\u7b80\u5355\uff0c\u4e0d\u7528@triton.autotune\u52a0BLOCK_SIZE = 1024\u3002<\/p>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/e759d62799a89af8f9e966087947c50ad3318d.webp\" data-type=\"block\" \/><\/p>\n<p><strong>(2) \u7f16\u8bd1\u573a\u666f\u4e0b RMSNorm\u7684Batch Invariance<\/strong><\/p>\n<p>\u5176\u5b9e\u8fd9\u91cc\u624d\u6700\u53cd\u76f4\u89c9\u7684\u5730\u65b9\uff0c\u4e0a\u9762\u8bf4\u4e86CUDA\u7684\u5b9e\u73b0\u3001Triton\u7684\u5b9e\u73b0\uff0c\u5b9e\u9645\u4e0a\u6700\u7ec8\u9ed8\u8ba4\u903b\u8f91\u8fd8\u662f\u8d70\u4e86torch.compile\u81ea\u52a8\u751f\u6210\u7684Triton kernel\u3002<\/p>\n<blockquote><p>&nbsp;<\/p>\n<p>When PyTorch Inductor is used, &#8216;none&#8217; is the default; otherwise &#8216;all&#8217;.<\/p>\n<p>&nbsp;<\/p><\/blockquote>\n<p>\u5728 TorchInductor \u7f16\u8bd1\u6a21\u5f0f\u4e0b,vLLM \u7684 CustomOp \u9ed8\u8ba4\u7981\u7528(custom_ops \u5728 Inductor \u540e\u7aef\u4e0b\u9ed8\u8ba4\u4e3a none),RMSNorm \u56e0\u6b64\u4e0d\u8d70\u81ea\u5b9a\u4e49 CUDA\/Triton \u5b9e\u73b0,\u800c\u662f\u5206\u53d1\u5230 forward_native,\u518d\u7ecfir.ops.rms_norm \/ fused_add_rms_norm \u7684 native provider(\u7f16\u8bd1\u6001\u4f18\u5148\u7ea7\u4e3a [&#8220;native&#8221;])\u3002\u8be5\u7eaf PyTorch \u53c2\u8003\u5b9e\u73b0\u7531 TorchDynamo \u6355\u83b7\u4e3a FX \u56fe,Inductor\u5bf9\u5176\u4e2d\u7684\u7c7b\u578b\u8f6c\u6362\u3001\u5e73\u65b9\u3001\u6309\u6700\u540e\u4e00\u7ef4\u6c42\u5747\u503c(\u5f52\u7ea6)\u4e0e\u6743\u91cd\u4e58\u6cd5\u6267\u884c\u7b97\u5b50\u878d\u5408,\u751f\u6210\u5355\u4e2a Triton \u5f52\u7ea6 kernel\u3002<\/p>\n<p>RMSNorm \u7684\u8f93\u5165\u5f62\u72b6\u4e3a [num_tokens, hidden_size],\u5f52\u7ea6\u6cbf\u6700\u540e\u4e00\u7ef4(hidden_size,\u6a21\u578b\u7684\u7f16\u8bd1\u671f\u9759\u6001\u5e38\u91cf)\u8fdb\u884c;num_tokens \u662f\u5e76\u884c\u7ef4,\u5404\u884c\u76f8\u4e92\u72ec\u7acb\u3001\u4e0d\u53c2\u4e0e\u5f7c\u6b64\u7684\u5f52\u7ea6\u3002\u56e0\u6b64 num_tokens\u7684\u53d6\u503c\u53ea\u51b3\u5b9a\u5e76\u884c\u884c\u6570,\u4e0d\u53c2\u4e0e\u5355\u884c\u5185\u5bf9 hidden_size \u7684\u6d6e\u70b9\u7d2f\u52a0\u987a\u5e8f\u3002\u8fd9\u662f\u8be5\u7b97\u5b50\u5728\u7f16\u8bd1\u6001\u53ef\u4fdd\u6301\u6279\u4e0d\u53d8\u7684\u7ed3\u6784\u524d\u63d0\u3002<\/p>\n<div>\n<div class=\"hljs-cto\">\n<div class=\"hljs-cto\"><button class=\"copy_btn disable\" data-clipboard-target=\"#code_id_12\">\u590d\u5236<\/button><\/p>\n<div class=\"code-toolbar\">\n<pre class=\"has-pre-numbering language-javascript\" tabindex=\"0\"><code class=\"language-javascript\">@register_op\r\ndef <span class=\"token function\">rms_norm<\/span><span class=\"token punctuation\">(<\/span>\r\n    <span class=\"token literal-property property\">x<\/span><span class=\"token operator\">:<\/span> Tensor<span class=\"token punctuation\">,<\/span> <span class=\"token literal-property property\">weight<\/span><span class=\"token operator\">:<\/span> Tensor <span class=\"token operator\">|<\/span> None<span class=\"token punctuation\">,<\/span> <span class=\"token literal-property property\">epsilon<\/span><span class=\"token operator\">:<\/span> float<span class=\"token punctuation\">,<\/span> <span class=\"token literal-property property\">variance_size<\/span><span class=\"token operator\">:<\/span> int <span class=\"token operator\">|<\/span> None <span class=\"token operator\">=<\/span> None\r\n<span class=\"token punctuation\">)<\/span> <span class=\"token operator\">-<\/span><span class=\"token operator\">&gt;<\/span> Tensor<span class=\"token operator\">:<\/span>\r\n    <span class=\"token string\">\"\"<\/span><span class=\"token string\">\"Weighted root-mean-square layer normalization\"<\/span><span class=\"token string\">\"\"<\/span>\r\n    orig_dtype <span class=\"token operator\">=<\/span> x<span class=\"token punctuation\">.<\/span>dtype\r\n    x <span class=\"token operator\">=<\/span> x<span class=\"token punctuation\">.<\/span><span class=\"token function\">to<\/span><span class=\"token punctuation\">(<\/span>torch<span class=\"token punctuation\">.<\/span>float32<span class=\"token punctuation\">)<\/span>\r\n    x_var <span class=\"token operator\">=<\/span> x <span class=\"token keyword\">if<\/span> variance_size is None <span class=\"token keyword\">else<\/span> x<span class=\"token punctuation\">[<\/span><span class=\"token operator\">...<\/span><span class=\"token punctuation\">,<\/span> <span class=\"token operator\">:<\/span>variance_size<span class=\"token punctuation\">]<\/span>\r\n    variance <span class=\"token operator\">=<\/span> x_var<span class=\"token punctuation\">.<\/span><span class=\"token function\">pow<\/span><span class=\"token punctuation\">(<\/span><span class=\"token number\">2<\/span><span class=\"token punctuation\">)<\/span><span class=\"token punctuation\">.<\/span><span class=\"token function\">mean<\/span><span class=\"token punctuation\">(<\/span>dim<span class=\"token operator\">=<\/span><span class=\"token operator\">-<\/span><span class=\"token number\">1<\/span><span class=\"token punctuation\">,<\/span> keepdim<span class=\"token operator\">=<\/span>True<span class=\"token punctuation\">)<\/span>\r\n    x <span class=\"token operator\">=<\/span> x <span class=\"token operator\">*<\/span> torch<span class=\"token punctuation\">.<\/span><span class=\"token function\">rsqrt<\/span><span class=\"token punctuation\">(<\/span>variance <span class=\"token operator\">+<\/span> epsilon<span class=\"token punctuation\">)<\/span>\r\n    <span class=\"token keyword\">if<\/span> weight is not None<span class=\"token operator\">:<\/span>\r\n        x <span class=\"token operator\">=<\/span> x<span class=\"token punctuation\">.<\/span><span class=\"token function\">to<\/span><span class=\"token punctuation\">(<\/span>weight<span class=\"token punctuation\">.<\/span>dtype<span class=\"token punctuation\">)<\/span> <span class=\"token operator\">*<\/span> weight\r\n    <span class=\"token keyword\">return<\/span> x<span class=\"token punctuation\">.<\/span><span class=\"token function\">to<\/span><span class=\"token punctuation\">(<\/span>orig_dtype<span class=\"token punctuation\">)<\/span><\/code><\/pre>\n<ul id=\"code_id_12\" class=\"pre-numbering\">\n<li>1.<\/li>\n<li>2.<\/li>\n<li>3.<\/li>\n<li>4.<\/li>\n<li>5.<\/li>\n<li>6.<\/li>\n<li>7.<\/li>\n<li>8.<\/li>\n<li>9.<\/li>\n<li>10.<\/li>\n<li>11.<\/li>\n<li>12.<\/li>\n<li>13.<\/li>\n<\/ul>\n<div class=\"toolbar\"><\/div>\n<\/div>\n<\/div>\n<\/div>\n<\/div>\n<p>\u7f16\u8bd1\u5668\u81ea\u52a8\u751f\u6210\u7684 Triton Kernel\uff0c\u5b83\u662f\u5982\u4f55\u7ef4\u6301 Batch Invariance \u7684\uff1f\u8fd9\u91cc\u7684\u6838\u5fc3\u95ee\u9898\u5728\u4e8e\u7f16\u8bd1\u5668\u7684\u00a0Split-Reduction\uff08\u62c6\u5206\u89c4\u7ea6\uff09\u00a0\u4f18\u5316\u3002\u7c7b\u4f3c\u4e8e GEMM \u4e2d\u7684 Split-K\uff0c\u5f53\u89c4\u7ea6\u7ef4\u5ea6\uff08hidden_size\uff09\u8f83\u957f\uff0c\u4e14\u5e76\u884c\u7ef4\u5ea6\uff08Batch Size\uff0c\u5373\u00a0num_tokens\uff09\u6781\u5c0f\u65f6\uff0c\u4e3a\u907f\u514d GPU \u7684 SM \u7a7a\u8f7d\uff0cInductor \u4f1a\u89e6\u53d1\u591a\u6bb5\u62c6\u5206\u903b\u8f91\uff1a\u5c06\u540c\u4e00\u884c\u7684\u00a0hidden_size\u00a0\u5f3a\u884c\u5207\u7247\uff0c\u5206\u53d1\u7ed9\u591a\u4e2a Thread Block \u5e76\u53d1\u8ba1\u7b97\u5c40\u90e8\u548c\uff0c\u6700\u7ec8\u518d\u6267\u884c\u5168\u5c40\u89c4\u7ea6\u5408\u5e76\u3002<\/p>\n<p>\u800c\u662f\u5426\u5bf9\u5f52\u7ea6\u505a\u591a\u6bb5\u62c6\u5206(split-reduction)\u7531 Reduction.num_splits \u51b3\u5b9a,\u5176\u4e2d:<\/p>\n<div>\n<div class=\"hljs-cto\">\n<div class=\"hljs-cto\"><button class=\"copy_btn disable\" data-clipboard-target=\"#code_id_13\">\u590d\u5236<\/button><\/p>\n<div class=\"code-toolbar\">\n<pre class=\"has-pre-numbering language-javascript\" tabindex=\"0\"><code class=\"language-javascript\"># torch<span class=\"token operator\">\/<\/span>_inductor<span class=\"token operator\">\/<\/span>ir<span class=\"token punctuation\">.<\/span>py<span class=\"token operator\">:<\/span><span class=\"token operator\">:<\/span>num_splits\r\nreduction_numel_hint <span class=\"token operator\">=<\/span> sizevars<span class=\"token punctuation\">.<\/span><span class=\"token function\">symbolic_hint<\/span><span class=\"token punctuation\">(<\/span>reduction_numel<span class=\"token punctuation\">)<\/span>   # hidden_size\r\nnumel_hint           <span class=\"token operator\">=<\/span> sizevars<span class=\"token punctuation\">.<\/span><span class=\"token function\">symbolic_hint<\/span><span class=\"token punctuation\">(<\/span><span class=\"token function\">sympy_product<\/span><span class=\"token punctuation\">(<\/span>ranges<span class=\"token punctuation\">)<\/span><span class=\"token punctuation\">)<\/span>  # num_tokens\r\n<span class=\"token keyword\">if<\/span> <span class=\"token function\">not<\/span> <span class=\"token punctuation\">(<\/span><span class=\"token function\">_is_static<\/span><span class=\"token punctuation\">(<\/span>reduction_numel_hint<span class=\"token punctuation\">)<\/span> and <span class=\"token function\">_is_static<\/span><span class=\"token punctuation\">(<\/span>numel_hint<span class=\"token punctuation\">)<\/span><span class=\"token punctuation\">)<\/span><span class=\"token operator\">:<\/span>\r\n    <span class=\"token keyword\">return<\/span> ReductionHint<span class=\"token punctuation\">.<\/span><span class=\"token constant\">DEFAULT<\/span><span class=\"token punctuation\">,<\/span> <span class=\"token number\">1<\/span>     # \u4ec5\u5bf9 unbacked symint \u6210\u7acb<\/code><\/pre>\n<ul id=\"code_id_13\" class=\"pre-numbering\">\n<li>1.<\/li>\n<li>2.<\/li>\n<li>3.<\/li>\n<li>4.<\/li>\n<li>5.<\/li>\n<\/ul>\n<div class=\"toolbar\"><\/div>\n<\/div>\n<\/div>\n<\/div>\n<\/div>\n<p>symbolic_hint \u7684\u884c\u4e3a\u662f:\u5bf9 backed \u7b26\u53f7\u7528\u4ee3\u8868\u503c(backed_var_to_val)\u66ff\u6362\u4e3a\u5177\u4f53\u6574\u6570,\u5bf9 unbacked \u7b26\u53f7\u4fdd\u7559\u7b26\u53f7\u3002<\/p>\n<p>vLLM \u9ed8\u8ba4\u4f7f\u7528\u00a0torch._dynamo.mark_dynamic(backed)\u3002\u56e0\u6b64num_tokens \u7ecf\u00a0symbolic_hint\u88ab\u66ff\u6362\u4e3a\u5176\u4ee3\u8868\u503c,_is_static(numel_hint)\u4e3a\u771f,\u4e0a\u8ff0\u63d0\u524d\u8fd4\u56de\u4e0d\u89e6\u53d1,\u540e\u7eed\u542f\u53d1\u5f0f\u6b63\u5e38\u6267\u884c,\u53ef\u80fd\u9009\u51fa split &gt; 1\u3002\u8be5\u63d0\u524d\u8fd4\u56de(\u5373\u6240\u8c13\u4fdd\u5b88\u8def\u5f84)\u4ec5\u5728\u4f7f\u7528 mark_unbacked \u65f6\u624d\u751f\u6548\u3002<\/p>\n<div>\n<div class=\"hljs-cto\">\n<div class=\"hljs-cto\"><button class=\"copy_btn disable\" data-clipboard-target=\"#code_id_14\">\u590d\u5236<\/button><\/p>\n<div class=\"code-toolbar\">\n<pre class=\"has-pre-numbering language-javascript\" tabindex=\"0\"><code class=\"language-javascript\">@<span class=\"token function\">support_torch_compile<\/span><span class=\"token punctuation\">(<\/span>\r\n    # mark_unbacked_dims<span class=\"token operator\">=<\/span><span class=\"token punctuation\">{<\/span><span class=\"token string-property property\">\"input_ids\"<\/span><span class=\"token operator\">:<\/span> <span class=\"token number\">0<\/span><span class=\"token punctuation\">}<\/span><span class=\"token punctuation\">,<\/span> \r\n    dynamic_arg_dims<span class=\"token operator\">=<\/span><span class=\"token punctuation\">{<\/span>\r\n        <span class=\"token string-property property\">\"input_ids\"<\/span><span class=\"token operator\">:<\/span> <span class=\"token punctuation\">{<\/span><span class=\"token number\">0<\/span><span class=\"token operator\">:<\/span> <span class=\"token string\">\"b\"<\/span><span class=\"token punctuation\">}<\/span><span class=\"token punctuation\">,<\/span>\r\n        <span class=\"token string-property property\">\"positions\"<\/span><span class=\"token operator\">:<\/span> <span class=\"token punctuation\">{<\/span><span class=\"token number\">0<\/span><span class=\"token operator\">:<\/span> <span class=\"token string\">\"b\"<\/span><span class=\"token punctuation\">}<\/span><span class=\"token punctuation\">,<\/span>\r\n        <span class=\"token string-property property\">\"intermediate_tensors\"<\/span><span class=\"token operator\">:<\/span> <span class=\"token punctuation\">{<\/span><span class=\"token number\">0<\/span><span class=\"token operator\">:<\/span> <span class=\"token string\">\"b\"<\/span><span class=\"token punctuation\">}<\/span><span class=\"token punctuation\">,<\/span>\r\n        <span class=\"token string-property property\">\"inputs_embeds\"<\/span><span class=\"token operator\">:<\/span> <span class=\"token punctuation\">{<\/span><span class=\"token number\">0<\/span><span class=\"token operator\">:<\/span> <span class=\"token string\">\"b\"<\/span><span class=\"token punctuation\">}<\/span><span class=\"token punctuation\">,<\/span>\r\n    <span class=\"token punctuation\">}<\/span><span class=\"token punctuation\">,<\/span>\r\n<span class=\"token punctuation\">)<\/span><\/code><\/pre>\n<ul id=\"code_id_14\" class=\"pre-numbering\">\n<li>1.<\/li>\n<li>2.<\/li>\n<li>3.<\/li>\n<li>4.<\/li>\n<li>5.<\/li>\n<li>6.<\/li>\n<li>7.<\/li>\n<li>8.<\/li>\n<li>9.<\/li>\n<\/ul>\n<div class=\"toolbar\"><\/div>\n<\/div>\n<\/div>\n<\/div>\n<\/div>\n<div>\n<div class=\"hljs-cto\">\n<div class=\"hljs-cto\"><button class=\"copy_btn disable\" data-clipboard-target=\"#code_id_15\">\u590d\u5236<\/button><\/p>\n<div class=\"code-toolbar\">\n<pre class=\"has-pre-numbering language-javascript\" tabindex=\"0\"><code class=\"language-javascript\">def <span class=\"token function\">mark_dynamic<\/span><span class=\"token punctuation\">(<\/span>arg<span class=\"token punctuation\">,<\/span> dim_shape_pairs<span class=\"token punctuation\">)<\/span><span class=\"token operator\">:<\/span>\r\n    <span class=\"token keyword\">if<\/span> ds_type <span class=\"token operator\">==<\/span> DynamicShapesType<span class=\"token punctuation\">.<\/span><span class=\"token constant\">UNBACKED<\/span><span class=\"token operator\">:<\/span>\r\n        <span class=\"token operator\">...<\/span>\r\n        torch<span class=\"token punctuation\">.<\/span>_dynamo<span class=\"token punctuation\">.<\/span>decorators<span class=\"token punctuation\">.<\/span><span class=\"token function\">mark_unbacked<\/span><span class=\"token punctuation\">(<\/span>arg<span class=\"token punctuation\">,<\/span> dim<span class=\"token punctuation\">,<\/span> hint_override<span class=\"token operator\">=<\/span><span class=\"token operator\">...<\/span><span class=\"token punctuation\">,<\/span> shape_id<span class=\"token operator\">=<\/span>shape_id<span class=\"token punctuation\">)<\/span>\r\n        <span class=\"token operator\">...<\/span>\r\n    <span class=\"token keyword\">else<\/span><span class=\"token operator\">:<\/span>\r\n        torch<span class=\"token punctuation\">.<\/span>_dynamo<span class=\"token punctuation\">.<\/span><span class=\"token function\">mark_dynamic<\/span><span class=\"token punctuation\">(<\/span>arg<span class=\"token punctuation\">,<\/span> dims<span class=\"token punctuation\">)<\/span>   # \u9ed8\u8ba4 <span class=\"token constant\">BACKED<\/span> \u8def\u5f84<\/code><\/pre>\n<ul id=\"code_id_15\" class=\"pre-numbering\">\n<li>1.<\/li>\n<li>2.<\/li>\n<li>3.<\/li>\n<li>4.<\/li>\n<li>5.<\/li>\n<li>6.<\/li>\n<li>7.<\/li>\n<\/ul>\n<div class=\"toolbar\"><\/div>\n<\/div>\n<\/div>\n<\/div>\n<\/div>\n<p>\u8fdb\u5165\u542f\u53d1\u5f0f\u540e,reduction_split_factor \u7684\u5224\u5b9a\u4e3a:<\/p>\n<div>\n<div class=\"hljs-cto\">\n<div class=\"hljs-cto\"><button class=\"copy_btn disable\" data-clipboard-target=\"#code_id_16\">\u590d\u5236<\/button><\/p>\n<div class=\"code-toolbar\">\n<pre class=\"has-pre-numbering language-javascript\" tabindex=\"0\"><code class=\"language-javascript\"># torch<span class=\"token operator\">\/<\/span>_inductor<span class=\"token operator\">\/<\/span>choices<span class=\"token punctuation\">.<\/span>py<span class=\"token operator\">:<\/span>reduction_split_factor\r\n<span class=\"token keyword\">if<\/span> numel_hint <span class=\"token operator\">&gt;=<\/span> <span class=\"token number\">2<\/span> <span class=\"token operator\">*<\/span> num_sm<span class=\"token operator\">:<\/span>      # \u5927 batch \u2192 \u4e0d split\r\n    <span class=\"token keyword\">return<\/span> <span class=\"token number\">1<\/span>\r\n<span class=\"token keyword\">if<\/span> reduction_numel_hint <span class=\"token operator\">&lt;=<\/span> <span class=\"token number\">8192<\/span><span class=\"token operator\">:<\/span>  # hidden \u2264 <span class=\"token number\">8192<\/span> \u2192 \u4e0d <span class=\"token function\">split<\/span><span class=\"token punctuation\">(<\/span>\u5bf9\u4efb\u610f num_tokens \u90fd\u6210\u7acb<span class=\"token punctuation\">)<\/span>\r\n    <span class=\"token keyword\">return<\/span> <span class=\"token number\">1<\/span>\r\n# \u4ec5 hidden <span class=\"token operator\">&gt;<\/span> <span class=\"token number\">8192<\/span> \u4e14 numel_hint <span class=\"token operator\">&lt;<\/span> <span class=\"token number\">2<\/span><span class=\"token operator\">*<\/span>num_sm \u624d split<span class=\"token operator\">&gt;<\/span><span class=\"token number\">1<\/span><span class=\"token punctuation\">,<\/span>\u4e14 split \u503c\u968f numel_hint \u53d8<\/code><\/pre>\n<ul id=\"code_id_16\" class=\"pre-numbering\">\n<li>1.<\/li>\n<li>2.<\/li>\n<li>3.<\/li>\n<li>4.<\/li>\n<li>5.<\/li>\n<li>6.<\/li>\n<\/ul>\n<div class=\"toolbar\"><\/div>\n<\/div>\n<\/div>\n<\/div>\n<\/div>\n<p>\u7531\u4e8e\u5f53\u4eca\u5927\u90e8\u5206\u5f00\u6e90\u6a21\u578b\u90fd\u6ee1\u8db3hidden_size \u2264 8192\uff0creduction_split_factor \u5bf9\u8be5\u6761\u4ef6\u6052\u8fd4\u56de 1,\u4e0e num_tokens \u4ee3\u8868\u503c\u65e0\u5173,\u5f52\u7ea6\u4fdd\u6301\u5355\u7a0b(single-pass)\u3002\u540c\u65f6vLLM\u501f\u52a9TorchCompileWithNoGuardsWrapper\u963b\u6b62\u7279\u5316\u4e0e\u91cd\u7f16\u8bd1\uff0ccompile_ranges_endpoints \u9ed8\u8ba4\u4e3a None,get_compile_ranges() \u8fd4\u56de\u7a7a,\u5373\u5bf9\u4efb\u610f batch size \u53ea\u7f16\u8bd1\u4e00\u4efd\u56fe\u3002\u5373\u4fbf\u914d\u7f6e\u591a\u4e2a compile_ranges \u751f\u6210\u591a\u5f20\u56fe,\u5bf9 hidden_size \u2264 8192,\u4e0a\u8ff0\u6761\u4ef6\u4f7f split \u6052\u4e3a 1\uff0c\u4ece\u800c\u5f52\u7ea6\u7ed3\u6784\u4e00\u81f4\uff0c\u4fdd\u6301 Batch Invariance\u3002<\/p>\n<p>\u7279\u522b\u6ce8\u610f\uff1a\u5728hidden size\u5927\u4e8e8192\u65f6\uff0c\u663e\u5f0f\u914d\u7f6e compile_ranges=[512] \u6216 compile_sizes=[N]\u3002\u6b64\u65f6\u7f16\u8bd1\u5668\u5c06\u9488\u5bf9\u7279\u5b9arange\/\u5c3a\u5bf8 N \u72ec\u7acb\u751f\u6210\u9759\u6001\u5f62\u72b6\u56fe\uff0cnumel_hint\u4e0ehidden_size\u7684\u4e58\u79ef\u5bfc\u81f4\u542f\u53d1\u5f0f\u9608\u503c\u5bfc\u81f4split\u53ef\u80fd\u53d1\u751f\u53d8\u5316\uff0c\u4ece\u800cInductor Split-Reduction \u62c6\u5206\u5f52\u7ea6\u53d1\u751f\u53d8\u5316\uff0c\u5bfc\u81f4\u8be5\u9759\u6001\u56fe\u4e0e\u52a8\u6001\u56fe\u7684\u7d2f\u52a0\u6811\u5f62\u6001\u4ea7\u751f\u7ed3\u6784\u6027\u7a81\u53d8\uff0c\u7834\u574f\u7b97\u5b50\u7ea7\u522b\u7684\u6570\u503c\u4e00\u81f4\u6027\u3002<\/p>\n<div>\n<div class=\"hljs-cto\">\n<div class=\"hljs-cto\"><button class=\"copy_btn disable\" data-clipboard-target=\"#code_id_17\">\u590d\u5236<\/button><\/p>\n<div class=\"code-toolbar\">\n<pre class=\"has-pre-numbering language-javascript\" tabindex=\"0\"><code class=\"language-javascript\">compiled <span class=\"token operator\">=<\/span> torch<span class=\"token punctuation\">.<\/span><span class=\"token function\">compile<\/span><span class=\"token punctuation\">(<\/span>fn<span class=\"token punctuation\">,<\/span> dynamic<span class=\"token operator\">=<\/span>False<span class=\"token punctuation\">)<\/span>   # \u6bcf\u4e2a shape \u5404\u7f16\u4e00\u4efd\r\ncompiled <span class=\"token operator\">=<\/span> torch<span class=\"token punctuation\">.<\/span><span class=\"token function\">compile<\/span><span class=\"token punctuation\">(<\/span>fn<span class=\"token punctuation\">,<\/span> dynamic<span class=\"token operator\">=<\/span>True<span class=\"token punctuation\">)<\/span>    # \u4e00\u4efd\u56fe\u901a\u5403\u6240\u6709 shape\r\ncompiled <span class=\"token operator\">=<\/span> torch<span class=\"token punctuation\">.<\/span><span class=\"token function\">compile<\/span><span class=\"token punctuation\">(<\/span>fn<span class=\"token punctuation\">,<\/span> dynamic<span class=\"token operator\">=<\/span>None<span class=\"token punctuation\">)<\/span>    # \u9ed8\u8ba4<span class=\"token operator\">:<\/span>\u5148\u7279\u5316<span class=\"token punctuation\">,<\/span>shape \u4e00\u53d8\u5c31\u8f6c\u6210\u52a8\u6001<\/code><\/pre>\n<ul id=\"code_id_17\" class=\"pre-numbering\">\n<li>1.<\/li>\n<li>2.<\/li>\n<li>3.<\/li>\n<\/ul>\n<div class=\"toolbar\"><\/div>\n<\/div>\n<\/div>\n<\/div>\n<\/div>\n<p>TorchCompileWithNoGuardsWrapper\uff1a \u4eceJIT\u5230AOT\u5728 PyTorch 2.11 \u7684 JIT\uff08\u5373\u65f6\u7f16\u8bd1\uff09\u673a\u5236\u4e0b\uff0c\u8ba1\u7b97\u56fe\u7684\u8fd0\u884c\u5f3a\u4f9d\u8d56 Shape Guard \u6765\u8fdb\u884c\u5408\u6cd5\u6027\u68c0\u67e5\u3002\u4f46\u5728 LLM \u63a8\u7406\u573a\u666f\u4e2d\uff0cBatch Size \u7684\u9891\u7e41\u6ce2\u52a8\u6613\u89e6\u78b0\u8fb9\u754c\u6761\u4ef6\uff08\u4f8b\u5982s0 &gt;= 2\u7684 Guard \u9a8c\u8bc1\u5931\u8d25\uff09\uff0c\u4ece\u800c\u89e6\u53d1 JIT \u5f15\u64ce\u7684\u5728\u7ebf\u91cd\u7f16\u8bd1\uff08Recompilation\uff09\uff0c\u5bfc\u81f4\u7aef\u5230\u7aef\u5ef6\u8fdf\u51fa\u73b0\u660e\u663e\u7684\u6bdb\u523a\u3002<\/p>\n<p>\u4e3a\u4e86\u5f7b\u5e95\u6d88\u9664 JIT \u5e26\u6765\u7684\u6027\u80fd\u9690\u60a3\uff0cvLLM \u5f15\u5165\u4e86\u00a0TorchCompileWithNoGuardsWrapper\u00a0\u7ed3\u5408 Piecewise Backend\uff0c\u5b9e\u73b0\u4e86\u7eaf AOT\uff08\u63d0\u524d\u7f16\u8bd1\uff09\u6d3e\u53d1\u3002\u5728\u9884\u70ed\u9636\u6bb5\uff0c\u7cfb\u7edf\u4f1a\u9488\u5bf9\u4e0d\u540c\u7684\u5f20\u91cf\u5f62\u72b6\u533a\u95f4\u63d0\u524d\u751f\u6210\u591a\u5f20\u9884\u7f16\u8bd1\u56fe\u3002\u8fdb\u5165\u8fd0\u884c\u65f6\u540e\uff0c\u8be5 Wrapper \u4f1a\u4e3b\u52a8\u62e6\u622a\u5e76\u4e22\u5f03\u5e95\u5c42\u6240\u6709\u7684 Dynamo Guard \u9a8c\u8bc1\uff0c\u7531 vLLM \u6846\u67b6\u5c42\u6839\u636e\u8f93\u5165\u5f62\u72b6\uff0c\u4ee5\u7eaf\u9759\u6001\u7684\u65b9\u5f0f\u76f4\u63a5\u8def\u7531\u5230\u5bf9\u5e94\u7684\u9884\u7f16\u8bd1\u56fe\u4e0a\u3002\u56e0\u6b64\uff0c\u8fd9\u4e00\u673a\u5236\u672c\u8d28\u4e0a\u662f\u7531 vLLM \u7edf\u7b79\u8c03\u5ea6\u7684\u9884\u7f16\u8bd1\u56fe\u6d3e\u53d1\uff08Dispatch\uff09\u7cfb\u7edf\uff0c\u800c\u975e\u7b80\u5355\u5730\u8ba9torch.compile\u53bb\u7ef4\u62a4\u5355\u5f20\u52a8\u6001\u56fe\u3002<\/p>\n<p>\u5728\u8fd9\u79cd\u5206\u6bb5\u7f16\u8bd1\u673a\u5236\u4e0b\uff0c\u5177\u4f53\u7b97\u5b50\u7684\u4ee3\u7801\u751f\u6210\uff08Codegen\uff09\u903b\u8f91\u9700\u89c6\u914d\u7f6e\u800c\u5b9a\u3002\u4ee5 RMSNorm \u4e3a\u4f8b\uff0c\u5f53\u00a0hidden_size &lt;= 8192\u00a0\u65f6\uff0c\u5176\u00a0reduction_split_factor\u00a0\u901a\u5e38\u4e3a 1\uff0c\u8fd9\u4f7f\u5f97\u5355\u884c\u5f52\u7ea6\uff08Inner Reduction\uff09\u80fd\u591f\u4fdd\u6301\u9ad8\u6548\u7684 Single-pass \u6267\u884c\u3002\u4f46\u9700\u8981\u6ce8\u610f\u7684\u662f\uff1a\u6700\u7ec8\u7cfb\u7edf\u4f1a\u5207\u5206\u51fa\u591a\u5c11\u5f20\u56fe\u3001\u6bcf\u5f20\u56fe\u91c7\u7528\u4ec0\u4e48\u4ee3\u8868\u503c\uff08Representative Value\uff09\uff0c\u4ee5\u53ca\u4e0d\u540c\u7684\u56fe\u662f\u5426\u4f1a\u547d\u4e2d\u76f8\u540c\u7684 Codegen \u7b56\u7565\uff0c\u90fd\u9ad8\u5ea6\u4f9d\u8d56\u4e8e vLLM \u7684 Compile Ranges\/Sizes \u8bbe\u5b9a\u3001\u52a8\u6001 Shape \u7684\u7c7b\u578b\u4ee5\u53ca PyTorch \u7684\u5177\u4f53\u7248\u672c\u3002\u56e0\u6b64\uff0c\u4efb\u4f55\u5173\u4e8e\u5e95\u5c42\u7b97\u5b50\u884c\u4e3a\u7684\u7edd\u5bf9\u4fdd\u8bc1\uff0c\u90fd\u5fc5\u987b\u4e25\u683c\u9650\u5b9a\u5728\u660e\u786e\u7684\u8f6f\u4ef6\u7248\u672c\u4e0e\u7f16\u8bd1\u914d\u7f6e\u73af\u5883\u5185\u3002<\/p>\n<h3>\u56db\u3001Attention\u7684Batch Invariance<\/h3>\n<h4>1. Attention\u8ba1\u7b97<\/h4>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/c250ce26874b81dad06801091dc26b7647b581.webp\" data-type=\"block\" \/><\/p>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/d8e28c308a05999266b41799230a697f79fe5e.webp\" data-type=\"block\" \/><\/p>\n<p><img decoding=\"async\" title=\"image-20260705144714859\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/3760f3a70d5ba6b1130929f2e64395c901e673.webp\" alt=\"image-20260705144714859\" data-type=\"inline\" \/><\/p>\n<p>\u8fd9\u91cc\u5c31\u4e0d\u8be6\u7ec6\u4ecb\u7ecdAttention\u7684\u57fa\u672c\u539f\u7406\u4e86\uff0c\u4e4b\u524d\u5199\u8fc7\u4e00\u7bc7\u6587\u7ae0\u300aAI Infra\u5165\u95e8\uff1a\u5927\u6a21\u578b\u662f\u5982\u4f55\u9ad8\u6548\u63a8\u7406\u7684 \u300b\uff0c\u6211\u4eec\u76f4\u63a5\u4eceFlashAttention\u5b9e\u73b0\u8bf4\u8d77\u3002<\/p>\n<table class=\"data-table\" data-transient-attributes=\"class\" data-width=\"676.997px\">\n<colgroup data-id=\"c7104f7d-fJE3A99m\">\n<col span=\"1\" width=\"169.236\" data-id=\"cd89ecb0-t5yQewIa\" \/>\n<col span=\"1\" width=\"169.236\" data-id=\"cd89ecb0-8zQPOr3w\" \/>\n<col span=\"1\" width=\"169.236\" data-id=\"cd89ecb0-CeOPvMcf\" \/>\n<col span=\"1\" width=\"169.288\" data-id=\"cd89ecb0-1R37f7pv\" \/><\/colgroup>\n<tbody data-id=\"t6d5e859-ranLrZ0o\">\n<tr data-id=\"t31e458f-WcfVLw5n\">\n<td data-id=\"t6267798-chijA1Wg\" data-transient-attributes=\"table-cell-selection\">\u64cd\u4f5c\u540d\u79f0<\/td>\n<td data-id=\"t6267798-HdHFPWFf\" data-transient-attributes=\"table-cell-selection\">\u6570\u5b66\u516c\u5f0f<\/td>\n<td data-id=\"t6267798-J0H7qT75\" data-transient-attributes=\"table-cell-selection\">\u5b8f\u89c2Shape\u53d8\u5316<\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-GnnfEPkV\" data-transient-attributes=\"table-cell-selection\">Tiling &#8211; \u5fae\u89c2<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-SON9CPus\">\n<td data-id=\"t6267798-spvqqq4I\" data-transient-attributes=\"table-cell-selection\">Q-K Dot Product<\/td>\n<td data-id=\"t6267798-GibwbcIL\" data-transient-attributes=\"table-cell-selection\">S\u00a0=\u00a0Q\u00a0\u00d7\u00a0KT<\/td>\n<td data-id=\"t6267798-iz3Ktyv6\" data-transient-attributes=\"table-cell-selection\">Q:\u00a0<code>[query_lens, num_heads, head_dim]<\/code><br \/>\nK:\u00a0<code>[seq_lens, num_kv_heads, head_dim]<\/code><br \/>\nS:\u00a0<code>[query_lens, num_heads, seq_lens]<\/code><\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-Y78678eR\" data-transient-attributes=\"table-cell-selection\">Q:\u00a0<code>[BLOCK_M, head_dim]<\/code><br \/>\nK:\u00a0<code>[BLOCK_N, head_dim]<\/code><br \/>\nS:\u00a0<code>[BLOCK_M, BLOCK_N]<\/code><\/td>\n<\/tr>\n<tr data-id=\"t31e458f-MQzyjawz\">\n<td data-id=\"t6267798-OuFMxZEY\" data-transient-attributes=\"table-cell-selection\">Scale<\/td>\n<td data-id=\"t6267798-ujaWWQO3\" data-transient-attributes=\"table-cell-selection\">S\u00a0=\u00a0S\u00a0\/ \u221a(d)<\/td>\n<td data-id=\"t6267798-rME3E2DT\" data-transient-attributes=\"table-cell-selection\">shape\u4e0d\u53d8<\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-NR2vBDad\" data-transient-attributes=\"table-cell-selection\">shape\u4e0d\u53d8<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-BvaLfci7\">\n<td data-id=\"t6267798-ZViv1rS8\" data-transient-attributes=\"table-cell-selection\">Mask<\/td>\n<td data-id=\"t6267798-PIxsPd1G\" data-transient-attributes=\"table-cell-selection\">S\u00a0=\u00a0S\u00a0+\u00a0M<\/td>\n<td data-id=\"t6267798-K9FZxbOT\" data-transient-attributes=\"table-cell-selection\">shape\u4e0d\u53d8<\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-aRwTwKyG\" data-transient-attributes=\"table-cell-selection\">shape\u4e0d\u53d8<\/td>\n<\/tr>\n<tr data-id=\"t31e458f-dLnioVD7\">\n<td data-id=\"t6267798-rkt6lokU\" data-transient-attributes=\"table-cell-selection\">Softmax<\/td>\n<td data-id=\"t6267798-rKu9WfzC\" data-transient-attributes=\"table-cell-selection\">P\u00a0= softmax(S)<\/td>\n<td data-id=\"t6267798-JuOTTNy7\" data-transient-attributes=\"table-cell-selection\">P:\u00a0<code>[query_lens, num_heads, seq_lens]<\/code><\/td>\n<td class=\"table-last-row\" data-id=\"t6267798-mWw7irm8\" data-transient-attributes=\"table-cell-selection\">P:\u00a0<code>[BLOCK_M, BLOCK_N]<\/code><\/td>\n<\/tr>\n<tr data-id=\"t31e458f-u3ioUDMV\">\n<td class=\"table-last-column\" data-id=\"t6267798-wf6a0ctF\" data-transient-attributes=\"table-cell-selection\">MatMul<\/td>\n<td class=\"table-last-column\" data-id=\"t6267798-A3aUx5Mf\" data-transient-attributes=\"table-cell-selection\">O\u00a0=\u00a0P\u00a0\u00d7\u00a0V<\/td>\n<td class=\"table-last-column\" data-id=\"t6267798-UppE8dNf\" data-transient-attributes=\"table-cell-selection\">P:\u00a0<code>[query_lens, num_heads, seq_lens]<\/code><br \/>\nV:\u00a0<code>[seq_lens, num_kv_heads, head_dim]<\/code><br \/>\nO:\u00a0<code>[query_lens, num_heads, head_dim]<\/code><\/td>\n<td class=\"table-last-column table-last-row\" data-id=\"t6267798-Fq4Gnpd5\" data-transient-attributes=\"table-cell-selection\">P:\u00a0<code>[BLOCK_M, BLOCK_N]<\/code><br \/>\nV:\u00a0<code>[BLOCK_N, head_dim]<\/code><br \/>\nO:\u00a0<code>[BLOCK_M, head_dim]<\/code><\/td>\n<\/tr>\n<\/tbody>\n<\/table>\n<p>Q TIle\u5bf9\u5e94\u7684\u5c31\u662f\u672c\u6b21\u4e00\u8d77\u5904\u7406\u7684Q\u7ef4\u5ea6Token\u5373[BLOCK_M, head_dim],\u672c\u8d28\u4e0a\u5c31\u662f\u5c06\u5b8f\u89c2\u7684Q\u00a0[query_lens, num_heads, head_dim]\u6cbf M \u8f74\u6309 BLOCK_M \u5c06query_len\u5207\u5206\u6210query_len \/ BLOCK_M\u4e2a Q tile\u3002 (num_heads\u4e5f\u53d8\u6210\u5e76\u53d1\u7ef4\u5ea6\u4e86)<\/p>\n<p>\u7136\u540e\u6bcf\u4e2aCTA\u8d1f\u8d23\u4e00\u4e2aQ Tile\uff0c\u6cbf K\/V \u65b9\u5411\u505a\u4e00\u6b21\u878d\u5408\u904d\u5386\u5904\u7406\u6240\u6709\u7684\u5386\u53f2KV\uff0c\u540c\u65f6\u5b8c\u6210\u00a0S=QKT\u3001online softmax\u3001\u4ee5\u53ca\u00a0O\u00a0+=\u00a0PV\u00a0\u4e09\u4ef6\u4e8b\u3002\u8ba9\u4e2d\u95f4\u77e9\u9635\u00a0S,\u00a0P\u00a0\u59cb\u7ec8\u6d3b\u5728\u5bc4\u5b58\u5668\u4e2d\u3001\u4e0d\u5199\u56de HBM\u2014\u2014\u6700\u540e\u53ea\u628a\u6bcf\u4e2a Q tile \u7684\u00a0O\u00a0\u548c LSE \u5199\u51fa\u53bb\u3002<\/p>\n<p>BLOCK_M\u53ea\u662f\u5f71\u54cd\u591a\u5c11\u4e2aQ\u7ef4\u5ea6\u7684token\u5e76\u53d1\uff08\u5373query_len\u7ef4\u5ea6\uff09\uff0c\u6240\u4ee5\u5b8c\u5168\u4e0d\u5e94\u5f71\u54cd\u6d6e\u70b9\u52a0\u6cd5\u7684\u987a\u5e8f\uff1b\u800cBLOCK_N\u662fKV\u7ef4\u5ea6\u7684token\u5e76\u53d1\uff08\u5373seq_len\u7ef4\u5ea6\uff09\uff0c\u662f\u89c4\u7ea6\u8f74\uff0c\u4f1a\u5f71\u54cd\u6d6e\u70b9\u52a0\u6cd5\u7684\u987a\u5e8f\u3002\u4e0d\u540c\u7684BLOCK_N\u5bfc\u81f4Online Softmax\u7f29\u653e\u7ef4\u5ea6\u4e0d\u4e00\u81f4\uff0c\u5373\u591a\u4e2aKV token\u8ba1\u7b97\u4e00\u4e2a\u5c40\u90e8\u7684mnew\uff0c \u540c\u65f6BLOCK_N\u5728O\u00a0=\u00a0P\u00a0\u00d7\u00a0V\u4e2d\u5176\u5b9e\u5c31\u662fGEMM\u4e2d\u7684BLOCK_K, \u4f1a\u89c4\u7ea6\u6811\u53d1\u751f\u53d8\u5316\uff0c\u4f46\u662fQK\u70b9\u79ef\u6ca1\u5f71\u54cd\uff0c\u56e0\u4e3aQK\u70b9\u79ef\u662fhead_dim\u7ef4\u5ea6\u89c4\u7ea6\u3002<\/p>\n<p>FlashAttention\u4e2d\uff0cBLOCK_M \u4e0e BLOCK_N \u90fd\u662f\u7f16\u8bd1\u671f\u5e38\u91cf,\u4e14\u5b83\u4eec\u7684\u53d6\u503c\u53ea\u7531 (head_dim, causal\/local, dtype, arch) \u51b3\u5b9a,\u4e0ebatch \u4e0e seqlen\u65e0\u5173\u2014\u2014\u5b83\u4eec\u51b3\u5b9a\u662f\u5355\u4e2a CTA \u5185\u90e8\u600e\u4e48\u7b97\u6700\u5feb (\u5373\u53d7\u6bcf\u4e2a SM \u56fa\u5b9a\u7684\u5bc4\u5b58\u5668\/SMEM \u9884\u7b97\u7ea6\u675f,head_dim \u662f\u628a tile \u6362\u7b97\u6210\u8d44\u6e90\u6d88\u8017\u7684\u57fa\u672c\u5355\u5143),\u800c\u8981\u8dd1\u591a\u5c11\u4e2a\u8fd9\u6837\u7684 CTA\u7531 grid \u4e0e num_splits \u8d1f\u8d23\u3002\u56e0\u6b64\u65e0\u8bba\u4e0d\u7ba1 batch \u591a\u5927\u3001seqlen \u591a\u957f,tile \u90fd\u662f\u540c\u4e00\u4e2a\u503c,\u89c4\u7ea6\u987a\u5e8f\u5728\u8be5 kernel \u5185\u5b8c\u5168\u56fa\u5b9a,\u7ed3\u679c\u9010 bit \u786e\u5b9a\u3002\u771f\u6b63\u6539\u53d8\u6d6e\u70b9\u52a0\u6cd5\u987a\u5e8f\u7684,\u662f\u5207\u5206\u89c4\u7ea6\u8f74(KV \u8f74)\u7684\u4f18\u5316\uff08SplitKV \/ num_splits\uff09:\u628a KV \u89c4\u7ea6\u8f74\u5207\u7ed9\u591a\u4e2a CTA \u518d\u7531 combine \u5408\u5e76,\u89c4\u7ea6\u6811\u7ed3\u6784\u6539\u53d8\u3002\u4e14\u5176\u81ea\u52a8\u542f\u53d1\u5f0f\u7684\u8f93\u5165 total_mblocks = batch_size \u00d7 num_head_kv \u00d7 num_m_blocks \u6b63\u6bd4\u4e8e batch,\u6240\u4ee5 batch \u53d8\u5316\u4f1a\u7ecf\u7531 num_splits \u7834\u574f batch \u4e0d\u53d8\u6027\u3002<\/p>\n<h4>2. chunked prefill<\/h4>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/460430345ed6eda72647747b2226c4e4a7df18.webp\" data-type=\"block\" \/><\/p>\n<p>FlashAttention\u8ba1\u7b97\u65f6\uff0cBLOCK_M\u662f\u5bf9\u4e8equery_len\u5fae\u89c2\u5c42\u9762\u7684\u5207\u5206\uff0c chunked prefill\u5219\u53ef\u4ee5\u7406\u89e3\u4e3a\u5bf9query_len\u7684\u5b8f\u89c2\u5c42\u9762\u4e0a\u7684\u5207\u5206\u3002\u4e24\u8005\u672c\u8d28\u76f8\u540c\u2014\u2014\u90fd\u6cbf query \u8f74(M \u8f74)\u5207\u3001\u5757\u95f4\u4e92\u4e0d\u5f52\u7ea6,\u6240\u4ee5\u90fd\u5929\u7136 batch-invariant-safe;\u4f46\u5207\u7684\u662f\u4e0d\u540c\u5c42\u7ea7\u7684query_len\u3002\u4e0a\u9762\u5df2\u7ecf\u770b\u5230\u4e86\uff0c\u5bf9\u4e8e\u76f8\u540c\u7684(head_dim, causal\/local, dtype, arch) \uff0cBLOCK_M\u662f\u786e\u5b9a\u7684\uff0c\u800cchunked prefill\u4e2dquery_len\u7684\u5207\u5206\u662f\u968fbatch\u4f1a\u53d1\u751f\u53d8\u5316\u7684\uff0c\u90a3\u8fd9\u91cc\u7684\u5207\u5206\u662f\u5426\u4f1a\u6539\u53d8\u6539\u53d8\u6d6e\u70b9\u52a0\u6cd5\u987a\u5e8f\u5462\uff1f<\/p>\n<p>chunked prefill\u662f\u8c03\u5ea6\u5c42\u7684\u5207\u5206\uff0cscheduler \u6309 token \u9884\u7b97\u51b3\u5b9a\u8fd9\u4e00 step \u5904\u7406\u54ea\u4e9b\u8bf7\u6c42\u7684\u54ea\u4e9b token;\u5f53\u67d0\u8bf7\u6c42\u7684 prompt \u4e00\u6b21\u7b97\u4e0d\u5b8c\u65f6,\u5b83\u5c31\u88ab\u8de8 step \u5207\u5f00,\u6bcf step \u5582\u8fdb kernel \u7684 query_len \u5373\u8be5 step \u5206\u5230\u7684 chunk \u5927\u5c0f(prompt\u77ed\u3001\u9884\u7b97\u591f\u65f6\u4e00\u6b65\u7b97\u5b8c,\u5e76\u4e0d\u4f1a\u88ab\u5207;\u540c\u4e00 step \u901a\u5e38\u8fd8\u6df7\u7740\u522b\u7684\u8bf7\u6c42\u7684 chunk \/ decode)\u3002<\/p>\n<p>BLOCK_M \u662f FA kernel \u5c42\u7684\u5207\u5206:kernel \u62ff\u5230\u672c step \u7684 query_len \u540e,\u6cbf M \u8f74\u6309 BLOCK_M \u5206\u6210 query_len \/ BLOCK_M\u4e2a Q tile\u3002<\/p>\n<p>&nbsp;<\/p>\n<p>chunked prefill\u7684\u8f93\u51fa\u5176\u5b9e\u5c31\u662fFA kernel\u7684\u8f93\u5165\u2014\u2014chunk \u5927\u5c0f\u5373 kernel \u62ff\u5230\u7684 query_len\u3002\u8fd9\u4e2a\u8f93\u5165\u4e0d\u4ec5\u5f71\u54cd\u51e0\u4e2a Q tile,\u8fd8\u7ecf num_m_blocks = \u2308query_len\/64\u2309 \u5f71\u54cd kernel \u8981\u4e0d\u8981 split-KV:chunk \u5927 \u2192 query_len \u5927 \u2192 \u5224&#8221;\u586b\u6ee1 SM&#8221;\u2192 \u4e0d\u5207 KV(prefill \u5e38\u6001,BI\u53cb\u597d);chunk \u9000\u5316\u5230 decode \u7684 1 \u2192 \u53ef\u80fd\u89e6\u53d1 split-KV\u3002<\/p>\n<p>&nbsp;<\/p>\n<p>\u4e4b\u6240\u4ee5query_len\u5207\u5206\u4e0d\u4f1a\u5f71\u54cd\u6d6e\u70b9\u52a0\u6cd5\u987a\u5e8f\u6709\u4e24\u4e2a\u539f\u56e0\uff1a<\/p>\n<ul data-id=\"u738a58b-g7Vfcsx3\">\n<li data-id=\"ld70c578-eGC1X9ab\">query_len\u7ef4\u5ea6(M \u8f74)\u4e0a\uff0cattention\u8ba1\u7b97\u662f\uff1a\u8bf7\u6c42\u5185\u4e0d\u540c query token \u4e4b\u95f4\u7684 attention \u8ba1\u7b97\u662f\u5e76\u884c\u7684\u3001\u5f7c\u6b64\u4e0d\u5b58\u5728\u5f52\u7ea6\uff0c\u8ba1\u7b97\u7684query_len\u7684\u6bcf\u4e2atoken \u4e0e seq_len \u4e2d\u7684 K\/V token \u7684\u5173\u7cfb\uff1b \u5207 query_len \u53ea\u51b3\u5b9a\u54ea\u4e9b token \u8fd9\u4e00\u6b65\u4e00\u8d77\u7b97, \u65e0\u8bba\u600e\u4e48\u5207,\u5b83\u5bf9\u5e94\u7684 n_block \u5207\u5206\u4e0e\u7d2f\u52a0\u987a\u5e8f\u9010\u5757\u5bf9\u9f50\u3001\u5b8c\u5168\u4e00\u81f4\u3002<\/li>\n<li data-id=\"ld70c578-MCvxqrRN\">causal mask \u4fdd\u8bc1\u6bcf\u4e2a token \u53ea\u4f9d\u8d56\u5b83\u81ea\u5df1\u548c\u5b83\u4e4b\u524d\u7684 token\u3002 K\/V \u662f\u9010 token \u72ec\u7acb\u751f\u6210\u7684(\u6bcf\u4e2a token \u7684 key\/value \u53ea\u7531\u5b83\u81ea\u5df1\u7684\u8f93\u5165\u7ecf\u6743\u91cd\u6295\u5f71\u5f97\u5230,\u8de8 token \u4e4b\u95f4\u6ca1\u6709\u8ba1\u7b97\u4f9d\u8d56);\u800c causal mask \u5728 attention \u8fd9\u4e00\u6b65\u628a\u6bcf\u4e2a token\u7684\u5f52\u7ea6\u8303\u56f4\u63a7\u5236\u5728\u5b83\u81ea\u5df1\u53ca\u4e4b\u524d\u7684\u6240\u6709 token\u3002\u56e0\u6b64Attention\u6700\u7ec8\u7684\u8f93\u51fa\u53d6\u51b3\u4e8equery \u5411\u91cf\u5bf9\u5b83\u81ea\u5df1\u53ca\u4e4b\u524d\u6240\u6709 token\u7684 K\/V \u505a softmax \u52a0\u6743\u548c\uff08\u5f52\u7ea6\u8f74\u662f KV \u65b9\u5411\uff09\uff0c\u4e0e query_len\u65e0\u5173\uff1bchunk \u600e\u4e48\u5207,\u5355\u4e2a token \u7684\u6570\u5b66\u7ed3\u679c\u4e0d\u53d8\u3002<\/li>\n<\/ul>\n<h4>3. FlashDecoding &#8211; split-kv<\/h4>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/1775b1759dd6c41bfeb0096469c4052f9d7665.webp\" data-type=\"block\" \/><\/p>\n<p>\u524d\u9762\u63d0\u5230\u4e86chunked prefill \u6cbf query \u8f74\uff08M \u8f74\uff09\u5207\uff0c\u5757\u95f4\u4e0d\u5f52\u7ea6\uff0c\u5929\u7136batch invariance\u3002\u4f46\u5230\u4e86 decode \u9636\u6bb5\uff0c\u60c5\u51b5\u6b63\u597d\u53cd\u8fc7\u6765\uff1a\u6bcf\u6b65 query_len \u9000\u5316\u4e3a 1\uff08M=1\uff09\uff0c\u800c seq_len (KV\uff09 \u5374\u53ef\u80fd\u6781\u957f\u3002\u5982\u679c\u8ba9\u4e00\u4e2a Q tile \u72ec\u5360\u4e00\u4e2a CTA\u3001\u6cbf KV \u4e32\u884c\u8ba1\u7b97\uff0c\u90a3\u4e48 batch \u5f88\u5c0f\u65f6\u5c31\u53ea\u80fd\u62c9\u8d77\u51e0\u4e2a CTA\uff0c\u7edd\u5927\u591a\u6570 SM \u7a7a\u8f6c\u2014\u2014\u8fd9\u91cc\u672c\u8d28\u4e0a\u8ddfGEMM\u4e2d\u7684Split-K\u8981\u89e3\u51b3\u7684\u95ee\u9898\u662f\u4e00\u81f4\u7684\uff0c\u53ea\u4e0d\u8fc7\u8fd9\u91cc\u88ab\u5207\u7684\u89c4\u7ea6\u8f74\u662f KV\u3002<\/p>\n<p>FlashDecoding \u7684\u505a\u6cd5\u5c31\u662f split-KV\uff1a\u6839\u636e\u00a0(batch, heads, seq_kv, num_sms)\u00a0\u52a8\u6001\u7684\uff0c\u628a\u4e00\u4e2a query \u7684 KV \u533a\u95f4\u5207\u6210num_splits\u6bb5\uff0c\u4ea4\u7ed9\u591a\u4e2a CTA \u5e76\u884c\uff0c\u5404\u81ea\u7b97\u51fa\u4e00\u6bb5\u7684\u5c40\u90e8\u00a0O\u00a0\u4e0e\u5c40\u90e8 LSE\uff08log-sum-exp\uff09\uff0c\u6700\u540e\u7531\u4e00\u4e2a combine\/reduce \u6b65\u9aa4\u6309 LSE \u628a\u8fd9\u4e9b\u5c40\u90e8\u7ed3\u679c\u5408\u5e76\u6210\u6700\u7ec8\u8f93\u51fa\u3002\u8fd9\u4e00\u6b65\u5728\u6570\u5b66\u4e0a\u7b49\u4ef7\u4e8e online softmax \u7684\u8de8\u6bb5\u5408\u5e76\u3002<\/p>\n<p>\u95ee\u9898\u5c31\u51fa\u5728\u5207\u51e0\u6bb5\u4e0a\uff1anum_splits\u662f\u6839\u636e batch\/seqlen \u52a8\u6001\u51b3\u5b9a\u7684\uff0c\u76ee\u6807\u662f\u586b\u6ee1 SM\u3002num_splits\u4e00\u53d8\uff0cLSE \u5408\u5e76\u65f6\u7684\u89c4\u7ea6\u62d3\u6251\u5c31\u53d8\uff0c\u6d6e\u70b9\u52a0\u6cd5\u987a\u5e8f\u968f\u4e4b\u6539\u53d8\u2014\u2014\u4e8e\u662f\u53c8\u4e22\u4e86 batch invariance\u3002\u540c\u6837\uff0c\u8fd9\u8ddfSplit-K\u7684\u903b\u8f91\u4e00\u81f4\uff1a\u56fa\u5b9a\u6bb5\u6570\u65f6 combine \u662f\u6709\u5e8f\u7684\u3001run-to-run \u786e\u5b9a\uff1b\u53ef\u4e00\u65e6 batch \u8ba9\u6bb5\u6570\u53d1\u751f\u53d8\u5316\uff0c\u7ed3\u679c\u7167\u6837\u4f1a\u53d8\u3002<\/p>\n<h4>4. vLLM\u4e2dAttention\u7684Batch Invariance\u652f\u6301<\/h4>\n<p>\u8fd9\u91cc\u5b9e\u73b0Batch Invarianc\u903b\u8f91\u4e5f\u5f88\u76f4\u63a5\uff0c\u5f3a\u5236\u4e0d\u62c6 KV\uff0c\u5373\u00a0num_splits = 1\u3002vLLM \u5728 FlashAttention \u540e\u7aef\u91cc\uff0c\u53ea\u8981\u5f00\u4e86\u00a0VLLM_BATCH_INVARIANT\u00a0\u5c31\u628a split\u5199\u6b7b\u4e3a1\uff1a<\/p>\n<div>\n<div class=\"hljs-cto\">\n<div class=\"hljs-cto\"><button class=\"copy_btn disable\" data-clipboard-target=\"#code_id_18\">\u590d\u5236<\/button><\/p>\n<div class=\"code-toolbar\">\n<pre class=\"has-pre-numbering language-javascript\" tabindex=\"0\"><code class=\"language-javascript\"># vllm<span class=\"token operator\">\/<\/span>v1<span class=\"token operator\">\/<\/span>attention<span class=\"token operator\">\/<\/span>backends<span class=\"token operator\">\/<\/span>flash_attn<span class=\"token punctuation\">.<\/span>py\r\n<span class=\"token keyword\">if<\/span> envs<span class=\"token punctuation\">.<\/span><span class=\"token constant\">VLLM_BATCH_INVARIANT<\/span><span class=\"token operator\">:<\/span>\r\n    max_num_splits <span class=\"token operator\">=<\/span> <span class=\"token number\">1<\/span>\r\n<span class=\"token operator\">...<\/span>\r\n# \u5404\u5904\u8c03\u7528 flash_attn \u65f6\u7edf\u4e00\uff1a\r\nnum_splits<span class=\"token operator\">=<\/span><span class=\"token number\">1<\/span> <span class=\"token keyword\">if<\/span> envs<span class=\"token punctuation\">.<\/span><span class=\"token constant\">VLLM_BATCH_INVARIANT<\/span> <span class=\"token keyword\">else<\/span> max_num_splits<\/code><\/pre>\n<ul id=\"code_id_18\" class=\"pre-numbering\">\n<li>1.<\/li>\n<li>2.<\/li>\n<li>3.<\/li>\n<li>4.<\/li>\n<li>5.<\/li>\n<li>6.<\/li>\n<\/ul>\n<div class=\"toolbar\"><\/div>\n<\/div>\n<\/div>\n<\/div>\n<\/div>\n<p>MLA \u7684 prefill\/decode \u8def\u5f84\u540c\u7406\uff08num_splits = 1\uff09\u3002\u4ee3\u4ef7\u662f decode \u5c0f batch \u4e0b SM \u5229\u7528\u7387\u4f1a\u4e0b\u964d\uff0c\u8fd9\u662f\u201c\u786e\u5b9a\u6027 \u2194 \u541e\u5410\u201d\u95f4 \u7684trade off\u5728 Attention \u4fa7\u7684\u53c8\u4e00\u6b21\u4f53\u73b0\uff1a\u548c GEMM \u4e00\u6837\uff0c\u6211\u4eec\u727a\u7272\u4e86\u4e00\u90e8\u5206\u5e76\u884c\u5ea6\uff0c\u6362\u56de\u4e86\u6d6e\u70b9\u52a0\u6cd5\u987a\u5e8f\u7684\u7a33\u5b9a\u3002<\/p>\n<h3>\u4e94\u3001NCCL\u7684Batch Invariance<\/h3>\n<p>\u524d\u9762\u6211\u4eec\u8ba8\u8bba\u7684\u90fd\u662f\u5355\u5361\u5185\u90e8\u5fae\u89c2\u7b97\u5b50\u7684\u8c03\u5ea6\u5bfc\u81f4\u7684\u7ed3\u679c\u504f\u5dee\u3002\u968f\u7740\u73b0\u4ee3\u5927\u6a21\u578b\u6a21\u578b\u53c2\u6570\u91cf\u7684\u7206\u70b8\uff0c\u5f53\u5355\u5361\u663e\u5b58\u65e0\u6cd5\u5bb9\u7eb3\u6574\u4e2a\u6a21\u578b\u65f6\uff0c\u7cfb\u7edf\u901a\u5e38\u9700\u8981\u5f15\u5165\u6a21\u578b\u5e76\u884c\u3002\u5176\u4e2d\uff0cTensor Parallelism (TP)\u662f\u4e00\u4e2a\u5e38\u7528\u7684\u65b9\u6848\uff1a\u5c06\u5355\u4e2a\u77e9\u9635\u4e58\u6cd5\u6cbf\u7279\u5b9a\u7ef4\u5ea6\u5207\u5206\u5230\u591a\u5f20 GPU \u4e0a\u3002\u4ece\u8fd9\u4e2a\u89d2\u5ea6\u770b\uff0cTP \u53ef\u4ee5\u88ab\u89c6\u4e3a\u4e00\u79cd\u8de8\u8bbe\u5907\u7684\u9ad8\u7ef4 Tiling\uff0c\u53ea\u4e0d\u8fc7 Tile \u4e4b\u95f4\u4e0d\u518d\u5171\u4eab\u540c\u4e00\u5757\u7247\u4e0a\u5b58\u50a8\uff0c\u800c\u662f\u901a\u8fc7\u5206\u5e03\u5f0f\u901a\u4fe1\u4ea4\u6362\u4e2d\u95f4\u7ed3\u679c\u3002<\/p>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/b90991999fc78c19c15783686b07aa66c1872c.webp\" data-type=\"block\" \/><\/p>\n<p>\u5206\u5e03\u5f0f\u8ba1\u7b97\u4e2d\uff0c\u8ba1\u7b97\u5f80\u5f80\u4e0d\u662f\u74f6\u9888\uff0c\u8de8\u8bbe\u5907\u901a\u4fe1\uff08\u5982 All-Reduce\uff09\u624d\u662f\u62d6\u6162\u6574\u4f53\u541e\u5410\u91cf\u7684\u5143\u51f6\u3002\u4e3a\u4e86\u51cf\u5c11\u901a\u4fe1\uff0c\u6211\u4eec\u9700\u8981\u4e00\u79cd\u77e9\u9635\u5207\u5206\u7b56\u7565\uff0c\u4f7f\u5f97\u591a\u4e2a\u8fde\u7eed\u7684\u7ebf\u6027\u5c42\u5728\u8ba1\u7b97\u8fc7\u7a0b\u4e2d\u5c3d\u53ef\u80fd\u4fdd\u6301\u72ec\u7acb\uff0c\u76f4\u5230\u6700\u540e\u4e00\u523b\u624d\u8fdb\u884c\u6570\u636e\u540c\u6b65\u3002\u8fd9\u5c31\u662f\u5728 MLP\u6216Attention \u4e2d\uff0c\u901a\u5e38\u4f1a\u91c7\u7528\u201c\u5217\u5e76\u884c\u00a0\u2192\u00a0\u884c\u5e76\u884c\u201d\u7684\u7ec4\u5408\u7b56\u7565\uff1a<\/p>\n<p>Column Parallelism (\u5217\u5207\u5206)\uff1a\u00a0Fused QKV \u548c FFN \u7684 Fused Gate\/Up \u90fd\u662f\u5217\u5207\u5206\u3002\u4f8b\u5982 QKV \u7684\u6743\u91cd\u7ef4\u5ea6\u5728\u5355\u5361\u4e0a\u4f1a\u53d8\u6210\u00a0[hidden_size, qkv_proj_size \/ TP_size]\uff0c\u8f93\u51fa\u4e5f\u76f8\u5e94\u53d8\u6210\u00a0[num_sched_tokens, qkv_proj_size \/ TP_size]\u3002<\/p>\n<p>Row Parallelism (\u884c\u5207\u5206)\uff1a\u00a0O_Proj \u548c FFN \u7684 Down_Proj \u662f\u884c\u5207\u5206\u3002\u6743\u91cd\u53d8\u4e3a\u00a0[hidden_size \/ TP_size, hidden_size]\u3002\u8ba1\u7b97\u5b8c\u6210\u540e[num_sched_tokens, hidden_size]\uff0c\u901a\u8fc7\u00a0AllReduce\u00a0\u7b97\u5b50\u6765\u805a\u5408\u8de8\u5361\u7ed3\u679c\u3002<\/p>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/23cc2e065a0d34c97be126560cb8585a6a366a.webp\" data-type=\"block\" \/><\/p>\n<p>\u5728\u4e0d\u8003\u8651 Sequence Parallel\u3001Late All-Reduce \u548c\u901a\u4fe1\u878d\u5408\u7684\u7ecf\u5178\u7684 Dense Transformer TP \u5b9e\u73b0\u4e2d\uff0c\u4e00\u4e2a Transformer Block \u901a\u5e38\u5305\u542b\u4e24\u6b21\u6b64\u7c7b\u540c\u6b65\uff1aAttention \u7684 O_Proj \u4e4b\u540e\u4e00\u6b21\uff0cMLP \u7684 Down_Proj \u4e4b\u540e\u4e00\u6b21\u3002\u6bcf\u4e2a Row-Parallel Linear \u672c\u8eab\u53ea\u9700\u8981\u4e00\u6b21 All-Reduce\u3002\u5177\u4f53\u800c\u8a00\uff0c\u5728\u884c\u5e76\u884c\u7ebf\u6027\u5c42\u8ba1\u7b97\u7ed3\u675f\u540e\uff0c\u6bcf\u5f20\u5361\u90fd\u4f1a\u4ea7\u751f\u4e00\u4e2a\u7ef4\u5ea6\u4e3a\u00a0[num_sched_tokens, hidden_size]\u00a0\u7684\u5c40\u90e8\u504f\u548c\uff08Partial Sum\uff09\u5f20\u91cf\u3002All-Reduce (Sum) \u64cd\u4f5c\u5c06\u8de8\u5361\u5bf9\u9f50\u89c4\u7ea6\u8fd9\u4e9b\u5f20\u91cf\uff0c\u64cd\u4f5c\u5b8c\u6210\u540e\uff0c\u96c6\u7fa4\u4e2d\u7684\u6bcf\u5f20\u5361\u90fd\u5c06\u62e5\u6709\u4e00\u4efd\u5b8c\u5168\u4e00\u81f4\u7684\u3001\u5168\u5c40\u6c42\u548c\u540e\u7684\u5b8c\u6574\u8f93\u51fa\u5f20\u91cf\uff08\u7ef4\u5ea6\u4fdd\u6301\u00a0[num_sched_tokens, hidden_size]\u00a0\u4e0d\u53d8\uff09\uff0c\u4ece\u800c\u4f7f\u5404\u5361\u53ef\u4ee5\u65e0\u901a\u4fe1\u5730\u72ec\u7acb\u6267\u884c\u540e\u7eed\u7684\u6b8b\u5dee\u8fde\u63a5\uff08Residual Add\uff09\u4e0e\u5f52\u4e00\u5316\uff08LayerNorm\uff09\u3002\u90a3\u4e48\u95ee\u9898\u6765\u4e86\uff0c\u8fd9\u91cc\u7684\u591a\u5361\u52a0\u6cd5\u7684\u987a\u5e8f\u662f\u6052\u5b9a\u7684\u5417\uff1f<\/p>\n<p>\u7b54\u6848\u662f\uff1a\u5728\u9ed8\u8ba4\u60c5\u51b5\u4e0b\uff0c\u4e0d\u4e00\u5b9a\u3002\u8fd9\u662f\u7531\u4e8eNCCL\u5e93\u5728\u541e\u5410\u548c\u5ef6\u8fdf\u95f4\u7684trade-off\u3002<\/p>\n<p>\u5bf9\u4e8e\u56fa\u5b9a\u6d88\u606f\u5927\u5c0f\uff0c\u5b83\u901a\u5e38\u662f\u7a33\u5b9a\u7684\uff1b\u4f46\u8de8\u4e0d\u540c Batch Size\uff0cNCCL \u5e76\u4e0d\u4fdd\u8bc1\u91c7\u7528\u76f8\u540c\u7684\u89c4\u7ea6\u7b97\u6cd5\u3001\u901a\u9053\u6570\u548c\u6570\u636e\u5207\u5206\u65b9\u5f0f\uff0c\u56e0\u6b64\u89c4\u7ea6\u987a\u5e8f\u4e0d\u4e00\u5b9a\u4fdd\u6301\u4e0d\u53d8\u3002<\/p>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/b82137e756dc0f5a022088528007449ef540a2.webp\" data-type=\"block\" \/><\/p>\n<p>\u4e1a\u754c\u7ecf\u5178\u7684\u00a0Ring All-Reduce\uff08\u00a0Reduce-Scatter\u00a0+\u00a0All-Gather\u00a0\uff09\u7684\u4f18\u52bf\u662f\u9ad8\u541e\u5410\uff0c\u4f46\u662f\u540c\u65f6\u5ef6\u8fdf\u4f1a\u968f\u8282\u70b9\u6570\u7ebf\u5f62\u589e\u957f\u3002\u4ece\u903b\u8f91\u4f9d\u8d56\u4e0a\u770b\uff0cRing All-Reduce \u9700\u8981 (N-1) \u4e2a Reduce-Scatter step \u548c (N-1) \u4e2a All-Gather step\u3002NCCL \u4f1a\u901a\u8fc7 chunk pipeline \u548c Multi-Channel \u5c06\u4e0d\u540c\u6570\u636e\u5757\u91cd\u53e0\u6267\u884c\uff0c\u4f46\u5355\u4e2a chunk \u7684\u4f9d\u8d56\u94fe\u957f\u5ea6\u4ecd\u968f rank \u6570\u8fd1\u4f3c\u7ebf\u6027\u589e\u957f\u3002\u56e0\u6b64\u5728\u5c0f\u6d88\u606f\u3001\u4f4e Batch \u7684 latency-bound \u573a\u666f\u4e0b\uff0cRing \u5f88\u96be\u5145\u5206\u53d1\u6325\u5176\u5e26\u5bbd\u4f18\u52bf\u3002<\/p>\n<p>&nbsp;<\/p>\n<p>\u503c\u5f97\u4e00\u63d0\u7684\u662f\uff0cBaidu Silicon Valley AI Lab \u8f83\u65e9\u5c06 HPC \u9886\u57df\u7684 Ring All-Reduce \u7cfb\u7edf\u6027\u5730\u5e94\u7528\u5e76\u63a8\u5e7f\u5230\u6df1\u5ea6\u5b66\u4e60\u5206\u5e03\u5f0f\u8bad\u7ec3\u4e2d\uff0c\u4f7f\u8fd9\u4e00\u7b97\u6cd5\u6210\u4e3a\u6570\u636e\u5e76\u884c\u8bad\u7ec3\u4e2d\u7684\u7ecf\u5178\u65b9\u6848\u3002NCCL \u4e5f\u957f\u671f\u5c06 Ring \u4f5c\u4e3a\u6838\u5fc3\u7b97\u6cd5\u4e4b\u4e00\u3002<\/p>\n<p>&nbsp;<\/p>\n<p><img decoding=\"async\" src=\"https:\/\/s2.51cto.com\/oss\/202608\/05\/f8ad12e970c9f6a079b71549764f2867768e64.webp\" data-type=\"block\" \/><\/p>\n<p>\u4e3a\u4e86\u5728\u4e0d\u540c\u8d1f\u8f7d\u4e0b\u90fd\u80fd\u538b\u69a8\u51fa\u6781\u9650\u6027\u80fd\uff0cNCCL \u5728 communicator \u521d\u59cb\u5316\u9636\u6bb5\u6839\u636e\u786c\u4ef6\u62d3\u6251\u6784\u9020\u591a\u79cd\u5019\u9009\u901a\u4fe1\u56fe\uff0c\u5e76\u5efa\u7acb latency\/bandwidth \u6027\u80fd\u6a21\u578b\uff1b\u5728\u6bcf\u6b21 collective\u5165\u961f\u65f6\uff0c\u518d\u6839\u636e\u6d88\u606f\u5927\u5c0f\u3001rank \u6570\u3001\u8282\u70b9\u6570\u548c\u53ef\u7528\u4f20\u8f93\u8def\u5f84\u9009\u62e9\u9884\u8ba1\u4ee3\u4ef7\u6700\u4f4e\u7684\u7b97\u6cd5\u3001\u534f\u8bae\u53ca channel \u914d\u7f6e\u3002<\/p>\n<p>\u4e0a\u9762\u5df2\u7ecf\u63d0\u5230\uff0c\u5728 LLM \u63a8\u7406\u7684\u52a8\u6001\u7ec4\u6279\uff08Dynamic Batching\uff09\u8fc7\u7a0b\u4e2d\uff0cAll-Reduce \u7684 Message Size \u662f\u6309\u00a0[num_tokens, hidden_size]\u00a0\u8ba1\u7b97\u7684\u3002\u8fd9\u610f\u5473\u7740\uff0cBatch Size \u7684\u6ce2\u52a8\uff0c\u4f1a\u76f4\u63a5\u5bfc\u81f4\u5355\u6b21\u901a\u4fe1\u8d1f\u8f7d\u7684\u5927\u5c0f\u53d1\u751f\u53d8\u5316\u3002\u5f53\u901a\u4fe1\u8f7d\u8377\u8de8\u8d8a\u67d0\u4e2a\u9608\u503c\u65f6\uff0cNCCL \u4f1a\u5728\u5e95\u5c42\u5207\u6362\u89c4\u7ea6\u62d3\u6251\uff1a<\/p>\n<ul data-id=\"u738a58b-n7qJuz91\">\n<li data-id=\"ld70c578-du5iab9e\">Ring \u7b97\u6cd5\uff1a \u91c7\u7528\u73af\u5f62\u4f20\u9012\uff0c\u52a0\u6cd5\u987a\u5e8f\u662f\u7ebf\u6027\u7684\u3002\u4f8b\u5982\u00a0GPU0\u00a0\u2192 GPU1\u00a0\u2192 GPU2\u00a0\u2192 GPU3\u3002<\/li>\n<li data-id=\"ld70c578-haqi9HAy\">Tree \/ Double Binary Tree \u7b97\u6cd5\uff1a \u91c7\u7528\u6811\u72b6\u5c42\u7ea7\u89c4\u7ea6\uff0c\u52a0\u6cd5\u987a\u5e8f\u662f\u5206\u6cbb\u7684\u3002\u4f8b\u5982\u00a0(GPU0\u00a0+ GPU1) + (GPU2\u00a0+ GPU3)\u3002<\/li>\n<\/ul>\n<p>\u4f46\u8fd9\u8fd8\u4e0d\u662f\u5168\u90e8\u3002\u9664\u4e86\u5b8f\u89c2\u7b97\u6cd5\u7684\u5207\u6362\uff0cNCCL\u5728\u5fae\u89c2\u6267\u884c\u94fe\u8def\u4e0a\u7684\u591a\u901a\u9053\u5e76\u53d1\uff08Multi-Channel\uff09\u5207\u5206\u673a\u5236\u4e5f\u4f1a\u5bfc\u81f4Batch Variance\u3002\u73b0\u4ee3 GPU \u8282\u70b9\u5185\u901a\u5e38\u5177\u6709\u591a\u6761\u7269\u7406\u4e92\u8054\u94fe\u8def\uff08\u5982\u591a\u6761 NVLink \u94fe\u8def\uff09\u3002\u4e3a\u4e86\u538b\u69a8\u53cc\u5411\u7269\u7406\u5e26\u5bbd\uff0cNCCL \u5f15\u5165\u4e86 Channel\u7684\u62bd\u8c61\u3002NCCL \u7684 Channel \u662f\u4e00\u6761\u72ec\u7acb\u7684\u903b\u8f91\u901a\u4fe1\u6d41\u6c34\u7ebf\u3002\u5728 GPU kernel \u4e2d\uff0c\u6bcf\u4e2a\u6d3b\u8dc3 Channel \u901a\u5e38\u7531\u4e00\u4e2a CUDA Thread Block \u8d1f\u8d23\uff1b\u6bcf\u4e2a Channel \u62e5\u6709\u81ea\u5df1\u7684 Ring\/Tree \u8fde\u63a5\u5173\u7cfb\u548cchunk \u6d41\u6c34\u7ebf\u3002\u5b83\u4e0d\u4e00\u5b9a\u4e0e\u67d0\u4e00\u6761\u7269\u7406 NVLink \u4e00\u4e00\u5bf9\u5e94\uff0c\u4f46\u591a\u4e2a Channel \u53ef\u4ee5\u5e2e\u52a9\u5e76\u884c\u5229\u7528\u591a\u6761\u53ef\u7528\u94fe\u8def\u3002<\/p>\n<p>\u5f53\u89e6\u53d1\u4e00\u4e2a\u8f83\u5927\u7684 All-Reduce \u4efb\u52a1\u65f6\uff0cNCCL \u4f1a\u5c06\u603b\u6570\u636e\u5757\u8fdb\u884c\u5207\u7247\uff0c\u5206\u53d1\u7ed9\u591a\u4e2a Channel \u5e76\u53d1\u6267\u884c\u89c4\u7ea6\u3002\u8fd9\u4e00\u673a\u5236\u5f15\u53d1 Batch Variance \u7684\u539f\u56e0\u5728\u4e8e\u4ee5\u4e0b\u4e24\u70b9\uff1a<\/p>\n<p>\u2460 \u62d3\u6251\u5f02\u6784\uff1a<\/p>\n<p>\u4e3a\u4e86\u5b9e\u73b0\u53cc\u5411\u94fe\u8def\u541e\u5410\u7684\u6700\u5927\u5316\uff0cNCCL \u4f1a\u4e3a\u4e0d\u540c\u7684 Channel \u5206\u914d\u7ed3\u6784\u6216\u65b9\u5411\u5b8c\u5168\u4e0d\u540c\u7684\u89c4\u7ea6\u62d3\u6251\u3002NCCL \u53ef\u4ee5\u4e3a\u4e0d\u540c Channel \u6784\u9020\u4e0d\u540c\u7684 Ring \u6216 Tree\u3002<\/p>\n<p>\u4f8b\u5982\uff0c\u5728\u67d0\u4e9b\u62d3\u6251\u4e0a\uff0c\u5206\u914d Channel 0 \u6267\u884c\u6b63\u5411 Ring\uff0c\u52a0\u6cd5\u987a\u5e8f\u4e3a\u00a0(((GPU0\u00a0+ GPU1) + GPU2) + GPU3)\uff1b\u540c\u65f6\u5206\u914d Channel 1 \u6267\u884c\u9006\u5411 Ring\uff0c\u52a0\u6cd5\u987a\u5e8f\u5219\u53d8\u4e3a\u00a0(((GPU0\u00a0+ GPU3) + GPU2) + GPU1)\u3002<\/p>\n<p>\u2461 \u6570\u636e\u5207\u5206\u8fb9\u754c\u7684\u52a8\u6001\u6ed1\u52a8\uff1a<\/p>\n<p>\u5728 LLM \u63a8\u7406\u7684\u52a8\u6001\u7ec4\u6279\u4e2d\uff0cBatch Size \u7684\u6ce2\u52a8\u4f1a\u5bfc\u81f4 All-Reduce \u7684\u603b Message Size \u53d1\u751f\u6539\u53d8\u3002\u4e00\u65e6\u603b\u6570\u636e\u91cf\u53d8\u5316\uff0cNCCL \u5185\u90e8\u7684\u542f\u53d1\u5f0f\u7b56\u7565\u5c31\u4f1a\u91cd\u65b0\u8ba1\u7b97\u5206\u914d\u7ed9\u5404\u4e2a Channel \u7684\u6570\u636e\u5207\u5206\u8fb9\u754c\u3002\u8fd9\u5c31\u5bfc\u81f4\u4e86\uff1a\u539f\u672c\u5728 Batch Size =\u00a0N\u00a0\u65f6\u88ab\u5206\u914d\u5230 Channel 0 \u8d1f\u8d23\u89c4\u7ea6\u7684\u67d0\u4e00\u6bb5\u5185\u5b58\u6570\u636e\uff0c\u5728 Batch Size =\u00a0N+1\u00a0\u65f6\uff0c\u53ef\u80fd\u56e0\u4e3a\u5207\u5206\u8fb9\u754c\u7684\u5fae\u8c03\u88ab\u5212\u5b9a\u5230\u4e86 Channel 1 \u7684\u8d1f\u8d23\u533a\u57df\u3002<\/p>\n<div>\n<div class=\"hljs-cto\">\n<div class=\"hljs-cto\"><button class=\"copy_btn disable\" data-clipboard-target=\"#code_id_19\">\u590d\u5236<\/button><\/p>\n<div class=\"code-toolbar\">\n<pre class=\"has-pre-numbering language-javascript\" tabindex=\"0\"><code class=\"language-javascript\"># vllm<span class=\"token operator\">\/<\/span>config<span class=\"token operator\">\/<\/span>parallel<span class=\"token punctuation\">.<\/span>py\r\n  <span class=\"token keyword\">if<\/span> envs<span class=\"token punctuation\">.<\/span><span class=\"token constant\">VLLM_BATCH_INVARIANT<\/span><span class=\"token operator\">:<\/span>\r\n      self<span class=\"token punctuation\">.<\/span>disable_custom_all_reduce <span class=\"token operator\">=<\/span> True<\/code><\/pre>\n<ul id=\"code_id_19\" class=\"pre-numbering\">\n<li>1.<\/li>\n<li>2.<\/li>\n<li>3.<\/li>\n<\/ul>\n<div class=\"toolbar\"><\/div>\n<\/div>\n<\/div>\n<\/div>\n<\/div>\n<p>vLLM \u7684 Custom All-Reduce \u662f\u8282\u70b9\u5185 TP \u7684\u4f4e\u5ef6\u8fdf\u6027\u80fd\u8def\u5f84\uff0c\u4f46\u5b83\u5185\u90e8\u540c\u6837\u5305\u542b\u968f\u6d88\u606f\u5927\u5c0f\u5207\u6362\u7684 1-stage\/2-stage \u7b97\u6cd5\uff0c\u56e0\u6b64\u4e0d\u80fd\u88ab\u89c6\u4e3a\u4e25\u683c Batch Invariant\u3002\u5f00\u542fVLLM_BATCH_INVARIANT=1 \u540e\uff0cvLLM \u9996\u5148\u7981\u7528 Custom All-Reduce \u548c symmetric-memory \u8def\u5f84\uff0c\u518d\u5c06 NCCL All-Reduce \u56fa\u5b9a\u5230 Tree + Simple\uff0c\u5e76\u5c06 Channel \u6570\u9501\u5b9a\u4e3a 1\uff0c\u540c\u65f6\u5173\u95ed NVLS\u3001CollNet \u7b49\u53ef\u80fd\u5f15\u5165\u5176\u4ed6\u89c4\u7ea6\u62d3\u6251\u7684\u901a\u4fe1\u8def\u5f84\u3002<\/p>\n<p>vLLM \u7684 Custom All-Reduce \u662f\u9ed8\u8ba4\u6027\u80fd\u8def\u5f84\u4e4b\u4e00\uff0c\u4f46\u5b83\u5e76\u4e0d\u662f\u5f53\u524d\u4e25\u683c Batch Invariance \u6a21\u5f0f\u7684\u89e3\u51b3\u65b9\u6848\u3002\u5b83\u5185\u90e8\u4e5f\u5305\u542b 1-stage \u548c 2-stage \u4e24\u5957\u7b97\u6cd5\uff1aTP4 \u4e0b\u901a\u5e38\u4ee5 512 KiB \u4e3a\u5207\u6362\u70b9\uff0cTP6\/TP8 \u4e0b\u901a\u5e38\u4ee5256 KiB \u4e3a\u5207\u6362\u70b9\u30021-stage \u4f1a\u5728\u6bcf\u4e2a rank \u4e0a\u6309\u7167\u56fa\u5b9a\u6307\u9488\u987a\u5e8f\u5b8c\u6574\u89c4\u7ea6\u6240\u6709\u5143\u7d20\uff1b2-stage \u5219\u5148\u7531\u4e0d\u540c rank \u8d1f\u8d23\u4e0d\u540c\u8f93\u51fa\u5206\u533a\uff0c\u518d\u6267\u884c All-Gather\uff0c\u800c\u4e14\u6bcf\u4e2a\u5206\u533a\u7684\u7d2f\u52a0\u987a\u5e8f\u4ee5\u5176 owner rank \u4e3a\u8d77\u70b9\u5faa\u73af\u5c55\u5f00\u3002\u5f53\u6d88\u606f\u5927\u5c0f\u6539\u53d8\u4e86\u7b97\u6cd5\u9009\u62e9\u6216\u5206\u533a owner \u65f6\uff0c\u540c\u4e00\u903b\u8f91\u5143\u7d20\u53ef\u80fd\u7ecf\u5386\u4e0d\u540c\u7684\u7d2f\u52a0\u987a\u5e8f\u3002\u56e0\u6b64 vLLM \u7684Batch Invariant \u6a21\u5f0f\u4f1a\u76f4\u63a5\u7981\u7528Custom All-Reduce\u3002<\/p>\n<p>Custom All-Reduce\u662fvLLM\u9488\u5bf9\u8282\u70b9\u5185 Tensor Parallel \u9ad8\u9891\u3001\u5c0f\u4e2d\u6d88\u606f All-Reduce \u6240\u589e\u52a0\u7684\u4e00\u6761\u4f4e\u5ef6\u8fdf\u5feb\u8def\u5f84\u3002\u5c06\u5355\u8282\u70b9 TP All-Reduce \u7b80\u5316\u6210\u4e00\u4e2a\u76f4\u63a5\u8bbf\u95ee\u6240\u6709 peer GPU \u663e\u5b58\u7684 CUDA kernel\u3002\u901a\u8fc7 CUDA IPC \u63d0\u524d\u83b7\u5f97 peer pointer\uff0c\u4f7f\u7528\u8f7b\u91cf\u7ea7 GPU barrier\u3001128-bit \u5411\u91cf\u5316\u8bbf\u5b58\u4ee5\u53ca1-stage\/2-stage \u4e13\u7528\u7b97\u6cd5\uff0c\u51cf\u5c11 NCCL \u901a\u7528\u8def\u5f84\u5728\u5c0f\u4e2d\u578b\u6d88\u606f\u4e0a\u7684\u56fa\u5b9a\u5ef6\u8fdf\u3002<\/p>\n<div>\n<div class=\"hljs-cto\">\n<div class=\"hljs-cto\"><button class=\"copy_btn disable\" data-clipboard-target=\"#code_id_20\">\u590d\u5236<\/button><\/p>\n<div class=\"code-toolbar\">\n<pre class=\"has-pre-numbering language-javascript\" tabindex=\"0\"><code class=\"language-javascript\">def <span class=\"token function\">override_envs_for_invariance<\/span><span class=\"token punctuation\">(<\/span><span class=\"token punctuation\">)<\/span><span class=\"token operator\">:<\/span>\r\n    os<span class=\"token punctuation\">.<\/span>environ<span class=\"token punctuation\">[<\/span><span class=\"token string\">\"VLLM_ALLREDUCE_USE_SYMM_MEM\"<\/span><span class=\"token punctuation\">]<\/span> <span class=\"token operator\">=<\/span> <span class=\"token string\">\"0\"<\/span> # \u5173\u95ed <span class=\"token constant\">NCCL<\/span><span class=\"token operator\">\/<\/span>Torch symmetric<span class=\"token operator\">-<\/span>memory \u5feb\u8def\u5f84\r\n\r\n    os<span class=\"token punctuation\">.<\/span>environ<span class=\"token punctuation\">[<\/span><span class=\"token string\">\"CUBLAS_WORKSPACE_CONFIG\"<\/span><span class=\"token punctuation\">]<\/span> <span class=\"token operator\">=<\/span> <span class=\"token string\">\":4096:8\"<\/span>\r\n\r\n    # <span class=\"token constant\">NCCL<\/span> determinism settings\r\n    os<span class=\"token punctuation\">.<\/span>environ<span class=\"token punctuation\">[<\/span><span class=\"token string\">\"NCCL_LAUNCH_MODE\"<\/span><span class=\"token punctuation\">]<\/span> <span class=\"token operator\">=<\/span> <span class=\"token string\">\"GROUP\"<\/span>\r\n    os<span class=\"token punctuation\">.<\/span>environ<span class=\"token punctuation\">[<\/span><span class=\"token string\">\"NCCL_COLLNET_ENABLE\"<\/span><span class=\"token punctuation\">]<\/span> <span class=\"token operator\">=<\/span> <span class=\"token string\">\"0\"<\/span>\r\n    os<span class=\"token punctuation\">.<\/span>environ<span class=\"token punctuation\">[<\/span><span class=\"token string\">\"NCCL_NVLS_ENABLE\"<\/span><span class=\"token punctuation\">]<\/span> <span class=\"token operator\">=<\/span> <span class=\"token string\">\"0\"<\/span>\r\n    os<span class=\"token punctuation\">.<\/span>environ<span class=\"token punctuation\">[<\/span><span class=\"token string\">\"NCCL_P2P_NET_DISABLE\"<\/span><span class=\"token punctuation\">]<\/span> <span class=\"token operator\">=<\/span> <span class=\"token string\">\"1\"<\/span>\r\n    os<span class=\"token punctuation\">.<\/span>environ<span class=\"token punctuation\">[<\/span><span class=\"token string\">\"NCCL_MIN_NCHANNELS\"<\/span><span class=\"token punctuation\">]<\/span> <span class=\"token operator\">=<\/span> <span class=\"token string\">\"1\"<\/span>\r\n    os<span class=\"token punctuation\">.<\/span>environ<span class=\"token punctuation\">[<\/span><span class=\"token string\">\"NCCL_MAX_NCHANNELS\"<\/span><span class=\"token punctuation\">]<\/span> <span class=\"token operator\">=<\/span> <span class=\"token string\">\"1\"<\/span>      # \u5f3a\u5236\u5355\u901a\u9053\u901a\u8baf\r\n    os<span class=\"token punctuation\">.<\/span>environ<span class=\"token punctuation\">[<\/span><span class=\"token string\">\"NCCL_PROTO\"<\/span><span class=\"token punctuation\">]<\/span> <span class=\"token operator\">=<\/span> <span class=\"token string\">\"Simple\"<\/span>\r\n    os<span class=\"token punctuation\">.<\/span>environ<span class=\"token punctuation\">[<\/span><span class=\"token string\">\"NCCL_ALGO\"<\/span><span class=\"token punctuation\">]<\/span> <span class=\"token operator\">=<\/span> <span class=\"token string\">\"allreduce:tree\"<\/span>  # \u5f3a\u5236\u4f7f\u7528tree\u7b97\u6cd5 <span class=\"token operator\">+<\/span> simple\u534f\u8bae\r\n    os<span class=\"token punctuation\">.<\/span>environ<span class=\"token punctuation\">[<\/span><span class=\"token string\">\"NCCL_NTHREADS\"<\/span><span class=\"token punctuation\">]<\/span> <span class=\"token operator\">=<\/span> <span class=\"token string\">\"1\"<\/span>\r\n    os<span class=\"token punctuation\">.<\/span>environ<span class=\"token punctuation\">[<\/span><span class=\"token string\">\"NCCL_SOCKET_NTHREADS\"<\/span><span class=\"token punctuation\">]<\/span> <span class=\"token operator\">=<\/span> <span class=\"token string\">\"1\"<\/span>\r\n\r\n    # torch<span class=\"token punctuation\">.<\/span>compile settings\r\n    os<span class=\"token punctuation\">.<\/span>environ<span class=\"token punctuation\">[<\/span><span class=\"token string\">\"VLLM_USE_AOT_COMPILE\"<\/span><span class=\"token punctuation\">]<\/span> <span class=\"token operator\">=<\/span> <span class=\"token string\">\"0\"<\/span><\/code><\/pre>\n<ul id=\"code_id_20\" class=\"pre-numbering\">\n<li>1.<\/li>\n<li>2.<\/li>\n<li>3.<\/li>\n<li>4.<\/li>\n<li>5.<\/li>\n<li>6.<\/li>\n<li>7.<\/li>\n<li>8.<\/li>\n<li>9.<\/li>\n<li>10.<\/li>\n<li>11.<\/li>\n<li>12.<\/li>\n<li>13.<\/li>\n<li>14.<\/li>\n<li>15.<\/li>\n<li>16.<\/li>\n<li>17.<\/li>\n<li>18.<\/li>\n<li>19.<\/li>\n<\/ul>\n<div class=\"toolbar\"><\/div>\n<\/div>\n<\/div>\n<\/div>\n<\/div>\n<p>VLLM_ALLREDUCE_USE_SYMM_MEM=0 \u5173\u95ed\u7684\u662f vLLM \u57fa\u4e8e PyTorch Symmetric Memory \u5b9e\u73b0\u7684 multimem\/two-shot All-Reduce \u540e\u7aef\uff0c\u907f\u514d\u8bf7\u6c42\u6839\u636e\u786c\u4ef6\u80fd\u529b\u8fdb\u5165\u53e6\u4e00\u5957\u89c4\u7ea6 kernel\u3002<\/p>\n<p>NCCL_COLLNET_ENABLE=0 \u7981\u6b62NCCL \u4f7f\u7528CollNetDirect\/CollNetChain \u7b49\u7f51\u7edc\u4fa7collective-offload \u7b97\u6cd5\uff1bNCCL_NVLS_ENABLE=0 \u5219\u7981\u6b62 Hopper\/Blackwell NVSwitch \u7cfb\u7edf\u4f7f\u7528NVLink SHARP multicast\/reduction \u8def\u5f84\u3002\u4e8c\u8005\u90fd\u662f\u4e3a\u4e86\u5c06\u5019\u9009\u89c4\u7ea6\u62d3\u6251\u6536\u7f29\u5230\u666e\u901a NCCL Tree\u3002\u5373\u4e0d\u5141\u8bb8\u7f51\u5361\/\u7f51\u7edc collective plugin\/NVSwitch \u53c2\u4e0e\u89c4\u7ea6\uff0c\u907f\u514d All-Reduce \u8d70 CollNet \u6216\u7f51\u7edc\u4fa7\u805a\u5408\u8def\u5f84\u3002<\/p>\n<p>\u8fd9\u4e0e\u5355\u5361\u7b97\u5b50\u7684\u89e3\u51b3\u65b9\u5f0f\u672c\u8d28\u4e00\u81f4\uff1aGEMM \u9700\u8981\u56fa\u5b9a Split-K \u548c K \u8f74\u5206\u5757\uff0cAttention \u9700\u8981\u56fa\u5b9a Split-KV\uff0c\u800c\u5206\u5e03\u5f0f All-Reduce \u5219\u9700\u8981\u56fa\u5b9a\u8de8 rank \u7684\u89c4\u7ea6\u56fe\u3001\u534f\u8bae\u548c\u6570\u636e\u5207\u5206\u65b9\u5f0f\u3002\u5b83\u4eec\u89e3\u51b3\u7684\u90fd\u662f\u540c\u4e00\u4e2a\u95ee\u9898\u2014\u2014\u963b\u6b62\u7cfb\u7edf\u6839\u636e Batch Shape \u52a8\u6001\u6539\u53d8 Reduction Topology\u3002<\/p>\n<h3>\u516d\u3001\u603b\u7ed3<\/h3>\n<p>\u4ee5\u4e0a\uff0c\u5927\u6a21\u578b\u5728\u63a8\u7406\u8fc7\u7a0b\u4e2d\u5bfc\u81f4Batch Variance\u7684\u6839\u672c\u539f\u56e0\uff0c\u5728\u4e8e\u5e95\u5c42\u7b97\u5b50\u4e3a\u6700\u5927\u5316\u786c\u4ef6\u8d44\u6e90\u5229\u7528\u7387\uff0c\u52a8\u6001\u8c03\u6574\u4e86\u89c4\u7ea6\u8f74\uff08\u5982 GEMM \u4e2d\u7684 Split-K \u6216 Attention \u4e2d\u7684 KV \u8f74\uff09\u7684\u5207\u5206\u7b56\u7565\uff0c\u8fdb\u800c\u6539\u53d8\u4e86\u6d6e\u70b9\u6570\u7d2f\u52a0\u6811\u7684\u62d3\u6251\u7ed3\u6784\u3002\u82e5\u8981\u5b9eBatch Invariance\uff0c\u5c31\u5fc5\u987b\u5728\u7b97\u5b50\u8c03\u5ea6\u5c42\u9762\u7ea6\u675f\u6b64\u7c7b\u52a8\u6001\u5207\u5206\u884c\u4e3a\u3002<\/p>\n<p>PS\uff1a\u53d7\u9650\u4e8e\u7bc7\u5e45\u540c\u65f6\u4fdd\u8bc1\u6587\u7ae0\u8d28\u91cf\uff0c\u672c\u6587\u8fd8\u67092\u4e2a\u5730\u65b9\u6ca1\u6709\u5c55\u5f00\u8bb2\uff1a\u4e00\u4e2a\u662f MoE\u67b6\u6784\u4e0b\u7279\u6709\u7684 Batch Invariance \u95ee\u9898\uff1b\u53e6\u4e00\u4e2a\u662f NCCL \u66f4\u6df1\u5c42\u7684\u901a\u4fe1\u62d3\u6251\u4e0e\u5e95\u5c42\u673a\u5236\u3002\u8fd9\u4e9b\u5185\u5bb9\u540e\u7eed\u518d\u5355\u5f00\u4e00\u7bc7\u586b\u5751\u3002<\/p>\n<p>\u7136\u800c\uff0c\u5de5\u7a0b\u5b9e\u73b0\u5374\u590d\u6742\u7684\u591a\u3002\u60f3\u8981\u505a\u5230\u77e5\u5176\u6240\u4ee5\u7136\uff0c\u6211\u4eec\u9700\u8981\u5411\u4e0b\u6df1\u5165 GPU \u7269\u7406\u5fae\u67b6\u6784\u4e0e\u7b97\u5b50\u6267\u884c\u8303\u5f0f\uff08Triton\/CUDA\uff09\uff0c\u5411\u4e0a\u89e3\u6790\u6846\u67b6\u5c42\u7684\u8ba1\u7b97\u56fe\u7f16\u8bd1\u903b\u8f91\uff08Inductor\/IR\uff09\u3002\u6240\u8c13\u7684 AI Infra\uff0c\u672c\u8d28\u4e0a\u662f\u5728\u8fd9\u591a\u5c42\u6280\u672f\u62bd\u8c61\u4e2d\uff0c\u5bfb\u6c42\u7cfb\u7edf\u6027\u80fd\u4e0e\u6570\u5b66\u7b49\u4ef7\u6027\u4e4b\u95f4\u7684\u6700\u4f18 Trade-off\u3002<\/p>\n<p>\u6587\u7ae0\u6765\u81ea\uff1a51CTO<\/p>\n<\/div>\n<\/div>\n<div class=\"pvc_clear\"><\/div>\n<p id=\"pvc_stats_29992\" class=\"pvc_stats total_only  \" data-element-id=\"29992\" style=\"\"><i class=\"pvc-stats-icon medium\" aria-hidden=\"true\"><svg xmlns=\"http:\/\/www.w3.org\/2000\/svg\" version=\"1.0\" viewBox=\"0 0 502 315\" preserveAspectRatio=\"xMidYMid meet\"><g transform=\"translate(0,332) scale(0.1,-0.1)\" fill=\"\" stroke=\"none\"><path d=\"M2394 3279 l-29 -30 -3 -207 c-2 -182 0 -211 15 -242 39 -76 157 -76 196 0 15 31 17 60 15 243 l-3 209 -33 29 c-26 23 -41 29 -80 29 -41 0 -53 -5 -78 -31z\"\/><path d=\"M3085 3251 c-45 -19 -58 -50 -96 -229 -47 -217 -49 -260 -13 -295 52 -53 146 -42 177 20 16 31 87 366 87 410 0 70 -86 122 -155 94z\"\/><path d=\"M1751 3234 c-13 -9 -29 -31 -37 -50 -12 -29 -10 -49 21 -204 19 -94 39 -189 45 -210 14 -50 54 -80 110 -80 34 0 48 6 76 34 21 21 34 44 34 59 0 14 -18 113 -40 219 -37 178 -43 195 -70 221 -36 32 -101 37 -139 11z\"\/><path d=\"M1163 3073 c-36 -7 -73 -59 -73 -102 0 -56 133 -378 171 -413 34 -32 83 -37 129 -13 70 36 67 87 -16 290 -86 209 -89 214 -129 231 -35 14 -42 15 -82 7z\"\/><path d=\"M3689 3066 c-15 -9 -33 -30 -42 -48 -48 -103 -147 -355 -147 -375 0 -98 131 -148 192 -74 13 15 57 108 97 206 80 196 84 226 37 273 -30 30 -99 39 -137 18z\"\/><path d=\"M583 2784 c-38 -19 -67 -74 -58 -113 9 -42 211 -354 242 -373 16 -10 45 -18 66 -18 51 0 107 52 107 100 0 39 -1 41 -124 234 -80 126 -108 162 -133 173 -41 17 -61 16 -100 -3z\"\/><path d=\"M4250 2784 c-14 -9 -74 -91 -133 -183 -95 -150 -107 -173 -107 -213 0 -55 33 -94 87 -104 67 -13 90 8 211 198 130 202 137 225 78 284 -27 27 -42 34 -72 34 -22 0 -50 -8 -64 -16z\"\/><path d=\"M2275 2693 c-553 -48 -1095 -270 -1585 -649 -135 -104 -459 -423 -483 -476 -23 -49 -22 -139 2 -186 73 -142 361 -457 571 -626 285 -228 642 -407 990 -497 242 -63 336 -73 660 -74 310 0 370 5 595 52 535 111 1045 392 1455 803 122 121 250 273 275 326 19 41 19 137 0 174 -41 79 -309 363 -465 492 -447 370 -946 591 -1479 653 -113 14 -422 18 -536 8z m395 -428 c171 -34 330 -124 456 -258 112 -119 167 -219 211 -378 27 -96 24 -300 -5 -401 -72 -255 -236 -447 -474 -557 -132 -62 -201 -76 -368 -76 -167 0 -236 14 -368 76 -213 98 -373 271 -451 485 -162 444 86 934 547 1084 153 49 292 57 452 25z m909 -232 c222 -123 408 -262 593 -441 76 -74 138 -139 138 -144 0 -16 -233 -242 -330 -319 -155 -123 -309 -223 -461 -299 l-81 -41 32 46 c18 26 49 83 70 128 143 306 141 649 -6 957 -25 52 -61 116 -79 142 l-34 47 45 -20 c26 -10 76 -36 113 -56z m-2057 25 c-40 -58 -105 -190 -130 -263 -110 -324 -59 -707 132 -981 25 -35 42 -64 37 -64 -19 0 -241 119 -326 174 -188 122 -406 314 -532 468 l-58 71 108 103 c185 178 428 349 672 473 66 33 121 60 123 61 2 0 -10 -19 -26 -42z\"\/><path d=\"M2375 1950 c-198 -44 -350 -190 -395 -379 -18 -76 -8 -221 19 -290 114 -284 457 -406 731 -260 98 52 188 154 231 260 27 69 37 214 19 290 -38 163 -166 304 -326 360 -67 23 -215 33 -279 19z\"\/><\/g><\/svg><\/i> <img loading=\"lazy\" decoding=\"async\" width=\"16\" height=\"16\" alt=\"Loading\" src=\"https:\/\/aif.amtbbs.org\/wp-content\/plugins\/page-views-count\/ajax-loader-2x.gif\" border=0 \/><\/p>\n<div class=\"pvc_clear\"><\/div>\n","protected":false},"excerpt":{"rendered":"<p>\u5927\u6a21\u578b\u5728\u63a8\u7406\u8fc7\u7a0b\u4e2d\u5bfc\u81f4Batch Variance\u7684\u6839\u672c\u539f\u56e0\uff0c\u5728\u4e8e\u5e95\u5c42\u7b97\u5b50\u4e3a\u6700\u5927\u5316\u786c\u4ef6\u8d44\u6e90\u5229\u7528\u7387\uff0c\u52a8\u6001\u8c03\u6574\u4e86\u89c4 [&hellip;]<\/p>\n<div class=\"pvc_clear\"><\/div>\n<p id=\"pvc_stats_29992\" class=\"pvc_stats total_only  \" data-element-id=\"29992\" style=\"\"><i class=\"pvc-stats-icon medium\" aria-hidden=\"true\"><svg xmlns=\"http:\/\/www.w3.org\/2000\/svg\" version=\"1.0\" viewBox=\"0 0 502 315\" preserveAspectRatio=\"xMidYMid meet\"><g transform=\"translate(0,332) scale(0.1,-0.1)\" fill=\"\" stroke=\"none\"><path d=\"M2394 3279 l-29 -30 -3 -207 c-2 -182 0 -211 15 -242 39 -76 157 -76 196 0 15 31 17 60 15 243 l-3 209 -33 29 c-26 23 -41 29 -80 29 -41 0 -53 -5 -78 -31z\"\/><path d=\"M3085 3251 c-45 -19 -58 -50 -96 -229 -47 -217 -49 -260 -13 -295 52 -53 146 -42 177 20 16 31 87 366 87 410 0 70 -86 122 -155 94z\"\/><path d=\"M1751 3234 c-13 -9 -29 -31 -37 -50 -12 -29 -10 -49 21 -204 19 -94 39 -189 45 -210 14 -50 54 -80 110 -80 34 0 48 6 76 34 21 21 34 44 34 59 0 14 -18 113 -40 219 -37 178 -43 195 -70 221 -36 32 -101 37 -139 11z\"\/><path d=\"M1163 3073 c-36 -7 -73 -59 -73 -102 0 -56 133 -378 171 -413 34 -32 83 -37 129 -13 70 36 67 87 -16 290 -86 209 -89 214 -129 231 -35 14 -42 15 -82 7z\"\/><path d=\"M3689 3066 c-15 -9 -33 -30 -42 -48 -48 -103 -147 -355 -147 -375 0 -98 131 -148 192 -74 13 15 57 108 97 206 80 196 84 226 37 273 -30 30 -99 39 -137 18z\"\/><path d=\"M583 2784 c-38 -19 -67 -74 -58 -113 9 -42 211 -354 242 -373 16 -10 45 -18 66 -18 51 0 107 52 107 100 0 39 -1 41 -124 234 -80 126 -108 162 -133 173 -41 17 -61 16 -100 -3z\"\/><path d=\"M4250 2784 c-14 -9 -74 -91 -133 -183 -95 -150 -107 -173 -107 -213 0 -55 33 -94 87 -104 67 -13 90 8 211 198 130 202 137 225 78 284 -27 27 -42 34 -72 34 -22 0 -50 -8 -64 -16z\"\/><path d=\"M2275 2693 c-553 -48 -1095 -270 -1585 -649 -135 -104 -459 -423 -483 -476 -23 -49 -22 -139 2 -186 73 -142 361 -457 571 -626 285 -228 642 -407 990 -497 242 -63 336 -73 660 -74 310 0 370 5 595 52 535 111 1045 392 1455 803 122 121 250 273 275 326 19 41 19 137 0 174 -41 79 -309 363 -465 492 -447 370 -946 591 -1479 653 -113 14 -422 18 -536 8z m395 -428 c171 -34 330 -124 456 -258 112 -119 167 -219 211 -378 27 -96 24 -300 -5 -401 -72 -255 -236 -447 -474 -557 -132 -62 -201 -76 -368 -76 -167 0 -236 14 -368 76 -213 98 -373 271 -451 485 -162 444 86 934 547 1084 153 49 292 57 452 25z m909 -232 c222 -123 408 -262 593 -441 76 -74 138 -139 138 -144 0 -16 -233 -242 -330 -319 -155 -123 -309 -223 -461 -299 l-81 -41 32 46 c18 26 49 83 70 128 143 306 141 649 -6 957 -25 52 -61 116 -79 142 l-34 47 45 -20 c26 -10 76 -36 113 -56z m-2057 25 c-40 -58 -105 -190 -130 -263 -110 -324 -59 -707 132 -981 25 -35 42 -64 37 -64 -19 0 -241 119 -326 174 -188 122 -406 314 -532 468 l-58 71 108 103 c185 178 428 349 672 473 66 33 121 60 123 61 2 0 -10 -19 -26 -42z\"\/><path d=\"M2375 1950 c-198 -44 -350 -190 -395 -379 -18 -76 -8 -221 19 -290 114 -284 457 -406 731 -260 98 52 188 154 231 260 27 69 37 214 19 290 -38 163 -166 304 -326 360 -67 23 -215 33 -279 19z\"\/><\/g><\/svg><\/i> <img loading=\"lazy\" decoding=\"async\" width=\"16\" height=\"16\" alt=\"Loading\" src=\"https:\/\/aif.amtbbs.org\/wp-content\/plugins\/page-views-count\/ajax-loader-2x.gif\" border=0 \/><\/p>\n<div class=\"pvc_clear\"><\/div>\n","protected":false},"author":56,"featured_media":29993,"comment_status":"open","ping_status":"open","sticky":false,"template":"","format":"standard","meta":{"_bbp_topic_count":0,"_bbp_reply_count":0,"_bbp_total_topic_count":0,"_bbp_total_reply_count":0,"_bbp_voice_count":0,"_bbp_anonymous_reply_count":0,"_bbp_topic_count_hidden":0,"_bbp_reply_count_hidden":0,"_bbp_forum_subforum_count":0,"footnotes":""},"categories":[3,23,20,80],"tags":[2266],"class_list":["post-29992","post","type-post","status-publish","format-standard","has-post-thumbnail","hentry","category-ai","category-23","category-20","category-80","tag-ai-infra-"],"_links":{"self":[{"href":"https:\/\/aif.amtbbs.org\/index.php\/wp-json\/wp\/v2\/posts\/29992","targetHints":{"allow":["GET"]}}],"collection":[{"href":"https:\/\/aif.amtbbs.org\/index.php\/wp-json\/wp\/v2\/posts"}],"about":[{"href":"https:\/\/aif.amtbbs.org\/index.php\/wp-json\/wp\/v2\/types\/post"}],"author":[{"embeddable":true,"href":"https:\/\/aif.amtbbs.org\/index.php\/wp-json\/wp\/v2\/users\/56"}],"replies":[{"embeddable":true,"href":"https:\/\/aif.amtbbs.org\/index.php\/wp-json\/wp\/v2\/comments?post=29992"}],"version-history":[{"count":1,"href":"https:\/\/aif.amtbbs.org\/index.php\/wp-json\/wp\/v2\/posts\/29992\/revisions"}],"predecessor-version":[{"id":29994,"href":"https:\/\/aif.amtbbs.org\/index.php\/wp-json\/wp\/v2\/posts\/29992\/revisions\/29994"}],"wp:featuredmedia":[{"embeddable":true,"href":"https:\/\/aif.amtbbs.org\/index.php\/wp-json\/wp\/v2\/media\/29993"}],"wp:attachment":[{"href":"https:\/\/aif.amtbbs.org\/index.php\/wp-json\/wp\/v2\/media?parent=29992"}],"wp:term":[{"taxonomy":"category","embeddable":true,"href":"https:\/\/aif.amtbbs.org\/index.php\/wp-json\/wp\/v2\/categories?post=29992"},{"taxonomy":"post_tag","embeddable":true,"href":"https:\/\/aif.amtbbs.org\/index.php\/wp-json\/wp\/v2\/tags?post=29992"}],"curies":[{"name":"wp","href":"https:\/\/api.w.org\/{rel}","templated":true}]}}