[Example] Add MXFP8 blockscaled grouped gemm examples with transB support - #2098
Conversation
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
📝 WalkthroughWalkthroughAdds optional Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant Quantize as Quantize & Pack\n(FP16→FP8, SF)
participant Kernel as Grouped GEMM\nKernel (2CTA / persistent)
participant Verify as Verify & Benchmark
User->>Quantize: Provide A, B
Quantize->>Quantize: Quantize to FP8\nCompute & pack SF per K-block
Quantize->>Kernel: FP8 tensors, packed SF, offsets
Kernel->>Kernel: TMA load A tiles
Kernel->>Kernel: TMA load B tiles\n(apply transpose_B layout)
Kernel->>Kernel: Load/transpose SF into TMEM\nExecute blockscaled GEMM per group
Kernel->>Verify: Output C
Verify->>Verify: Unpack SF\nReference GEMM (fp32)\nCompute similarity & errors
Verify->>User: Verification results & perf
sequenceDiagram
participant Input as Input B
participant Check as Transpose Check
participant Layout as Shared-Mem Layout
participant TMA as TMA Copy
participant Compute as Compute Pipeline
Input->>Check: B tile
Check->>Check: transpose_B ?
alt transpose_B = true
Check->>Layout: Treat B as [N, K]\nreshape buffers
Layout->>TMA: Adjust TMA indexing for transposed tiles
else transpose_B = false
Check->>Layout: Treat B as [K, N]
Layout->>TMA: Standard TMA indexing
end
TMA->>Compute: Loaded B tile in\nconsistent compute layout
Compute->>Compute: Perform blockscaled MMA
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
examples/blockscaled_gemm_sm100/grouped_gemm_mxfp8_blockscaled_1d1d.py (1)
562-562: Unused unpacked variablem_total.
m_totalfroma.shapeisn't referenced inrun_grouped_mxfp8_blockscaled_gemm. Rename to_to match the pattern already used for the leading_on the next two lines and silence theRUF059hint.♻️ Proposed change
- m_total, k = a.shape + _, k = a.shape if transpose_B: _, n, k2 = b.shape else: _, k2, n = b.shape🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/blockscaled_gemm_sm100/grouped_gemm_mxfp8_blockscaled_1d1d.py` at line 562, In run_grouped_mxfp8_blockscaled_gemm the tuple unpacking uses m_total, k = a.shape but m_total is unused; change the unpack to use _ instead (i.e., _, k = a.shape) to silence the unused-variable warning (RUF059) and match the existing underscore usage elsewhere in the function.examples/blockscaled_gemm_sm100/gemm_mxfp8_blockscaled_1d1d.py (1)
26-26: Default oftranspose_Bdiffers from the new grouped example.Here
transpose_Bdefaults toFalse(and the same formxfp8_blockscaled_gemm_2cta/_2cta_persistent), but ingrouped_gemm_mxfp8_blockscaled_1d1d.pythe corresponding kernels default toTrue. The CLI parsers default toFalsein both files, so the discrepancy is hidden when invoked viamain(), but a direct call to the grouped kernel withouttranspose_B=will silently behave differently from the non-grouped kernel. Consider aligning the defaults (e.g., bothFalseto match the non-grouped, "B is[K, N]" convention) to reduce surprise.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/blockscaled_gemm_sm100/gemm_mxfp8_blockscaled_1d1d.py` at line 26, The grouped kernel defaults for transpose_B are inconsistent with the non-grouped kernels: change the default value of transpose_B in the grouped kernel definitions (the constructors / function signatures for the grouped mxfp8 blockscaled kernels such as the grouped variant of mxfp8_blockscaled_gemm_1d1d and any grouped versions of mxfp8_blockscaled_gemm_2cta / _2cta_persistent) to match the non‑grouped kernels (use transpose_B=False) so both grouped_gemm_mxfp8_blockscaled_1d1d and the non‑grouped mxfp8_blockscaled_gemm_* kernels follow the same "B is [K, N]" convention and avoid silent behavioral differences when callers omit the transpose_B argument.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@examples/blockscaled_gemm_sm100/grouped_gemm_mxfp8_blockscaled_1d1d.py`:
- Around line 92-97: When an expert's m_size is zero or clamped_pid_m would
point past the expert's actual blocks the kernel still issues TMA loads; add an
explicit runtime guard to skip the TMA+MMA pipeline for such empty/fully-clamped
tiles: compute m_size = offsets[eid+1]-offsets[eid], expert_m_blocks =
T.ceildiv(m_size, block_M) and set a boolean like use_tma = (m_size > 0) &&
(clamped_pid_m < expert_m_blocks) (or equivalently check tile_m < end_m) and
early-bail or route to a no-op path when false so the TMA descriptor and MMA are
not invoked for empty/partial groups (apply the same change at the symmetric
block around lines 292-299).
- Line 56: The dtype annotations use the string literal "int32" which is
inconsistent and causes F821; replace occurrences of the quoted type with the
typed reference T.int32 in the annotations (e.g., update the type for the
offsets tensor declaration and the matching annotation later in the file) so
they match other annotations like T.uint32 and silence the linter.
---
Nitpick comments:
In `@examples/blockscaled_gemm_sm100/gemm_mxfp8_blockscaled_1d1d.py`:
- Line 26: The grouped kernel defaults for transpose_B are inconsistent with the
non-grouped kernels: change the default value of transpose_B in the grouped
kernel definitions (the constructors / function signatures for the grouped mxfp8
blockscaled kernels such as the grouped variant of mxfp8_blockscaled_gemm_1d1d
and any grouped versions of mxfp8_blockscaled_gemm_2cta / _2cta_persistent) to
match the non‑grouped kernels (use transpose_B=False) so both
grouped_gemm_mxfp8_blockscaled_1d1d and the non‑grouped mxfp8_blockscaled_gemm_*
kernels follow the same "B is [K, N]" convention and avoid silent behavioral
differences when callers omit the transpose_B argument.
In `@examples/blockscaled_gemm_sm100/grouped_gemm_mxfp8_blockscaled_1d1d.py`:
- Line 562: In run_grouped_mxfp8_blockscaled_gemm the tuple unpacking uses
m_total, k = a.shape but m_total is unused; change the unpack to use _ instead
(i.e., _, k = a.shape) to silence the unused-variable warning (RUF059) and match
the existing underscore usage elsewhere in the function.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 9a3b230d-afa1-4d61-8c54-b7e5b97fc347
📒 Files selected for processing (2)
examples/blockscaled_gemm_sm100/gemm_mxfp8_blockscaled_1d1d.pyexamples/blockscaled_gemm_sm100/grouped_gemm_mxfp8_blockscaled_1d1d.py
| start_m = offsets[eid] | ||
| end_m = offsets[eid + 1] | ||
| m_size = end_m - start_m | ||
| expert_m_blocks = T.ceildiv(m_size, block_M) | ||
| clamped_pid_m = T.min(pid_m, T.max(expert_m_blocks, 1) - 1) | ||
| tile_m = start_m + clamped_pid_m * block_M |
There was a problem hiding this comment.
Empty / partial groups: TMA still issues loads at clamped tile_m.
For experts with m_size == 0 (or when pid_m_cluster * 2 + cta_id exceeds the expert's blocks), clamped_pid_m/safe_pid_m is forced to 0 and the kernel still runs the full TMA + MMA pipeline; only the epilogue's pid_m * block_M < m_size guard prevents stores. That's correct for the global-A bounds (TMA descriptor handles OOB) and avoids stale writes, but it does waste cycles for empty groups and reads cross-group rows for the last partial tile. The default batch sizes in main() are all multiples of block_M=128, so this is benign for the example, but worth documenting (or skipping the TMA path) for users who plug in arbitrary offsets.
Also applies to: 292-299
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@examples/blockscaled_gemm_sm100/grouped_gemm_mxfp8_blockscaled_1d1d.py`
around lines 92 - 97, When an expert's m_size is zero or clamped_pid_m would
point past the expert's actual blocks the kernel still issues TMA loads; add an
explicit runtime guard to skip the TMA+MMA pipeline for such empty/fully-clamped
tiles: compute m_size = offsets[eid+1]-offsets[eid], expert_m_blocks =
T.ceildiv(m_size, block_M) and set a boolean like use_tma = (m_size > 0) &&
(clamped_pid_m < expert_m_blocks) (or equivalently check tile_m < end_m) and
early-bail or route to a no-op path when false so the TMA descriptor and MMA are
not invoked for empty/partial groups (apply the same change at the symmetric
block around lines 292-299).
There was a problem hiding this comment.
🧹 Nitpick comments (3)
examples/blockscaled_gemm_sm100/mxfp8_illustrated.md (3)
8-8: Consider hyphenating "K-scale blocks" for clarity.The phrase "K scale blocks" should be hyphenated as "K-scale blocks" when used as a compound modifier.
📝 Proposed fix
-and B column per 128 K elements. Four adjacent K scale blocks are packed into +and B column per 128 K elements. Four adjacent K-scale blocks are packed into🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/blockscaled_gemm_sm100/mxfp8_illustrated.md` at line 8, Replace the unhyphenated compound "K scale blocks" with the hyphenated form "K-scale blocks" in the sentence containing "Four adjacent K scale blocks are packed into" (update the phrase to "Four adjacent K-scale blocks are packed into") so the compound modifier is clear and consistent.
54-54: Clarify the K=32 example reference.The text states "Take K=32, SFA as an example" but the actual kernels use
sf_granularity_k=128andblock_K=128. If this refers to a simplified example in the figure, consider clarifying that the figure shows a reduced example for illustration purposes to avoid confusion with the actual kernel parameters.📝 Suggested clarification
-Blackwell blockscaled tcgen05 MMA instructions require special layout for scale factors in TMEM. Take K=32, SFA as an example: +Blackwell blockscaled tcgen05 MMA instructions require special layout for scale factors in TMEM. The figure below shows a simplified example with K=32 for SFA:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/blockscaled_gemm_sm100/mxfp8_illustrated.md` at line 54, The sentence "Take K=32, SFA as an example" is misleading because the real kernels use sf_granularity_k=128 and block_K=128; update the phrasing in mxfp8_illustrated.md to state that the figure shows a simplified toy example (K=32) for illustration only and explicitly note the actual kernel parameters (sf_granularity_k=128, block_K=128) used in the implementation, e.g., change the sentence to clarify that K=32 is reduced for the diagram and point readers to sf_granularity_k and block_K for the real values.
104-104: Clarify the "128-column" terminology in the comment.The comment "two 128-column N chunks" is ambiguous. Consider clarifying to "two 128-word chunks" or "two chunks, each covering 128 N elements" to match the actual structure:
block_N=256yields 2 chunks of 128 words, with each chunk mapped to 4 TMEM columns (totaling 8 columns).📝 Suggested clarification
-SFB_tmem: [128 lanes, 8 columns] # two 128-column N chunks +SFB_tmem: [128 lanes, 8 columns] # two 128-word chunks (N=256 split into 2×128)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/blockscaled_gemm_sm100/mxfp8_illustrated.md` at line 104, Update the ambiguous comment "two 128-column N chunks" to explicitly state the unit and mapping — e.g., change SFB_tmem: [128 lanes, 8 columns] # two 128-word chunks (each covering 128 N elements; block_N=256 → 2 chunks of 128 words, each chunk mapped to 4 TMEM columns for a total of 8 columns) — so readers understand "128" refers to words/elements and how it maps to TMEM columns; locate and modify the comment near the SFB_tmem line to match this wording.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@examples/blockscaled_gemm_sm100/mxfp8_illustrated.md`:
- Line 8: Replace the unhyphenated compound "K scale blocks" with the hyphenated
form "K-scale blocks" in the sentence containing "Four adjacent K scale blocks
are packed into" (update the phrase to "Four adjacent K-scale blocks are packed
into") so the compound modifier is clear and consistent.
- Line 54: The sentence "Take K=32, SFA as an example" is misleading because the
real kernels use sf_granularity_k=128 and block_K=128; update the phrasing in
mxfp8_illustrated.md to state that the figure shows a simplified toy example
(K=32) for illustration only and explicitly note the actual kernel parameters
(sf_granularity_k=128, block_K=128) used in the implementation, e.g., change the
sentence to clarify that K=32 is reduced for the diagram and point readers to
sf_granularity_k and block_K for the real values.
- Line 104: Update the ambiguous comment "two 128-column N chunks" to explicitly
state the unit and mapping — e.g., change SFB_tmem: [128 lanes, 8 columns] #
two 128-word chunks (each covering 128 N elements; block_N=256 → 2 chunks of 128
words, each chunk mapped to 4 TMEM columns for a total of 8 columns) — so
readers understand "128" refers to words/elements and how it maps to TMEM
columns; locate and modify the comment near the SFB_tmem line to match this
wording.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 38a6a3f7-66ab-4de4-a720-3b4957f40217
⛔ Files ignored due to path filters (5)
examples/blockscaled_gemm_sm100/figures/blockscaled_data_path.svgis excluded by!**/*.svgexamples/blockscaled_gemm_sm100/figures/blockscaled_sf_layout.svgis excluded by!**/*.svgexamples/blockscaled_gemm_sm100/figures/blockscaled_variants.svgis excluded by!**/*.svgexamples/blockscaled_gemm_sm100/figures/blockscaled_warp_specialization.svgis excluded by!**/*.svgexamples/blockscaled_gemm_sm100/figures/sfa.pngis excluded by!**/*.png
📒 Files selected for processing (1)
examples/blockscaled_gemm_sm100/mxfp8_illustrated.md
There was a problem hiding this comment.
🧹 Nitpick comments (2)
examples/blockscaled_gemm_sm100/grouped_gemm_mxfp8_blockscaled_1d1d.py (2)
549-549: Drop the unusedm_totalbinding.
m_totalis unpacked but never used inrun_grouped_mxfp8_blockscaled_gemm(the kernel infers shapes itself). Renaming silences RuffRUF059.♻️ Proposed change
- m_total, k = a.shape + _, k = a.shape🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/blockscaled_gemm_sm100/grouped_gemm_mxfp8_blockscaled_1d1d.py` at line 549, The tuple unpacking "m_total, k = a.shape" in run_grouped_mxfp8_blockscaled_gemm binds an unused m_total; change the unpacking to ignore that value (e.g., use "_ , k = a.shape" or assign k via a.shape[1]) so m_total is not bound and the RUF059 warning is silenced while preserving k for subsequent use.
76-77: Optional: unify SF-TMEM shape expressions across the two kernels.The non-persistent kernel hardcodes
[block_M, 4]/[block_M, 8], while the persistent variant on lines 257–258 usesblock_M // 128 * 4/block_N // 128 * 4. Aligning them makes the relationshipcols = (tile / 128) * 4explicit in both places and slightly future-proofs against any tile-size relaxation of the existingblock_M == 128/block_N == 256asserts.♻️ Proposed change
- SFA_tmem = T.alloc_tmem([block_M, 4], "uint32") - SFB_tmem = T.alloc_tmem([block_M, 8], "uint32") + SFA_tmem = T.alloc_tmem([block_M, block_M // 128 * 4], "uint32") + SFB_tmem = T.alloc_tmem([block_M, block_N // 128 * 4], "uint32")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/blockscaled_gemm_sm100/grouped_gemm_mxfp8_blockscaled_1d1d.py` around lines 76 - 77, SFA_tmem and SFB_tmem use hardcoded shapes ([block_M, 4] / [block_M, 8]) while the persistent kernel computes shapes as block_M // 128 * 4 and block_N // 128 * 4; change the non-persistent alloc_tmem calls (SFA_tmem = T.alloc_tmem(...) and SFB_tmem = T.alloc_tmem(...)) to use the same expression pattern so cols = (tile / 128) * 4 is explicit (e.g., replace 4/8 with (block_M // 128) * 4 and (block_N // 128) * 4 or equivalent using block_M/block_N variables), keeping the expressions consistent with the persistent kernel and preserving existing assertions that block_M == 128 / block_N == 256.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@examples/blockscaled_gemm_sm100/grouped_gemm_mxfp8_blockscaled_1d1d.py`:
- Line 549: The tuple unpacking "m_total, k = a.shape" in
run_grouped_mxfp8_blockscaled_gemm binds an unused m_total; change the unpacking
to ignore that value (e.g., use "_ , k = a.shape" or assign k via a.shape[1]) so
m_total is not bound and the RUF059 warning is silenced while preserving k for
subsequent use.
- Around line 76-77: SFA_tmem and SFB_tmem use hardcoded shapes ([block_M, 4] /
[block_M, 8]) while the persistent kernel computes shapes as block_M // 128 * 4
and block_N // 128 * 4; change the non-persistent alloc_tmem calls (SFA_tmem =
T.alloc_tmem(...) and SFB_tmem = T.alloc_tmem(...)) to use the same expression
pattern so cols = (tile / 128) * 4 is explicit (e.g., replace 4/8 with (block_M
// 128) * 4 and (block_N // 128) * 4 or equivalent using block_M/block_N
variables), keeping the expressions consistent with the persistent kernel and
preserving existing assertions that block_M == 128 / block_N == 256.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 9128cdea-fd6a-42d5-a01b-54ae8af38e1f
📒 Files selected for processing (2)
examples/blockscaled_gemm_sm100/grouped_gemm_mxfp8_blockscaled_1d1d.pyexamples/blockscaled_gemm_sm100/mxfp8_illustrated.md
Summary by CodeRabbit
New Features
Documentation