[Blitz Decode] Integrate Embedding with H2D (#37913)

### Ticket
[Link to Github
Issue](#37662)

### Problem description
- We need to add the embedding op for Blitz Decode:
  - `[vocab_size, embedding_dim] = [129280, 7168]`
  - DRAM Interleaved with `page_size = embedding_dim * sizeof(bfloat16)`
- For this implementation, we decided to fuse the embedding lookup with
the H2D receiver

### What's changed
- Modify `HostInterface` to accept an embedding tensor as an optional
input
- If specified, the interface will be run using the new
`fused_h2d_receiver_embedding.cpp` kernel, which does a direct embedding
lookup from input tokens and sends the embedding to the downstream
receiver
- Bugfix: Support xfer sizes > `NOC_MAX_BURST_SIZE` in `HostInterface`.
The original issue was that the NOC read/write APIs used for interfacing
with PCIe only supported up to a single NOC packet. Utility APIs for
support arbitrary data sizes are added to
`models/demos/deepseek_v3_b1/micro_ops/host_io/kernels/pcie_noc_utils.h`
- Tested main use case for Blitz Decode and additional embedding tensor
shapes

### Checklist

- [ ] [![All post-commit
tests](https://github.com/tenstorrent/tt-metal/actions/workflows/all-post-commit-workflows.yaml/badge.svg?branch=asaigal/embedding_rebased)](https://github.com/tenstorrent/tt-metal/actions/workflows/all-post-commit-workflows.yaml?query=branch:asaigal/embedding_rebased)
- [ ] [![Blackhole Post
commit](https://github.com/tenstorrent/tt-metal/actions/workflows/blackhole-post-commit.yaml/badge.svg?branch=asaigal/embedding_rebased)](https://github.com/tenstorrent/tt-metal/actions/workflows/blackhole-post-commit.yaml?query=branch:asaigal/embedding_rebased)
- [ ]
[![cpp-unit-tests](https://github.com/tenstorrent/tt-metal/actions/workflows/tt-metal-l2-nightly.yaml/badge.svg?branch=asaigal/embedding_rebased)](https://github.com/tenstorrent/tt-metal/actions/workflows/tt-metal-l2-nightly.yaml?query=branch:asaigal/embedding_rebased)
- [ ] New/Existing tests provide coverage for changes


#### Model tests

If your changes cover model-related code, you should run tests
corresponding to affected models and platforms (Single card, T3K,
Galaxy). "Choose your pipeline" workflows facilitate running multiple
kinds of tests in a single run. Each offers `models-mandatory` and
`models-extended` presets.
The former includes a minimal set of tests, to be run always. The latter
extends that with additional ones - use your best judgement in deciding
which is the most appropriate for your PR.

- [ ] [![(Single) Choose your
pipeline](https://github.com/tenstorrent/tt-metal/actions/workflows/pipeline-select.yaml/badge.svg?branch=asaigal/embedding_rebased)](https://github.com/tenstorrent/tt-metal/actions/workflows/pipeline-select.yaml?query=branch:asaigal/embedding_rebased)
- [ ] `models-mandatory` preset (runs: [Device perf
regressions](https://github.com/tenstorrent/tt-metal/actions/workflows/perf-device-models.yaml)
and [Frequent model and ttnn
tests](https://github.com/tenstorrent/tt-metal/actions/workflows/fast-dispatch-full-regressions-and-models.yaml))
- [ ] `models-extended` preset (runs: the mandatory tests, plus
[Demo](https://github.com/tenstorrent/tt-metal/actions/workflows/single-card-demo-tests.yaml)
and [Model
perf](https://github.com/tenstorrent/tt-metal/actions/workflows/perf-models.yaml)
tests)
  - [ ] other selection - specify runs

- [ ] [![(T3K) Choose your
pipeline](https://github.com/tenstorrent/tt-metal/actions/workflows/pipeline-select-t3k.yaml/badge.svg?branch=asaigal/embedding_rebased)](https://github.com/tenstorrent/tt-metal/actions/workflows/pipeline-select-t3k.yaml?query=branch:asaigal/embedding_rebased)
- [ ] `models-mandatory` preset (runs: [Unit
tests](https://github.com/tenstorrent/tt-metal/actions/workflows/t3000-unit-tests.yaml))
- [ ] `models-extended` preset (runs: the mandatory tests, plus
[Demo](https://github.com/tenstorrent/tt-metal/actions/workflows/t3000-demo-tests.yaml)
and [Model
perf](https://github.com/tenstorrent/tt-metal/actions/workflows/t3000-model-perf-tests.yaml)
tests)
  - [ ] other selection - specify runs

- [ ] [![(Galaxy) Choose your
pipeline](https://github.com/tenstorrent/tt-metal/actions/workflows/pipeline-select-galaxy.yaml/badge.svg?branch=asaigal/embedding_rebased)](https://github.com/tenstorrent/tt-metal/actions/workflows/pipeline-select-galaxy.yaml?query=branch:asaigal/embedding_rebased)
- [ ] `models-mandatory` preset (runs: [Quick
tests](https://github.com/tenstorrent/tt-metal/actions/workflows/galaxy-quick.yaml))
- [ ] `models-extended` preset (runs: the mandatory tests, plus
[Demo](https://github.com/tenstorrent/tt-metal/actions/workflows/galaxy-demo-tests.yaml)
and [Model
perf](https://github.com/tenstorrent/tt-metal/actions/workflows/galaxy-perf-tests.yaml)
tests)
  - [ ] other selection - specify runs

Author: tt-asaigal

models/demos/deepseek_v3_b1/micro_ops/host_io/kernels/d2h_sender.cpp

@@ -4,7 +4,7 @@
 #include <cstdint>
 #include "api/dataflow/dataflow_api.h"
 #include "api/socket_api.h"
-#include "api/debug/dprint.h"
+#include "pcie_noc_utils.h"
 
 FORCE_INLINE bool socket_wait_for_pages_with_termination(
     const SocketReceiverInterface& socket, uint32_t num_pages, volatile tt_l1_ptr uint32_t* termination_semaphore) {
@@ -63,14 +63,13 @@ void kernel_main() {
                 break;
             }
             uint32_t read_addr = get_read_ptr(upstream_interface_index);
-            noc_wwrite_with_state<noc_mode, write_cmd_buf, CQ_NOC_SNDL, CQ_NOC_SEND, CQ_NOC_WAIT, true, false>(
+            noc_async_wide_write_any_len_with_state(
                 NOC_INDEX,
                 read_addr,
                 pcie_xy_enc,
                 ((static_cast<uint64_t>(write_addr_hi) << 32) | sender_socket.downstream_fifo_addr) +
                     sender_socket.write_ptr,
-                page_size,
-                1);
+                page_size);
             noc_async_writes_flushed();
             cb_pop_front(upstream_interface_index, 1);
         } else {
@@ -79,14 +78,13 @@ void kernel_main() {
                 break;
             }
             uint32_t read_addr = receiver_socket.read_ptr;
-            noc_wwrite_with_state<noc_mode, write_cmd_buf, CQ_NOC_SNDL, CQ_NOC_SEND, CQ_NOC_WAIT, true, false>(
+            noc_async_wide_write_any_len_with_state(
                 NOC_INDEX,
                 read_addr,
                 pcie_xy_enc,
                 ((static_cast<uint64_t>(write_addr_hi) << 32) | sender_socket.downstream_fifo_addr) +
                     sender_socket.write_ptr,
-                page_size,
-                1);
+                page_size);
             socket_pop_pages(receiver_socket, 1);
             noc_async_writes_flushed();
             socket_notify_sender(receiver_socket);
@@ -102,5 +100,4 @@ void kernel_main() {
 
     noc_async_write_barrier();
     noc_async_read_barrier();
-    DPRINT << "End D2H Main Loop" << ENDL();
 }

models/demos/deepseek_v3_b1/micro_ops/host_io/kernels/fused_h2d_receiver_embedding.cpp

@@ -0,0 +1,116 @@
+// SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+//
+// SPDX-License-Identifier: Apache-2.0
+#include <cstdint>
+#include "api/dataflow/dataflow_api.h"
+#include "api/socket_api.h"
+#include "api/tensor/tensor_accessor.h"
+#include "pcie_noc_utils.h"
+
+FORCE_INLINE bool socket_wait_for_pages_with_termination(
+    const SocketReceiverInterface& socket, uint32_t num_pages, volatile tt_l1_ptr uint32_t* termination_semaphore) {
+    constexpr uint32_t termination_value = 1;
+    while (!socket_wait_for_pages(socket, num_pages, 1000)) {
+        invalidate_l1_cache();
+        if (termination_semaphore[0] == termination_value) {
+            return false;
+        }
+    }
+    return true;
+}
+
+void kernel_main() {
+    // Get this value from MeshSocket struct on host
+    constexpr uint32_t recv_socket_config_addr = get_compile_time_arg_val(0);
+    constexpr uint32_t termination_semaphore_addr = get_compile_time_arg_val(1);
+    constexpr uint32_t token_page_size = get_compile_time_arg_val(2);
+    constexpr bool pull_from_host = get_compile_time_arg_val(3);
+    constexpr bool loopback_mode = get_compile_time_arg_val(4);
+    constexpr uint32_t downstream_interface_index = get_compile_time_arg_val(5);
+    constexpr uint32_t embedding_cb_index = get_compile_time_arg_val(6);
+    constexpr uint32_t embedding_page_size = get_compile_time_arg_val(7);
+    constexpr uint32_t embedding_addr = get_compile_time_arg_val(8);
+    // TensorAccessorArgs for embedding tensor at CT arg index 9
+    constexpr auto embedding_args = TensorAccessorArgs<9>();
+
+    auto embedding_accessor = TensorAccessor(embedding_args, embedding_addr, embedding_page_size);
+
+    SocketReceiverInterface receiver_socket = create_receiver_socket_interface(recv_socket_config_addr);
+    SocketSenderInterface sender_socket = {};
+
+    if constexpr (!loopback_mode) {
+        sender_socket = create_sender_socket_interface(downstream_interface_index);
+        set_sender_socket_page_size(sender_socket, embedding_page_size);
+    }
+    set_receiver_socket_page_size(receiver_socket, token_page_size);
+
+    // Read first page of embedding tensor from DRAM into CB
+
+    uint32_t read_addr_hi = receiver_socket.h2d.data_addr_hi;
+    uint32_t read_addr_lo = receiver_socket.h2d.data_addr_lo;
+    uint32_t pcie_xy_enc = receiver_socket.h2d.pcie_xy_enc;
+
+    volatile tt_l1_ptr uint32_t* termination_semaphore =
+        reinterpret_cast<volatile tt_l1_ptr uint32_t*>(termination_semaphore_addr);
+    while (true) {
+        // Wait for pages in H2D socket
+        if (!socket_wait_for_pages_with_termination(receiver_socket, 1, termination_semaphore)) {
+            break;
+        }
+        if constexpr (pull_from_host) {
+            // Pages available in H2D socket - read over PCIe
+            noc_async_wide_read_any_len_with_state(
+                NOC_INDEX,
+                pcie_xy_enc,
+                ((static_cast<uint64_t>(read_addr_hi) << 32) | read_addr_lo) + receiver_socket.read_ptr -
+                    receiver_socket.fifo_addr,
+                receiver_socket.read_ptr,
+                token_page_size);
+            noc_async_read_barrier();
+        }
+
+        // TODO: Add and assert that token id is within vocab size
+        volatile tt_l1_ptr uint32_t* token_id_ptr =
+            reinterpret_cast<volatile tt_l1_ptr uint32_t*>(receiver_socket.read_ptr);
+        // Embedding CB is a scratch pad for now. We only read into the first slot of the CB.
+        // TODO: Setup separate reader to pipeline reads.
+        uint32_t l1_write_addr = get_write_ptr(embedding_cb_index);
+        uint64_t noc_addr = embedding_accessor.get_noc_addr(*token_id_ptr);
+        noc_async_read(noc_addr, l1_write_addr, embedding_page_size);
+        noc_async_read_barrier();
+
+        if constexpr (loopback_mode) {
+            cb_reserve_back(downstream_interface_index, 1);
+            noc_async_write(
+                get_noc_addr(get_read_ptr(embedding_cb_index)),
+                get_noc_addr(get_write_ptr(downstream_interface_index)),
+                embedding_page_size);
+            noc_async_write_barrier();
+            cb_push_back(downstream_interface_index, 1);
+        } else {
+            sender_downstream_encoding downstream_enc = get_downstream_encoding(sender_socket, 0);
+            socket_reserve_pages(sender_socket, 1);
+            noc_async_write(
+                get_noc_addr(get_read_ptr(embedding_cb_index)),
+                get_noc_addr(
+                    downstream_enc.d2d.downstream_noc_x,
+                    downstream_enc.d2d.downstream_noc_y,
+                    sender_socket.write_ptr + sender_socket.downstream_fifo_addr),
+                embedding_page_size);
+            socket_push_pages(sender_socket, 1);
+            socket_notify_receiver(sender_socket);
+            noc_async_writes_flushed();
+        }
+        socket_pop_pages(receiver_socket, 1);
+        socket_notify_sender(receiver_socket);
+        invalidate_l1_cache();
+    }
+
+    update_socket_config(receiver_socket);
+    if constexpr (!loopback_mode) {
+        socket_barrier(sender_socket);
+    }
+
+    noc_async_write_barrier();
+    noc_async_read_barrier();
+}

models/demos/deepseek_v3_b1/micro_ops/host_io/kernels/h2d_receiver.cpp

@@ -4,6 +4,7 @@
 #include <cstdint>
 #include "api/dataflow/dataflow_api.h"
 #include "api/socket_api.h"
+#include "pcie_noc_utils.h"
 
 FORCE_INLINE bool socket_wait_for_pages_with_termination(
     const SocketReceiverInterface& socket, uint32_t num_pages, volatile tt_l1_ptr uint32_t* termination_semaphore) {
@@ -50,7 +51,7 @@ void kernel_main() {
         }
         if constexpr (pull_from_host) {
             // Pages available in H2D socket - read over PCIe
-            noc_read_with_state<noc_mode, read_cmd_buf, CQ_NOC_SNDL, CQ_NOC_SEND, CQ_NOC_WAIT>(
+            noc_async_wide_read_any_len_with_state(
                 NOC_INDEX,
                 pcie_xy_enc,
                 ((static_cast<uint64_t>(read_addr_hi) << 32) | read_addr_lo) + receiver_socket.read_ptr -
@@ -59,15 +60,16 @@ void kernel_main() {
                 page_size);
             noc_async_read_barrier();
         }
+
         if constexpr (loopback_mode) {
             cb_reserve_back(downstream_interface_index, 1);
             noc_async_write(
                 receiver_socket.read_ptr, get_noc_addr(get_write_ptr(downstream_interface_index)), page_size);
             noc_async_write_barrier();
             cb_push_back(downstream_interface_index, 1);
         } else {
-            socket_reserve_pages(sender_socket, 1);
             sender_downstream_encoding downstream_enc = get_downstream_encoding(sender_socket, 0);
+            socket_reserve_pages(sender_socket, 1);
             noc_async_write(
                 receiver_socket.read_ptr,
                 get_noc_addr(
@@ -92,5 +94,4 @@ void kernel_main() {
 
     noc_async_write_barrier();
     noc_async_read_barrier();
-    DPRINT << "End H2D Main Loop" << ENDL();
 }

models/demos/deepseek_v3_b1/micro_ops/host_io/kernels/pcie_noc_utils.h

@@ -0,0 +1,40 @@
+// SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+//
+// SPDX-License-Identifier: Apache-2.0
+#pragma once
+
+#include <cstdint>
+#include "api/dataflow/dataflow_api.h"
+
+// Contains utility functions to perform IO operations of variable length
+// over PCIe.
+
+// This implementation is currently not optimized to minimize RISC cycles.
+// APIs can be made more stateful, especially for the HostIO op, since the PCIe
+// NOC encoding is constant.
+
+FORCE_INLINE void noc_async_wide_write_any_len_with_state(
+    uint32_t noc, uint32_t src_addr, uint32_t dst_noc_addr, uint64_t dst_addr, uint32_t len_bytes) {
+    while (len_bytes > NOC_MAX_BURST_SIZE) {
+        noc_wwrite_with_state<noc_mode, write_cmd_buf, CQ_NOC_SNDL, CQ_NOC_SEND, CQ_NOC_WAIT, true, false>(
+            noc, src_addr, dst_noc_addr, dst_addr, NOC_MAX_BURST_SIZE, 1);
+        len_bytes -= NOC_MAX_BURST_SIZE;
+        src_addr += NOC_MAX_BURST_SIZE;
+        dst_addr += NOC_MAX_BURST_SIZE;
+    }
+    noc_wwrite_with_state<noc_mode, write_cmd_buf, CQ_NOC_SNDL, CQ_NOC_SEND, CQ_NOC_WAIT, true, false>(
+        noc, src_addr, dst_noc_addr, dst_addr, len_bytes, 1);
+}
+
+FORCE_INLINE void noc_async_wide_read_any_len_with_state(
+    uint32_t noc, uint64_t src_noc_encoding, uint64_t src_addr, uint32_t dst_addr, uint32_t len_bytes) {
+    while (len_bytes > NOC_MAX_BURST_SIZE) {
+        noc_read_with_state<noc_mode, read_cmd_buf, CQ_NOC_SNDL, CQ_NOC_SEND, CQ_NOC_WAIT>(
+            noc, src_noc_encoding, src_addr, dst_addr, NOC_MAX_BURST_SIZE);
+        len_bytes -= NOC_MAX_BURST_SIZE;
+        src_addr += NOC_MAX_BURST_SIZE;
+        dst_addr += NOC_MAX_BURST_SIZE;
+    }
+    noc_read_with_state<noc_mode, read_cmd_buf, CQ_NOC_SNDL, CQ_NOC_SEND, CQ_NOC_WAIT>(
+        noc, src_noc_encoding, src_addr, dst_addr, len_bytes);
+}